triton dispach fix

Author: BlackSamorez

inference_lib/setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="fp_quant",
-    version="0.1.5",
+    version="0.1.6",
     packages=find_packages(where="src"),
     package_dir={"": "src"},
     author="Andrei Panferov",
@@ -20,7 +20,6 @@
     install_requires=[
         "torch>=2.7.0",
         "scipy>=1.13.0",
-        "qutlass>=0.0.1",
         "triton>=3.3.0",
     ],
 )

inference_lib/src/fp_quant/module/triton/pseudoquant.py

@@ -215,18 +215,19 @@ def mxfp4_forward_kernel_wrapper(
     grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
 
     # Launch optimized kernel
-    mxfp4_forward_kernel[grid](
-        x_ptr=x,
-        hadamard_matrix_ptr=hadamard_matrix,
-        output_ptr=output,
-        clip_mask_ptr=clip_mask,
-        n_elements=n_elements,
-        hadamard_dim=hadamard_matrix.shape[-1],
-        group_size=32,
-        gaussian_scale=gaussian_scale,
-        stochastic_round=stochastic_round,
-        seed=seed,
-        quest=quest,
-    )
+    with torch.cuda.device(x.device):
+        mxfp4_forward_kernel[grid](
+            x_ptr=x,
+            hadamard_matrix_ptr=hadamard_matrix,
+            output_ptr=output,
+            clip_mask_ptr=clip_mask,
+            n_elements=n_elements,
+            hadamard_dim=hadamard_matrix.shape[-1],
+            group_size=32,
+            gaussian_scale=gaussian_scale,
+            stochastic_round=stochastic_round,
+            seed=seed,
+            quest=quest,
+        )
 
     return output, clip_mask