minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2 (#30272)

Browse Source
This commit is contained in:
Gabriel Wu
2026-07-18 03:01:06 -07:00
committed by GitHub
parent 6c6175fabd
commit faf6894093
18 changed files with 502 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h
+2 -5
View File
@@ -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/kernels/ops/attention/dsa/tilelang_kernel.py
+7 -4
View File
@@ -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)
python/sglang/kernels/ops/attention/flash_mla_sm120.py
+32 -22
View File
@@ -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)
python/sglang/srt/arg_groups/overrides.py
+5 -6
View File
@@ -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 {}
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+8 -3
View File
@@ -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,
python/sglang/srt/layers/attention/dsv4/indexer.py
+19
View File
@@ -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
python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py
+2 -1
View File
@@ -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
+45 -22
View File
@@ -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,
python/sglang/srt/layers/quantization/fp8.py
+5 -4
View File
@@ -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,
python/sglang/srt/layers/quantization/marlin_utils_fp4.py
+5
View File
@@ -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:
python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py
+82 -129
View File
@@ -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)
python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py
+6 -4
View File
@@ -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,
)
python/sglang/srt/server_args.py
+2
View File
@@ -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)