diff --git a/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx b/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx
index c524dab44..30348ff22 100644
--- a/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx
+++ b/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx
@@ -149,7 +149,7 @@ import { Playground } from "/src/snippets/_playground.jsx";
| DeepSeek-V4-Flash |
284B |
13B |
- single-node serving on B200 / B300 / GB200 / GB300 / H200 (TP=4); H100 (TP=8) |
+ single-node serving on B200 / B300 / GB200 / GB300 / H200 (TP=4); RTX PRO 6000 (TP=2); H100 (TP=8) |
| DeepSeek-V4-Pro |
@@ -293,9 +293,8 @@ TCP, which can lead to garbled KV transfer on large checkpoints.
**RTX PRO 6000 (SM120 / Blackwell Desktop) note**
-RTX PRO 6000 (96 GB) runs **Flash only** — V4-Pro doesn't fit on 8× 96 GB. It uses the
-**low-latency / TP-only** recipe (TP=4, single node) with the **Marlin** W4A16 MoE runner and
-`--mem-fraction-static 0.70`; the Deploy panel greys out the other recipes for this card.
+RTX PRO 6000 (96 GB) runs **Flash only** with the FlashInfer MXFP4 MoE runner.
+V4-Pro doesn't fit on 8× 96 GB; the Deploy panel greys out unsupported recipes.
HiCache and MegaMoE are **not** supported on RTX PRO 6000.
**AMD (MI300X / MI355X) note**
diff --git a/docs_new/src/snippets/configs/deepseek-ai/deepseek-v4.jsx b/docs_new/src/snippets/configs/deepseek-ai/deepseek-v4.jsx
index 1f56cdfb0..df4b109ef 100644
--- a/docs_new/src/snippets/configs/deepseek-ai/deepseek-v4.jsx
+++ b/docs_new/src/snippets/configs/deepseek-ai/deepseek-v4.jsx
@@ -1421,7 +1421,6 @@ sgl-eval run aime25 \\
// ====================================================================
// RTX PRO 6000 (SM120 / Blackwell Desktop) — Flash + low-latency only
- // (V4-Pro doesn't fit on 8× 96 GB); TP-only, Marlin MoE runner.
// ====================================================================
{
match: { hw: "rtx6000", variant: "flash", quant: "fp4", strategy: "low-latency", nodes: "single" },
@@ -1430,9 +1429,9 @@ sgl-eval run aime25 \\
flags: [
"--trust-remote-code",
"--model-path {{MODEL_NAME}}",
- "--tp 4",
- "--moe-runner-backend marlin",
- "--mem-fraction-static 0.70",
+ "--tp 2",
+ "--moe-runner-backend flashinfer_mxfp4",
+ "--mem-fraction-static 0.92",
"--cuda-graph-max-bs-decode 32",
"--host {{HOST_IP}}",
"--port {{PORT}}",
diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h b/python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h
index a06054990..cc2d15986 100644
--- a/python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h
+++ b/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;
diff --git a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py
index 2b6e1ad07..d2dc85132 100644
--- a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py
+++ b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py
@@ -1443,11 +1443,14 @@ def fp8_paged_mqa_logits_kernel(
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)
diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py
index 5bca1372a..81d9bfc17 100644
--- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py
+++ b/python/sglang/kernels/ops/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)
+ # 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_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,
+ _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)
diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py
index 1ef908905..c9de6dca6 100644
--- a/python/sglang/srt/arg_groups/overrides.py
+++ b/python/sglang/srt/arg_groups/overrides.py
@@ -1434,16 +1434,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 {}
diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index 6467c682d..306e60acc 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -1688,9 +1688,14 @@ class DeepseekV4AttnBackend(
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,
diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py
index 4fe15baf8..0ecd103ae 100644
--- a/python/sglang/srt/layers/attention/dsv4/indexer.py
+++ b/python/sglang/srt/layers/attention/dsv4/indexer.py
@@ -174,6 +174,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
diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py
index adf2e653d..7189b63fe 100644
--- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py
+++ b/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,
diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
index a57b5593b..cca9331c3 100644
--- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
+++ b/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,6 +332,33 @@ 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()):
@@ -342,15 +366,13 @@ def fused_experts_none_to_flashinfer_mxfp4(
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,
diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py
index fc920d84e..33c2f16f9 100644
--- a/python/sglang/srt/layers/quantization/fp8.py
+++ b/python/sglang/srt/layers/quantization/fp8.py
@@ -383,9 +383,8 @@ class Fp8Config(QuantizationConfig):
return Mxfp4HummingMoEMethod(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() or is_sm120_supported():
from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import (
Mxfp4FlashinferCutlassMoEMethod,
)
@@ -1053,6 +1052,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
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
@@ -1188,7 +1188,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
# 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,
diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py
index 44cd5181e..f1463ae23 100644
--- a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py
+++ b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py
@@ -399,6 +399,11 @@ def prepare_moe_mxfp4_layer_for_marlin(layer: torch.nn.Module) -> 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:
diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py
index 0cc330d8a..b79cb7cb5 100644
--- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py
+++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py
@@ -1,19 +1,7 @@
-"""DeepSeek-V4 MXFP4 expert backend backed by FlashInfer's SM90 cutlass
-mixed-input MoE GEMM (FlashInfer PR #3084).
+"""DeepSeek-V4 MXFP4 expert backend backed by FlashInfer CUTLASS MoE.
-Sibling of :class:`Mxfp4MarlinMoEMethod` and :class:`Mxfp4FlashinferTrtllmMoEMethod`.
-Wired into :func:`Fp8MoEConfig.get_quant_method` when
-``is_fp4_experts=True`` and ``--moe-runner-backend flashinfer_mxfp4`` is
-selected on a Hopper (SM90) device. SM100 still routes to
-:class:`Mxfp4FlashinferTrtllmMoEMethod` (trtllm-gen).
-
-Performance trade-off vs Marlin (kernel-level on H100, GPT-OSS-like body):
- - decode (M <= 64) : Marlin +12-15 %
- - tie (M ~= 256)
- - prefill (M >= 1024) : FlashInfer +24-36 %
-
-PD-disaggregated prefill workers are the natural fit; decode workers should
-keep the Marlin default.
+``Fp8Config`` selects this backend for SM90 and SM120; SM100 uses the
+TRT-LLM implementation.
"""
from __future__ import annotations
@@ -26,29 +14,12 @@ import torch
from torch.nn import Module
from torch.nn.parameter import Parameter
-from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.utils import is_flashinfer_available, log_info_on_rank0
+from sglang.srt.utils.common import is_sm120_supported
-# Silence the TRT-LLM cutlass autotune trace embedded inside FlashInfer's
-# cutlass_fused_moe. Its C++ logger reads TLLM_LOG_LEVEL on first kernel launch;
-# setdefault preserves any explicit user override.
+# Suppress TRT-LLM CUTLASS trace logs without overriding user configuration.
os.environ.setdefault("TLLM_LOG_LEVEL", "INFO")
-if is_flashinfer_available():
- try:
- from flashinfer.fused_moe import (
- interleave_moe_scales_for_sm90_mixed_gemm,
- interleave_moe_weights_for_sm90_mixed_gemm,
- )
-
- _FI_HAS_SM90_CUTLASS_MXFP4 = True
- except ImportError:
- interleave_moe_scales_for_sm90_mixed_gemm = None
- interleave_moe_weights_for_sm90_mixed_gemm = None
- _FI_HAS_SM90_CUTLASS_MXFP4 = False
-else:
- _FI_HAS_SM90_CUTLASS_MXFP4 = False
-
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
@@ -59,26 +30,21 @@ _GROUP_SIZE = 32
class Mxfp4FlashinferCutlassMoEMethod:
- """DeepSeek-V4 W4A16 MXFP4 MoE via FlashInfer's SM90 mixed-input cutlass
- grouped GEMM. The fused kernel does GEMM1 + clamped SwiGLU + GEMM2 in one
- call after a one-shot weight/scale interleave at load time."""
+ """FlashInfer MXFP4 MoE: W4A16 on SM90 and W4A8 on SM120."""
def __init__(self, fp8_method, prefix: str):
- if not _FI_HAS_SM90_CUTLASS_MXFP4:
- raise RuntimeError(
- "Mxfp4FlashinferCutlassMoEMethod requires FlashInfer >= 0.6.11 "
- "(PR #3084 SM90 mixed-input helpers). Older builds lack "
- "interleave_moe_{weights,scales}_for_sm90_mixed_gemm; "
- "either upgrade flashinfer-python or fall back to "
- "--moe-runner-backend marlin."
- )
+ if not is_flashinfer_available():
+ raise RuntimeError("Mxfp4FlashinferCutlassMoEMethod requires FlashInfer.")
+ self._use_mxfp8_act_scaling = is_sm120_supported()
self._fp8 = fp8_method
self.prefix = prefix
- self._swiglu_alpha_tensor: torch.Tensor | None = None
- self._swiglu_beta_tensor: torch.Tensor | None = None
self._swiglu_limit_tensor: torch.Tensor | None = None
+ self._mxfp4_weight_global_scale_tensor: torch.Tensor | None = None
- # --- Lifecycle ---------------------------------------------------------
+ @property
+ def load_up_proj_weight_first(self) -> bool:
+ """Load W13 directly as ``[up; gate]`` for FlashInfer CUTLASS."""
+ return True
def create_weights(
self,
@@ -89,11 +55,7 @@ class Mxfp4FlashinferCutlassMoEMethod:
params_dtype,
**extra_weight_attrs,
):
- # SM90 mixed-input GEMM: contraction dim K must be a multiple of 128
- # (interleave factor = 128 / group_size = 4). For DSv4 (hidden=7168,
- # inter=2048) both are already multiples of 128; we assert rather than
- # silently pad here, since padding the FP8-base buffers in-place would
- # require deeper changes.
+ # Both CUTLASS paths require dimensions aligned to 128.
if hidden_size % 128 != 0 or intermediate_size_per_partition % 128 != 0:
raise ValueError(
"Mxfp4FlashinferCutlassMoEMethod requires hidden_size and "
@@ -101,14 +63,14 @@ class Mxfp4FlashinferCutlassMoEMethod:
f"(got hidden={hidden_size}, "
f"intermediate={intermediate_size_per_partition})."
)
- # Raw weight shapes match what the fp8 base method allocates for fp4
- # experts (uint8 4-bit packed weights, fp32 E8M0 scales). Delegate.
+ # Keep checkpoint scales in native E8M0 instead of staging them as FP32.
self._fp8.create_weights(
layer,
num_experts,
hidden_size,
intermediate_size_per_partition,
params_dtype,
+ fp4_scale_dtype=torch.float8_e8m0fnu,
**extra_weight_attrs,
)
@@ -118,29 +80,21 @@ class Mxfp4FlashinferCutlassMoEMethod:
self.moe_runner_config = moe_runner_config
- # DSv4 uses standard SwiGLU plus a config-driven activation clamp.
- # We pass all three (alpha, beta, limit) as explicit per-expert tensors
- # rather than mixing tensors with None: the cutlass SwiGLU kernel
- # branches on whether each is None, and partial-None inputs land in
- # less-tested code paths. ``alpha=1.0``, ``beta=0.0`` reproduce plain
- # ``silu(gate) * up``; ``limit`` enforces the activation clamp the
- # checkpoint was trained with.
+ E = layer.num_local_experts
+ device = layer.w13_weight.device
+ if self._use_mxfp8_act_scaling:
+ # FlashInfer's MXFP4 ABI requires a neutral per-expert global scale.
+ self._mxfp4_weight_global_scale_tensor = torch.ones(
+ E, dtype=torch.float32, device=device
+ )
+
+ # FlashInfer defaults alpha/beta to 1/0, so DSv4 only supplies its clamp.
swiglu_limit = getattr(moe_runner_config, "swiglu_limit", None)
if swiglu_limit is not None:
- E = layer.num_local_experts
- device = layer.w13_weight.device
- self._swiglu_alpha_tensor = torch.ones(
- E, dtype=torch.float32, device=device
- )
- self._swiglu_beta_tensor = torch.zeros(
- E, dtype=torch.float32, device=device
- )
self._swiglu_limit_tensor = torch.full(
(E,), float(swiglu_limit), dtype=torch.float32, device=device
)
else:
- self._swiglu_alpha_tensor = None
- self._swiglu_beta_tensor = None
self._swiglu_limit_tensor = None
# Register the fused func at runner construction so the FusedOpPool
@@ -150,67 +104,71 @@ class Mxfp4FlashinferCutlassMoEMethod:
self.runner = MoeRunner(MoeRunnerBackend.FLASHINFER_MXFP4, moe_runner_config)
def process_weights_after_loading(self, layer: Module) -> None:
- from sglang.srt.layers.quantization.utils import reorder_w1w3_to_w3w1
-
- # Run the fp8 base hook first (ROCm normalization, mxfp8 requant, ...).
+ # Preserve the base FP4 post-load handling.
self._fp8.process_weights_after_loading(layer)
if getattr(layer, "_mega_moe_weights_built", False):
return
- # cutlass_fused_moe expects fc1 in [w3; w1] = [up; gate] order, just
- # like the trtllm-gen path. The HF / FP8 loader emits [w1; w3].
- w13, w13_s = reorder_w1w3_to_w3w1(
- layer.w13_weight.data, layer.w13_weight_scale_inv.data
- )
- layer.w13_weight = Parameter(w13, requires_grad=False)
- layer.w13_weight_scale_inv = Parameter(w13_s, requires_grad=False)
-
+ arch = "SM120" if self._use_mxfp8_act_scaling else "SM90"
log_info_on_rank0(
logger,
- f"Preparing DSv4 MXFP4 experts for FlashInfer SM90 cutlass "
+ f"Preparing DSv4 MXFP4 experts for FlashInfer {arch} CUTLASS "
f"(layer: {self.prefix})...",
)
- # FP8 base stores scales as fp32 numerical values (= 2**e). The
- # FlashInfer SM90 helper reads raw E8M0 bytes (uint8 with the
- # exponent + 127 bias). Cast through float8_e8m0fnu to extract the
- # raw byte without losing the exponent.
- w13_scale_u8 = (
- layer.w13_weight_scale_inv.data.to(torch.float8_e8m0fnu)
- .view(torch.uint8)
- .contiguous()
- )
- w2_scale_u8 = (
- layer.w2_weight_scale_inv.data.to(torch.float8_e8m0fnu)
- .view(torch.uint8)
- .contiguous()
- )
+ # FlashInfer consumes the raw bytes of the checkpoint's E8M0 scales.
+ for name in ("w13_weight_scale_inv", "w2_weight_scale_inv"):
+ scale = getattr(layer, name)
+ if scale.dtype != torch.float8_e8m0fnu:
+ raise TypeError(
+ f"{name} must remain native E8M0 for FlashInfer MXFP4, "
+ f"got {scale.dtype}."
+ )
+ w13_scale_u8 = layer.w13_weight_scale_inv.data.view(torch.uint8)
+ w2_scale_u8 = layer.w2_weight_scale_inv.data.view(torch.uint8)
- # C++ byte interleave on packed 4-bit weights.
- w13_il = interleave_moe_weights_for_sm90_mixed_gemm(
- layer.w13_weight.data.view(torch.uint8).contiguous(), "fp4"
- )
- w2_il = interleave_moe_weights_for_sm90_mixed_gemm(
- layer.w2_weight.data.view(torch.uint8).contiguous(), "fp4"
- )
- # Pure-PyTorch reshape+permute on E8M0 block scales.
- w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
- w13_scale_u8, group_size=_GROUP_SIZE
- )
- w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
- w2_scale_u8, group_size=_GROUP_SIZE
- )
+ if self._use_mxfp8_act_scaling:
+ from flashinfer import block_scale_interleave
- layer.w13_weight = Parameter(w13_il, requires_grad=False)
- layer.w2_weight = Parameter(w2_il, requires_grad=False)
- layer.w13_weight_scale_inv = Parameter(w13_s_il, requires_grad=False)
- layer.w2_weight_scale_inv = Parameter(w2_s_il, requires_grad=False)
+ if (
+ not layer.w13_weight.is_contiguous()
+ or not layer.w2_weight.is_contiguous()
+ ):
+ raise ValueError("SM120 FlashInfer MXFP4 weights must be contiguous.")
+ for scale_u8 in (w13_scale_u8, w2_scale_u8):
+ scale_u8.copy_(block_scale_interleave(scale_u8).reshape_as(scale_u8))
+ else:
+ from flashinfer.fused_moe import (
+ interleave_moe_scales_for_sm90_mixed_gemm,
+ interleave_moe_weights_for_sm90_mixed_gemm,
+ )
- layer._dsv4_mxfp4_backend = "flashinfer_cutlass_sm90"
- torch.cuda.empty_cache()
+ w13_il = interleave_moe_weights_for_sm90_mixed_gemm(
+ layer.w13_weight.data.view(torch.uint8).contiguous(), "fp4"
+ )
+ w2_il = interleave_moe_weights_for_sm90_mixed_gemm(
+ layer.w2_weight.data.view(torch.uint8).contiguous(), "fp4"
+ )
+ w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
+ w13_scale_u8, group_size=_GROUP_SIZE
+ )
+ w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
+ w2_scale_u8, group_size=_GROUP_SIZE
+ )
+ layer.w13_weight = Parameter(w13_il, requires_grad=False)
+ layer.w2_weight = Parameter(w2_il, requires_grad=False)
+ layer.w13_weight_scale_inv = Parameter(w13_s_il, requires_grad=False)
+ layer.w2_weight_scale_inv = Parameter(w2_s_il, requires_grad=False)
- # --- Forward -----------------------------------------------------------
+ layer._dsv4_mxfp4_backend = (
+ "flashinfer_cutlass_sm120"
+ if self._use_mxfp8_act_scaling
+ else "flashinfer_cutlass_sm90"
+ )
+ # SM90 creates full-size interleaved copies; release old layouts per layer.
+ if not self._use_mxfp8_act_scaling:
+ torch.cuda.empty_cache()
def apply(
self,
@@ -221,26 +179,21 @@ class Mxfp4FlashinferCutlassMoEMethod:
FlashInferCutlassMxfp4MoeQuantInfo,
)
- # DSv4 always feeds StandardDispatchOutput; the fused func tolerates
- # bypassed too but we keep the strict check here as a contract guard.
- topk_output = dispatch_output.topk_output
- if not TopKOutputChecker.format_is_standard(topk_output):
- raise ValueError(f"Unsupported topk output format: {topk_output.format}")
-
quant_info = FlashInferCutlassMxfp4MoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
w13_weight_scale=layer.w13_weight_scale_inv,
w2_weight_scale=layer.w2_weight_scale_inv,
- w13_bias=None, # DSv4 has no MoE expert bias.
+ mxfp4_weight_global_scale=self._mxfp4_weight_global_scale_tensor,
+ w13_bias=None,
w2_bias=None,
- swiglu_alpha=self._swiglu_alpha_tensor, # ones: standard SiLU gate
- swiglu_beta=self._swiglu_beta_tensor, # zeros: standard up
+ swiglu_alpha=None,
+ swiglu_beta=None,
swiglu_limit=self._swiglu_limit_tensor,
moe_tp_size=layer.moe_tp_size,
moe_tp_rank=layer.moe_tp_rank,
moe_ep_size=layer.moe_ep_size,
moe_ep_rank=layer.moe_ep_rank,
- padded_hidden=None, # DSv4 hidden_size is already a multiple of 128.
+ padded_hidden=None,
)
return self.runner.run(dispatch_output, quant_info)
diff --git a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py
index b0391fd27..3cdf7cbdc 100644
--- a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py
+++ b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py
@@ -71,21 +71,23 @@ class Mxfp4MarlinMoEMethod:
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
+ # Store loader scales in E8M0; uint8 127 encodes 1.0.
+ def _e8m0_ones(*shape: int) -> torch.Tensor:
+ return torch.full(shape, 127, dtype=torch.uint8).view(torch.float8_e8m0fnu)
+
w13_weight_scale = torch.nn.Parameter(
- torch.ones(
+ _e8m0_ones(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // fp4_block_k,
- dtype=torch.float32,
),
requires_grad=False,
)
w2_weight_scale = torch.nn.Parameter(
- torch.ones(
+ _e8m0_ones(
num_experts,
hidden_size,
intermediate_size_per_partition // fp4_block_k,
- dtype=torch.float32,
),
requires_grad=False,
)
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index d31d2f419..a260605de 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -4470,6 +4470,8 @@ class ServerArgs:
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.set(False)
envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False)
envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.set(True)
+ # Prefer TileLang over the Torch fallback.
+ envs.SGLANG_OPT_USE_TILELANG_INDEXER.set(True)
elif is_hip():
envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False)
envs.SGLANG_OPT_USE_FUSED_COMPRESS.set(True)
diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
new file mode 100644
index 000000000..38079a140
--- /dev/null
+++ b/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
@@ -0,0 +1,251 @@
+"""SM120 FlashInfer MXFP8-by-MXFP4 MoE integration test."""
+
+from __future__ import annotations
+
+import builtins
+import importlib
+import sys
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
+
+
+def _random_weights(num_experts: int, hidden: int, intermediate: int):
+ generator = torch.Generator(device="cuda").manual_seed(0)
+ w13 = torch.randint(
+ -128,
+ 128,
+ (num_experts, 2 * intermediate, hidden // 2),
+ dtype=torch.int8,
+ device="cuda",
+ generator=generator,
+ )
+ w2 = torch.randint(
+ -128,
+ 128,
+ (num_experts, hidden, intermediate // 2),
+ dtype=torch.int8,
+ device="cuda",
+ generator=generator,
+ )
+ w13_scale_u8 = torch.randint(
+ 125,
+ 130,
+ (num_experts, 2 * intermediate, hidden // 32),
+ dtype=torch.uint8,
+ device="cuda",
+ generator=generator,
+ )
+ w2_scale_u8 = torch.randint(
+ 125,
+ 130,
+ (num_experts, hidden, intermediate // 32),
+ dtype=torch.uint8,
+ device="cuda",
+ generator=generator,
+ )
+ return (
+ w13,
+ w2,
+ w13_scale_u8.view(torch.float8_e8m0fnu),
+ w2_scale_u8.view(torch.float8_e8m0fnu),
+ )
+
+
+def test_cutlass_adapter_import_does_not_require_flashinfer(monkeypatch):
+ module_name = "sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe"
+ # Load the package before blocking FlashInfer so this test isolates the
+ # adapter import exercised by non-CUDA backends.
+ importlib.import_module("sglang.srt.layers.quantization")
+ cached_module = sys.modules.pop(module_name, None)
+ real_import = builtins.__import__
+
+ def import_without_flashinfer(name, *args, **kwargs):
+ if name == "flashinfer" or name.startswith("flashinfer."):
+ raise ModuleNotFoundError("No module named 'flashinfer'")
+ return real_import(name, *args, **kwargs)
+
+ monkeypatch.setattr(builtins, "__import__", import_without_flashinfer)
+ try:
+ module = importlib.import_module(module_name)
+ assert hasattr(module, "Mxfp4FlashinferCutlassMoEMethod")
+ finally:
+ sys.modules.pop(module_name, None)
+ if cached_module is not None:
+ sys.modules[module_name] = cached_module
+
+
+def test_dsv4_sm120_load_contract(monkeypatch):
+ import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as adapter_module
+
+ monkeypatch.setattr(adapter_module, "is_sm120_supported", lambda: True)
+
+ captured = {}
+
+ class _Fp8Method:
+ def create_weights(self, *args, **kwargs):
+ captured.update(kwargs)
+
+ method = adapter_module.Mxfp4FlashinferCutlassMoEMethod(_Fp8Method(), "test")
+ method.create_weights(
+ SimpleNamespace(),
+ num_experts=4,
+ hidden_size=256,
+ intermediate_size_per_partition=256,
+ params_dtype=torch.bfloat16,
+ )
+
+ assert method.load_up_proj_weight_first
+ assert captured["fp4_scale_dtype"] == torch.float8_e8m0fnu
+
+
+def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
+ if not torch.cuda.is_available():
+ pytest.skip("CUDA required")
+ if torch.cuda.get_device_capability()[0] != 12:
+ pytest.skip("SM120 required")
+ pytest.importorskip("flashinfer.fused_moe")
+
+ from flashinfer import block_scale_interleave, mxfp8_quantize
+ from flashinfer.fused_moe import cutlass_fused_moe
+ from flashinfer.fused_moe.core import ActivationType
+
+ import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as runner_module
+ from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+ from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
+ from sglang.srt.layers.moe.topk import StandardTopKOutput
+ from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import (
+ Mxfp4FlashinferCutlassMoEMethod,
+ )
+
+ monkeypatch.setattr(
+ runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext()
+ )
+ monkeypatch.setattr(runner_module, "is_allocation_symmetric", lambda: False)
+ monkeypatch.setattr(runner_module, "get_tp_group", lambda: None)
+
+ num_experts, hidden, intermediate = 4, 256, 256
+ w13, w2, w13_scale, w2_scale = _random_weights(num_experts, hidden, intermediate)
+ w1, w3 = w13.chunk(2, dim=1)
+ w1_scale, w3_scale = w13_scale.chunk(2, dim=1)
+ # Simulate FusedMoE's ``load_up_proj_weight_first`` loader contract.
+ w31 = torch.cat((w3, w1), dim=1)
+ w31_scale = torch.cat(
+ (w3_scale.view(torch.uint8), w1_scale.view(torch.uint8)),
+ dim=1,
+ ).view(torch.float8_e8m0fnu)
+ layer = SimpleNamespace(
+ w13_weight=torch.nn.Parameter(w31.clone(), requires_grad=False),
+ w2_weight=torch.nn.Parameter(w2.clone(), requires_grad=False),
+ w13_weight_scale_inv=torch.nn.Parameter(w31_scale.clone(), requires_grad=False),
+ w2_weight_scale_inv=torch.nn.Parameter(w2_scale.clone(), requires_grad=False),
+ num_local_experts=num_experts,
+ moe_tp_size=1,
+ moe_tp_rank=0,
+ moe_ep_size=1,
+ moe_ep_rank=0,
+ )
+
+ method = Mxfp4FlashinferCutlassMoEMethod(
+ SimpleNamespace(process_weights_after_loading=lambda layer: None), "test"
+ )
+ config = MoeRunnerConfig(
+ num_experts=num_experts,
+ num_local_experts=num_experts,
+ hidden_size=hidden,
+ intermediate_size_per_partition=intermediate,
+ top_k=2,
+ activation="silu",
+ is_gated=True,
+ swiglu_limit=10,
+ )
+ method.create_moe_runner(layer, config)
+
+ w13_parameter = layer.w13_weight
+ w2_parameter = layer.w2_weight
+ w13_scale_parameter = layer.w13_weight_scale_inv
+ w2_scale_parameter = layer.w2_weight_scale_inv
+ method.process_weights_after_loading(layer)
+
+ expected_w13_scale = block_scale_interleave(w31_scale.view(torch.uint8)).reshape_as(
+ w31_scale
+ )
+ expected_w2_scale = block_scale_interleave(w2_scale.view(torch.uint8)).reshape_as(
+ w2_scale
+ )
+ assert layer.w13_weight is w13_parameter
+ assert layer.w2_weight is w2_parameter
+ assert layer.w13_weight_scale_inv is w13_scale_parameter
+ assert layer.w2_weight_scale_inv is w2_scale_parameter
+ assert torch.equal(layer.w13_weight_scale_inv.view(torch.uint8), expected_w13_scale)
+ assert torch.equal(layer.w2_weight_scale_inv.view(torch.uint8), expected_w2_scale)
+
+ generator = torch.Generator(device="cuda").manual_seed(1)
+ x = (
+ torch.randn(
+ 8,
+ hidden,
+ dtype=torch.bfloat16,
+ device="cuda",
+ generator=generator,
+ )
+ * 0.1
+ )
+ logits = torch.randn(
+ 8,
+ num_experts,
+ dtype=torch.float32,
+ device="cuda",
+ generator=generator,
+ )
+ topk_weights, topk_ids = torch.topk(torch.softmax(logits, dim=-1), 2, dim=-1)
+ topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
+ topk = StandardTopKOutput(topk_weights, topk_ids.to(torch.int32), logits)
+ dispatch_output = StandardDispatchOutput(x, None, topk)
+
+ actual = method.apply(layer, dispatch_output).hidden_states
+
+ x_quant, x_scale = mxfp8_quantize(
+ x,
+ is_sf_swizzled_layout=True,
+ alignment=32,
+ )
+ global_scale = torch.ones(num_experts, dtype=torch.float32, device="cuda")
+ swiglu_limit = torch.full((num_experts,), 10.0, dtype=torch.float32, device="cuda")
+ expected = torch.empty_like(x)
+ cutlass_fused_moe(
+ input=x_quant,
+ token_selected_experts=topk_ids.to(torch.int32),
+ token_final_scales=topk_weights,
+ fc1_expert_weights=layer.w13_weight.view(torch.int64),
+ fc2_expert_weights=layer.w2_weight.view(torch.int64),
+ output_dtype=torch.bfloat16,
+ quant_scales=[
+ layer.w13_weight_scale_inv.view(torch.int32),
+ global_scale,
+ layer.w2_weight_scale_inv.view(torch.int32),
+ global_scale,
+ ],
+ input_sf=x_scale,
+ # Compare the adapter's implicit defaults against the old explicit
+ # alpha=1/beta=0 representation.
+ swiglu_alpha=torch.ones(num_experts, dtype=torch.float32, device="cuda"),
+ swiglu_beta=torch.zeros(num_experts, dtype=torch.float32, device="cuda"),
+ swiglu_limit=swiglu_limit,
+ use_mxfp8_act_scaling=True,
+ activation_type=ActivationType.Swiglu,
+ tune_max_num_tokens=8,
+ output=expected,
+ )
+
+ assert torch.equal(actual, expected)
+
+
+if __name__ == "__main__":
+ sys.exit(pytest.main([__file__, "-v"]))
diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py
index 8463e95ff..2fd5cf637 100644
--- a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py
+++ b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py
@@ -404,8 +404,7 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
def _make_random_dsv4_mxfp4(num_experts, hidden, inter, seed=0):
- """Mirrors the fp8 base method's allocation for fp4 experts: int8-packed
- 4-bit weights, fp32 scales (containing 2**e values, not raw E8M0 bytes)."""
+ """Create native checkpoint-style packed MXFP4 weights and E8M0 scales."""
g = torch.Generator(device="cuda").manual_seed(seed)
# int8 storage (signed) -- matches Fp8MoEMethod.create_weights for fp4_experts.
w13 = torch.randint(
@@ -424,8 +423,7 @@ def _make_random_dsv4_mxfp4(num_experts, hidden, inter, seed=0):
device="cuda",
generator=g,
)
- # fp32 scales whose bit pattern after .to(float8_e8m0fnu).view(uint8) lands
- # in a sane E8M0 band -- generate exponents around 0 (= 2**0).
+ # Native E8M0 scales with exponents around 0 (= 2**0).
raw_e = torch.randint(
125,
130,
@@ -442,8 +440,8 @@ def _make_random_dsv4_mxfp4(num_experts, hidden, inter, seed=0):
device="cuda",
generator=g,
)
- w13_s = raw_e.view(torch.float8_e8m0fnu).to(torch.float32)
- w2_s = raw_e2.view(torch.float8_e8m0fnu).to(torch.float32)
+ w13_s = raw_e.view(torch.float8_e8m0fnu)
+ w2_s = raw_e2.view(torch.float8_e8m0fnu)
return w13, w2, w13_s, w2_s
@@ -460,12 +458,11 @@ def test_dsv4_apply_matches_flashinfer_direct(
):
"""End-to-end: SGLang's DSv4 ``Mxfp4FlashinferCutlassMoEMethod.apply``
output must match a direct FlashInfer ``cutlass_fused_moe`` call with
- the equivalent reorder + scale-cast + interleave applied manually."""
+ the equivalent native E8M0 scale/weight interleave applied manually."""
from types import SimpleNamespace
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as fi_cutlass_mod
import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as ds_mod
- from sglang.srt.layers.quantization.utils import reorder_w1w3_to_w3w1
# Bypass symmetric-memory / TP-group stack in the new fused-func module
# (where DSv4 ``apply`` now dispatches the kernel call through).
@@ -476,29 +473,31 @@ def test_dsv4_apply_matches_flashinfer_direct(
monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
w13, w2, w13_s, w2_s = _make_random_dsv4_mxfp4(num_experts, hidden, inter)
+ w1, w3 = w13.chunk(2, dim=1)
+ w1_s, w3_s = w13_s.chunk(2, dim=1)
+ # Simulate FusedMoE's ``load_up_proj_weight_first`` loader contract.
+ w31 = torch.cat((w3, w1), dim=1)
+ w31_s = torch.cat(
+ (w3_s.view(torch.uint8), w1_s.view(torch.uint8)),
+ dim=1,
+ ).view(torch.float8_e8m0fnu)
x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1
topk_w, topk_i = _make_topk(tokens, num_experts, top_k)
# ---- SGLang DSv4 path ----
- method = ds_mod.Mxfp4FlashinferCutlassMoEMethod.__new__(
- ds_mod.Mxfp4FlashinferCutlassMoEMethod
- )
- method._fp8 = SimpleNamespace(
- process_weights_after_loading=lambda layer: None,
- )
- method.prefix = "test"
# plain SiLU * up — all three SwiGLU scalars None (no clamp configured).
- method._swiglu_alpha_tensor = None
- method._swiglu_beta_tensor = None
- method._swiglu_limit_tensor = None
+ method = ds_mod.Mxfp4FlashinferCutlassMoEMethod(
+ SimpleNamespace(process_weights_after_loading=lambda layer: None),
+ "test",
+ )
# Wire the unified MoeRunner -> flashinfer_mxfp4 fused func that
# ``apply`` now dispatches through.
method.runner = _build_flashinfer_mxfp4_runner(num_experts, hidden, inter)
layer = _MockLayer()
- layer.w13_weight = torch.nn.Parameter(w13.clone(), requires_grad=False)
+ layer.w13_weight = torch.nn.Parameter(w31.clone(), requires_grad=False)
layer.w2_weight = torch.nn.Parameter(w2.clone(), requires_grad=False)
- layer.w13_weight_scale_inv = torch.nn.Parameter(w13_s.clone(), requires_grad=False)
+ layer.w13_weight_scale_inv = torch.nn.Parameter(w31_s.clone(), requires_grad=False)
layer.w2_weight_scale_inv = torch.nn.Parameter(w2_s.clone(), requires_grad=False)
layer.num_local_experts = num_experts
layer.moe_tp_size = 1
@@ -513,11 +512,10 @@ def test_dsv4_apply_matches_flashinfer_direct(
).hidden_states
# ---- Direct FlashInfer reference ----
- w13_re, w13_s_re = reorder_w1w3_to_w3w1(w13, w13_s)
- w13_s_u8 = w13_s_re.to(torch.float8_e8m0fnu).view(torch.uint8).contiguous()
- w2_s_u8 = w2_s.to(torch.float8_e8m0fnu).view(torch.uint8).contiguous()
+ w13_s_u8 = w31_s.view(torch.uint8)
+ w2_s_u8 = w2_s.view(torch.uint8)
ref_w13 = interleave_moe_weights_for_sm90_mixed_gemm(
- w13_re.view(torch.uint8).contiguous(), "fp4"
+ w31.view(torch.uint8).contiguous(), "fp4"
)
ref_w2 = interleave_moe_weights_for_sm90_mixed_gemm(
w2.view(torch.uint8).contiguous(), "fp4"
diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py
index 471b64884..eb162a8af 100644
--- a/test/registered/unit/test_model_overrides.py
+++ b/test/registered/unit/test_model_overrides.py
@@ -521,6 +521,7 @@ class TestGoldenModelOverrides(_IsolatedPublish):
"llama",
dllm_algorithm="LowConfidence",
disable_radix_cache=True,
+ attention_backend="triton",
)
self.assertEqual(sa.attention_backend, "flashinfer") # materialized
self.assertIn(
@@ -821,7 +822,8 @@ class TestGoldenModelOverrides(_IsolatedPublish):
with patch.object(overrides_module, "is_sm120_supported", return_value=True):
self.assertEqual(
- _deepseek_v4_sm120_moe(_view()), {"moe_runner_backend": "marlin"}
+ _deepseek_v4_sm120_moe(_view()),
+ {"moe_runner_backend": "flashinfer_mxfp4"},
)
self.assertEqual(
_deepseek_v4_sm120_moe(_view(moe_runner_backend="triton")), {}