minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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:] = (
test/registered/kernels/ops/moe/test_moe_fused_gate.py
+36
View File
@@ -521,5 +521,41 @@ def test_grouped_dispatch_flag_matches_default(
)
@pytest.mark.parametrize("bias_dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("scoring_func", ["sigmoid", "sqrtsoftplus"])
def test_moe_fused_gate_accepts_non_fp32_bias(
bias_dtype: torch.dtype, scoring_func: str
) -> None:
"""A non-fp32 correction bias must route identically to its fp32 copy.
DeepSeek-V4 stores ``e_score_correction_bias`` in bf16 and topk.py now passes
it straight through (no host-side fp32 copy per MoE invocation). The kernel
upcasts the bias to fp32 on load, so a bf16/fp16 bias must be bit-identical to
casting it to fp32 first -- this guards both the relaxed dtype assertion and
the removed per-call cast. (top-6 exercises the DSV4 non-pow2 routing width.)
"""
M, num_experts, topk = 256, 128, 6
torch.manual_seed(num_experts * 11 + topk)
gating = torch.randn(M, num_experts, dtype=torch.float32, device=DEVICE) * 2.0
bias_lowp = torch.randn(num_experts, dtype=bias_dtype, device=DEVICE) * 0.5
bias_fp32 = bias_lowp.to(torch.float32)
lowp_w, lowp_i = moe_fused_gate(
gating, bias_lowp, topk=topk, scoring_func=scoring_func, renormalize=True
)
fp32_w, fp32_i = moe_fused_gate(
gating, bias_fp32, topk=topk, scoring_func=scoring_func, renormalize=True
)
torch.cuda.synchronize()
# Widening bf16/fp16 -> fp32 is exact, so the two routings must match exactly.
torch.testing.assert_close(
_scatter_by_expert(lowp_w, lowp_i, num_experts),
_scatter_by_expert(fp32_w, fp32_i, num_experts),
rtol=0,
atol=0,
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py
+69 -3
View File
@@ -34,7 +34,7 @@ def _reference_append_remap(
):
"""Pure-torch golden reference mirroring the kernel's documented contract.
Routed IDs: e -> e + e // num_local_routed
Routed IDs: e -> e + (e // num_local_routed) * s
Shared IDs: shared_id_base + arange(s)
Routed wgt: passthrough
Shared wgt: scale_factor
@@ -44,7 +44,7 @@ def _reference_append_remap(
out_w = torch.empty(
(m, k + s), dtype=topk_weights.dtype, device=topk_weights.device
)
out_ids[:, :k] = topk_ids + topk_ids // num_local_routed
out_ids[:, :k] = topk_ids + (topk_ids // num_local_routed) * s
out_w[:, :k] = topk_weights
shared = shared_id_base + torch.arange(s, device=topk_ids.device)
out_ids[:, k:] = shared.to(topk_ids.dtype)
@@ -57,12 +57,16 @@ def _reference_append_remap(
)
class TestFusedAppendRemapPerRankSharedSlots(CustomTestCase):
# (m, k, num_physical_routed, ep_size, ep_rank, num_fused_shared_experts).
# k and num_fused_shared_experts are kept powers of two (tl.arange constraint).
# Includes non-power-of-two k and num_fused_shared_experts (DeepSeek-V4 routes
# top-6): the kernel blocks over next_power_of_2 and masks, so these must work.
CASES = [
(1, 8, 256, 8, 0, 1),
(4, 8, 256, 8, 7, 1),
(17, 8, 264, 8, 3, 1),
(128, 16, 128, 4, 2, 2),
(1, 6, 258, 6, 0, 1), # DSV4: k=6 (non-pow2), s=1
(13, 6, 258, 6, 5, 1), # DSV4: k=6 (non-pow2), non-zero ep_rank
(32, 6, 264, 4, 2, 3), # non-pow2 k=6 and non-pow2 s=3 together
]
def _make_inputs(self, m, k, num_physical_routed, ids_dtype=torch.int64):
@@ -111,6 +115,68 @@ class TestFusedAppendRemapPerRankSharedSlots(CustomTestCase):
self.assertTrue(torch.equal(got_ids, exp_ids))
self.assertTrue(torch.allclose(got_w, exp_w))
def test_no_routed_shared_collision_across_ranks(self):
"""Remapped routed ids never land on any rank's shared slots (S > 1).
Independent of the gap-insertion math the kernel/eager path use: the
per-rank layout is, by definition, ep_size contiguous blocks of width
num_local_experts == num_local_routed + S, each block being
[num_local_routed routed ids ... S shared ids]. So a physical routed id
``e`` must map to ``rank * num_local_experts + local`` where
``rank = e // num_local_routed`` and ``local = e % num_local_routed`` --
derived from the block layout, not from ``e + (e // nlr) * S``.
This is the regression guard for the S > 1 bug: the old
``e + e // num_local_routed`` shifts by a single slot, so e.g. the first
routed id of rank 1 (e == num_local_routed) mapped to
``num_local_routed + 1``, colliding with rank 0's shared slots when
S > 1. The check asserts (a) the kernel matches the block-derived ids and
(b) no remapped routed id intersects the shared-slot id set of ANY rank.
"""
# Every config here uses S > 1 and spans all ep_size ranks (npr == m*k
# feeds each physical routed id exactly once) so rank boundaries are hit.
# (m, k, num_physical_routed, ep_size, num_fused_shared_experts).
CASES = [
(44, 6, 264, 4, 3), # DSV4-shaped: non-pow2 k=6, non-pow2 S=3
(32, 8, 256, 8, 2), # pow2 k, S=2
(43, 6, 258, 6, 4), # non-pow2 npr/rank boundaries, S=4
]
for m, k, npr, ep_size, s in CASES:
with self.subTest(m=m, k=k, npr=npr, ep_size=ep_size, s=s):
self.assertEqual(m * k, npr) # cover each physical id once
num_local_routed = npr // ep_size
num_local_experts = num_local_routed + s
device = get_device()
# Feed every physical routed id [0, npr) through the kernel.
all_ids = torch.arange(npr, device=device, dtype=torch.int64).view(m, k)
weights = torch.ones((m, k), dtype=torch.float32, device=device)
# shared_id_base / ep_rank only affect the appended shared columns,
# not the routed remap under test; ep_rank 0 is fine here.
shared_id_base = num_local_routed
got_ids, _ = fused_append_remap_shared_experts_deepep(
all_ids, weights, s, 1.0, shared_id_base, num_local_routed
)
routed_out = got_ids[:, :k].reshape(-1)
# (a) Independent block-derived expectation.
e = torch.arange(npr, device=device, dtype=torch.int64)
rank = e // num_local_routed
local = e % num_local_routed
expected = rank * num_local_experts + local
self.assertTrue(torch.equal(routed_out, expected))
# (b) No remapped routed id hits any rank's shared slots.
shared_slots = set()
for r in range(ep_size):
base = r * num_local_experts + num_local_routed
shared_slots.update(range(base, base + s))
routed_set = set(routed_out.tolist())
self.assertEqual(routed_set & shared_slots, set())
# Routed ids stay unique and inside the global id space.
self.assertEqual(len(routed_set), npr)
self.assertLess(max(routed_set), ep_size * num_local_experts)
def test_equivalence_with_eager_append_then_remap(self):
"""Fused kernel == append shared experts + per-rank shared-slot remap.
test/registered/moe/test_fused_append_shared_experts_top6.py
+115
View File
@@ -0,0 +1,115 @@
"""Unit tests for ``fused_append_shared_experts`` with a non-power-of-two topk.
``_fused_append_shared_experts_kernel`` previously indexed the routed and shared
lanes with ``tl.arange(0, K)`` / ``tl.arange(0, S)``, which Triton only accepts
for power-of-two ranges. DeepSeek-V4 routes top-6 (K=6), so enabling shared
experts fusion (``--enforce-shared-experts-fusion``) crashed with
``ValueError: arange's range must be a power of 2``. The kernel now iterates
over ``next_power_of_2`` blocks with masking, so these tests specifically cover
non-power-of-two K and S. The kernel is GPU-only (Triton), so the tests are
skipped when no accelerator is present.
"""
import unittest
import torch
from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
fused_append_shared_experts,
)
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
register_amd_ci(est_time=20, suite="stage-b-test-1-gpu-small-amd")
def _reference_append(topk_ids, topk_weights, s, scale_factor, n_base):
"""Pure-torch golden reference: routed lanes pass through, shared lanes are
appended as ``n_base + arange(s)`` with weight ``scale_factor``."""
m, k = topk_ids.shape
out_ids = torch.empty((m, k + s), dtype=topk_ids.dtype, device=topk_ids.device)
out_w = torch.empty(
(m, k + s), dtype=topk_weights.dtype, device=topk_weights.device
)
out_ids[:, :k] = topk_ids
out_w[:, :k] = topk_weights
shared = n_base + torch.arange(s, device=topk_ids.device)
out_ids[:, k:] = shared.to(topk_ids.dtype)
out_w[:, k:] = scale_factor
return out_ids, out_w
@unittest.skipUnless(
torch.cuda.is_available(), "fused_append_shared_experts kernel requires a GPU"
)
class TestFusedAppendSharedExpertsTop6(CustomTestCase):
# (m, k, s). k and/or s are deliberately NON power-of-two -- the case that
# used to crash. k=6 is the DeepSeek-V4 top-6 routing width.
CASES = [
(1, 6, 1), # DSV4 top-6, single shared expert (the original crash)
(4, 6, 1),
(17, 6, 1),
(128, 6, 1),
(8, 6, 3), # non-pow2 K and non-pow2 S together
(33, 5, 2), # non-pow2 K and non-pow2 S, odd M
(4, 8, 1), # power-of-two K still correct (regression guard)
]
N_BASE = 256 # shared-expert base id (num routed experts)
def _make_inputs(self, m, k, ids_dtype=torch.int64):
device = get_device()
g = torch.Generator(device="cpu").manual_seed(m * 1000 + k * 7 + 1)
topk_ids = torch.randint(
0, self.N_BASE, (m, k), generator=g, dtype=ids_dtype
).to(device)
topk_weights = torch.rand((m, k), generator=g, dtype=torch.float32).to(device)
return topk_ids, topk_weights
def test_matches_golden_reference(self):
"""Kernel output equals routed-passthrough + shared-append, incl. K=6."""
scale_factor = 0.5
for m, k, s in self.CASES:
with self.subTest(m=m, k=k, s=s):
topk_ids, topk_weights = self._make_inputs(m, k)
got_ids, got_w = fused_append_shared_experts(
topk_ids.clone(),
topk_weights.clone(),
s,
scale_factor,
N=self.N_BASE,
)
exp_ids, exp_w = _reference_append(
topk_ids, topk_weights, s, scale_factor, self.N_BASE
)
self.assertEqual(tuple(got_ids.shape), (m, k + s))
self.assertEqual(tuple(got_w.shape), (m, k + s))
self.assertTrue(torch.equal(got_ids, exp_ids))
self.assertTrue(torch.allclose(got_w, exp_w))
def test_routed_lanes_unmodified(self):
"""The first K columns must be the original routed ids/weights verbatim."""
m, k, s = 16, 6, 1
topk_ids, topk_weights = self._make_inputs(m, k)
got_ids, got_w = fused_append_shared_experts(
topk_ids.clone(), topk_weights.clone(), s, 1.0, N=self.N_BASE
)
self.assertTrue(torch.equal(got_ids[:, :k], topk_ids))
self.assertTrue(torch.allclose(got_w[:, :k], topk_weights))
def test_no_shared_experts_is_noop(self):
"""s == 0 returns the inputs untouched (no kernel launch)."""
topk_ids, topk_weights = self._make_inputs(4, 6)
got_ids, got_w = fused_append_shared_experts(
topk_ids, topk_weights, 0, 1.0, N=self.N_BASE
)
self.assertTrue(torch.equal(got_ids, topk_ids))
self.assertTrue(torch.equal(got_w, topk_weights))
if __name__ == "__main__":
unittest.main()