diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py index 9650ed9c3..e727e0f6d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py @@ -1236,6 +1236,92 @@ def fused_append_shared_experts( return out_ids, out_weights +@triton.jit +def _fused_append_remap_shared_experts_deepep_kernel( + topk_ids_ptr, + topk_weights_ptr, + out_ids_ptr, + out_weights_ptr, + shared_id_base, # runtime scalar: ep_rank * num_local_experts + num_local_routed + num_local_routed, # runtime scalar: routed experts per rank (for gap-insertion) + scale_factor, # runtime scalar: shared-expert weight + K: tl.constexpr, + S: tl.constexpr, +): + """Append shared experts AND apply the DeepEP interleaved remap in one pass. + + Equivalent to fused_append_shared_experts() immediately followed by + topk._remap_topk_for_deepep(), but the remap math runs on the rows already + 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) + Shared IDs: shared_id_base + arange(S) (one id per shared slot) + Shared wgt: scale_factor (1.0 on aiter; 1/rsf otherwise) + """ + pid = tl.program_id(0) + + 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) + + # 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 + + tl.store(out_ids_ptr + out_ids_row_ptr + offs_k, ids) + tl.store(out_weights_ptr + out_ids_row_ptr + offs_k, ws) + + offs_s = tl.arange(0, S) + shared_ids = tl.cast(shared_id_base + offs_s, ids.dtype) + shared_ws = tl.full([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) + + +def fused_append_remap_shared_experts_deepep( + topk_ids, + topk_weights, + num_fused_shared_experts, + scale_factor, + shared_id_base, + num_local_routed, +): + """Fused append + DeepEP remap (see kernel docstring). + + Replaces the fused_append_shared_experts() + _remap_topk_for_deepep() pair on + the aiter/DeepEP-class path. Host computes the scalar remap params so the + kernel stays branch-free. + """ + m, k = topk_ids.shape + s = int(num_fused_shared_experts) + if s <= 0: + return topk_ids, topk_weights + + out_ids = torch.empty((m, k + s), dtype=topk_ids.dtype, device=topk_ids.device) + out_weights = torch.empty( + (m, k + s), dtype=topk_weights.dtype, device=topk_weights.device + ) + + _fused_append_remap_shared_experts_deepep_kernel[(m,)]( + topk_ids, + topk_weights, + out_ids, + out_weights, + shared_id_base, + num_local_routed, + scale_factor, + K=k, + S=s, + num_warps=1, + ) + return out_ids, out_weights + + @triton.jit def _fused_append_shared_experts_with_weights_kernel( topk_ids_ptr, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 40fe623f1..a97b12bf6 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -139,6 +139,7 @@ _is_musa = is_musa() # an accuracy run before becoming the default. _skip_hip_pad_mask = get_bool_env_var("SGLANG_MORI_NO_PAD_MASK", "False") + if _is_cuda: from sgl_kernel import moe_fused_gate @@ -1234,6 +1235,74 @@ def _eplb_remap_enabled() -> bool: ) +@triton.jit +def _fill_padded_rows_kernel( + out_ptr, + num_token_non_padded_ptr, + n_cols, + fill_value, + stride_row, + BLOCK_COLS: tl.constexpr, +): + row = tl.program_id(0) + n_valid = tl.load(num_token_non_padded_ptr) + if row >= n_valid: + cols = tl.arange(0, BLOCK_COLS) + mask = cols < n_cols + ptrs = out_ptr + row * stride_row + cols + fill = tl.full((BLOCK_COLS,), fill_value, dtype=out_ptr.dtype.element_ty) + tl.store(ptrs, fill, mask=mask) + + +def _can_fuse_padded_region(x: torch.Tensor) -> bool: + # The fused kernel uses one program per row and assumes a row-major 2D + # tensor (columns contiguous); fall back to eager for anything else. + return x.dim() == 2 and x.stride(1) == 1 + + +def _fill_padded_rows( + x: torch.Tensor, + num_token_non_padded: torch.Tensor, + fill_value, +) -> None: + """Set ``x[row, :] = fill_value`` for every padded row (row index + ``>= num_token_non_padded``) using a single Triton launch. + + Replaces the eager ``arange + (>=) + boolean index_put_`` sequence, which + issues several launch-latency-bound kernels per call. The grid is static + (one program per row) and the pad count is read from device memory inside + the kernel, so this is safe to capture inside a CUDA/HIP graph. + """ + # Metadata-only checks (no device sync): the kernel reads a single scalar + # routing count from device memory, so it must be a 1-element integer tensor + # on the same device as ``x``. Use explicit raises (not asserts) so the + # checks survive ``python -O`` and invalid inputs fail loudly instead of + # turning into opaque Triton/memory errors. + if not isinstance(num_token_non_padded, torch.Tensor): + raise TypeError("num_token_non_padded must be a torch.Tensor") + if num_token_non_padded.numel() != 1: + raise ValueError( + "num_token_non_padded must be a single-element tensor, got shape " + f"{tuple(num_token_non_padded.shape)}" + ) + if num_token_non_padded.dtype.is_floating_point: + raise TypeError( + "num_token_non_padded must be an integer tensor, got " + f"{num_token_non_padded.dtype}" + ) + if num_token_non_padded.device != x.device: + raise ValueError("num_token_non_padded and x must be on the same device") + n_rows, n_cols = x.shape + _fill_padded_rows_kernel[(n_rows,)]( + x, + num_token_non_padded, + n_cols, + fill_value, + x.stride(0), + BLOCK_COLS=triton.next_power_of_2(n_cols), + ) + + def _mask_topk_ids_padded_region( topk_ids: torch.Tensor, num_token_non_padded: Optional[torch.Tensor] = None, @@ -1244,6 +1313,8 @@ def _mask_topk_ids_padded_region( # TODO: let the kernel support other dtypes if _is_cuda and topk_ids.dtype == torch.int32 and fill_value == -1: mask_topk_ids(topk_ids, num_token_non_padded) + elif _can_fuse_padded_region(topk_ids): + _fill_padded_rows(topk_ids, num_token_non_padded, fill_value) elif _is_npu: # On NPU, bool-indexed scatter `topk_ids[bool_mask, :] = -1` lowers # to aclnnNonzeroV2 and can trigger an aicore timeout under long @@ -1698,7 +1769,48 @@ def _post_process_topk_ids( if recorder_topk_ids is None: recorder_topk_ids = topk_ids - if num_fused_shared_experts > 0 and _use_aiter: + _aiter_append = num_fused_shared_experts > 0 and _use_aiter + _deepep_remap = num_fused_shared_experts > 0 and is_deepep_class_backend() + + if _aiter_append and _deepep_remap: + # Fused path: append shared experts AND apply the DeepEP interleaved + # remap in a single Triton kernel. This replaces the original + # fused_append_shared_experts() + eager _remap_topk_for_deepep() pair, + # collapsing ~6 launch-bound elementwise kernels/layer (div_floor / add / + # arange / fill / copy) into the one append kernel that already runs. + # + # Shared weight is 1.0 here because this branch is aiter-only: + # aiter_biased_grouped_topk folds routed_scaling_factor into the routed + # weights and forward_deepep skips the post-MoE multiply for _use_aiter, + # so the always-on shared expert must contribute 1.0x. (The eager + # _remap_topk_for_deepep instead sets shared weight to + # 1/routed_scaling_factor to compensate a post-MoE scale that the aiter + # path does not apply; see PR #28237.) + num_physical_routed_experts = ( + expert_location_dispatch_info.num_physical_experts + if expert_location_dispatch_info is not None + else router_logits.shape[1] + ) + ep_size = get_parallel().moe_ep_size + ep_rank = get_parallel().moe_ep_rank + num_local_routed = num_physical_routed_experts // ep_size + num_local_experts = num_local_routed + num_fused_shared_experts + shared_id_base = ep_rank * num_local_experts + num_local_routed + + # Lazy import to avoid circular-import issues + from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import ( + fused_append_remap_shared_experts_deepep, + ) + + topk_ids, topk_weights = fused_append_remap_shared_experts_deepep( + topk_ids, + topk_weights, + num_fused_shared_experts, + 1.0, # shared-expert weight on the aiter path + shared_id_base, + num_local_routed, + ) + elif _aiter_append: M, N = router_logits.shape scale_factor = ( 1.0 @@ -1718,10 +1830,9 @@ def _post_process_topk_ids( scale_factor, N, # base id for shared experts ) - - # DeepEP: remap to interleaved expert layout where each rank's shared - # expert has a unique ID for dispatch routing. - if num_fused_shared_experts > 0 and is_deepep_class_backend(): + elif _deepep_remap: + # DeepEP: remap to interleaved expert layout where each rank's shared + # expert has a unique ID for dispatch routing. num_physical_routed_experts = ( expert_location_dispatch_info.num_physical_experts if expert_location_dispatch_info is not None diff --git a/test/registered/moe/test_fused_append_remap_deepep.py b/test/registered/moe/test_fused_append_remap_deepep.py new file mode 100644 index 000000000..ca291cd04 --- /dev/null +++ b/test/registered/moe/test_fused_append_remap_deepep.py @@ -0,0 +1,181 @@ +"""Unit tests for the fused append + DeepEP-remap shared-experts Triton kernel. + +Covers ``fused_append_remap_shared_experts_deepep``, which collapses +``fused_append_shared_experts()`` followed by ``_remap_topk_for_deepep()`` into a +single Triton launch on the aiter/DeepEP-class path. The kernel is GPU-only +(Triton), so these tests are skipped when no accelerator is present. +""" + +import unittest + +import torch + +from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import ( + fused_append_remap_shared_experts_deepep, + fused_append_shared_experts, +) +from sglang.srt.layers.moe.topk import TopKConfig, _remap_topk_for_deepep +from sglang.srt.runtime_context import get_parallel +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_remap( + topk_ids, topk_weights, s, scale_factor, shared_id_base, num_local_routed +): + """Pure-torch golden reference mirroring the kernel's documented contract. + + Routed IDs: e -> e + e // num_local_routed + Shared IDs: shared_id_base + arange(s) + Routed wgt: passthrough + Shared wgt: 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 + topk_ids // num_local_routed + out_w[:, :k] = topk_weights + shared = shared_id_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+remap kernel requires a GPU" +) +class TestFusedAppendRemapDeepEP(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). + 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), + ] + + def _make_inputs(self, m, k, num_physical_routed, ids_dtype=torch.int64): + device = get_device() + g = torch.Generator(device="cpu").manual_seed(m * 1000 + k * 7 + 1) + topk_ids = torch.randint( + 0, num_physical_routed, (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 + + @staticmethod + def _shared_id_base(num_physical_routed, ep_size, ep_rank, s): + num_local_routed = num_physical_routed // ep_size + num_local_experts = num_local_routed + s + return ep_rank * num_local_experts + num_local_routed, num_local_routed + + def test_matches_golden_reference(self): + """Kernel output equals the documented routed-remap + shared-append math.""" + for m, k, npr, ep_size, ep_rank, s in self.CASES: + with self.subTest(m=m, k=k, npr=npr, ep_rank=ep_rank, s=s): + shared_id_base, num_local_routed = self._shared_id_base( + npr, ep_size, ep_rank, s + ) + scale_factor = 1.0 + topk_ids, topk_weights = self._make_inputs(m, k, npr) + + got_ids, got_w = fused_append_remap_shared_experts_deepep( + topk_ids, + topk_weights, + s, + scale_factor, + shared_id_base, + num_local_routed, + ) + exp_ids, exp_w = _reference_append_remap( + topk_ids, + topk_weights, + s, + scale_factor, + shared_id_base, + num_local_routed, + ) + + self.assertEqual(tuple(got_ids.shape), (m, k + s)) + self.assertTrue(torch.equal(got_ids, exp_ids)) + self.assertTrue(torch.allclose(got_w, exp_w)) + + def test_equivalence_with_eager_append_then_remap(self): + """Fused kernel == fused_append_shared_experts() + _remap_topk_for_deepep(). + + The eager remap overwrites the shared weight with 1/routed_scaling_factor, + so the fused kernel is invoked with that same value to make the two paths + bit-identical (ids are identical regardless of the scaling factor). + """ + rsf = 2.5 + scale_factor = 1.0 / rsf + for m, k, npr, ep_size, ep_rank, s in self.CASES: + with self.subTest(m=m, k=k, npr=npr, ep_rank=ep_rank, s=s): + shared_id_base, num_local_routed = self._shared_id_base( + npr, ep_size, ep_rank, s + ) + topk_ids, topk_weights = self._make_inputs(m, k, npr) + + fused_ids, fused_w = fused_append_remap_shared_experts_deepep( + topk_ids.clone(), + topk_weights.clone(), + s, + scale_factor, + shared_id_base, + num_local_routed, + ) + + with get_parallel().override(moe_ep_size=ep_size, moe_ep_rank=ep_rank): + eager_ids, eager_w = fused_append_shared_experts( + topk_ids.clone(), + topk_weights.clone(), + s, + scale_factor, + npr, # shared-expert base id (overwritten by the remap) + ) + eager_ids, eager_w = _remap_topk_for_deepep( + eager_ids, + eager_w, + s, + npr, + TopKConfig( + top_k=k, + num_fused_shared_experts=s, + routed_scaling_factor=rsf, + ), + ) + + self.assertTrue(torch.equal(fused_ids, eager_ids)) + self.assertTrue(torch.allclose(fused_w, eager_w)) + + def test_shared_weight_is_one_on_aiter_path(self): + """On the aiter path the always-on shared expert must contribute 1.0x.""" + m, k, npr, ep_size, ep_rank, s = 8, 8, 256, 8, 1, 1 + shared_id_base, num_local_routed = self._shared_id_base( + npr, ep_size, ep_rank, s + ) + topk_ids, topk_weights = self._make_inputs(m, k, npr) + + _, got_w = fused_append_remap_shared_experts_deepep( + topk_ids, topk_weights, s, 1.0, shared_id_base, num_local_routed + ) + self.assertTrue(torch.all(got_w[:, -s:] == 1.0)) + + 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, 8, 256) + got_ids, got_w = fused_append_remap_shared_experts_deepep( + topk_ids, topk_weights, 0, 1.0, 0, 32 + ) + self.assertTrue(torch.equal(got_ids, topk_ids)) + self.assertTrue(torch.equal(got_w, topk_weights)) + + +if __name__ == "__main__": + unittest.main()