[Kernel Slimming] Migrate NVFP4 kernels to JIT (#19437)

python/sglang/srt/layers/moe/cutlass_moe.py

@@ -11,17 +11,20 @@ _is_cuda = is_cuda()
 if _is_cuda:
     from sgl_kernel import (
         apply_shuffle_mul_sum,
-        cutlass_fp4_group_mm,
         es_fp8_blockwise_scaled_grouped_mm,
         es_sm100_mxfp8_blockscaled_grouped_mm,
         es_sm100_mxfp8_blockscaled_grouped_quant,
         fp8_blockwise_scaled_grouped_mm,
         prepare_moe_input,
-        scaled_fp4_experts_quant,
         shuffle_rows,
         silu_and_mul,
     )
 
+    from sglang.jit_kernel.nvfp4 import (
+        cutlass_fp4_group_mm,
+        scaled_fp4_experts_quant,
+    )
+
 
 def cutlass_fused_experts_fp8(
     a: torch.Tensor,

python/sglang/srt/layers/moe/cutlass_moe_params.py

@@ -72,6 +72,11 @@ class CutlassMoEParams:
     # b_scales_ptrs: [e] dtype: int64
     a_scales_ptrs: torch.Tensor
     b_scales_ptrs: torch.Tensor
+    # Pointers for per-expert alpha values
+    alpha_ptrs: torch.Tensor
+    # CUTLASS blockscale layouts for A and B operands
+    layout_sfa: torch.Tensor
+    layout_sfb: torch.Tensor
 
     # Offsets that mark at which token index each expert begins its computation
     # The number of tokens computed with expert E is expert_offsets[E + 1] - expert_offsets[E]
@@ -139,6 +144,13 @@ class CutlassMoEParams:
         self.b_scales_ptrs = torch.empty(
             (self.e,), dtype=torch.int64, device=self.device
         )
+        self.alpha_ptrs = torch.empty((self.e,), dtype=torch.int64, device=self.device)
+        self.layout_sfa = torch.empty(
+            (self.e, 5), dtype=torch.int64, device=self.device
+        )
+        self.layout_sfb = torch.empty(
+            (self.e, 5), dtype=torch.int64, device=self.device
+        )
 
     def to_gemm1_args(self) -> dict:
         return {
@@ -147,11 +159,14 @@ class CutlassMoEParams:
             "problem_sizes": self.problem_sizes1,
             "expert_offsets": self.expert_offsets[:-1],
             "blockscale_offsets": self.blockscale_offsets[:-1],
-            #    "a_ptrs": self.a_ptrs,
-            #    "b_ptrs": self.b_ptrs,
-            #    "out_ptrs": self.out_ptrs,
-            #    "a_scales_ptrs": self.a_scales_ptrs,
-            #    "b_scales_ptrs": self.b_scales_ptrs,
+            "a_ptrs": self.a_ptrs,
+            "b_ptrs": self.b_ptrs,
+            "out_ptrs": self.out_ptrs,
+            "a_scales_ptrs": self.a_scales_ptrs,
+            "b_scales_ptrs": self.b_scales_ptrs,
+            "alpha_ptrs": self.alpha_ptrs,
+            "layout_sfa": self.layout_sfa,
+            "layout_sfb": self.layout_sfb,
         }
 
     def to_gemm2_args(self) -> dict:
@@ -161,9 +176,12 @@ class CutlassMoEParams:
             "problem_sizes": self.problem_sizes2,
             "expert_offsets": self.expert_offsets[:-1],
             "blockscale_offsets": self.blockscale_offsets[:-1],
-            #    "a_ptrs": self.a_ptrs,
-            #    "b_ptrs": self.b_ptrs,
-            #    "out_ptrs": self.out_ptrs,
-            #    "a_scales_ptrs": self.a_scales_ptrs,
-            #    "b_scales_ptrs": self.b_scales_ptrs,
+            "a_ptrs": self.a_ptrs,
+            "b_ptrs": self.b_ptrs,
+            "out_ptrs": self.out_ptrs,
+            "a_scales_ptrs": self.a_scales_ptrs,
+            "b_scales_ptrs": self.b_scales_ptrs,
+            "alpha_ptrs": self.alpha_ptrs,
+            "layout_sfa": self.layout_sfa,
+            "layout_sfb": self.layout_sfb,
         }

python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py

@@ -38,7 +38,7 @@ if TYPE_CHECKING:
 if is_flashinfer_available() and is_sm120_supported():
     from flashinfer import fp4_quantize
 elif is_cuda_alike():
-    from sgl_kernel import scaled_fp4_quant as fp4_quantize
+    from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
 else:
     fp4_quantize = None
 

python/sglang/srt/layers/moe/token_dispatcher/standard.py

@@ -43,7 +43,7 @@ try:
     if is_sm120_supported():
         from flashinfer import fp4_quantize
     else:
-        from sgl_kernel import scaled_fp4_quant as fp4_quantize
+        from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
 
     from flashinfer import fp4_quantize as fp4_quantize_flashinfer
 except ImportError:

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -74,9 +74,9 @@ try:
         try:
             from flashinfer import fp4_quantize
         except ImportError:
-            from sgl_kernel import scaled_fp4_quant as fp4_quantize
+            from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
     else:
-        from sgl_kernel import scaled_fp4_quant as fp4_quantize
+        from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
 except ImportError:
     fp4_quantize = None
 
@@ -87,7 +87,7 @@ try:
     enable_flashinfer_fp4_gemm = True
 except ImportError:
     if is_cuda():
-        from sgl_kernel import cutlass_scaled_fp4_mm as cutlass_fp4_gemm
+        from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm as cutlass_fp4_gemm
     enable_flashinfer_fp4_gemm = False
     reorder_rows_for_gated_act_gemm = None
     shuffle_matrix_a = None