Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2

python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h

@@ -375,14 +375,11 @@ __global__ void Marlin(
       is_zp_float ? prob_n * prob_k / group_size / 8 : prob_n * prob_k / group_size / (pack_factor * 4);
   const int b_bias_expert_stride = prob_n / 8;
 
-  // parallel: num valid moe blocks
   int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
   int parallel = num_tokens_past_padded / moe_block_size;
   int num_valid_blocks = parallel;
-  if (is_ep) {
-    for (int i = 0; i < parallel; i++) {
-      if (expert_ids_ptr[i] == -1) num_valid_blocks--;
-    }
+  for (int i = 0; i < parallel; i++) {
+    if (expert_ids_ptr[i] == -1) num_valid_blocks--;
   }
   int num_invalid_blocks = parallel - num_valid_blocks;
   parallel = num_valid_blocks;

python/sglang/srt/arg_groups/overrides.py

@@ -1349,16 +1349,15 @@ def _deepseek_v4_kv_cache_dtype(view: Any) -> dict:
 
 @register_post_process
 def _deepseek_v4_sm120_moe(view: Any) -> dict:
-    """Slot pass in the DeepSeek V4 validation branch: SM120 lacks
-    tcgen05/TMEM, fall back to the marlin MoE runner (reads the
-    mid-resolution moe_runner_backend, after the dispatch-time nvfp4
-    default)."""
+    """Default DeepSeek V4 MXFP4 experts to FlashInfer CUTLASS on SM120."""
     hf_config = view.get_model_config().hf_config
     if hf_config.architectures[0] != "DeepseekV4ForCausalLM":
         return {}
     if is_sm120_supported() and view.moe_runner_backend == "auto":
-        logger.info("Use marlin as MoE runner backend on SM120 for DeepseekV4")
-        return {"moe_runner_backend": "marlin"}
+        logger.info(
+            "Use flashinfer_mxfp4 as MoE runner backend on SM120 for DeepseekV4"
+        )
+        return {"moe_runner_backend": "flashinfer_mxfp4"}
     return {}
 
 

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -1411,9 +1411,14 @@ def match_num_queries(x, value):
                     extra_indices.shape[-1] % 64 == 0
                 ), f"{extra_indices.shape=}'s last dimension is not aligned to 64"
 
-            if forward_batch.forward_mode.is_extend_without_speculative() and (
-                q.shape[0] > _LARGE_INDEXER_QUERY_THRESHOLD
-                or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get()
+            # sparse_prefill_fwd does not support SM120.
+            if (
+                forward_batch.forward_mode.is_extend_without_speculative()
+                and not _is_sm120
+                and (
+                    q.shape[0] > _LARGE_INDEXER_QUERY_THRESHOLD
+                    or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get()
+                )
             ):
                 return self._forward_prefill_sparse(
                     q=q,

python/sglang/srt/layers/attention/dsa/tilelang_kernel.py

@@ -1443,11 +1443,14 @@ def fp8_paged_mqa_logits(
             for j in T.Pipelined(n_iters, num_stages=2):
                 i = i_start + j
                 page = page_table[bx, i]
-                k_smem_u8 = T.alloc_shared((B * D,), UINT8)
-                T.copy(kvcache_u8[page, 0:SCALE_OFFSET], k_smem_u8)
+                k_smem_u8 = T.alloc_shared((1, B * D), UINT8)
+                T.copy(kvcache_u8[page : page + 1, 0:SCALE_OFFSET], k_smem_u8)
                 k_smem = T.view(k_smem_u8, (B, D), FP8)
-                k_s_smem_u8 = T.alloc_shared((B * 4,), UINT8)
-                T.copy(kvcache_u8[page, SCALE_OFFSET:BLOCK_BYTES], k_s_smem_u8)
+                k_s_smem_u8 = T.alloc_shared((1, B * 4), UINT8)
+                T.copy(
+                    kvcache_u8[page : page + 1, SCALE_OFFSET:BLOCK_BYTES],
+                    k_s_smem_u8,
+                )
                 k_s_smem = T.view(k_s_smem_u8, (B,), FP32)
                 k_s_frag = T.alloc_fragment((B,), FP32)
                 T.copy(k_s_smem, k_s_frag)

python/sglang/srt/layers/attention/dsv4/indexer.py

@@ -176,6 +176,25 @@ def fp8_paged_mqa_logits_torch_sm120(
     block_size = kvcache_fp8.shape[1]
     device = q_fp8.device
 
+    _QUERY_CHUNK = 1024
+    if batch_size > _QUERY_CHUNK:
+        return torch.cat(
+            [
+                fp8_paged_mqa_logits_torch_sm120(
+                    q_fp8[start : start + _QUERY_CHUNK],
+                    kvcache_fp8,
+                    weight[start : start + _QUERY_CHUNK],
+                    seq_lens[start : start + _QUERY_CHUNK],
+                    page_table[start : start + _QUERY_CHUNK],
+                    deep_gemm_metadata,
+                    max_seq_len,
+                    clean_logits=clean_logits,
+                )
+                for start in range(0, batch_size, _QUERY_CHUNK)
+            ],
+            dim=0,
+        )
+
     assert head_dim == 128, "Vectorized torch impl hardcodes DSV4 indexer head_dim=128"
     assert (
         block_size == 64

python/sglang/srt/layers/attention/flash_mla_sm120.py

@@ -400,14 +400,19 @@ def _flash_mla_flashinfer(
     extra_indices,
     extra_topk_length,
 ):
-    """FlashInfer SM120 sparse MLA via sparse_mla_sm120_decode_dsv4.
+    """FlashInfer SM120 sparse MLA via the paged-attention dispatcher.
 
     SGLang SWA pool uses page_size=256 (footer format: 256*576 bytes data + 256*8 bytes scale).
     FlashInfer decode_dsv4 fast path requires page_block_size=64 (footer: 64*576 + 64*8).
     We split 256-token pages into 4 virtual 64-token pages.
     Token indices are invariant under page-split (identity mapping).
     """
-    from flashinfer.mla._sparse_mla_sm120 import sparse_mla_sm120_decode_dsv4
+    from flashinfer.mla._sparse_mla_sm120 import (
+        _DECODE_MAX_TOKENS as _FI_DECODE_MAX_TOKENS,
+    )
+    from flashinfer.mla._sparse_mla_sm120 import (
+        _sparse_mla_sm120_paged_attention,
+    )
 
     B, _, H, D = q.shape  # (batch, 1, num_heads, head_dim)
     dev = q.device
@@ -435,32 +440,37 @@ def _flash_mla_flashinfer(
     output = torch.empty(B, H, head_dim_v, dtype=torch.bfloat16, device=dev)
     out_lse = torch.empty(B, H, dtype=torch.float32, device=dev)
 
-    # Pre-allocate split-K scratch for decode-dsv4 fast path.
-    topk = idx.shape[-1]
-    extra_topk = extra_idx.shape[-1] if extra_idx is not None else 0
-    _BI = 64
-    num_splits = (topk + _BI - 1) // _BI + (
-        (extra_topk + _BI - 1) // _BI if extra_topk > 0 else 0
-    )
-    mid_out = torch.empty(
-        B, H, num_splits, head_dim_v, dtype=torch.bfloat16, device=dev
-    )
-    mid_lse = torch.empty(B, H, num_splits, dtype=torch.float32, device=dev)
-
-    sparse_mla_sm120_decode_dsv4(
-        q=q.squeeze(1) if q.ndim == 4 else q,
-        kv_cache=kv_64,
-        indices=idx,
-        mid_out=mid_out,
-        mid_lse=mid_lse,
-        output=output,
-        out_lse=out_lse,
-        sm_scale=softmax_scale,
+    # Use split-K for decode-sized batches and paged attention otherwise.
+    if B <= _FI_DECODE_MAX_TOKENS:
+        topk = idx.shape[-1]
+        extra_topk = extra_idx.shape[-1] if extra_idx is not None else 0
+        _BI = 64
+        num_splits = (topk + _BI - 1) // _BI + (
+            (extra_topk + _BI - 1) // _BI if extra_topk > 0 else 0
+        )
+        mid_out = torch.empty(
+            B, H, num_splits, head_dim_v, dtype=torch.bfloat16, device=dev
+        )
+        mid_lse = torch.empty(B, H, num_splits, dtype=torch.float32, device=dev)
+    else:
+        mid_out = None
+        mid_lse = None
+
+    _sparse_mla_sm120_paged_attention(
+        q.squeeze(1) if q.ndim == 4 else q,
+        kv_64,
+        idx,
+        output,
+        out_lse,
+        softmax_scale,
+        d_v=head_dim_v,
         topk_length=topk_length,
         attn_sink=attn_sink,
         extra_kv_cache=extra_kv_64,
         extra_indices=extra_idx,
         extra_topk_length=extra_topk_length,
+        mid_out=mid_out,
+        mid_lse=mid_lse,
     )
 
     return (output.unsqueeze(1), None)

python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py

@@ -204,7 +204,8 @@ def fused_marlin_moe(
         device=hidden_states.device,
         dtype=hidden_states.dtype,
     )
-    intermediate_cache13 = torch.empty(
+    # Marlin skips masked expert rows, so their shared cache must start at zero.
+    intermediate_cache13 = torch.zeros(
         (M * topk_ids.shape[1] * max(gemm1_n, K),),
         device=hidden_states.device,
         dtype=hidden_states.dtype,

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

@@ -1,7 +1,7 @@
 """FlashInfer CUTLASS MoE fused funcs.
 
 This module owns the FlashInfer ``cutlass_fused_moe`` calls used by the
-unquantized, ModelOpt FP8, ModelOpt NVFP4, and SM90 MXFP4 MoE paths.
+unquantized, ModelOpt FP8, ModelOpt NVFP4, and MXFP4 MoE paths.
 Quantization methods prepare a small quant_info payload and route through
 ``MoeRunner``.
 """
@@ -62,27 +62,28 @@ class FlashInferCutlassMoeQuantInfo(MoeQuantInfo):
 
 @dataclass
 class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo):
-    """Quantization payload for the SM90 CUTLASS W4A16 MXFP4 MoE path.
+    """Quantization payload for CUTLASS MXFP4 MoE.
 
-    Weights and scales are pre-interleaved at load time via
-    ``interleave_moe_{weights,scales}_for_sm90_mixed_gemm``; this dataclass
-    only carries references plus the per-call routing/topology fields.
+    SM90 consumes W4A16-interleaved weights and scales. SM120 consumes packed
+    MXFP4 weights and block-interleaved scales with MXFP8 activations.
     """
 
-    # Pre-interleaved weights (uint8, packed FP4)
+    # SM90 weights are interleaved; SM120 weights remain checkpoint-packed.
     w13_weight: torch.Tensor  # [E, 2*N, K/2]
     w2_weight: torch.Tensor  # [E, K, N/2]
 
-    # Pre-interleaved E8M0 block scales (uint8; viewed as int32 at call time)
+    # E8M0 block scales in the layout selected by the quantization method.
     w13_weight_scale: torch.Tensor  # [E, 2*N, K/32]
     w2_weight_scale: torch.Tensor  # [E, K, N/32]
 
+    # A non-None global scale selects the SM120 MXFP8 activation path.
+    mxfp4_weight_global_scale: Optional[torch.Tensor] = None
+
     # Per-expert bias. GPT-OSS has both; DSv4 leaves both None.
     w13_bias: Optional[torch.Tensor] = None  # bf16 [E, 2*N]
     w2_bias: Optional[torch.Tensor] = None  # bf16 [E, K]
 
-    # Per-expert SwiGLU scalars (fp32 [E]). Either all three are present
-    # (clamped SwiGLU) or all three are None (kernel default SwiGLU).
+    # Optional per-expert SwiGLU overrides, fp32 [E].
     swiglu_alpha: Optional[torch.Tensor] = None
     swiglu_beta: Optional[torch.Tensor] = None
     swiglu_limit: Optional[torch.Tensor] = None
@@ -297,11 +298,7 @@ def fused_experts_none_to_flashinfer_mxfp4(
     quant_info: MoeQuantInfo,
     runner_config: MoeRunnerConfig,
 ) -> StandardCombineInput:
-    """SM90 W4A16 MXFP4 fused expert forward pass.
-
-    This preserves the ``flashinfer_mxfp4`` runner backend registration while
-    centralizing the CUTLASS execution in this module.
-    """
+    """Run the FlashInfer CUTLASS MXFP4 fused experts."""
     from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
     from sglang.srt.layers.moe.topk import TopKOutputChecker
 
@@ -335,22 +332,47 @@ def fused_experts_none_to_flashinfer_mxfp4(
             value=0.0,
         )
 
+    weight_global_scale = quant_info.mxfp4_weight_global_scale
+    use_mxfp8_act_scaling = weight_global_scale is not None
+    input_sf = None
+    fc1_expert_weights = quant_info.w13_weight
+    fc2_expert_weights = quant_info.w2_weight
+    if weight_global_scale is not None:
+        from flashinfer import mxfp8_quantize
+
+        x, input_sf = mxfp8_quantize(
+            x,
+            is_sf_swizzled_layout=True,
+            alignment=32,
+        )
+        fc1_expert_weights = fc1_expert_weights.view(torch.int64)
+        fc2_expert_weights = fc2_expert_weights.view(torch.int64)
+        quant_scales = [
+            quant_info.w13_weight_scale.view(torch.int32),
+            weight_global_scale,
+            quant_info.w2_weight_scale.view(torch.int32),
+            weight_global_scale,
+        ]
+    else:
+        quant_scales = [
+            quant_info.w13_weight_scale.view(torch.int32),
+            quant_info.w2_weight_scale.view(torch.int32),
+        ]
+
     out_hidden = padded_hidden if do_pad else origin_hidden
     output_dtype = torch.bfloat16
     with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()):
         out = torch.empty(x.shape[0], out_hidden, dtype=output_dtype, device=x.device)
 
     flashinfer_cutlass_fused_moe(
         input=x,
-        token_selected_experts=topk_ids.to(torch.int),
+        token_selected_experts=topk_ids.to(torch.int32),
         token_final_scales=topk_weights,
-        fc1_expert_weights=quant_info.w13_weight,
-        fc2_expert_weights=quant_info.w2_weight,
+        fc1_expert_weights=fc1_expert_weights,
+        fc2_expert_weights=fc2_expert_weights,
         output_dtype=output_dtype,
-        quant_scales=[
-            quant_info.w13_weight_scale.view(torch.int32),
-            quant_info.w2_weight_scale.view(torch.int32),
-        ],
+        quant_scales=quant_scales,
+        input_sf=input_sf,
         fc1_expert_biases=quant_info.w13_bias,
         fc2_expert_biases=quant_info.w2_bias,
         swiglu_alpha=quant_info.swiglu_alpha,
@@ -360,7 +382,8 @@ def fused_experts_none_to_flashinfer_mxfp4(
         tp_rank=quant_info.moe_tp_rank,
         ep_size=quant_info.moe_ep_size,
         ep_rank=quant_info.moe_ep_rank,
-        use_w4_group_scaling=True,
+        use_w4_group_scaling=not use_mxfp8_act_scaling,
+        use_mxfp8_act_scaling=use_mxfp8_act_scaling,
         activation_type=ActivationType.Swiglu,
         tune_max_num_tokens=next_power_of_2(x.shape[0]),
         output=out,

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

@@ -364,9 +364,10 @@ def get_quant_method(
                 return Mxfp4MarlinMoEMethod(fp8_method, prefix=prefix)
 
             if self.is_fp4_experts and get_moe_runner_backend().is_flashinfer_mxfp4():
-                # SM100 (Blackwell) -> trtllm-gen path.
-                # SM90  (Hopper)    -> cutlass mixed-input path (FlashInfer #3084).
-                if is_sm90_supported() and not is_sm100_supported():
+                # SM100 uses TRT-LLM; SM90 uses W4A16 and SM120 uses MXFP8xMXFP4.
+                if (
+                    is_sm90_supported() and not is_sm100_supported()
+                ) or is_sm120_supported():
                     from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import (
                         Mxfp4FlashinferCutlassMoEMethod,
                     )
@@ -970,6 +971,7 @@ def create_weights(
         intermediate_size_per_partition: int,
         params_dtype: torch.dtype,
         with_bias: bool = False,
+        fp4_scale_dtype: Optional[torch.dtype] = None,
         **extra_weight_attrs,
     ):
         self.with_bias = with_bias
@@ -1105,7 +1107,8 @@ def create_weights(
         # WEIGHT_SCALES
         if self.is_fp4_expert:
             fp4_block_k = 32
-            fp4_scale_dtype = torch.float8_e8m0fnu if _use_aiter else torch.float32
+            if fp4_scale_dtype is None:
+                fp4_scale_dtype = torch.float8_e8m0fnu if _use_aiter else torch.float32
             w13_weight_scale = torch.nn.Parameter(
                 torch.ones(
                     num_experts,

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

@@ -399,6 +399,11 @@ def _permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None:
             _permute_bias(w2_bias_data), requires_grad=False
         )
 
+    # Marlin uses the repacked scales; release the loader-format parameters.
+    for stale in ("w13_weight_scale_inv", "w2_weight_scale_inv"):
+        if hasattr(layer, stale):
+            delattr(layer, stale)
+
 
 def prepare_moe_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
     if layer.quant_config.group_size != 16: