diff --git a/python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu b/python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu index ae1fd0dcd..0050dc0b1 100644 --- a/python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu +++ b/python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu @@ -478,10 +478,10 @@ struct MoeAlignBlockSizeKernel { int64_t max_num_tokens_padded = sorted_token_ids.size(0); // num_experts from Python is actual_num_experts + 1 (for EP offset convention). - // The v2 kernel (>1024 experts) uses 1024 threads with EXPERTS_PER_THREAD=4, - // covering at most 4096 expert indices, so num_experts (including the +1 - // offset bucket) must be <= 4096. This means up to 4095 real experts. - RuntimeCheck(num_experts <= 4096, "moe_align_block_size: num_experts must be <= 4096, got ", num_experts); + // The v2 kernel (>1024 experts) uses 1024 threads with EXPERTS_PER_THREAD up + // to 8, covering at most 8192 expert indices. This supports up to 8191 real + // experts, sufficient for LoRA virtual experts (num_moe_experts * max_loras). + RuntimeCheck(num_experts <= 8192, "moe_align_block_size: num_experts must be <= 8192, got ", num_experts); const scalar_t* topk_ids_ptr = static_cast(topk_ids.data_ptr()); int32_t* sorted_token_ids_ptr = static_cast(sorted_token_ids.data_ptr()); @@ -561,8 +561,10 @@ struct MoeAlignBlockSizeKernel { if (padded_num_experts <= 2048) { launch_v2(std::integral_constant{}); - } else { + } else if (padded_num_experts <= 4096) { launch_v2(std::integral_constant{}); + } else { + launch_v2(std::integral_constant{}); } const int block_threads = std::min(256, threads); diff --git a/python/sglang/srt/lora/triton_ops/virtual_experts.py b/python/sglang/srt/lora/triton_ops/virtual_experts.py index e3467fba5..4781dfe50 100644 --- a/python/sglang/srt/lora/triton_ops/virtual_experts.py +++ b/python/sglang/srt/lora/triton_ops/virtual_experts.py @@ -9,6 +9,8 @@ import torch import triton import triton.language as tl +from sglang.jit_kernel.moe_align import moe_align_block_size as jit_moe_align_block_size + @triton.jit def _fused_virtual_topk_ids_kernel( @@ -298,6 +300,90 @@ def _invoke_moe_lora_shrink_splitk( ) +def _align_block_size_jit( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """CUDA JIT align_block_size for num_experts > 1024 (up to 8191). + + Uses the v2 kernel from moe_align_kernel.cu which supports large expert + counts via per-thread multi-expert processing and a two-level warp scan, + replacing the previous pure-PyTorch fallback that had excessive CPU overhead + from 15+ individual kernel launches and torch.argsort. + + The JIT kernel uses a +1 offset convention: topk_ids are shifted by +1 so + that the EP sentinel value (-1) maps to bucket 0. The kernel internally + handles histogram, padded prefix-sum, expert_ids assignment, and token + scattering in just 2–3 CUDA kernel launches. + """ + assert num_experts <= 8191, ( + f"_align_block_size_jit supports at most 8191 experts " + f"(num_moe_experts * max_loras), got {num_experts}" + ) + + device = topk_ids.device + flat_topk_ids = topk_ids.reshape(-1) + if flat_topk_ids.dtype == torch.int64: + flat_topk_ids = flat_topk_ids.to(torch.int32) + num_total_tokens = flat_topk_ids.numel() + + if num_total_tokens == 0: + empty = torch.empty(0, dtype=torch.int32, device=device) + return empty, empty, torch.zeros(1, dtype=torch.int32, device=device) + + # JIT kernel uses +1 offset convention: -1 -> bucket 0 (sentinel), + # expert i -> bucket i+1. So pass num_experts + 1 as the bucket count. + jit_num_experts = num_experts + 1 + + if num_total_tokens < jit_num_experts: + max_num_tokens_padded = num_total_tokens * block_size + else: + max_num_tokens_padded = num_total_tokens + jit_num_experts * (block_size - 1) + + # Align every sub-buffer offset to a multiple of 4 (VEC_SIZE). The CUDA + # kernel fills sorted_token_ids with vectorized int4 writes whose last + # store can spill up to 3 int32s past the logical end. With a fused + # allocation the spill would corrupt the adjacent sub-buffer. + _A4 = lambda n: (n + 3) & ~3 # noqa: E731 + max_num_tokens_padded = _A4(max_num_tokens_padded) + max_num_m_blocks = (max_num_tokens_padded + block_size - 1) // block_size + max_num_m_blocks_padded = _A4(max_num_m_blocks) + num_post_pad_size = _A4(1) # 1 element, padded to 4 + cumsum_size = _A4(jit_num_experts + 1) + + # Single allocation sliced into 4 views (zero-copy) to avoid + # per-call Python overhead of 4 separate torch.empty calls. + total_buf = ( + max_num_tokens_padded + + max_num_m_blocks_padded + + num_post_pad_size + + cumsum_size + ) + buf = torch.empty(total_buf, dtype=torch.int32, device=device) + off = 0 + sorted_token_ids = buf[off : off + max_num_tokens_padded] + off += max_num_tokens_padded + expert_ids = buf[off : off + max_num_m_blocks] + off += max_num_m_blocks_padded + num_tokens_post_padded = buf[off : off + 1] + off += num_post_pad_size + cumsum_buffer = buf[off : off + jit_num_experts + 1] + + jit_moe_align_block_size( + flat_topk_ids, + jit_num_experts, + block_size, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + cumsum_buffer, + True, # pad_sorted_token_ids + ) + + return sorted_token_ids, expert_ids, num_tokens_post_padded + + @torch.compile(dynamic=True) def _align_block_size_torch( topk_ids: torch.Tensor, @@ -306,6 +392,8 @@ def _align_block_size_torch( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Pure-PyTorch align_block_size for num_experts > 1024, compiled via torch.compile. + Fallback for platforms where the CUDA JIT kernel is unavailable (e.g. AMD/ROCm). + Out-of-range topk_ids (negative sentinels left by EP dispatch, or virtual- expert IDs >= num_experts produced when those sentinels are combined with a per-adapter offset) are routed into a dedicated sentinel bucket. Without @@ -317,9 +405,6 @@ def _align_block_size_torch( flat_topk_ids = topk_ids.reshape(-1).to(torch.int64) num_total_tokens = flat_topk_ids.numel() - # Map every invalid id to the sentinel bucket (`num_experts`). The bucket - # itself is allocated below via `bucket_count = num_experts + 1` and is - # excluded from block→expert assignment so its blocks stay marked -1. sentinel = num_experts valid_mask = (flat_topk_ids >= 0) & (flat_topk_ids < num_experts) safe_topk_ids = torch.where( @@ -373,8 +458,6 @@ def _align_block_size_torch( sorted_order.to(torch.int32), ) - # Drop the sentinel bucket from the block→expert assignment so its blocks - # remain -1 instead of getting a real expert id from `searchsorted`. block_counts = padded_counts // block_size real_block_counts = block_counts.clone() real_block_counts[sentinel] = 0 @@ -399,7 +482,18 @@ def _align_block_size_torch( return sorted_token_ids, expert_ids, total_padded_tokens -_align_block_size_large = _align_block_size_torch +def _align_block_size_large( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Dispatch to the CUDA JIT kernel when available, otherwise fall back to + the pure-PyTorch torch.compile path (needed on AMD/ROCm or when the JIT + module fails to load).""" + try: + return _align_block_size_jit(topk_ids, block_size, num_experts) + except Exception: + return _align_block_size_torch(topk_ids, block_size, num_experts) def _merged_experts_fused_moe_lora_add_fake( diff --git a/test/registered/lora/test_virtual_experts_kernels.py b/test/registered/lora/test_virtual_experts_kernels.py index 0730e57c3..c42457f06 100644 --- a/test/registered/lora/test_virtual_experts_kernels.py +++ b/test/registered/lora/test_virtual_experts_kernels.py @@ -9,16 +9,15 @@ Covers two regression bugs that surface only with `--lora-use-virtual-experts` them onto a real virtual-expert slot belonging to another adapter and triggered OOB loads in downstream LoRA kernels. -- `_align_block_size_torch` (the `>= 1024`-expert torch.compile fallback) - must route `-1` and `>= num_experts` IDs into a sentinel bucket so they - don't OOB-index `padded_offsets[sorted_expert_ids]` (negative wrap, or - past-end) and don't get assigned to a real expert in the consumer-block - table. +- `_align_block_size_torch` / `_align_block_size_jit` (the `>= 1024`-expert + fallback paths) must route `-1` and `>= num_experts` IDs into a sentinel + bucket so they don't OOB-index `padded_offsets[sorted_expert_ids]` (negative + wrap, or past-end) and don't get assigned to a real expert in the + consumer-block table. -Both kernels run on CUDA. The torch-compile fallback is gated on -`virtual_num_experts > 1024` in production, but we exercise it directly -here at smaller sizes for cheaper iteration; one test sticks to the >1024 -regime to mirror the production trigger. +Both kernels run on CUDA. The fallback is gated on `virtual_num_experts >= 1024` +in production, but we exercise it directly here at smaller sizes for cheaper +iteration; one test sticks to the >1024 regime to mirror the production trigger. Usage: python -m pytest test/registered/lora/test_virtual_experts_kernels.py -v @@ -34,8 +33,10 @@ from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=15, suite="stage-b-test-1-gpu-small") from sglang.srt.lora.triton_ops.virtual_experts import ( + _align_block_size_jit, _align_block_size_torch, _fused_virtual_topk_ids, + fused_sanitize_expert_ids, ) @@ -130,18 +131,23 @@ class TestFusedVirtualTopkIdsPreservesSentinels(CustomTestCase): self.assertFalse(bool(mask[0].item())) -class TestAlignBlockSizeTorchSentinelBucket(CustomTestCase): - """Item C regression: invalid `topk_ids` must not OOB-index - `padded_offsets[sorted_expert_ids]`, must not be assigned to any real - expert in the consumer-block table, and the function must remain - correct on legitimate (all-valid) inputs.""" +class _AlignBlockSizeSentinelBucketBase(CustomTestCase): + """Shared tests for both the torch.compile and JIT align_block_size paths. + + Subclasses override ``_align`` to select the concrete implementation. + """ @classmethod def setUpClass(cls): if not torch.cuda.is_available(): raise unittest.SkipTest("CUDA required") + if cls is _AlignBlockSizeSentinelBucketBase: + raise unittest.SkipTest("Base class") cls.device = "cuda:0" + def _align(self, topk_ids, block_size, num_experts): + raise NotImplementedError + @staticmethod def _assigned_experts(expert_ids: torch.Tensor) -> list: """Return the list of real expert ids assigned to blocks (filtering @@ -166,12 +172,10 @@ class TestAlignBlockSizeTorchSentinelBucket(CustomTestCase): device=self.device, ) - _, expert_ids, _ = _align_block_size_torch(topk_ids, block_size, num_experts) + _, expert_ids, _ = self._align(topk_ids, block_size, num_experts) self._assert_only_real_or_sentinel(expert_ids, num_experts) assigned = set(self._assigned_experts(expert_ids)) - # Every distinct real input expert must be assigned to at least one - # block. self.assertEqual(assigned, set(range(num_experts))) def test_negative_ids_routed_to_sentinel(self): @@ -185,7 +189,7 @@ class TestAlignBlockSizeTorchSentinelBucket(CustomTestCase): device=self.device, ) - _, expert_ids, _ = _align_block_size_torch(topk_ids, block_size, num_experts) + _, expert_ids, _ = self._align(topk_ids, block_size, num_experts) self._assert_only_real_or_sentinel(expert_ids, num_experts) assigned = self._assigned_experts(expert_ids) @@ -204,20 +208,19 @@ class TestAlignBlockSizeTorchSentinelBucket(CustomTestCase): device=self.device, ) - _, expert_ids, _ = _align_block_size_torch(topk_ids, block_size, num_experts) + _, expert_ids, _ = self._align(topk_ids, block_size, num_experts) self._assert_only_real_or_sentinel(expert_ids, num_experts) assigned = self._assigned_experts(expert_ids) for valid_eid in (0, 1, 7, 50): - # 50 is OOR (>= 8), should NOT be in assigned. if valid_eid >= num_experts: self.assertNotIn(valid_eid, assigned) else: self.assertIn(valid_eid, assigned) def test_mixed_invalid_at_production_size(self): - """Mirror the production trigger: `num_experts > 1024` (only path - where `_align_block_size_torch` is invoked instead of the native + """Mirror the production trigger: `num_experts >= 1024` (only path + where the large-expert fallback is invoked instead of the native align kernel).""" num_experts = 1500 block_size = 16 @@ -232,7 +235,7 @@ class TestAlignBlockSizeTorchSentinelBucket(CustomTestCase): device=self.device, ) - _, expert_ids, _ = _align_block_size_torch(topk_ids, block_size, num_experts) + _, expert_ids, _ = self._align(topk_ids, block_size, num_experts) self._assert_only_real_or_sentinel(expert_ids, num_experts) assigned = self._assigned_experts(expert_ids) @@ -246,14 +249,32 @@ class TestAlignBlockSizeTorchSentinelBucket(CustomTestCase): block_size = 16 topk_ids = torch.empty((0, 2), dtype=torch.int32, device=self.device) - sorted_token_ids, expert_ids, num_post_padded = _align_block_size_torch( + sorted_token_ids, expert_ids, num_post_padded = self._align( topk_ids, block_size, num_experts ) self.assertEqual(num_post_padded.item(), 0) - # Whatever expert_ids contains, it must be sentinel only. self.assertEqual(self._assigned_experts(expert_ids), []) +class TestAlignBlockSizeTorchSentinelBucket(_AlignBlockSizeSentinelBucketBase): + """Test the pure-PyTorch torch.compile fallback path (AMD/ROCm compatible).""" + + def _align(self, topk_ids, block_size, num_experts): + return _align_block_size_torch(topk_ids, block_size, num_experts) + + +class TestAlignBlockSizeJitSentinelBucket(_AlignBlockSizeSentinelBucketBase): + """Test the CUDA JIT kernel path (with fused_sanitize_expert_ids, as in + production).""" + + def _align(self, topk_ids, block_size, num_experts): + sorted_token_ids, expert_ids, num_tokens_post_padded = _align_block_size_jit( + topk_ids, block_size, num_experts + ) + expert_ids = fused_sanitize_expert_ids(expert_ids, num_experts) + return sorted_token_ids, expert_ids, num_tokens_post_padded + + if __name__ == "__main__": unittest.main(verbosity=2)