(3/n - prefill optimize)[LoRA][MoE] Optimize virtual experts: remove CPU-GPU sync & multi-block CUDA JIT histogram (#24262)

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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<const scalar_t*>(topk_ids.data_ptr());
     int32_t* sorted_token_ids_ptr = static_cast<int32_t*>(sorted_token_ids.data_ptr());
@@ -561,8 +561,10 @@ struct MoeAlignBlockSizeKernel {
 
       if (padded_num_experts <= 2048) {
         launch_v2(std::integral_constant<int, 2>{});
-      } else {
+      } else if (padded_num_experts <= 4096) {
         launch_v2(std::integral_constant<int, 4>{});
+      } else {
+        launch_v2(std::integral_constant<int, 8>{});
       }
 
       const int block_threads = std::min(256, threads);

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(

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)