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

[AMD] Fuse topk padded-token masking into a single Triton kernel (#28084)

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
karverma-amd
2026-06-23 11:15:12 -07:00
committed by GitHub
co-authored by Cursor
parent ed26a109ee
commit e0dc8b7137
2 changed files with 278 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/topk.py
+84 -4
View File
@@ -30,6 +30,8 @@ from typing import (
import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from sglang.srt.runtime_context import get_parallel
@@ -130,6 +132,13 @@ _is_xpu = is_xpu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_is_musa = is_musa()
# Experimental: skip the HIP padded-token routing-weight masking entirely.
# Padded (CUDA-graph) rows are discarded downstream and the MoE combine is
# per-token, so zeroing their weights is in principle unnecessary. Gated off by
# default because it is a numerics-affecting change that must be validated with
# 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
@@ -1134,6 +1143,71 @@ def is_power_of_two(n):
return n > 0 and math.log2(n).is_integer()
@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``.
assert isinstance(
num_token_non_padded, torch.Tensor
), "num_token_non_padded must be a torch.Tensor"
assert num_token_non_padded.numel() == 1, (
"num_token_non_padded must be a single-element tensor, got shape "
f"{tuple(num_token_non_padded.shape)}"
)
assert (
not num_token_non_padded.dtype.is_floating_point
), f"num_token_non_padded must be an integer tensor, got {num_token_non_padded.dtype}"
assert (
num_token_non_padded.device == x.device
), "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 _eplb_remap_enabled() -> bool:
# A real logical->physical mapping only exists when EPLB is enabled, the
# initial expert placement is non-trivial, or there are redundant physical
@@ -1172,6 +1246,8 @@ def _mask_topk_ids_padded_region(
indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device)
mask = (indices >= num_token_non_padded).unsqueeze(-1)
topk_ids = torch.where(mask, torch.full_like(topk_ids, -1), topk_ids)
elif _can_fuse_padded_region(topk_ids):
_fill_padded_rows(topk_ids, num_token_non_padded, fill_value)
else:
indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device)
topk_ids[indices >= num_token_non_padded, :] = fill_value
@@ -1183,6 +1259,9 @@ def _zero_topk_weights_padded_region(
):
if num_token_non_padded is None:
return
if _can_fuse_padded_region(topk_weights):
_fill_padded_rows(topk_weights, num_token_non_padded, 0.0)
return
indices = torch.arange(0, topk_weights.shape[0], device=topk_weights.device)
topk_weights[indices >= num_token_non_padded, :] = 0.0
@@ -1591,9 +1670,10 @@ def _post_process_topk_ids(
topk_ids = topk_ids_logical_to_physical(
topk_ids, expert_location_dispatch_info
)
# On AMD HIP the aiter MoE kernels do not handle topk_ids=-1 safely, so
# padded tokens are neutralized by zeroing their routing weights.
_zero_topk_weights_padded_region(topk_weights, num_token_non_padded)
# NOTE (HIP): padded-token routing-weight zeroing is deferred to the
# single pass at the end of this function (gated by SGLANG_MORI_NO_PAD_MASK).
# That final pass re-zeros after any shared-expert append/remap, so a
# second zeroing here would be redundant (zeroing is idempotent).
if recorder_topk_ids is None:
recorder_topk_ids = topk_ids
@@ -1635,7 +1715,7 @@ def _post_process_topk_ids(
topk_config,
)
if _is_hip:
if _is_hip and not _skip_hip_pad_mask:
# Shared-expert append/remap can introduce non-zero weights after the
# initial HIP padding mask above. Ensure padded tokens leave this helper
# with all expert weights zeroed.
test/registered/moe/test_topk_padded_region.py
+194
View File
@@ -0,0 +1,194 @@
import unittest
import torch
import sglang.srt.layers.moe.topk as topk_mod
from sglang.srt.layers.moe.topk import (
TopKConfig,
_can_fuse_padded_region,
_fill_padded_rows,
_mask_topk_ids_padded_region,
_post_process_topk_ids,
_zero_topk_weights_padded_region,
)
from sglang.srt.utils import is_hip
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=60, stage="base-b", runner_config="1-gpu-large")
register_amd_ci(est_time=60, stage="stage-b", runner_config="1-gpu-small-amd")
_IS_HIP = is_hip()
torch.manual_seed(1234)
def _eager_fill_padded_rows(x, num_token_non_padded, fill_value):
out = x.clone()
indices = torch.arange(0, x.shape[0], device=x.device)
out[indices >= num_token_non_padded, :] = fill_value
return out
@unittest.skipUnless(
torch.cuda.is_available(), "fused padded-region kernel needs a GPU"
)
class TestTopkPaddedRegion(CustomTestCase):
DEVICE = "cuda"
def test_matches_eager_across_shapes(self):
configs = [
# (n_tokens, topk, dtype, fill, helper)
(37, 9, torch.float32, 0.0, _zero_topk_weights_padded_region),
(256, 8, torch.float32, 0.0, _zero_topk_weights_padded_region),
(1, 9, torch.float32, 0.0, _zero_topk_weights_padded_region),
(37, 9, torch.int32, -1, _mask_topk_ids_padded_region),
(512, 16, torch.int32, -1, _mask_topk_ids_padded_region),
]
for n, k, dtype, fill, helper in configs:
for n_valid in (0, 1, 5, n - 1, n):
with self.subTest(n=n, k=k, dtype=dtype, n_valid=n_valid):
if dtype.is_floating_point:
x = torch.rand((n, k), device=self.DEVICE, dtype=dtype) + 0.5
else:
x = torch.randint(
0, 100, (n, k), device=self.DEVICE, dtype=dtype
)
self.assertTrue(_can_fuse_padded_region(x))
num_token_non_padded = torch.tensor(
n_valid, device=self.DEVICE, dtype=torch.int32
)
expected = _eager_fill_padded_rows(x, num_token_non_padded, fill)
fused = x.clone()
helper(fused, num_token_non_padded)
self.assertTrue(torch.equal(fused, expected))
def test_none_pad_count_is_noop(self):
x = torch.rand((16, 8), device=self.DEVICE, dtype=torch.float32) + 0.5
ref = x.clone()
_zero_topk_weights_padded_region(x, None)
self.assertTrue(torch.equal(x, ref))
def test_non_contiguous_falls_back_to_eager(self):
# A column slice is not row-major contiguous, so the fused path must be
# skipped while still producing the correct result via the eager branch.
base = torch.rand((32, 16), device=self.DEVICE, dtype=torch.float32) + 0.5
view = base[:, ::2]
self.assertFalse(_can_fuse_padded_region(view))
num_token_non_padded = torch.tensor(5, device=self.DEVICE, dtype=torch.int32)
expected = _eager_fill_padded_rows(view, num_token_non_padded, 0.0)
_zero_topk_weights_padded_region(view, num_token_non_padded)
self.assertTrue(torch.equal(view, expected))
def test_cuda_graph_capture_and_replay(self):
n, k = 256, 9
weights = torch.rand((n, k), device=self.DEVICE, dtype=torch.float32) + 0.5
num_token_non_padded = torch.tensor(n, device=self.DEVICE, dtype=torch.int32)
# Warmup on a side stream before capture.
side = torch.cuda.Stream()
side.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(side):
for _ in range(3):
tmp = weights.clone()
_zero_topk_weights_padded_region(tmp, num_token_non_padded)
torch.cuda.current_stream().wait_stream(side)
work = weights.clone()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
_zero_topk_weights_padded_region(work, num_token_non_padded)
for n_valid in (n, 5, 0, 100):
work.copy_(weights)
num_token_non_padded.fill_(n_valid)
graph.replay()
torch.cuda.synchronize()
expected = _eager_fill_padded_rows(weights, num_token_non_padded, 0.0)
self.assertTrue(torch.equal(work, expected))
def test_invalid_pad_count_tensor_raises(self):
x = torch.rand((8, 8), device=self.DEVICE, dtype=torch.float32)
with self.assertRaises(AssertionError):
_fill_padded_rows(x, 4, 0.0) # python int, not a tensor
with self.assertRaises(AssertionError):
_fill_padded_rows(
x,
torch.tensor([1, 2], device=self.DEVICE, dtype=torch.int32),
0.0,
)
with self.assertRaises(AssertionError):
_fill_padded_rows(
x,
torch.tensor(4.0, device=self.DEVICE, dtype=torch.float32),
0.0,
)
@unittest.skipUnless(torch.cuda.is_available(), "padded-region masking needs a GPU")
class TestZeroPaddedRegionIdempotent(CustomTestCase):
"""The HIP post-process keeps a single padded-row zeroing pass. Removing the
earlier (redundant) pass is only safe if zeroing is idempotent."""
DEVICE = "cuda"
def test_zeroing_twice_equals_once(self):
for n, k in [(37, 9), (256, 8), (512, 16)]:
for n_valid in (0, 5, n - 1, n):
with self.subTest(n=n, k=k, n_valid=n_valid):
base = torch.rand((n, k), device=self.DEVICE) + 0.5
pad = torch.tensor(n_valid, device=self.DEVICE, dtype=torch.int32)
once = base.clone()
_zero_topk_weights_padded_region(once, pad)
twice = base.clone()
_zero_topk_weights_padded_region(twice, pad)
_zero_topk_weights_padded_region(twice, pad)
self.assertTrue(torch.equal(once, twice))
@unittest.skipUnless(
_IS_HIP and torch.cuda.is_available(),
"_post_process_topk_ids padded masking is HIP-only",
)
class TestPostProcessPaddedMaskingHip(CustomTestCase):
DEVICE = "cuda"
def _run(self, n=256, k=8, n_valid=5):
topk_weights = torch.rand((n, k), device=self.DEVICE, dtype=torch.float32) + 0.5
topk_ids = torch.randint(0, 64, (n, k), device=self.DEVICE, dtype=torch.int32)
router_logits = torch.rand((n, 64), device=self.DEVICE, dtype=torch.float32)
pad = torch.tensor(n_valid, device=self.DEVICE, dtype=torch.int32)
cfg = TopKConfig(top_k=k, num_fused_shared_experts=0)
_, out_weights, _ = _post_process_topk_ids(
topk_ids,
topk_weights,
cfg,
router_logits,
layer_id=0,
num_token_non_padded=pad,
)
return out_weights, n_valid
def test_padded_rows_zeroed_by_default(self):
# Flag off (default): padded rows must be fully zeroed, valid rows kept.
self.assertFalse(topk_mod._skip_hip_pad_mask)
out, n_valid = self._run()
self.assertTrue(torch.all(out[n_valid:] == 0.0))
self.assertTrue(torch.all(out[:n_valid] > 0.0))
def test_flag_skips_masking(self):
# Flag on: padded rows are left untouched (kept non-zero here).
orig = topk_mod._skip_hip_pad_mask
topk_mod._skip_hip_pad_mask = True
try:
out, n_valid = self._run()
self.assertTrue(torch.all(out[n_valid:] > 0.0))
finally:
topk_mod._skip_hip_pad_mask = orig
if __name__ == "__main__":
unittest.main()