diff --git a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py index 409789759..7e895045a 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py @@ -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, ) diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index 5e9c8b23a..cf1c56fda 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -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" diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index e5511d15f..b9393e388 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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:] = ( diff --git a/test/registered/kernels/ops/moe/test_moe_fused_gate.py b/test/registered/kernels/ops/moe/test_moe_fused_gate.py index 52f239955..49614597d 100644 --- a/test/registered/kernels/ops/moe/test_moe_fused_gate.py +++ b/test/registered/kernels/ops/moe/test_moe_fused_gate.py @@ -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"])) diff --git a/test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py b/test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py index 2bf85a904..efafba563 100644 --- a/test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py +++ b/test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py @@ -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. diff --git a/test/registered/moe/test_fused_append_shared_experts_top6.py b/test/registered/moe/test_fused_append_shared_experts_top6.py new file mode 100644 index 000000000..d6edc8614 --- /dev/null +++ b/test/registered/moe/test_fused_append_shared_experts_top6.py @@ -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()