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

Amd/dsv4 shared experts fusion top6 (#32340)

Browse Source 
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
This commit is contained in:
karverma-amd
2026-08-19 22:57:58 -07:00
committed by GitHub
co-authored by Claude Opus 4.8 kk Thomas Wang
parent 02b93e7e01
commit 50dae2d99d
6 changed files with 291 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py
+49 -23
View File
@@ -1259,6 +1259,8 @@ def _fused_append_shared_experts_kernel(
scale_factor, # runtime scalar
K: tl.constexpr,
S: tl.constexpr,
BLOCK_K: tl.constexpr,
BLOCK_S: tl.constexpr,
):
"""
for m in range(M):
@@ -1276,20 +1278,25 @@ def _fused_append_shared_experts_kernel(
out_ids_row_ptr = pid * (K + S)
out_w_row_ptr = pid * (K + S)
offs_k = tl.arange(0, K)
ids = tl.load(topk_ids_ptr + ids_row_ptr + offs_k)
ws = tl.load(topk_weights_ptr + w_row_ptr + offs_k)
# tl.arange requires a power-of-2 range, but K (topk) and S (num shared
# experts) need not be pow2 -- DeepSeek-V4 uses top-6. Iterate over the
# next-pow2 block and mask the tail (mirrors the _with_weights sibling).
offs_k = tl.arange(0, BLOCK_K)
mask_k = offs_k < K
ids = tl.load(topk_ids_ptr + ids_row_ptr + offs_k, mask=mask_k)
ws = tl.load(topk_weights_ptr + w_row_ptr + offs_k, mask=mask_k)
tl.store(out_ids_ptr + out_ids_row_ptr + offs_k, ids)
tl.store(out_weights_ptr + out_w_row_ptr + offs_k, ws)
tl.store(out_ids_ptr + out_ids_row_ptr + offs_k, ids, mask=mask_k)
tl.store(out_weights_ptr + out_w_row_ptr + offs_k, ws, mask=mask_k)
offs_s = tl.arange(0, S)
offs_s = tl.arange(0, BLOCK_S)
mask_s = offs_s < S
shared_ids = tl.cast(N_BASE + offs_s, ids.dtype)
shared_ws = tl.full([S], scale_factor, dtype=ws.dtype)
shared_ws = tl.full([BLOCK_S], scale_factor, dtype=ws.dtype)
tl.store(out_ids_ptr + out_ids_row_ptr + K + offs_s, shared_ids)
tl.store(out_weights_ptr + out_w_row_ptr + K + offs_s, shared_ws)
tl.store(out_ids_ptr + out_ids_row_ptr + K + offs_s, shared_ids, mask=mask_s)
tl.store(out_weights_ptr + out_w_row_ptr + K + offs_s, shared_ws, mask=mask_s)
def fused_append_shared_experts(
@@ -1315,6 +1322,8 @@ def fused_append_shared_experts(
scale_factor=scale_factor,
K=k,
S=s,
BLOCK_K=triton.next_power_of_2(k),
BLOCK_S=triton.next_power_of_2(s),
num_warps=1,
)
return out_ids, out_weights
@@ -1333,6 +1342,8 @@ def _fused_append_remap_shared_experts_deepep_kernel(
pad_fill_id, # runtime scalar: routed-id fill for padded rows
K: tl.constexpr,
S: tl.constexpr,
BLOCK_K: tl.constexpr,
BLOCK_S: tl.constexpr,
HAS_PADDING: tl.constexpr,
):
"""Append shared experts AND apply the DeepEP interleaved remap in one pass.
@@ -1342,7 +1353,8 @@ def _fused_append_remap_shared_experts_deepep_kernel(
loaded into registers, so it costs a few ALU ops instead of ~6 extra eager
kernel launches (div_floor / add / arange / fill / copy) per MoE layer.
Routed IDs: e -> e + e // num_local_routed (insert gaps for shared slots)
Routed IDs: e -> e + (e // num_local_routed) * S (insert S-wide gaps for
the shared slots that precede this id's rank)
Shared IDs: shared_id_base + arange(S) (one id per shared slot)
Shared wgt: scale_factor (1.0 on aiter; 1/rsf otherwise)
"""
@@ -1351,31 +1363,43 @@ def _fused_append_remap_shared_experts_deepep_kernel(
ids_row_ptr = pid * K
out_ids_row_ptr = pid * (K + S)
offs_k = tl.arange(0, K)
ids = tl.load(topk_ids_ptr + ids_row_ptr + offs_k)
ws = tl.load(topk_weights_ptr + ids_row_ptr + offs_k)
# tl.arange requires a power-of-2 range, but K (topk) and S (num shared
# experts) need not be pow2 -- DeepSeek-V4 uses top-6. Iterate over the
# next-pow2 block and mask the tail (mirrors the _append/_with_weights
# siblings), otherwise K=6 fails with "arange's range must be a power of 2".
offs_k = tl.arange(0, BLOCK_K)
mask_k = offs_k < K
ids = tl.load(topk_ids_ptr + ids_row_ptr + offs_k, mask=mask_k)
ws = tl.load(topk_weights_ptr + ids_row_ptr + offs_k, mask=mask_k)
# DeepEP interleaved layout: shift each routed id past the shared slots that
# precede it. Matches `routed + routed // num_local_routed` exactly.
ids = ids + ids // num_local_routed
# DeepEP interleaved layout: shift each routed id past ALL shared slots that
# precede its rank. Rank r == id // num_local_routed contributes r*S shared
# slots ahead of the id, so the gap is (id // num_local_routed) * S -- not a
# single slot. With S == 1 this reduces to the old `id // num_local_routed`,
# but S > 1 (e.g. multiple fused shared experts) needs the full S-wide gap or
# routed ids collide with an earlier rank's shared slots.
ids = ids + (ids // num_local_routed) * S
if HAS_PADDING:
# Fold the padded-topk_ids fill (previously a separate _fill_padded_rows
# launch): rows >= num_token_non_padded get pad_fill_id in every routed
# slot. Matches the old fill(topk_ids=0) -> remap(0)=0 when pad_fill_id==0.
# ids is a BLOCK_K-wide register tile (K need not be pow2), so fill the
# whole tile and let the masked store below drop the tail.
n_valid = tl.load(num_token_non_padded_ptr)
if pid >= n_valid:
ids = tl.full((K,), pad_fill_id, dtype=ids.dtype)
ids = tl.full((BLOCK_K,), pad_fill_id, dtype=ids.dtype)
tl.store(out_ids_ptr + out_ids_row_ptr + offs_k, ids)
tl.store(out_weights_ptr + out_ids_row_ptr + offs_k, ws)
tl.store(out_ids_ptr + out_ids_row_ptr + offs_k, ids, mask=mask_k)
tl.store(out_weights_ptr + out_ids_row_ptr + offs_k, ws, mask=mask_k)
offs_s = tl.arange(0, S)
offs_s = tl.arange(0, BLOCK_S)
mask_s = offs_s < S
shared_ids = tl.cast(shared_id_base + offs_s, ids.dtype)
shared_ws = tl.full([S], scale_factor, dtype=ws.dtype)
shared_ws = tl.full([BLOCK_S], scale_factor, dtype=ws.dtype)
tl.store(out_ids_ptr + out_ids_row_ptr + K + offs_s, shared_ids)
tl.store(out_weights_ptr + out_ids_row_ptr + K + offs_s, shared_ws)
tl.store(out_ids_ptr + out_ids_row_ptr + K + offs_s, shared_ids, mask=mask_s)
tl.store(out_weights_ptr + out_ids_row_ptr + K + offs_s, shared_ws, mask=mask_s)
def fused_append_remap_shared_experts_deepep(
@@ -1419,6 +1443,8 @@ def fused_append_remap_shared_experts_deepep(
pad_fill_id,
K=k,
S=s,
BLOCK_K=triton.next_power_of_2(k),
BLOCK_S=triton.next_power_of_2(s),
HAS_PADDING=has_padding,
num_warps=1,
)
python/sglang/kernels/ops/moe/moe_fused_gate.py
+9 -2
View File
@@ -89,7 +89,7 @@ def moe_fused_gate_jit(
@triton.jit
def _router_triton_kernel(
scores_ptr, # [M, N] fp32, GEMM output (raw logits)
bias_ptr, # [N] fp32
bias_ptr, # [N] fp32/fp16/bf16 (upcast to fp32 on load)
out_weights_ptr, # [M, K] fp32
out_indices_ptr, # [M, K] int32
M,
@@ -282,7 +282,14 @@ def moe_fused_gate(
torch.float16,
torch.bfloat16,
), "scores must be float32/float16/bfloat16"
assert bias.dtype == torch.float32, "bias must be float32"
# The kernel loads the bias and upcasts it to fp32 in-register (see
# _router_triton_kernel), so a non-fp32 bias (DeepSeek-V4 stores the
# correction bias in bf16) needs no host-side cast/copy.
assert bias.dtype in (
torch.float32,
torch.float16,
torch.bfloat16,
), "bias must be float32/float16/bfloat16"
assert scores.ndim == 2, "scores must be 2D"
assert bias.ndim == 1, "bias must be 1D"
assert scores.size(1) == bias.size(0), "scores and bias must have same num_experts"
python/sglang/srt/layers/moe/topk.py
+13 -3
View File
@@ -1271,6 +1271,10 @@ def biased_topk_jit_kernel_impl(
else:
from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate
# DeepSeek-V4 stores e_score_correction_bias in bf16 (for the aiter
# sqrtsoftplus topk path). moe_fused_gate upcasts the bias to fp32
# in-register, so pass it through directly rather than allocating a fresh
# fp32 copy of this static routing bias on every MoE invocation.
topk_weights, topk_ids = moe_fused_gate(
gating_output,
correction_bias,
@@ -1839,7 +1843,7 @@ def remap_topk_for_per_rank_shared_slots(
so tokens route to the correct rank. The layout is ordered by rank:
[rank0 routed..., rank0 shared, rank1 routed..., rank1 shared, ...].
Routed IDs: e -> e + e // num_local_routed
Routed IDs: e -> e + (e // num_local_routed) * num_fused_shared_experts
Shared IDs: ep_rank * num_local_experts + num_local_routed
Shared weight: 1.0 on the aiter path, else 1/routed_scaling_factor (see below).
"""
@@ -1854,9 +1858,15 @@ def remap_topk_for_per_rank_shared_slots(
num_local_routed = num_physical_routed_experts // ep_size
num_local_experts = num_local_routed + num_fused_shared_experts
# Remap routed IDs: insert gaps for shared expert slots (single fused op)
# Remap routed IDs: insert gaps for shared expert slots (single fused op).
# Each rank r == e // num_local_routed is preceded by r shared-slot blocks of
# width num_fused_shared_experts, so shift by (e // num_local_routed) * S --
# a single-slot shift (S == 1) would let routed ids collide with an earlier
# rank's shared slots once S > 1.
routed = topk_ids[:, :-num_fused_shared_experts]
topk_ids[:, :-num_fused_shared_experts] = routed + routed // num_local_routed
topk_ids[:, :-num_fused_shared_experts] = (
routed + (routed // num_local_routed) * num_fused_shared_experts
)
# Set shared expert IDs to route to home rank (vectorized)
topk_ids[:, -num_fused_shared_experts:] = (