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

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(