Use sgl-kernel sgl_per_token_group_quant_int8 (#4971)

2025-04-26 20:19:42 -07:00
parent bdbe5f816b
commit a086a11305
2 changed files with 39 additions and 2 deletions
@@ -755,6 +755,9 @@ def invoke_fused_moe_kernel(
        from sglang.srt.layers.quantization.fp8_kernel import (
            sglang_per_token_group_quant_fp8,
        )
+        from sglang.srt.layers.quantization.int8_kernel import (
+            sglang_per_token_group_quant_int8,
+        )
    else:
        from sglang.srt.layers.quantization.fp8_kernel import per_token_group_quant_fp8

@@ -794,7 +797,10 @@ def invoke_fused_moe_kernel(
            # activation block-wise int8 quantization
            assert len(block_shape) == 2
            block_n, block_k = block_shape[0], block_shape[1]
-            A, A_scale = per_token_group_quant_int8(A, block_k)
+            if _is_cuda:
+                A, A_scale = sglang_per_token_group_quant_int8(A, block_k)
+            else:
+                A, A_scale = per_token_group_quant_int8(A, block_k)
            assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
            assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2]
            assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1]