From 0d9a2a9de378e0d9ed7d968cd49100d4ea0cb637 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Sat, 30 May 2026 17:02:56 +0800 Subject: [PATCH] [MoE Refactor] Migrate SM90 Cutlass W4A16 to MoeRunner (#26489) Co-authored-by: luoyuan.luo --- .../layers/moe/moe_runner/flashinfer_mxfp4.py | 174 ++++++++++++++++++ .../srt/layers/moe/moe_runner/runner.py | 2 + .../sglang/srt/layers/quantization/mxfp4.py | 97 ++++------ .../mxfp4_flashinfer_cutlass_moe.py | 73 +++----- .../quantization/test_mxfp4_sm90_cutlass.py | 51 ++++- 5 files changed, 281 insertions(+), 116 deletions(-) create mode 100644 python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py new file mode 100644 index 000000000..ec00efe17 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py @@ -0,0 +1,174 @@ +"""FlashInfer SM90 cutlass mixed-input W4A16 MXFP4 MoE fused func. + +Registered for ``("none", "flashinfer_mxfp4")``. Drives FlashInfer's +``cutlass_fused_moe(use_w4_group_scaling=True)`` (PR #3084 in flashinfer, +SM90 only). Quant methods build the quant_info each forward and call +``MoeRunner.run(dispatch_output, quant_info)``. + +Two production call sites share this fused func: + - GPT-OSS via :class:`Mxfp4MoEMethod` (input pad/output trim + per-expert + SwiGLU scalars + per-expert bias) + - DSv4 via :class:`Mxfp4FlashinferCutlassMoEMethod` (no bias, optional + SwiGLU scalars, no padding) + +The SM100 trtllm-gen path also lives under ``MoeRunnerBackend.FLASHINFER_MXFP4`` +but is intentionally left in the legacy bypass path for now; migrating it is a +follow-up. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.distributed import get_tp_group +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.layers.dp_attention import is_allocation_symmetric +from sglang.srt.layers.moe.moe_runner.base import ( + MoeQuantInfo, + MoeRunnerConfig, + register_fused_func, +) +from sglang.srt.utils import is_flashinfer_available +from sglang.srt.utils.common import next_power_of_2 + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + + +@dataclass +class FlashInferMxfp4CutlassMoeQuantInfo(MoeQuantInfo): + """Quantization payload for the SM90 cutlass W4A16 MXFP4 MoE path. + + 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. + """ + + # Pre-interleaved weights (uint8, packed FP4) + 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) + w13_weight_scale: torch.Tensor # [E, 2*N, K/32] + w2_weight_scale: torch.Tensor # [E, K, N/32] + + # 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). + swiglu_alpha: Optional[torch.Tensor] = None + swiglu_beta: Optional[torch.Tensor] = None + swiglu_limit: Optional[torch.Tensor] = None + + # TP/EP topology (forwarded to the FlashInfer kernel) + moe_tp_size: int = 1 + moe_tp_rank: int = 0 + moe_ep_size: int = 1 + moe_ep_rank: int = 0 + + # GPT-OSS pads its input hidden dim up to the (pre-padded) loaded weight + # width and trims the output back. DSv4 leaves this as ``None`` (no pad). + padded_hidden: Optional[int] = None + + +def _flashinfer_cutlass_fused_moe(): + """Lazy import — keeps non-flashinfer wheels importable.""" + if not is_flashinfer_available(): + raise RuntimeError( + "flashinfer_mxfp4 runner backend requires flashinfer to be installed." + ) + from flashinfer.fused_moe import cutlass_fused_moe + from flashinfer.fused_moe.core import ActivationType + + return cutlass_fused_moe, ActivationType + + +@register_fused_func("none", "flashinfer_mxfp4") +def fused_experts_none_to_flashinfer_mxfp4( + dispatch_output: "StandardDispatchOutput", + quant_info: MoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> "StandardCombineInput": + """SM90 W4A16 MXFP4 fused expert forward pass. + + Mirrors the legacy ``Mxfp4MoEMethod._apply_sm90_cutlass`` and DSv4's + ``Mxfp4FlashinferCutlassMoEMethod.apply`` exactly; difference vs those is + that all per-layer state arrives via ``quant_info`` rather than via the + layer module, so this function is layer-agnostic. + """ + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + from sglang.srt.layers.moe.topk import TopKOutputChecker + + assert isinstance( + quant_info, FlashInferMxfp4CutlassMoeQuantInfo + ), f"Unexpected quant_info type for flashinfer_mxfp4: {type(quant_info)}" + + flashinfer_cutlass_fused_moe, ActivationType = _flashinfer_cutlass_fused_moe() + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + # Under ``--moe-runner-backend flashinfer_mxfp4`` topk may be in bypassed + # form (the SM100 trtllm-gen path does routing internally). The cutlass + # SM90 path needs explicit topk_ids / topk_weights; materialize here. + if TopKOutputChecker.format_is_bypassed(topk_output): + topk_output = topk_output.to_standard() + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + + # GPT-OSS: pad input hidden dim up to the loaded weight width. DSv4 + # leaves padded_hidden as None (or equal to origin_hidden), no pad. + origin_hidden = x.shape[-1] + padded_hidden = quant_info.padded_hidden + do_pad = padded_hidden is not None and padded_hidden != origin_hidden + if do_pad: + x = torch.nn.functional.pad( + x, + (0, padded_hidden - origin_hidden), + mode="constant", + value=0.0, + ) + + 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_final_scales=topk_weights, + fc1_expert_weights=quant_info.w13_weight, + fc2_expert_weights=quant_info.w2_weight, + output_dtype=output_dtype, + quant_scales=[ + quant_info.w13_weight_scale.view(torch.int32), + quant_info.w2_weight_scale.view(torch.int32), + ], + fc1_expert_biases=quant_info.w13_bias, + fc2_expert_biases=quant_info.w2_bias, + swiglu_alpha=quant_info.swiglu_alpha, + swiglu_beta=quant_info.swiglu_beta, + swiglu_limit=quant_info.swiglu_limit, + tp_size=quant_info.moe_tp_size, + 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, + activation_type=ActivationType.Swiglu, + tune_max_num_tokens=next_power_of_2(x.shape[0]), + output=out, + ) + + if do_pad: + out = out[:, :origin_hidden].contiguous() + + return StandardCombineInput(hidden_states=out) diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 392534517..1e4aa56d3 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -61,6 +61,8 @@ class MoeRunner: self.runner_core = None # FlashInfer TRT-LLM only supports fused path elif runner_backend.is_flashinfer_cutedsl(): self.runner_core = None # FlashInfer CuteDSL only supports fused path + elif runner_backend.is_flashinfer_mxfp4(): + self.runner_core = None # FlashInfer MXFP4 only supports fused path else: raise NotImplementedError(f"Unsupported runner backend: {runner_backend}") diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index c45d44cc8..6ac793d02 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -77,9 +77,7 @@ if is_flashinfer_available(): nvfp4_block_scale_interleave, trtllm_fp4_block_scale_moe, ) - from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe from flashinfer.fused_moe.core import ( - ActivationType, get_w2_permute_indices_with_cache, ) @@ -1056,78 +1054,47 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): or moe_runner_backend.is_triton() or moe_runner_backend.is_marlin() ): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + elif ( + moe_runner_backend.is_flashinfer_mxfp4() + and self._fi_kernel == "cutlass_sm90" + ): + # Register the fused func at runner construction so the FusedOpPool + # lookup at `MoeRunner.__init__` finds it. + import sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 # noqa: F401 + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) else: - # TODO(cwan): refactor other backends + # Legacy bypass path (e.g. SM100 trtllm-gen under flashinfer_mxfp4) + # routes through `apply` without a MoeRunner. TODO(cwan): migrate. pass - def _apply_sm90_cutlass(self, layer, x, topk_output): + def _apply_sm90_cutlass(self, layer, dispatch_output): """SM90 (Hopper) MXFP4 x BF16 MoE via FlashInfer's cutlass mixed-input - path (PR #3084). The fused kernel does GEMM1 + SwiGLU + GEMM2 in one - call; weights/scales were pre-interleaved at load time.""" - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - from sglang.srt.layers.moe.topk import TopKOutputChecker + path (PR #3084). Routed through the unified ``MoeRunner`` -- this + helper only builds the quant_info; the actual kernel call lives in + :mod:`sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4`.""" + from sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 import ( + FlashInferMxfp4CutlassMoeQuantInfo, + ) - # Under ``--moe-runner-backend flashinfer_mxfp4`` the SGLang TopK layer - # emits BypassedTopKOutput by default (the SM100 trtllm-gen kernel does - # routing internally). The cutlass kernel needs explicit topk_ids / - # topk_weights, so materialize them here when bypassed. - if TopKOutputChecker.format_is_bypassed(topk_output): - topk_output = topk_output.to_standard() - topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids - - # Pad input hidden dim to the (already-padded) loaded weight width. - origin_hidden = x.shape[-1] - padded_hidden = self._padded_hidden - if padded_hidden != origin_hidden: - x = torch.nn.functional.pad( - x, - (0, padded_hidden - origin_hidden), - mode="constant", - value=0.0, - ) - - output_dtype = torch.bfloat16 - # Output is allocated at padded width (kernel writes padded_hidden - # columns), then trimmed back to origin_hidden before returning. - with use_symmetric_memory( - get_tp_group(), disabled=not is_allocation_symmetric() - ): - out_padded = torch.empty( - x.shape[0], padded_hidden, dtype=output_dtype, device=x.device - ) - - flashinfer_cutlass_fused_moe( - input=x, - token_selected_experts=topk_ids.to(torch.int), - token_final_scales=topk_weights, - fc1_expert_weights=layer.w13_weight, # uint8 [E, 2*N, K/2] interleaved - fc2_expert_weights=layer.w2_weight, # uint8 [E, K, N/2] interleaved - output_dtype=output_dtype, - quant_scales=[ - layer.w13_weight_scale.view(torch.int32), - layer.w2_weight_scale.view(torch.int32), - ], - fc1_expert_biases=layer.w13_weight_bias, # bf16 [E, 2*N] - fc2_expert_biases=layer.w2_weight_bias, # bf16 [E, K] + quant_info = FlashInferMxfp4CutlassMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + w13_bias=layer.w13_weight_bias, + w2_bias=layer.w2_weight_bias, swiglu_alpha=layer.swiglu_alpha, swiglu_beta=layer.swiglu_beta, swiglu_limit=layer.swiglu_limit, - tp_size=layer.moe_tp_size, - tp_rank=layer.moe_tp_rank, - ep_size=layer.moe_ep_size, - ep_rank=layer.moe_ep_rank, - use_w4_group_scaling=True, - activation_type=ActivationType.Swiglu, - tune_max_num_tokens=next_power_of_2(x.shape[0]), - output=out_padded, + 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=self._padded_hidden, ) - - if padded_hidden != origin_hidden: - out = out_padded[:, :origin_hidden].contiguous() - else: - out = out_padded - return StandardCombineInput(hidden_states=out) + return self.runner.run(dispatch_output, quant_info) def apply( self, @@ -1183,7 +1150,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): return self.runner.run(dispatch_output, quant_info) if self._fi_kernel == "cutlass_sm90": - return self._apply_sm90_cutlass(layer, x, topk_output) + return self._apply_sm90_cutlass(layer, dispatch_output) if self.use_flashinfer: # When bf16 mode is enabled, we don't need to quantize the input, # TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations, 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 7fce478e1..31552cdfe 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py @@ -26,15 +26,8 @@ import torch from torch.nn import Module from torch.nn.parameter import Parameter -from sglang.srt.distributed import get_tp_group -from sglang.srt.distributed.device_communicators.pynccl_allocator import ( - use_symmetric_memory, -) -from sglang.srt.layers.dp_attention import is_allocation_symmetric -from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput 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 next_power_of_2 # 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; @@ -42,9 +35,6 @@ from sglang.srt.utils.common import next_power_of_2 os.environ.setdefault("TLLM_LOG_LEVEL", "INFO") if is_flashinfer_available(): - from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe - from flashinfer.fused_moe.core import ActivationType - try: from flashinfer.fused_moe import ( interleave_moe_scales_for_sm90_mixed_gemm, @@ -123,6 +113,9 @@ class Mxfp4FlashinferCutlassMoEMethod: ) def create_moe_runner(self, layer: Module, moe_runner_config) -> None: + from sglang.srt.layers.moe.moe_runner.runner import MoeRunner + from sglang.srt.layers.moe.utils import MoeRunnerBackend + self.moe_runner_config = moe_runner_config # DSv4 uses standard SwiGLU plus a config-driven activation clamp. @@ -150,6 +143,12 @@ class Mxfp4FlashinferCutlassMoEMethod: self._swiglu_beta_tensor = None self._swiglu_limit_tensor = None + # Register the fused func at runner construction so the FusedOpPool + # lookup at `MoeRunner.__init__` finds it. + import sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 # noqa: F401 + + 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 @@ -218,46 +217,30 @@ class Mxfp4FlashinferCutlassMoEMethod: layer: Module, dispatch_output: "DispatchOutput", ) -> "CombineInput": + from sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 import ( + FlashInferMxfp4CutlassMoeQuantInfo, + ) + + # 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}") - x = dispatch_output.hidden_states - topk_weights = topk_output.topk_weights - topk_ids = topk_output.topk_ids - - output_dtype = torch.bfloat16 - with use_symmetric_memory( - get_tp_group(), disabled=not is_allocation_symmetric() - ): - out = torch.empty( - x.shape[0], x.shape[-1], dtype=output_dtype, device=x.device - ) - - flashinfer_cutlass_fused_moe( - input=x, - token_selected_experts=topk_ids.to(torch.int), - token_final_scales=topk_weights, - fc1_expert_weights=layer.w13_weight, - fc2_expert_weights=layer.w2_weight, - output_dtype=output_dtype, - quant_scales=[ - layer.w13_weight_scale_inv.view(torch.int32), - layer.w2_weight_scale_inv.view(torch.int32), - ], - fc1_expert_biases=None, # DSv4 has no MoE expert bias. - fc2_expert_biases=None, + quant_info = FlashInferMxfp4CutlassMoeQuantInfo( + 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. + w2_bias=None, swiglu_alpha=self._swiglu_alpha_tensor, # ones: standard SiLU gate swiglu_beta=self._swiglu_beta_tensor, # zeros: standard up swiglu_limit=self._swiglu_limit_tensor, - tp_size=layer.moe_tp_size, - tp_rank=layer.moe_tp_rank, - ep_size=layer.moe_ep_size, - ep_rank=layer.moe_ep_rank, - use_w4_group_scaling=True, - activation_type=ActivationType.Swiglu, - tune_max_num_tokens=next_power_of_2(x.shape[0]), - output=out, + 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. ) - - return StandardCombineInput(hidden_states=out) + return self.runner.run(dispatch_output, quant_info) 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 c48270b0f..be54ca7bc 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py @@ -160,9 +160,34 @@ def _build_method(num_experts, hidden, inter): method._padded_hidden = _round_up(hidden, 128) method._padded_intermediate = _round_up(inter, 128) method.use_flashinfer = True + method.runner = _build_flashinfer_mxfp4_runner(num_experts, hidden, inter) return method +def _build_flashinfer_mxfp4_runner(num_experts, hidden, inter): + """Construct a real MoeRunner bound to the flashinfer_mxfp4 fused func. + + Bypasses ``create_moe_runner`` (which needs a live server arg context) + and wires the runner with a minimal MoeRunnerConfig sufficient for the + cutlass SM90 fused func, which only reads dispatch_output / quant_info. + """ + import sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 # noqa: F401 + from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig + from sglang.srt.layers.moe.moe_runner.runner import MoeRunner + from sglang.srt.layers.moe.utils import MoeRunnerBackend + + cfg = MoeRunnerConfig( + num_experts=num_experts, + num_local_experts=num_experts, + hidden_size=hidden, + intermediate_size_per_partition=inter, + top_k=None, + activation="silu", + is_gated=True, + ) + return MoeRunner(MoeRunnerBackend.FLASHINFER_MXFP4, cfg) + + def _expected_w13_processed(w13_un, w13_s_un, w13_b_un, N_pad, K_pad, group_size): """Replicate ``_process_weights_for_sm90_cutlass`` for w13: de-interleave HF's pair-wise ``[g_0, u_0, g_1, u_1, ...]`` layout into halved @@ -300,14 +325,21 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct( covered separately by ``test_process_weights_matches_direct_interleave``; here we just verify that ``apply`` calls the kernel with the right arguments (incl. input padding + output trim).""" + import sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 as fi_mxfp4_mod import sglang.srt.layers.quantization.mxfp4 as mxfp4_mod - # Bypass symmetric-memory / TP-group: not relevant to numerics. + # Bypass symmetric-memory / TP-group in both the legacy quant_method and + # the new fused-func module (where the kernel call now lives). monkeypatch.setattr( mxfp4_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext() ) monkeypatch.setattr(mxfp4_mod, "is_allocation_symmetric", lambda: False) monkeypatch.setattr(mxfp4_mod, "get_tp_group", lambda: None) + monkeypatch.setattr( + fi_mxfp4_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext() + ) + monkeypatch.setattr(fi_mxfp4_mod, "is_allocation_symmetric", lambda: False) + monkeypatch.setattr(fi_mxfp4_mod, "get_tp_group", lambda: None) w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter) x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1 @@ -321,7 +353,7 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct( method._process_weights_for_sm90_cutlass(layer) out_sglang = method._apply_sm90_cutlass( - layer, x.clone(), _MockTopKOutput(topk_w, topk_i) + layer, _MockDispatchOutput(x.clone(), topk_w, topk_i) ).hidden_states # ---- FlashInfer-direct reference using the same processed weights ---- @@ -431,13 +463,17 @@ def test_dsv4_apply_matches_flashinfer_direct( the equivalent reorder + scale-cast + interleave applied manually.""" from types import SimpleNamespace + import sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 as fi_mxfp4_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 -- not relevant to numerics. - monkeypatch.setattr(ds_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()) - monkeypatch.setattr(ds_mod, "is_allocation_symmetric", lambda: False) - monkeypatch.setattr(ds_mod, "get_tp_group", lambda: None) + # Bypass symmetric-memory / TP-group stack in the new fused-func module + # (where DSv4 ``apply`` now dispatches the kernel call through). + monkeypatch.setattr( + fi_mxfp4_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext() + ) + monkeypatch.setattr(fi_mxfp4_mod, "is_allocation_symmetric", lambda: False) + monkeypatch.setattr(fi_mxfp4_mod, "get_tp_group", lambda: None) w13, w2, w13_s, w2_s = _make_random_dsv4_mxfp4(num_experts, hidden, inter) x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1 @@ -455,6 +491,9 @@ def test_dsv4_apply_matches_flashinfer_direct( method._swiglu_alpha_tensor = None method._swiglu_beta_tensor = None method._swiglu_limit_tensor = None + # 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)