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

[MoE Refactor] Migrate SM90 Cutlass W4A16 to MoeRunner (#26489)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-05-30 02:02:56 -07:00
committed by GitHub
co-authored by luoyuan.luo
parent c93a559e5a
commit 0d9a2a9de3
5 changed files with 281 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py
+174
View File
@@ -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)
python/sglang/srt/layers/moe/moe_runner/runner.py
+2
View File
@@ -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}")
python/sglang/srt/layers/quantization/mxfp4.py
+31 -64
View File
@@ -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
# 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,
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,
)
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,
python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py
+28 -45
View File
@@ -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)
test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py
+45 -6
View File
@@ -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)