minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity
Files
8ac19cc19f8ade51ed203478f17c9a09c706d730
sglang/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py
T

910 lines
34 KiB
Python
Raw Blame History

"""Unit test for the SM90 cutlass MXFP4 path in :class:`Mxfp4MoEMethod`.
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 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:
python -m pytest test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py -v
"""
from __future__ import annotations
from contextlib import nullcontext
import pytest
import torch
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=11, 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",
allow_module_level=True,
)
if not torch.cuda.is_available():
pytest.skip("CUDA required", allow_module_level=True)
from sglang.srt.utils import is_sm90_supported, is_sm100_supported
if not is_sm90_supported() or is_sm100_supported():
pytest.skip(
"SM90-only path; require Hopper without SM100 promotion",
allow_module_level=True,
)
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 sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
GROUP_SIZE = 32 # MXFP4 block size
@pytest.fixture
def stated_tp_group():
"""Provide a TP-group placeholder for kernels with mocked symmetric memory."""
from sglang.srt.runtime_context import get_parallel
with get_parallel().override(tp_group=None):
yield
class _MockLayer:
"""Stand-in for ``FusedMoE`` carrying the attributes the SM90 helpers read.
We construct one by hand so the test stays out of SGLang's distributed init
path (``get_tp_group`` etc.).
"""
def __init__(self):
# The SM90 weight-processing path reads the runner config for the
# gate/up row layout (``gate_up_interleaved``) and the activation. A
# real ``FusedMoE`` always carries one, so the stand-in does too.
self.moe_runner_config = MoeRunnerConfig()
class _MockTopKOutput:
def __init__(self, weights, ids):
self.topk_weights = weights
self.topk_ids = ids
def _make_random_mxfp4(num_experts, hidden, inter, seed=0):
g = torch.Generator(device="cuda").manual_seed(seed)
w13 = torch.randint(
0,
256,
(num_experts, 2 * inter, hidden // 2),
dtype=torch.uint8,
device="cuda",
generator=g,
)
w2 = torch.randint(
0,
256,
(num_experts, hidden, inter // 2),
dtype=torch.uint8,
device="cuda",
generator=g,
)
# E8M0 scales centered around 127 (= 2^0); narrow band keeps dequant values
# in a sane range so SwiGLU clamp doesn't dominate.
w13_s = torch.randint(
125,
130,
(num_experts, 2 * inter, hidden // GROUP_SIZE),
dtype=torch.uint8,
device="cuda",
generator=g,
)
w2_s = torch.randint(
125,
130,
(num_experts, hidden, inter // GROUP_SIZE),
dtype=torch.uint8,
device="cuda",
generator=g,
)
w13_b = (
torch.randn(
num_experts, 2 * inter, dtype=torch.float32, device="cuda", generator=g
).to(torch.bfloat16)
* 0.01
)
w2_b = (
torch.randn(
num_experts, hidden, dtype=torch.float32, device="cuda", generator=g
).to(torch.bfloat16)
* 0.01
)
return w13, w2, w13_s, w2_s, w13_b, w2_b
def _make_topk(tokens, num_experts, top_k, seed=1):
g = torch.Generator(device="cuda").manual_seed(seed)
logits = torch.randn(
tokens, num_experts, dtype=torch.float32, device="cuda", generator=g
)
weights, ids = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1)
weights = weights / weights.sum(dim=-1, keepdim=True)
return weights.to(torch.float32), ids.to(torch.int32)
def _build_mock_layer(num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b):
layer = _MockLayer()
layer.w13_weight = torch.nn.Parameter(w13.clone(), requires_grad=False)
layer.w2_weight = torch.nn.Parameter(w2.clone(), requires_grad=False)
layer.w13_weight_scale = torch.nn.Parameter(w13_s.clone(), requires_grad=False)
layer.w2_weight_scale = torch.nn.Parameter(w2_s.clone(), requires_grad=False)
layer.w13_weight_bias = torch.nn.Parameter(w13_b.clone(), requires_grad=False)
layer.w2_weight_bias = torch.nn.Parameter(w2_b.clone(), requires_grad=False)
layer.num_local_experts = num_experts # tests run with EP size = 1
layer.moe_tp_size = 1
layer.moe_tp_rank = 0
layer.moe_ep_size = 1
layer.moe_ep_rank = 0
return layer
def _round_up(x, base):
return ((x + base - 1) // base) * base
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
# values to mirror what ``create_weights`` records.
method.hidden_size = hidden
method.intermediate_size_per_partition = 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_cutlass # 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
``[up; gate]``, pad each half along its row dim from ``N_un -> N_pad``
and last dim from ``K_un -> K_pad`` with zeros, then run the FlashInfer
SM90 byte / scale interleave helpers."""
E, two_n_un, last_un_w = w13_un.shape
N_un = two_n_un // 2
K_un = last_un_w * 2 # packed 4-bit -> *2 for raw K
def _split_and_pad(unpadded, last_pad, last_un, dtype):
gate = unpadded[:, 0::2, :]
up = unpadded[:, 1::2, :]
out = torch.zeros(E, 2 * N_pad, last_pad, dtype=dtype, device=unpadded.device)
out[:, :N_un, :last_un] = up
out[:, N_pad : N_pad + N_un, :last_un] = gate
return out
w13_pad = _split_and_pad(
w13_un.view(torch.uint8), K_pad // 2, K_un // 2, w13_un.dtype
)
w13_s_pad = _split_and_pad(
w13_s_un, K_pad // group_size, K_un // group_size, w13_s_un.dtype
)
gate_b = w13_b_un[:, 0::2]
up_b = w13_b_un[:, 1::2]
w13_b_pad = torch.zeros(E, 2 * N_pad, dtype=w13_b_un.dtype, device=w13_b_un.device)
w13_b_pad[:, :N_un] = up_b
w13_b_pad[:, N_pad : N_pad + N_un] = gate_b
w13_il = interleave_moe_weights_for_sm90_mixed_gemm(w13_pad, "fp4")
w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm(
w13_s_pad, group_size=group_size
)
return w13_il, w13_s_il, w13_b_pad
def _expected_w2_processed(w2_un, w2_s_un, w2_b_un, N_pad, K_pad, group_size):
"""w2 needs padding only (no halving / no de-interleave)."""
E, K_un, last_un_w = w2_un.shape
N_un = last_un_w * 2
def _pad(unpadded, last_pad, last_un):
out = torch.zeros(
E, K_pad, last_pad, dtype=unpadded.dtype, device=unpadded.device
)
out[:, :K_un, :last_un] = unpadded
return out
w2_pad = _pad(w2_un.view(torch.uint8), N_pad // 2, N_un // 2)
w2_s_pad = _pad(w2_s_un, N_pad // group_size, N_un // group_size)
w2_b_pad = torch.zeros(E, K_pad, dtype=w2_b_un.dtype, device=w2_b_un.device)
w2_b_pad[:, :K_un] = w2_b_un
w2_il = interleave_moe_weights_for_sm90_mixed_gemm(w2_pad, "fp4")
w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm(w2_s_pad, group_size=group_size)
return w2_il, w2_s_il, w2_b_pad
@pytest.mark.parametrize(
"num_experts,hidden,inter",
[
# Aligned shapes (no padding needed).
(4, 256, 256),
(8, 768, 384),
(8, 1024, 1024),
# Non-aligned shapes (exercise the de-interleave + pad path).
# 192 % 128 = 64, so N_pad = K_pad = 256 (round_up(192, 128)).
(4, 192, 192),
# GPT-OSS-20B-like: hidden=2880, inter=2880 -> padded to 2944.
# Use smaller E to keep memory bounded.
(4, 2880, 2880),
],
)
def test_process_weights_matches_direct_interleave(num_experts, hidden, inter):
"""``_process_weights_for_sm90_cutlass`` must produce the same bytes as
a manual de-interleave + pad + halved-swap + interleave reference."""
w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter)
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)
method._process_weights_for_sm90_cutlass(layer)
N_pad = _round_up(inter, 128)
K_pad = _round_up(hidden, 128)
ref_w13, ref_w13_s, ref_w13_b = _expected_w13_processed(
w13, w13_s, w13_b, N_pad, K_pad, GROUP_SIZE
)
ref_w2, ref_w2_s, ref_w2_b = _expected_w2_processed(
w2, w2_s, w2_b, N_pad, K_pad, GROUP_SIZE
)
assert torch.equal(layer.w13_weight.data, ref_w13)
assert torch.equal(layer.w2_weight.data, ref_w2)
assert torch.equal(layer.w13_weight_scale.data, ref_w13_s)
assert torch.equal(layer.w2_weight_scale.data, ref_w2_s)
assert torch.equal(layer.w13_weight_bias.data, ref_w13_b)
assert torch.equal(layer.w2_weight_bias.data, ref_w2_b)
# SwiGLU per-expert scalars seeded with GPT-OSS defaults.
assert torch.allclose(
layer.swiglu_alpha,
torch.full((num_experts,), 1.702, dtype=torch.float32, device="cuda"),
)
assert torch.allclose(
layer.swiglu_beta,
torch.full((num_experts,), 1.0, dtype=torch.float32, device="cuda"),
)
assert torch.allclose(
layer.swiglu_limit,
torch.full((num_experts,), 7.0, dtype=torch.float32, device="cuda"),
)
@pytest.mark.parametrize(
"tokens,num_experts,hidden,inter,top_k",
[
# Aligned shapes (no padding).
(4, 4, 256, 256, 2),
(16, 8, 768, 384, 2),
(32, 8, 1024, 1024, 4),
# Non-aligned (exercises pad x + trim output).
(8, 4, 192, 192, 2),
],
)
def test_apply_sm90_cutlass_matches_flashinfer_direct(
tokens, num_experts, hidden, inter, top_k, monkeypatch, stated_tp_group
):
"""End-to-end: SGLang's ``_apply_sm90_cutlass`` must produce the same
output as a direct FlashInfer ``cutlass_fused_moe`` call fed with the
same processed weights / scales / biases. The processing pipeline is
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_cutlass as fi_cutlass_mod
# Bypass symmetric-memory / TP-group in the fused-func module, which is where
# the kernel call lives.
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.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
topk_w, topk_i = _make_topk(tokens, num_experts, top_k)
# ---- SGLang path ----
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)
method._process_weights_for_sm90_cutlass(layer)
out_sglang = method._apply_sm90_cutlass(
layer, _MockDispatchOutput(x.clone(), topk_w, topk_i)
).hidden_states
# ---- FlashInfer-direct reference using the same processed weights ----
K_pad = method._padded_hidden
if K_pad != hidden:
x_padded = torch.nn.functional.pad(
x.clone(), (0, K_pad - hidden), mode="constant", value=0.0
)
else:
x_padded = x.clone()
out_ref_padded = torch.empty(tokens, K_pad, dtype=torch.bfloat16, device="cuda")
cutlass_fused_moe(
input=x_padded,
token_selected_experts=topk_i.to(torch.int),
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.w2_weight_scale.view(torch.int32),
],
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,
use_w4_group_scaling=True,
activation_type=ActivationType.Swiglu,
use_fused_finalize=False,
output=out_ref_padded,
)
out_ref = (
out_ref_padded[:, :hidden].contiguous() if K_pad != hidden else out_ref_padded
)
assert torch.equal(out_sglang, out_ref), (
f"SGLang vs FlashInfer-direct mismatch; "
f"max abs diff = {(out_sglang.float() - out_ref.float()).abs().max().item():.4g}"
)
@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, stated_tp_group
):
"""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.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 +
# is_flashinfer_mxfp4 + is_fp4_experts.
# =============================================================================
def _make_random_dsv4_mxfp4(num_experts, hidden, inter, seed=0):
"""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(
-128,
128,
(num_experts, 2 * inter, hidden // 2),
dtype=torch.int8,
device="cuda",
generator=g,
)
w2 = torch.randint(
-128,
128,
(num_experts, hidden, inter // 2),
dtype=torch.int8,
device="cuda",
generator=g,
)
# Native E8M0 scales with exponents around 0 (= 2**0).
raw_e = torch.randint(
125,
130,
(num_experts, 2 * inter, hidden // GROUP_SIZE),
dtype=torch.uint8,
device="cuda",
generator=g,
)
raw_e2 = torch.randint(
125,
130,
(num_experts, hidden, inter // GROUP_SIZE),
dtype=torch.uint8,
device="cuda",
generator=g,
)
w13_s = raw_e.view(torch.float8_e8m0fnu)
w2_s = raw_e2.view(torch.float8_e8m0fnu)
return w13, w2, w13_s, w2_s
@pytest.mark.parametrize(
"tokens,num_experts,hidden,inter,top_k",
[
(4, 4, 256, 256, 2),
(16, 8, 768, 384, 2),
(256, 8, 1024, 1024, 4),
],
)
def test_dsv4_apply_matches_flashinfer_direct(
tokens, num_experts, hidden, inter, top_k, monkeypatch, stated_tp_group
):
"""End-to-end: SGLang's DSv4 ``Mxfp4FlashinferCutlassMoEMethod.apply``
output must match a direct FlashInfer ``cutlass_fused_moe`` call with
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
# Bypass symmetric-memory / TP-group stack in the new fused-func module
# (where DSv4 ``apply`` now dispatches the kernel call through).
monkeypatch.setattr(
fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
)
monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
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 ----
# plain SiLU * up — all three SwiGLU scalars None (no clamp configured).
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)
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
layer.moe_tp_size = 1
layer.moe_tp_rank = 0
layer.moe_ep_size = 1
layer.moe_ep_rank = 0
method.process_weights_after_loading(layer)
out_sglang = method.apply(
layer, _MockDispatchOutput(x.clone(), topk_w, topk_i)
).hidden_states
# ---- Direct FlashInfer reference ----
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(
w31.view(torch.uint8).contiguous(), "fp4"
)
ref_w2 = interleave_moe_weights_for_sm90_mixed_gemm(
w2.view(torch.uint8).contiguous(), "fp4"
)
ref_w13_s = interleave_moe_scales_for_sm90_mixed_gemm(
w13_s_u8, group_size=GROUP_SIZE
)
ref_w2_s = interleave_moe_scales_for_sm90_mixed_gemm(w2_s_u8, group_size=GROUP_SIZE)
out_ref = torch.empty(tokens, hidden, dtype=torch.bfloat16, device="cuda")
cutlass_fused_moe(
input=x.clone(),
token_selected_experts=topk_i,
token_final_scales=topk_w,
fc1_expert_weights=ref_w13,
fc2_expert_weights=ref_w2,
output_dtype=torch.bfloat16,
quant_scales=[ref_w13_s.view(torch.int32), ref_w2_s.view(torch.int32)],
fc1_expert_biases=None,
fc2_expert_biases=None,
swiglu_alpha=None,
swiglu_beta=None,
swiglu_limit=None,
use_w4_group_scaling=True,
activation_type=ActivationType.Swiglu,
output=out_ref,
)
assert torch.equal(out_sglang, out_ref), (
f"DSv4 SGLang vs FlashInfer-direct mismatch; "
f"max abs diff = "
f"{(out_sglang.float() - out_ref.float()).abs().max().item():.4g}"
)
@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``
(an isinstance check) returns True without distributed init."""
def __init__(self, hidden_states, topk_weights, topk_ids):
from sglang.srt.layers.moe.topk import StandardTopKOutput
self.hidden_states = hidden_states
# router_logits is unused by Mxfp4FlashinferCutlassMoEMethod.apply;
# supply a placeholder of the right shape to keep the NamedTuple happy.
router_logits = torch.zeros(
topk_ids.shape[0],
int(topk_ids.max().item()) + 1 if topk_ids.numel() else 1,
dtype=torch.float32,
device=topk_ids.device,
)
self.topk_output = StandardTopKOutput(
topk_weights=topk_weights,
topk_ids=topk_ids,
router_logits=router_logits,
)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v"]))
Reference in New Issue View Git Blame Copy Permalink