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

Support DeepGEMM for standard MoE dispatch (#33128)

Browse Source 
Co-authored-by: Sam Li <lsam@nvidia.com>
This commit is contained in:
YAMY
2026-08-02 21:48:13 -07:00
committed by GitHub
co-authored by Sam Li
parent f5f021672a
commit 5fe97637df
5 changed files with 485 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/moe/test_minimax_quant_scatter.py
+240 -1
View File
@@ -1,14 +1,31 @@
import random
import sys
from contextlib import nullcontext
import pytest
import torch
from sglang.kernels.ops.moe.ep_moe_kernels import fill_gateup_input_triton_kernel
import sglang.srt.layers.moe.moe_runner.deep_gemm as deep_gemm_runner
from sglang.kernels.ops.moe.ep_moe_kernels import (
fill_gateup_input_triton_kernel,
moe_ep_deepgemm_preprocess,
)
from sglang.kernels.ops.quantization.minimax_quant_ue8m0 import (
per_token_quant_fp8_ue8m0,
per_token_quant_fp8_ue8m0_scatter,
)
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.deep_gemm import (
DeepGemmMoeQuantInfo,
DeepGemmRunnerCore,
post_permute_deep_gemm_to_standard,
pre_permute_standard_to_deep_gemm,
)
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
from sglang.srt.layers.quantization.fp8_utils import (
quant_weight_ue8m0,
transform_scale_ue8m0,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
@@ -80,5 +97,227 @@ def test_quant_scatter_matches_quant_plus_fill(num_tokens, topk, hidden, group):
), f"scale mismatch token={t} slot={j} expert={e}"
def test_standard_deepgemm_preprocess_quantizes_with_ue8m0_scale():
arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
if arch_major <= 9:
pytest.skip("UE8M0 fusion is Blackwell-only")
num_tokens, topk, hidden, group, num_experts = 7, 4, 2048, 128, 8
torch.manual_seed(1234)
x = torch.randn(num_tokens, hidden, device=dev, dtype=torch.bfloat16)
topk_ids = torch.stack(
[
(torch.arange(topk, device=dev, dtype=torch.int32) + token) % num_experts
for token in range(num_tokens)
]
)
_, _, src2dst, grouped_x, grouped_scale = moe_ep_deepgemm_preprocess(
topk_ids=topk_ids,
num_local_experts=num_experts,
hidden_states=x,
top_k=topk,
block_shape=[group, group],
output_dtype=torch.float8_e4m3fn,
use_mxfp8=False,
)
direct_x, direct_scale = per_token_quant_fp8_ue8m0(x, group)
assert grouped_scale.dtype == torch.int32
for token in range(num_tokens):
for slot in range(topk):
dst = int(src2dst[token * topk + slot])
expert, row = divmod(dst, grouped_x.shape[1])
assert torch.equal(
grouped_x[expert, row].view(torch.uint8),
direct_x[token].view(torch.uint8),
)
assert torch.equal(
grouped_scale[expert, row],
direct_scale[token],
)
@pytest.mark.parametrize(
"num_assignments,num_experts,expected",
[
(14, 2, 256),
(10, 512, 1280),
(20, 512, 2560),
(320, 512, 40960),
(640, 512, 65664),
(1280, 512, 66304),
],
)
def test_compact_all_tokens_uses_tight_routing_independent_bound(
num_assignments, num_experts, expected
):
assert (
deep_gemm_runner._get_compact_all_tokens(num_assignments, num_experts)
== expected
)
@pytest.mark.parametrize("weight_dtype", ["fp8", "bf16"])
def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_dtype):
"""Exercise both production grouped GEMMs through the standard path."""
arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
if arch_major <= 9:
pytest.skip("DeepGEMM UE8M0 is Blackwell-only")
# This kernel test runs outside a model-parallel process. Bypass only the
# symmetric-allocation context; all pre-permute, DeepGEMM, activation,
# quantization, down-GEMM, and post-permute kernels remain real.
monkeypatch.setattr(deep_gemm_runner, "get_tp_group", lambda: None)
monkeypatch.setattr(
deep_gemm_runner,
"use_symmetric_memory",
lambda *args, **kwargs: nullcontext(),
)
# UE8M0 packs four 128-wide scale groups into each int32. Use the smallest
# legal K for both the gate/up and down GEMMs.
num_tokens, hidden, intermediate, topk, num_local_experts = 7, 512, 512, 2, 2
torch.manual_seed(20260730)
hidden_states = torch.randn(num_tokens, hidden, device=dev, dtype=torch.bfloat16)
topk_ids = torch.tensor(
[
[0, -1],
[1, -1],
[0, 1],
[-1, 1],
[0, -1],
[1, 0],
[-1, 1],
],
device=dev,
dtype=torch.int32,
)
topk_weights = torch.tensor(
[
[0.8, 0.2],
[0.7, 0.3],
[0.6, 0.4],
[0.1, 0.9],
[0.75, 0.25],
[0.55, 0.45],
[0.35, 0.65],
],
device=dev,
dtype=torch.float32,
)
weight_std = hidden**-0.5
w13_bf16 = (
torch.randn(
num_local_experts,
2 * intermediate,
hidden,
device=dev,
dtype=torch.bfloat16,
)
* weight_std
)
w2_bf16 = (
torch.randn(
num_local_experts,
hidden,
intermediate,
device=dev,
dtype=torch.bfloat16,
)
* weight_std
)
if weight_dtype == "fp8":
w13, w13_scale = quant_weight_ue8m0(w13_bf16, [128, 128])
w2, w2_scale = quant_weight_ue8m0(w2_bf16, [128, 128])
quant_info = DeepGemmMoeQuantInfo(
w13_weight=w13,
w2_weight=w2,
use_fp8=True,
w13_scale=transform_scale_ue8m0(w13_scale, mn=w13.shape[-2]),
w2_scale=transform_scale_ue8m0(w2_scale, mn=w2.shape[-2]),
block_shape=[128, 128],
)
else:
quant_info = DeepGemmMoeQuantInfo(
w13_weight=w13_bf16,
w2_weight=w2_bf16,
use_fp8=False,
)
dispatch_output = StandardDispatchOutput(
hidden_states=hidden_states,
hidden_states_scale=None,
topk_output=(topk_weights, topk_ids, None),
)
def run_with_num_experts(num_experts):
config = MoeRunnerConfig(
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden,
intermediate_size_per_partition=intermediate,
top_k=topk,
activation="silu",
is_gated=True,
inplace=False,
)
running_state = {}
runner_input = pre_permute_standard_to_deep_gemm(
dispatch_output,
quant_info,
config,
running_state,
)
runner_output = DeepGemmRunnerCore(config).run(
runner_input,
quant_info,
running_state,
)
return (
runner_input.use_masked_gemm,
running_state.get("all_tokens"),
runner_input.m_indices,
post_permute_deep_gemm_to_standard(
runner_output,
quant_info,
config,
running_state,
).hidden_states,
)
compact_is_masked, compact_all_tokens, compact_m_indices, compact_output = (
run_with_num_experts(num_local_experts)
)
masked_is_masked, masked_all_tokens, masked_m_indices, masked_output = (
run_with_num_experts(8)
)
torch.cuda.synchronize()
assert not compact_is_masked
assert masked_is_masked
assert compact_all_tokens == 256
assert masked_all_tokens is None
assert masked_m_indices is None
valid_assignments = topk_ids[topk_ids >= 0]
assert torch.equal(
torch.bincount(
compact_m_indices[compact_m_indices >= 0],
minlength=num_local_experts,
),
torch.bincount(valid_assignments, minlength=num_local_experts),
)
assert (compact_m_indices == -1).sum() == compact_all_tokens - len(
valid_assignments
)
torch.testing.assert_close(
masked_output,
compact_output,
rtol=5e-2,
atol=5e-2,
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-x"]))
test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py
+56 -1
View File
@@ -5,13 +5,14 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest
from unittest.mock import patch
from unittest.mock import call, patch
import torch
from compressed_tensors.quantization import QuantizationStrategy
import sglang.srt.layers.quantization.fp8_utils as fp8_utils
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.quantization import fp8 as fp8_quant
from sglang.srt.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 import (
CompressedTensorsW8A8Fp8,
)
@@ -162,6 +163,60 @@ class TestDeepGemmUE8M0Requant(CustomTestCase):
self.assertTrue(layer.weight_scale.format_ue8m0)
requant.assert_called_once()
def test_fp8_moe_requants_standard_layer_for_deepgemm(self):
method = fp8_quant.Fp8MoEMethod.__new__(fp8_quant.Fp8MoEMethod)
method.convert_mxfp8_to_block = False
method.use_mxfp8 = False
method.is_fp4_expert = False
method.dequant_fp4_to_fp8 = False
method.quant_config = unittest.mock.Mock(weight_block_size=BLOCK_SIZE)
layer = torch.nn.Module()
layer.w13_weight, layer.w13_weight_scale_inv = _make_params()
layer.w2_weight, layer.w2_weight_scale_inv = _make_params()
def _mark_ue8m0(weight, weight_scale, *args, **kwargs):
weight_scale.format_ue8m0 = True
return True
with patch.multiple(
fp8_quant,
_is_cpu=False,
_is_fp8_fnuz=False,
_use_aiter=False,
), patch.object(
method, "is_deepgemm_moe_runner_backend_enabled", return_value=True
), patch.object(
fp8_quant,
"requant_block_scale_ue8m0_for_deepgemm",
side_effect=_mark_ue8m0,
) as requant:
method.process_weights_after_loading_block_quant(layer)
self.assertEqual(
requant.call_args_list,
[
call(
layer.w13_weight,
layer.w13_weight_scale_inv,
BLOCK_SIZE,
use_deepgemm_runner=True,
output_dtype=torch.bfloat16,
weight_shape=layer.w13_weight.shape[-2:],
),
call(
layer.w2_weight,
layer.w2_weight_scale_inv,
BLOCK_SIZE,
use_deepgemm_runner=True,
output_dtype=torch.bfloat16,
weight_shape=layer.w2_weight.shape[-2:],
),
],
)
self.assertTrue(layer.w13_weight_scale_inv.format_ue8m0)
self.assertTrue(layer.w2_weight_scale_inv.format_ue8m0)
if __name__ == "__main__":
unittest.main(verbosity=3)