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

[MoE] Add FlashInfer SM90 MXFP4 W4A8 CUTLASS MoE (#34967)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-08-31 20:04:41 -07:00
committed by GitHub
co-authored by luoyuan.luo
parent 22337e9c56
commit 5b04408784
13 changed files with 662 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/layers/moe/bench_mxfp4_sm90_kernels.py
+81 -7
View File
@@ -1,10 +1,10 @@
"""Benchmark MXFP4 MoE kernels on H100/H200: SGLang Marlin vs FlashInfer cutlass.
"""Benchmark MXFP4 MoE kernels on H100/H200: Marlin vs FlashInfer CUTLASS.
Compares per-call latency of:
* Marlin path : ``fused_marlin_moe(...)`` after Marlin weight repack
* FlashInfer : ``cutlass_fused_moe(use_w4_group_scaling=True, ...)``
(PR #3084's SM90 mixed-input path)
* Marlin path : ``fused_marlin_moe(...)`` after Marlin weight repack
* FlashInfer W4A16 : PR #3084's SM90 mixed-input path
* FlashInfer W4A8 : PR #3738/#4431's corrected Humming path, when available
Both run on the same random MXFP4 weights/scales (semantics differ slightly --
Marlin uses a scalar swiglu clamp + no bias, FlashInfer fuses per-expert
@@ -15,7 +15,7 @@ Run on H100/H200:
cd /sgl-workspace/sglang_dev3 && \\
PYTHONPATH=python:/sgl-workspace/flashinfer FLASHINFER_DISABLE_VERSION_CHECK=1 \\
python python/sglang/test/bench_mxfp4_sm90_kernels.py
python test/manual/layers/moe/bench_mxfp4_sm90_kernels.py
"""
from __future__ import annotations
@@ -25,15 +25,29 @@ from dataclasses import dataclass
from typing import Callable, List, Tuple
import torch
from flashinfer.autotuner import autotune
# ---- FlashInfer ----
from flashinfer import __version__ as flashinfer_version
from flashinfer.autotuner import autotune
from flashinfer.fused_moe import (
cutlass_fused_moe,
interleave_moe_scales_for_sm90_mixed_gemm,
interleave_moe_weights_for_sm90_mixed_gemm,
)
from flashinfer.fused_moe.core import ActivationType
from packaging.version import Version
try:
from flashinfer.fused_moe import (
preprocess_moe_weights_for_sm90_mixed_gemm_humming,
)
except ImportError:
preprocess_moe_weights_for_sm90_mixed_gemm_humming = None
_fi_release = Version(flashinfer_version).release
_fi_release = _fi_release + (0,) * (3 - len(_fi_release))
if _fi_release[:3] < (0, 6, 18):
preprocess_moe_weights_for_sm90_mixed_gemm_humming = None
# ---- SGLang Marlin ----
from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
@@ -162,14 +176,50 @@ def build_flashinfer_inputs(shape: Shape, w13, w2, w13_s, w2_s, w13_b, w2_b):
}
def build_flashinfer_humming_inputs(shape: Shape, w13, w2, w13_s, w2_s, w13_b, w2_b):
if preprocess_moe_weights_for_sm90_mixed_gemm_humming is None:
raise RuntimeError("FlashInfer does not provide the corrected Humming API.")
w13_il, w13_s_il, w13_residual = preprocess_moe_weights_for_sm90_mixed_gemm_humming(
w13, w13_s
)
w2_il, w2_s_il, w2_residual = preprocess_moe_weights_for_sm90_mixed_gemm_humming(
w2, w2_s
)
e = shape.num_experts
return {
"w13": w13_il,
"w2": w2_il,
"quant_scales": [
w13_s_il.view(torch.int32),
(w13_residual * 64.0).contiguous(),
torch.ones((), dtype=torch.float32, device="cuda"),
w2_s_il.view(torch.int32),
(w2_residual * 64.0).contiguous(),
],
"w13_b": w13_b,
"w2_b": w2_b,
"swiglu_alpha": torch.full((e,), 1.702, dtype=torch.float32, device="cuda"),
"swiglu_beta": torch.full((e,), 1.0, dtype=torch.float32, device="cuda"),
"swiglu_limit": torch.full((e,), 7.0, dtype=torch.float32, device="cuda"),
}
def make_flashinfer_runner(
shape: Shape, prep, x, topk_w, topk_i, autotuned: bool, with_bias: bool = True
shape: Shape,
prep,
x,
topk_w,
topk_i,
autotuned: bool,
with_bias: bool = True,
use_humming: bool = False,
):
out = torch.empty(shape.tokens, shape.hidden, dtype=torch.bfloat16, device="cuda")
fc1_b = prep["w13_b"] if with_bias else None
fc2_b = prep["w2_b"] if with_bias else None
def _call():
humming_kwargs = {"use_wfp4afp8_humming": True} if use_humming else {}
cutlass_fused_moe(
input=x,
token_selected_experts=topk_i,
@@ -186,6 +236,7 @@ def make_flashinfer_runner(
use_w4_group_scaling=True,
activation_type=ActivationType.Swiglu,
output=out,
**humming_kwargs,
)
if autotuned:
@@ -328,6 +379,29 @@ def run_one_shape(shape: Shape, run_marlin: bool):
)
fi_med = fi_at_med # alias for downstream speedup print
if preprocess_moe_weights_for_sm90_mixed_gemm_humming is not None:
humming_prep = build_flashinfer_humming_inputs(
shape, w13, w2, w13_s, w2_s, w13_b, w2_b
)
humming_call = make_flashinfer_runner(
shape,
humming_prep,
x,
topk_w,
topk_i,
autotuned=True,
with_bias=True,
use_humming=True,
)
humming_med, humming_min = time_call(humming_call)
print(
f" FlashInfer Humming W4A8: median={humming_med:.3f} ms "
f"min={humming_min:.3f} ms"
)
print(f" speedup (FI W4A16 / FI W4A8): {fi_med / humming_med:.2f}x")
else:
print(" FlashInfer Humming W4A8: SKIPPED (requires >= 0.6.18)")
# Marlin
if run_marlin:
try:
test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
+8 -4
View File
@@ -84,6 +84,7 @@ def test_cutlass_adapter_import_does_not_require_flashinfer(monkeypatch):
def test_dsv4_sm120_load_contract(monkeypatch, request):
import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as adapter_module
from sglang.srt.runtime_context import get_context
platform = override_platform(is_sm120=True)
platform.install()
@@ -95,7 +96,8 @@ def test_dsv4_sm120_load_contract(monkeypatch, request):
def create_weights(self, *args, **kwargs):
captured.update(kwargs)
method = adapter_module.Mxfp4FlashinferCutlassMoEMethod(_Fp8Method(), "test")
with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"):
method = adapter_module.Mxfp4FlashinferCutlassMoEMethod(_Fp8Method(), "test")
method.create_weights(
SimpleNamespace(),
num_experts=4,
@@ -126,6 +128,7 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import (
Mxfp4FlashinferCutlassMoEMethod,
)
from sglang.srt.runtime_context import get_context
monkeypatch.setattr(
runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext()
@@ -155,9 +158,10 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
moe_ep_rank=0,
)
method = Mxfp4FlashinferCutlassMoEMethod(
SimpleNamespace(process_weights_after_loading=lambda layer: None), "test"
)
with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"):
method = Mxfp4FlashinferCutlassMoEMethod(
SimpleNamespace(process_weights_after_loading=lambda layer: None), "test"
)
config = MoeRunnerConfig(
num_experts=num_experts,
num_local_experts=num_experts,
test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py
+328 -8
View File
@@ -2,9 +2,9 @@
Builds a single-layer GPT-OSS-style MoE with random MXFP4 weights, drives the
SGLang plumbing (``_process_weights_for_sm90_cutlass`` + ``_apply_sm90_cutlass``)
and compares against a direct FlashInfer ``cutlass_fused_moe`` call with the
same inputs. Both paths invoke the same SM90 kernel from FlashInfer PR #3084,
so outputs must be bit-exact.
and compares against direct FlashInfer ``cutlass_fused_moe`` calls. It covers
both PR #3084's W4A16 path and PR #3738/#4431's corrected Humming W4A8 path;
outputs must be bit-exact within each path.
Run on H100/H200:
@@ -24,6 +24,16 @@ register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
flashinfer_fused_moe = pytest.importorskip("flashinfer.fused_moe")
HAS_CORRECTED_HUMMING_API = hasattr(
flashinfer_fused_moe,
"preprocess_moe_weights_for_sm90_mixed_gemm_humming",
)
preprocess_humming = getattr(
flashinfer_fused_moe,
"preprocess_moe_weights_for_sm90_mixed_gemm_humming",
None,
)
if not hasattr(flashinfer_fused_moe, "interleave_moe_weights_for_sm90_mixed_gemm"):
pytest.skip(
"FlashInfer build does not include PR #3084 SM90 mixed-input helpers",
@@ -154,11 +164,12 @@ def _round_up(x, base):
return ((x + base - 1) // base) * base
def _build_method(num_experts, hidden, inter):
def _build_method(num_experts, hidden, inter, *, use_humming=False):
from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod)
method._fi_kernel = "cutlass_sm90"
method._use_sm90_humming = use_humming
method.num_experts = num_experts
# The new SM90 cutlass path tracks padded sizes in dedicated attrs;
# ``hidden_size`` / ``intermediate_size_per_partition`` keep the unpadded
@@ -348,6 +359,11 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
)
monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
monkeypatch.setattr(
fi_cutlass_mod.envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE,
"get",
lambda: False,
)
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
@@ -392,6 +408,7 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
swiglu_limit=layer.swiglu_limit,
use_w4_group_scaling=True,
activation_type=ActivationType.Swiglu,
use_fused_finalize=False,
output=out_ref_padded,
)
out_ref = (
@@ -404,6 +421,258 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
)
@pytest.mark.skipif(
not HAS_CORRECTED_HUMMING_API,
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
)
def test_process_weights_humming_matches_flashinfer_direct():
"""The SM90 fp8 option must use #3738's preprocessing and retain #4431's
per-local-expert residual contract."""
num_experts, hidden, inter = 4, 256, 256
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter)
# Make each expert's residual distinct so an accidental scalar/broadcast
# contract cannot pass this check.
expert_offsets = torch.arange(num_experts, dtype=torch.uint8, device="cuda").view(
-1, 1, 1
)
w13_s = w13_s + expert_offsets
w2_s = w2_s + expert_offsets + 1
layer = _build_mock_layer(
num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b
)
method = _build_method(num_experts, hidden, inter, use_humming=True)
method._process_weights_for_sm90_cutlass(layer)
# GPT-OSS loads pair-wise [gate, up]; FlashInfer consumes halved [up; gate].
ref_w13 = torch.cat((w13[:, 1::2], w13[:, 0::2]), dim=1).contiguous()
ref_w13_s = torch.cat((w13_s[:, 1::2], w13_s[:, 0::2]), dim=1).contiguous()
expected_w13, expected_w13_s, expected_w13_residual = preprocess_humming(
ref_w13, ref_w13_s
)
expected_w2, expected_w2_s, expected_w2_residual = preprocess_humming(w2, w2_s)
assert torch.equal(layer.w13_weight, expected_w13)
assert torch.equal(layer.w2_weight, expected_w2)
assert torch.equal(layer.w13_weight_scale, expected_w13_s)
assert torch.equal(layer.w2_weight_scale, expected_w2_s)
assert torch.equal(layer.w13_humming_residual_scale, expected_w13_residual * 64.0)
assert torch.equal(layer.w2_humming_residual_scale, expected_w2_residual * 64.0)
assert layer.w13_humming_residual_scale.shape == (num_experts,)
assert layer.w2_humming_residual_scale.shape == (num_experts,)
assert layer.humming_fc2_act_scale.shape == ()
@pytest.mark.skipif(
not HAS_CORRECTED_HUMMING_API,
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
)
def test_humming_padding_preserves_per_expert_residual():
"""Synthetic alignment padding must not change an expert's E8M0 range."""
num_experts, hidden, inter = 4, 192, 192
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter)
ref_w13 = torch.cat((w13[:, 1::2], w13[:, 0::2]), dim=1).contiguous()
ref_w13_s = torch.cat((w13_s[:, 1::2], w13_s[:, 0::2]), dim=1).contiguous()
_, _, expected_w13_residual = preprocess_humming(
ref_w13, ref_w13_s, interleave=False
)
_, _, expected_w2_residual = preprocess_humming(w2, w2_s, interleave=False)
layer = _build_mock_layer(
num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b
)
method = _build_method(num_experts, hidden, inter, use_humming=True)
method._process_weights_for_sm90_cutlass(layer)
assert torch.equal(layer.w13_humming_residual_scale, expected_w13_residual * 64.0)
assert torch.equal(layer.w2_humming_residual_scale, expected_w2_residual * 64.0)
def _build_prerounded_case(E, hidden_real, hidden_rounded, inter, tail_fill, seed=0):
"""Weights as ``create_weights`` leaves them when FusedMoE pre-rounds hidden.
Buffers are allocated at ``hidden_rounded``; the loader only ever writes the
first ``hidden_real`` columns, so the tail keeps whatever the buffer was
filled with (``_UE8M0_ONE`` in production). Real scales sit well BELOW 2^0
so a leaked tail moves the per-expert max.
"""
g = torch.Generator(device="cuda").manual_seed(seed)
kr_bytes = hidden_real // 2
kr_grp = hidden_real // GROUP_SIZE
w13 = torch.zeros(
(E, 2 * inter, hidden_rounded // 2), dtype=torch.uint8, device="cuda"
)
w13[:, :, :kr_bytes] = torch.randint(
0, 256, (E, 2 * inter, kr_bytes), dtype=torch.uint8, device="cuda", generator=g
)
w2 = torch.zeros((E, hidden_rounded, inter // 2), dtype=torch.uint8, device="cuda")
w2[:, :hidden_real, :] = torch.randint(
0,
256,
(E, hidden_real, inter // 2),
dtype=torch.uint8,
device="cuda",
generator=g,
)
w13_s = torch.full(
(E, 2 * inter, hidden_rounded // GROUP_SIZE),
tail_fill,
dtype=torch.uint8,
device="cuda",
)
w13_s[:, :, :kr_grp] = torch.randint(
100, 110, (E, 2 * inter, kr_grp), dtype=torch.uint8, device="cuda", generator=g
)
w2_s = torch.full(
(E, hidden_rounded, inter // GROUP_SIZE),
tail_fill,
dtype=torch.uint8,
device="cuda",
)
w2_s[:, :hidden_real, :] = torch.randint(
100,
110,
(E, hidden_real, inter // GROUP_SIZE),
dtype=torch.uint8,
device="cuda",
generator=g,
)
w13_b = torch.zeros((E, 2 * inter), dtype=torch.bfloat16, device="cuda")
w2_b = torch.zeros((E, hidden_rounded), dtype=torch.bfloat16, device="cuda")
return w13, w2, w13_s, w2_s, w13_b, w2_b
@pytest.mark.skipif(
not HAS_CORRECTED_HUMMING_API,
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
)
def test_humming_range_ignores_prerounded_hidden_tail():
"""FusedMoE rounds GPT-OSS hidden 2880 -> 3072 BEFORE ``create_weights``, so
the trailing scale columns keep the ``_UE8M0_ONE`` buffer fill. Those bytes
are 2^0 -- above any real per-expert max -- and must not reach Humming's
min/max, or the residual shifts and perturbs the real weights.
Invariant: the residual must not depend on what the never-written tail holds.
"""
from sglang.srt.layers.quantization.mxfp4 import _UE8M0_ONE
E, hidden_real, hidden_rounded, inter = 4, 2880, 3072, 256
residuals = []
for tail_fill in (_UE8M0_ONE, 105): # 105 sits inside the real 100..110 band
w13, w2, w13_s, w2_s, w13_b, w2_b = _build_prerounded_case(
E, hidden_real, hidden_rounded, inter, tail_fill
)
layer = _build_mock_layer(
E, hidden_rounded, inter, w13, w2, w13_s, w2_s, w13_b, w2_b
)
method = _build_method(E, hidden_rounded, inter, use_humming=True)
# What create_weights records from layer.hidden_size_unpadded.
method._unpadded_hidden = hidden_real
method._process_weights_for_sm90_cutlass(layer)
residuals.append(
(
layer.w13_humming_residual_scale.clone(),
layer.w2_humming_residual_scale.clone(),
)
)
assert torch.equal(residuals[0][0], residuals[1][0]), (
"w13 Humming residual changed with the never-written hidden tail; "
"the _UE8M0_ONE fill leaked into the per-expert E8M0 range"
)
assert torch.equal(
residuals[0][1], residuals[1][1]
), "w2 Humming residual changed with the never-written hidden tail"
@pytest.mark.skipif(
not HAS_CORRECTED_HUMMING_API,
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
)
@pytest.mark.parametrize(
"tokens,hidden,inter,ep_size,ep_rank",
[(8, 256, 256, 1, 0), (8, 192, 192, 1, 0), (8, 256, 256, 2, 1)],
)
def test_apply_sm90_humming_matches_flashinfer_direct(
tokens, hidden, inter, ep_size, ep_rank, monkeypatch
):
"""SGLang must forward the five Humming scales and enable the new kernel."""
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as fi_cutlass_mod
monkeypatch.setattr(
fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
)
monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
monkeypatch.setattr(
fi_cutlass_mod.envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE,
"get",
lambda: False,
)
num_experts, top_k = 4, 2
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
topk_w, topk_i = _make_topk(tokens, num_experts, top_k)
topk_i = topk_i + ep_rank * num_experts
layer = _build_mock_layer(
num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b
)
layer.moe_ep_size = ep_size
layer.moe_ep_rank = ep_rank
method = _build_method(num_experts, hidden, inter, use_humming=True)
method._process_weights_for_sm90_cutlass(layer)
out_sglang = method._apply_sm90_cutlass(
layer, _MockDispatchOutput(x.clone(), topk_w, topk_i)
).hidden_states
padded_hidden = method._padded_hidden
x_ref = (
torch.nn.functional.pad(x, (0, padded_hidden - hidden))
if padded_hidden != hidden
else x
)
out_ref_padded = torch.empty(
tokens, padded_hidden, dtype=torch.bfloat16, device="cuda"
)
cutlass_fused_moe(
input=x_ref,
token_selected_experts=topk_i,
token_final_scales=topk_w,
fc1_expert_weights=layer.w13_weight,
fc2_expert_weights=layer.w2_weight,
output_dtype=torch.bfloat16,
quant_scales=[
layer.w13_weight_scale.view(torch.int32),
layer.w13_humming_residual_scale,
layer.humming_fc2_act_scale,
layer.w2_weight_scale.view(torch.int32),
layer.w2_humming_residual_scale,
],
fc1_expert_biases=layer.w13_weight_bias,
fc2_expert_biases=layer.w2_weight_bias,
swiglu_alpha=layer.swiglu_alpha,
swiglu_beta=layer.swiglu_beta,
swiglu_limit=layer.swiglu_limit,
ep_size=ep_size,
ep_rank=ep_rank,
use_w4_group_scaling=True,
use_wfp4afp8_humming=True,
activation_type=ActivationType.Swiglu,
tune_max_num_tokens=tokens,
use_fused_finalize=False,
output=out_ref_padded,
)
out_ref = out_ref_padded[:, :hidden].contiguous()
assert torch.equal(out_sglang, out_ref)
# =============================================================================
# DeepSeek-V4 path: Mxfp4FlashinferCutlassMoEMethod (sibling of Marlin /
# trtllm-gen). Wired into fp8.py's get_quant_method when SM90 +
@@ -494,10 +763,13 @@ def test_dsv4_apply_matches_flashinfer_direct(
# ---- SGLang DSv4 path ----
# plain SiLU * up — all three SwiGLU scalars None (no clamp configured).
method = ds_mod.Mxfp4FlashinferCutlassMoEMethod(
SimpleNamespace(process_weights_after_loading=lambda layer: None),
"test",
)
from sglang.srt.runtime_context import get_context
with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"):
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)
@@ -559,6 +831,54 @@ def test_dsv4_apply_matches_flashinfer_direct(
)
@pytest.mark.skipif(
not HAS_CORRECTED_HUMMING_API,
reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18",
)
def test_dsv4_process_weights_humming_matches_flashinfer_direct():
"""DSv4's native [up; gate] layout must use the same #3738 transform."""
from types import SimpleNamespace
import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as ds_mod
from sglang.srt.runtime_context import get_context
num_experts, hidden, inter = 4, 256, 256
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)
w31 = torch.cat((w3, w1), dim=1).contiguous()
w31_s = torch.cat((w3_s.view(torch.uint8), w1_s.view(torch.uint8)), dim=1).view(
torch.float8_e8m0fnu
)
with get_context().override_server_args(flashinfer_mxfp4_moe_precision="fp8"):
method = ds_mod.Mxfp4FlashinferCutlassMoEMethod(
SimpleNamespace(process_weights_after_loading=lambda layer: None),
"test",
)
layer = _MockLayer()
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(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
method.process_weights_after_loading(layer)
ref_w13, ref_w13_s, ref_w13_residual = preprocess_humming(
w31.view(torch.uint8), w31_s.view(torch.uint8)
)
ref_w2, ref_w2_s, ref_w2_residual = preprocess_humming(
w2.view(torch.uint8), w2_s.view(torch.uint8)
)
assert torch.equal(layer.w13_weight, ref_w13)
assert torch.equal(layer.w2_weight, ref_w2)
assert torch.equal(layer.w13_weight_scale_inv, ref_w13_s)
assert torch.equal(layer.w2_weight_scale_inv, ref_w2_s)
assert torch.equal(layer.w13_humming_residual_scale, ref_w13_residual * 64.0)
assert torch.equal(layer.w2_humming_residual_scale, ref_w2_residual * 64.0)
class _MockDispatchOutput:
"""Stand-in for StandardDispatchOutput. ``topk_output`` is a real
``StandardTopKOutput`` so ``TopKOutputChecker.format_is_standard``