Use Cute-DSL NVFP4 quantization kernels (#23745)

Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>

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

@@ -313,10 +313,9 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
     quant_info: CuteDslFp4MoeQuantInfo,
     runner_config: MoeRunnerConfig,
 ) -> StandardCombineInput:
-    from flashinfer import fp4_quantize
-
     from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
     from sglang.srt.layers.moe.topk import TopKOutputChecker
+    from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
 
     assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
 

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

@@ -49,7 +49,7 @@ if TYPE_CHECKING:
     )
 
 if is_flashinfer_available():
-    from flashinfer import fp4_quantize
+    from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
 elif is_cuda_alike():
     from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
 else:

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

@@ -25,11 +25,13 @@ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import get_int_env_var
 
 try:
-    from flashinfer import fp4_quantize, nvfp4_block_scale_interleave
+    from flashinfer import nvfp4_block_scale_interleave
     from flashinfer.comm import MoeAlltoAll, moe_a2a_get_workspace_size_per_rank
     from flashinfer.comm.mapping import Mapping
     from flashinfer.comm.mnnvl import MnnvlConfig
 
+    from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
+
     use_flashinfer = True
 except ImportError:
     use_flashinfer = False

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

@@ -44,10 +44,13 @@ if TYPE_CHECKING:
 
 
 try:
-    from flashinfer import fp4_quantize as fp4_quantize_flashinfer
     from flashinfer import (
         nvfp4_block_scale_interleave as nvfp4_block_scale_interleave_flashinfer,
     )
+
+    from sglang.srt.layers.quantization.modelopt_quant import (
+        fp4_quantize as fp4_quantize_flashinfer,
+    )
 except ImportError:
     fp4_quantize_flashinfer = None
     nvfp4_block_scale_interleave_flashinfer = None

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py

@@ -304,7 +304,9 @@ class CompressedTensorsW4A4Nvfp4MoE(CompressedTensorsMoEScheme):
         topk_output = dispatch_output.topk_output
 
         if self.use_flashinfer_trtllm:
-            from flashinfer import fp4_quantize, trtllm_fp4_block_scale_moe
+            from flashinfer import trtllm_fp4_block_scale_moe
+
+            from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
 
             router_logits = topk_output.router_logits
             topk_config = topk_output.topk_config

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

@@ -2,9 +2,12 @@ from __future__ import annotations
 
 import logging
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
+
+import torch
 
 from sglang.srt.utils.common import is_sm100_supported, is_sm120_supported
+from sglang.srt.utils.custom_op import register_custom_op_from_extern
 
 if TYPE_CHECKING:
     from sglang.srt.server_args import ServerArgs
@@ -12,6 +15,77 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+fp4_quantize = None
+try:
+    from flashinfer import fp4_quantize as _flashinfer_fp4_quantize
+
+    _flashinfer_fp4_quantize_backend = "cute-dsl" if is_sm100_supported() else "cuda"
+
+    def _round_up(x: int, y: int) -> int:
+        return ((x + y - 1) // y) * y
+
+    def _flashinfer_fp4_quantize_impl(
+        input: torch.Tensor,
+        global_scale: Optional[torch.Tensor] = None,
+        sf_vec_size: int = 16,
+        sf_use_ue8m0: bool = False,
+        is_sf_swizzled_layout: bool = True,
+        is_sf_8x4_layout: bool = False,
+        enable_pdl: Optional[bool] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return _flashinfer_fp4_quantize(
+            input,
+            global_scale,
+            sf_vec_size,
+            sf_use_ue8m0,
+            is_sf_swizzled_layout,
+            is_sf_8x4_layout,
+            enable_pdl,
+            backend=_flashinfer_fp4_quantize_backend,
+        )
+
+    def _flashinfer_fp4_quantize_fake(
+        input: torch.Tensor,
+        global_scale: Optional[torch.Tensor] = None,
+        sf_vec_size: int = 16,
+        sf_use_ue8m0: bool = False,
+        is_sf_swizzled_layout: bool = True,
+        is_sf_8x4_layout: bool = False,
+        enable_pdl: Optional[bool] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        is_column_major = input.stride(-2) == 1
+        if is_column_major:
+            m = input.shape[-1]
+            K = input.shape[-2]
+        else:
+            m = input.numel() // input.shape[-1]
+            K = input.shape[-1]
+        if is_column_major:
+            x_q = input.new_empty((*input.shape[:-2], K // 2, m), dtype=torch.uint8)
+        else:
+            x_q = input.new_empty((*input.shape[:-1], K // 2), dtype=torch.uint8)
+        if is_sf_swizzled_layout:
+            row_size = 8 if is_sf_8x4_layout else 128
+            sf_rows = _round_up(m, row_size)
+            sf_cols = _round_up(K // sf_vec_size, 4)
+        else:
+            sf_rows = m
+            sf_cols = K // sf_vec_size
+        if is_column_major:
+            sf = input.new_empty((sf_cols, sf_rows), dtype=torch.uint8)
+        else:
+            sf = input.new_empty((sf_rows, sf_cols), dtype=torch.uint8)
+        return x_q, sf
+
+    fp4_quantize = register_custom_op_from_extern(
+        _flashinfer_fp4_quantize_impl,
+        op_name="flashinfer_fp4_quantize",
+        fake_impl=_flashinfer_fp4_quantize_fake,
+    )
+except ImportError:
+    fp4_quantize = None
+
+
 class Fp4GemmRunnerBackend(Enum):
     """Enum for FP4 GEMM runner backend selection."""
 

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

@@ -35,7 +35,10 @@ from sglang.srt.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
-from sglang.srt.layers.quantization.fp4_utils import get_fp4_gemm_runner_backend
+from sglang.srt.layers.quantization.fp4_utils import (
+    fp4_quantize,
+    get_fp4_gemm_runner_backend,
+)
 from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.layers.quantization.fp8_utils import (
     apply_fp8_linear,
@@ -70,18 +73,6 @@ if TYPE_CHECKING:
     )
     from sglang.srt.models.utils import WeightsMapper
 
-fp4_quantize = None
-try:
-    if is_sm120_supported():
-        try:
-            from flashinfer import fp4_quantize
-        except ImportError:
-            from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
-    else:
-        from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize
-except ImportError:
-    fp4_quantize = None
-
 try:
     from flashinfer import mm_fp4 as flashinfer_fp4_gemm
     from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_sf_a