evan/ai-toolkit
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Publish libdevice as a module global so older triton builds without closure capture can compile the convrot kernels

Browse Source
This commit is contained in:
Jaret Burkett committed 2026-10-17 14:54:32 -06:00
1 parent 9d408eb64d
commit 577004a9cd
1 file changed
+12 -13
Download Patch File Download Diff File Expand all files Collapse all files
toolkit/util/convrot_quant.py
+12 -13
View File
@@ -293,13 +293,16 @@ def _import_triton():
variables through ``fn.__globals__`` only (no closure capture), so ``tl`` variables through ``fn.__globals__`` only (no closure capture), so ``tl``
referenced inside a kernel defined in a builder function must exist at referenced inside a kernel defined in a builder function must exist at
module scope or those versions die with ``NameError: 'tl' is not defined`` 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 as _triton
import triton.language as _tl import triton.language as _tl
from triton.language.extra import libdevice as _libdevice
globals()["triton"] = _triton globals()["triton"] = _triton
globals()["tl"] = _tl globals()["tl"] = _tl
return _triton, _tl globals()["libdevice"] = _libdevice
return _triton, _tl, _libdevice
_kernel = None _kernel = None
@@ -309,7 +312,7 @@ def _get_kernel():
global _kernel global _kernel
if _kernel is not None: if _kernel is not None:
return _kernel return _kernel
triton, tl = _import_triton() triton, tl, _ = _import_triton()
@triton.jit @triton.jit
def nvfp4_act_quant_kernel( def nvfp4_act_quant_kernel(
@@ -481,7 +484,7 @@ def _get_dequant_kernel():
global _dequant_kernel global _dequant_kernel
if _dequant_kernel is not None: if _dequant_kernel is not None:
return _dequant_kernel return _dequant_kernel
triton, tl = _import_triton() triton, tl, _ = _import_triton()
@triton.jit @triton.jit
def nvfp4_dequant_kernel( def nvfp4_dequant_kernel(
@@ -848,8 +851,7 @@ def _get_int8_kernels():
global _int8_kernels global _int8_kernels
if _int8_kernels is not None: if _int8_kernels is not None:
return _int8_kernels return _int8_kernels
triton, tl = _import_triton() triton, tl, libdevice = _import_triton()
from triton.language.extra import libdevice
@triton.jit @triton.jit
def int8_act_quant_kernel( def int8_act_quant_kernel(
@@ -1686,7 +1688,7 @@ def _get_intn_kernel():
global _intn_kernel global _intn_kernel
if _intn_kernel is not None: if _intn_kernel is not None:
return _intn_kernel return _intn_kernel
triton, tl = _import_triton() triton, tl, _ = _import_triton()
@triton.jit @triton.jit
def intn_unpack_kernel( def intn_unpack_kernel(
@@ -1729,8 +1731,7 @@ def _get_intn_grouped_kernel():
global _intn_grouped_kernel global _intn_grouped_kernel
if _intn_grouped_kernel is not None: if _intn_grouped_kernel is not None:
return _intn_grouped_kernel return _intn_grouped_kernel
triton, tl = _import_triton() triton, tl, libdevice = _import_triton()
from triton.language.extra import libdevice
@triton.jit @triton.jit
def intn_unpack_grouped_kernel( def intn_unpack_grouped_kernel(
@@ -1781,8 +1782,7 @@ def _get_bitnet_kernel():
global _bitnet_kernel global _bitnet_kernel
if _bitnet_kernel is not None: if _bitnet_kernel is not None:
return _bitnet_kernel return _bitnet_kernel
triton, tl = _import_triton() triton, tl, libdevice = _import_triton()
from triton.language.extra import libdevice
@triton.jit @triton.jit
def bitnet_unpack_kernel( def bitnet_unpack_kernel(
@@ -1912,8 +1912,7 @@ def _get_int_gemv_kernel():
global _int_gemv_kernel global _int_gemv_kernel
if _int_gemv_kernel is not None: if _int_gemv_kernel is not None:
return _int_gemv_kernel return _int_gemv_kernel
triton, tl = _import_triton() triton, tl, libdevice = _import_triton()
from triton.language.extra import libdevice
@triton.jit @triton.jit
def _unpack_lane( def _unpack_lane(