Add vLLM grouped gemm backend for MoE inference

megatron/core/inference/moe/__init__.py

@@ -3,6 +3,7 @@
 import enum
 
 from .fused_moe import ActivationType, mcore_fused_moe
+from .vllm_fused_moe import vllm_fused_moe
 
 
 class InferenceGroupedGemmBackend(enum.Enum):
@@ -14,3 +15,4 @@ class InferenceGroupedGemmBackend(enum.Enum):
 
     FLASHINFER = "flashinfer"
     TORCH = "torch"
+    VLLM = "vllm"

megatron/core/inference/moe/permute.py

@@ -8,6 +8,7 @@
 - Unpermute expert outputs back to original token order
 """
 
+from typing import Optional
 from unittest.mock import MagicMock
 
 import torch
@@ -28,6 +29,16 @@
     tl = MagicMock()
 
 
+_NUM_SMS: Optional[int] = None
+
+
+def _get_num_sms(device: torch.device) -> int:
+    global _NUM_SMS
+    if _NUM_SMS is None:
+        _NUM_SMS = torch.cuda.get_device_properties(device).multi_processor_count
+    return _NUM_SMS
+
+
 def _ceil_div(a, b):
     return (a + b - 1) // b
 
@@ -57,17 +68,53 @@ def _count_local_tokens_kernel(
     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = offsets < valid_pairs
     expert_ids = tl.load(routing_map_ptr + offsets, mask=mask, other=-1)
-    # Map global expert IDs to local indices; non-local experts become negative
     local_ids = expert_ids - local_expert_start
     is_local = (local_ids >= 0) & (local_ids < num_local_experts) & mask
     tl.atomic_add(tokens_per_expert_ptr + local_ids, 1, mask=is_local)
 
 
+@triton.jit
+def _count_local_tokens_kernel_persistent(
+    routing_map_ptr,  # [max_tokens, topk] flattened expert assignments
+    tokens_per_expert_ptr,  # [num_local_experts] output counters (zeroed by caller)
+    valid_tokens_ptr,  # scalar int32 CUDA tensor: number of valid tokens this iteration
+    topk,  # number of expert choices per token
+    local_expert_start,  # first global expert index owned by this rank
+    num_local_experts: tl.constexpr,  # number of experts on this rank
+    num_sms,  # number of SMs (grid size for persistent kernel)
+    BLOCK_SIZE: tl.constexpr,  # number of pairs processed per iteration
+):
+    """Count tokens routed to local experts using a persistent grid.
+
+    Launches num_sms CTAs. Each CTA loops over its share of BLOCK_SIZE-sized
+    chunks, with total work determined device-side from valid_tokens.
+    """
+    pid = tl.program_id(0)
+    valid_tokens = tl.load(valid_tokens_ptr)
+    valid_pairs = valid_tokens * topk
+
+    total_blocks = tl.cdiv(valid_pairs, BLOCK_SIZE)
+    blocks_per_cta = tl.cdiv(total_blocks, num_sms)
+    block_start = pid * blocks_per_cta
+
+    if block_start < total_blocks:
+        block_end = tl.minimum(block_start + blocks_per_cta, total_blocks)
+
+        for block_id in tl.range(block_start, block_end):
+            offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+            mask = offsets < valid_pairs
+            expert_ids = tl.load(routing_map_ptr + offsets, mask=mask, other=-1)
+            local_ids = expert_ids - local_expert_start
+            is_local = (local_ids >= 0) & (local_ids < num_local_experts) & mask
+            tl.atomic_add(tokens_per_expert_ptr + local_ids, 1, mask=is_local)
+
+
 def compute_local_tokens_per_expert(
     routing_map: torch.Tensor,
     local_expert_start: int,
     num_local_experts: int,
     valid_tokens: torch.Tensor,
+    persistent: bool = False,
 ) -> torch.Tensor:
     """Count tokens routed to each local expert.
 
@@ -78,20 +125,34 @@ def compute_local_tokens_per_expert(
         num_local_experts: number of experts on this rank.
         valid_tokens: scalar int32 CUDA tensor with the number of valid tokens
             this iteration. Fixed address; value updated each step before graph replay.
+        persistent: use persistent-grid kernel variant (fewer CTAs, looped).
     """
     max_pairs = routing_map.numel()
     topk = routing_map.shape[1]
     tokens_per_expert = torch.zeros(num_local_experts, dtype=torch.int32, device=routing_map.device)
     BLOCK = 1024
-    _count_local_tokens_kernel[(_ceil_div(max_pairs, BLOCK),)](
-        routing_map,
-        tokens_per_expert,
-        valid_tokens,
-        topk,
-        local_expert_start,
-        num_local_experts,
-        BLOCK_SIZE=BLOCK,
-    )
+    if persistent:
+        num_sms = _get_num_sms(routing_map.device)
+        _count_local_tokens_kernel_persistent[(num_sms,)](
+            routing_map,
+            tokens_per_expert,
+            valid_tokens,
+            topk,
+            local_expert_start,
+            num_local_experts,
+            num_sms,
+            BLOCK_SIZE=BLOCK,
+        )
+    else:
+        _count_local_tokens_kernel[(_ceil_div(max_pairs, BLOCK),)](
+            routing_map,
+            tokens_per_expert,
+            valid_tokens,
+            topk,
+            local_expert_start,
+            num_local_experts,
+            BLOCK_SIZE=BLOCK,
+        )
     return tokens_per_expert