diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 4f504e2a2..f8d9df3a2 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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. diff --git a/test/registered/moe/test_topk_padded_region.py b/test/registered/moe/test_topk_padded_region.py new file mode 100644 index 000000000..03d258a3a --- /dev/null +++ b/test/registered/moe/test_topk_padded_region.py @@ -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()