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
python/sglang/kernels/ops/moe/ep_moe_kernels.py
+14 -5
View File
@@ -1019,6 +1019,7 @@ def post_reorder_triton_kernel(
@triton.jit
def _fwd_kernel_ep_scatter_1(
num_recv_tokens_per_expert,
num_valid_tokens_per_expert,
expert_start_loc,
m_indices,
num_experts: tl.constexpr,
@@ -1037,15 +1038,17 @@ def _fwd_kernel_ep_scatter_1(
tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts)
cur_expert_start = tl.load(expert_start_loc + cur_expert)
cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert)
cur_expert_padded_token_num = tl.load(num_recv_tokens_per_expert + cur_expert)
cur_expert_valid_token_num = tl.load(num_valid_tokens_per_expert + cur_expert)
m_indices_start_ptr = m_indices + cur_expert_start
off_expert = tl.arange(0, BLOCK_E)
for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4):
for start_m in tl.range(0, cur_expert_padded_token_num, BLOCK_E, num_stages=4):
offsets = start_m + off_expert
tl.store(
m_indices_start_ptr + start_m + off_expert,
cur_expert,
m_indices_start_ptr + offsets,
tl.where(offsets < cur_expert_valid_token_num, cur_expert, -1),
)
@@ -1137,6 +1140,7 @@ def ep_scatter(
recv_x_scale: torch.Tensor,
recv_topk: torch.Tensor,
num_recv_tokens_per_expert: torch.Tensor,
num_valid_tokens_per_expert: torch.Tensor,
expert_start_loc: torch.Tensor,
output_tensor: torch.Tensor,
output_tensor_scale: torch.Tensor,
@@ -1172,6 +1176,7 @@ def ep_scatter(
_fwd_kernel_ep_scatter_1[(grid,)](
num_recv_tokens_per_expert,
num_valid_tokens_per_expert,
expert_start_loc,
m_indices,
num_experts=num_experts,
@@ -1575,7 +1580,11 @@ def moe_ep_deepgemm_preprocess(
assert len(block_shape) == 2
block_n, block_k = block_shape[0], block_shape[1]
is_fp8 = output_dtype == torch.float8_e4m3fn
if is_fp8 and use_mxfp8:
# Quantize FP8 values with the UE8M0 scale directly. Rounding only the
# scale afterward can change the represented activation by up to 2x.
from sglang.srt.layers import deep_gemm_wrapper
if is_fp8 and (use_mxfp8 or deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0):
from sglang.kernels.ops.quantization.minimax_quant_ue8m0 import (
per_token_quant_fp8_ue8m0_scatter,
)
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+167 -18
View File
@@ -98,6 +98,24 @@ def copy_list_to_gpu_no_ce(arr: List[int]):
return tensor_gpu
def _should_use_masked_standard_layout(runner_config: MoeRunnerConfig) -> bool:
"""Use masked GEMM when expert parallelism keeps its buffer small."""
return (
runner_config.num_experts > runner_config.num_local_experts
and runner_config.num_local_experts <= 32
)
def _get_compact_all_tokens(
num_assignments: int, num_experts: int, block_e: int = 128
) -> int:
"""Return the maximum padded rows over all routings of the assignments."""
max_nonempty_experts = min(num_assignments, num_experts)
return block_e * (
max_nonempty_experts + (num_assignments - max_nonempty_experts) // block_e
)
@dataclass
class DeepGemmRunnerInput(RunnerInput):
hidden_states: torch.Tensor
@@ -670,7 +688,11 @@ def pre_permute_standard_to_deep_gemm(
runner_config: MoeRunnerConfig,
running_state: dict,
) -> DeepGemmRunnerInput:
from sglang.kernels.ops.moe.ep_moe_kernels import moe_ep_deepgemm_preprocess
from sglang.kernels.ops.moe.ep_moe_kernels import (
ep_scatter,
fused_moe_dispatch_index,
moe_ep_deepgemm_preprocess,
)
hidden_states, topk_output = (
dispatch_output.hidden_states,
@@ -685,25 +707,151 @@ def pre_permute_standard_to_deep_gemm(
topk_weights, topk_ids = topk_weights, topk_ids
# PreReorder
if _should_use_masked_standard_layout(runner_config):
output_dtype = (
torch.bfloat16
if quant_info.w13_weight.dtype == torch.bfloat16
else torch.float8_e4m3fn
)
masked_m, _, src2dst, hidden_states, hidden_states_scale = (
moe_ep_deepgemm_preprocess(
topk_ids,
runner_config.num_local_experts,
hidden_states,
runner_config.top_k,
quant_info.block_shape,
output_dtype=output_dtype,
use_mxfp8=quant_info.use_mxfp8,
)
)
# Use the global expert count because expected_m is a tuning hint, not
# the per-rank buffer capacity.
expected_m = max(
1,
ceil_div(
hidden_states_shape[0] * runner_config.top_k,
runner_config.num_experts,
),
)
if runner_config.inplace:
dispose_tensor(hidden_states_ref)
running_state["topk_ids"] = topk_ids
running_state["topk_weights"] = topk_weights
running_state["hidden_states_shape"] = hidden_states_shape
running_state["hidden_states_dtype"] = hidden_states_dtype
running_state["hidden_states_device"] = hidden_states_device
running_state["src2dst"] = src2dst
running_state["mxfp8_act_gran_k"] = (
quant_info.block_shape[1] if quant_info.block_shape else 128
)
return DeepGemmRunnerInput(
hidden_states=hidden_states,
hidden_states_scale=hidden_states_scale,
use_masked_gemm=True,
masked_m=masked_m,
expected_m=expected_m,
)
# The compact layout avoids scaling masked buffers with the expert count.
# Scatter and post-permute skip non-local experts mapped to -1.
block_e = 128
num_experts = runner_config.num_local_experts
num_assignments = topk_ids.numel()
all_tokens = _get_compact_all_tokens(num_assignments, num_experts, block_e)
tokens_per_expert, unused_masked_dst = fused_moe_dispatch_index(
topk_ids, num_experts, 1
)
dispose_tensor(unused_masked_dst)
valid_tokens_per_expert = tokens_per_expert
tokens_per_expert = (ceil_div(tokens_per_expert, block_e) * block_e).to(torch.int32)
# Keep graph-static shapes by appending padding to the final segment.
# Its m_indices stay -1, so DeepGEMM skips those rows.
tokens_per_expert[-1].add_(all_tokens - tokens_per_expert.sum())
k = hidden_states.size(1)
output_dtype = (
torch.bfloat16
if quant_info.w13_weight.dtype == torch.bfloat16
else torch.float8_e4m3fn
)
masked_m, expected_m, src2dst, hidden_states, hidden_states_scale = (
moe_ep_deepgemm_preprocess(
topk_ids,
runner_config.num_local_experts,
hidden_states,
runner_config.top_k,
quant_info.block_shape,
output_dtype=output_dtype,
use_mxfp8=quant_info.use_mxfp8,
if output_dtype == torch.bfloat16:
packed_input_source = hidden_states
packed_input_source_scale = None
packed_input = torch.empty(
(all_tokens, k), device=hidden_states_device, dtype=torch.bfloat16
)
# ep_scatter ignores scales for BF16, but a real tensor keeps its
# Triton signature uniform across the existing DeepEP caller.
packed_input_scale = torch.empty(
(all_tokens, 1), device=hidden_states_device, dtype=torch.float32
)
else:
from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_fp8,
)
)
dispose_tensor(hidden_states_ref)
block_k = quant_info.block_shape[1] if quant_info.block_shape else 128
packed_input_source, packed_input_source_scale = (
sglang_per_token_group_quant_fp8(
hidden_states,
block_k,
column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
)
)
packed_input = torch.zeros(
(all_tokens, k),
device=hidden_states_device,
dtype=torch.float8_e4m3fn,
)
scale_width = k // block_k
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
scale_width = ceil_div(scale_width, 4)
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
packed_input_scale = torch.zeros(
(scale_width, all_tokens),
device=hidden_states_device,
dtype=packed_input_source_scale.dtype,
).transpose(0, 1)
else:
packed_input_scale = torch.zeros(
(all_tokens, scale_width),
device=hidden_states_device,
dtype=packed_input_source_scale.dtype,
)
expert_start_loc = torch.empty(
num_experts, device=hidden_states_device, dtype=torch.int32
)
m_indices = torch.empty(all_tokens, device=hidden_states_device, dtype=torch.int32)
src2dst = torch.empty_like(topk_ids, dtype=torch.int32)
ep_scatter(
packed_input_source,
packed_input_source_scale,
topk_ids,
tokens_per_expert,
valid_tokens_per_expert,
expert_start_loc,
packed_input,
packed_input_scale,
m_indices,
src2dst,
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
quant_block_size=(quant_info.block_shape[1] if quant_info.block_shape else 128),
)
if packed_input_source is not hidden_states:
dispose_tensor(packed_input_source)
if packed_input_source_scale is not None:
dispose_tensor(packed_input_source_scale)
# Preserve the input when a shared expert or its gate may still use it.
if runner_config.inplace:
dispose_tensor(hidden_states_ref)
running_state["topk_ids"] = topk_ids
running_state["topk_weights"] = topk_weights
@@ -711,16 +859,16 @@ def pre_permute_standard_to_deep_gemm(
running_state["hidden_states_dtype"] = hidden_states_dtype
running_state["hidden_states_device"] = hidden_states_device
running_state["src2dst"] = src2dst
running_state["all_tokens"] = all_tokens
running_state["mxfp8_act_gran_k"] = (
quant_info.block_shape[1] if quant_info.block_shape else 128
)
return DeepGemmRunnerInput(
hidden_states=hidden_states,
hidden_states_scale=hidden_states_scale,
use_masked_gemm=True,
masked_m=masked_m,
expected_m=expected_m,
hidden_states=packed_input,
hidden_states_scale=packed_input_scale,
use_masked_gemm=False,
m_indices=m_indices,
)
@@ -891,6 +1039,7 @@ def pre_permute_deepep_normal_to_deep_gemm(
hidden_states_scale,
topk_ids,
num_recv_tokens_per_expert_gpu,
num_recv_tokens_per_expert_gpu,
expert_start_loc,
input_tensor,
input_tensor_scale,
python/sglang/srt/layers/quantization/fp8.py
+8 -12
View File
@@ -1603,7 +1603,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
else:
# For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE
# Check if MoE will actually use DeepGEMM runner
will_use_deepgemm = self.is_deepgemm_moe_runner_backend_enabled()
@@ -1677,20 +1676,17 @@ class Fp8MoEMethod(FusedMoEMethodBase):
if not self.is_fp4_expert:
weight_block_size = self.quant_config.weight_block_size
if requant_block_scale_ue8m0_for_deepgemm(
layer.w13_weight,
layer.w13_weight_scale_inv,
weight_block_size,
use_deepgemm_runner=will_use_deepgemm,
for weight, weight_scale in (
(layer.w13_weight, layer.w13_weight_scale_inv),
(layer.w2_weight, layer.w2_weight_scale_inv),
):
assert isinstance(
layer, DeepEPMoE
), "DeepGemm MoE is only supported with DeepEPMoE"
requant_block_scale_ue8m0_for_deepgemm(
layer.w2_weight,
layer.w2_weight_scale_inv,
weight,
weight_scale,
weight_block_size,
use_deepgemm_runner=True,
use_deepgemm_runner=will_use_deepgemm,
output_dtype=torch.bfloat16,
weight_shape=weight.shape[-2:],
)
def _convert_mxfp8_moe_to_block_fp8(self, layer: Module) -> None:
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)