diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index 231d24e..6e116d6 100644 --- a/toolkit/util/convrot_quant.py +++ b/toolkit/util/convrot_quant.py @@ -293,13 +293,16 @@ def _import_triton(): variables through ``fn.__globals__`` only (no closure capture), so ``tl`` referenced inside a kernel defined in a builder function must exist at module scope or those versions die with ``NameError: 'tl' is not defined`` - at first compile.""" + at first compile. ``libdevice`` is published the same way for the kernels + that call ``libdevice.rint``.""" import triton as _triton import triton.language as _tl + from triton.language.extra import libdevice as _libdevice globals()["triton"] = _triton globals()["tl"] = _tl - return _triton, _tl + globals()["libdevice"] = _libdevice + return _triton, _tl, _libdevice _kernel = None @@ -309,7 +312,7 @@ def _get_kernel(): global _kernel if _kernel is not None: return _kernel - triton, tl = _import_triton() + triton, tl, _ = _import_triton() @triton.jit def nvfp4_act_quant_kernel( @@ -481,7 +484,7 @@ def _get_dequant_kernel(): global _dequant_kernel if _dequant_kernel is not None: return _dequant_kernel - triton, tl = _import_triton() + triton, tl, _ = _import_triton() @triton.jit def nvfp4_dequant_kernel( @@ -848,8 +851,7 @@ def _get_int8_kernels(): global _int8_kernels if _int8_kernels is not None: return _int8_kernels - triton, tl = _import_triton() - from triton.language.extra import libdevice + triton, tl, libdevice = _import_triton() @triton.jit def int8_act_quant_kernel( @@ -1686,7 +1688,7 @@ def _get_intn_kernel(): global _intn_kernel if _intn_kernel is not None: return _intn_kernel - triton, tl = _import_triton() + triton, tl, _ = _import_triton() @triton.jit def intn_unpack_kernel( @@ -1729,8 +1731,7 @@ def _get_intn_grouped_kernel(): global _intn_grouped_kernel if _intn_grouped_kernel is not None: return _intn_grouped_kernel - triton, tl = _import_triton() - from triton.language.extra import libdevice + triton, tl, libdevice = _import_triton() @triton.jit def intn_unpack_grouped_kernel( @@ -1781,8 +1782,7 @@ def _get_bitnet_kernel(): global _bitnet_kernel if _bitnet_kernel is not None: return _bitnet_kernel - triton, tl = _import_triton() - from triton.language.extra import libdevice + triton, tl, libdevice = _import_triton() @triton.jit def bitnet_unpack_kernel( @@ -1912,8 +1912,7 @@ def _get_int_gemv_kernel(): global _int_gemv_kernel if _int_gemv_kernel is not None: return _int_gemv_kernel - triton, tl = _import_triton() - from triton.language.extra import libdevice + triton, tl, libdevice = _import_triton() @triton.jit def _unpack_lane(