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: