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

[AMD] Fuse shared-expert append + DeepEP remap into one Triton kernel (#28450)

Browse Source 
Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
This commit is contained in:
Rita Brugarolas
2026-06-25 01:22:57 -07:00
committed by GitHub
parent 4a8200565e
commit de2d01c8da
3 changed files with 383 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py
+86
View File
@@ -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,
python/sglang/srt/layers/moe/topk.py
+116 -5
View File
@@ -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
test/registered/moe/test_fused_append_remap_deepep.py
+181
View File
@@ -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()