Batch-invariance generation with vLLM fused-MoE backend at improved performance

megatron/core/inference/batch_dimensions_utils.py

@@ -16,6 +16,14 @@
 
 from megatron.core.utils import get_pg_size, round_up_to_nearest_multiple
 
+# Canonical token-count alignment multiple for dynamic inference. The eager
+# path's DynamicInferenceContext.TOKEN_ROUNDER references this constant, and
+# batch-invariant mode aligns every CUDA-graph bucket to it (norm kernels
+# select reduction codepaths by M % 32 alignment class; see
+# _batch_invariant_token_align). Single source of truth: do not restate the
+# literal elsewhere.
+TOKEN_ROUNDER = 64
+
 
 @dataclass(order=True, frozen=True)
 class InferenceBatchDimensions:
@@ -175,7 +183,7 @@ def adjust_batch_dims_for_expert_parallelism(
         if ep_zmq_communicator is not None:
             # CPU-only sync via ZMQ: avoids a NCCL AllReduce kernel on the
             # compute stream plus the H2D/D2H pair that sandwiches it.
-            (max_token_count, max_is_non_decode) = ep_zmq_communicator.sync_all_reduce_max(
+            max_token_count, max_is_non_decode = ep_zmq_communicator.sync_all_reduce_max(
                 local_batch_dims.token_count, int(is_non_decode)
             )
         else:
@@ -205,6 +213,30 @@ def adjust_batch_dims_for_expert_parallelism(
         return adjusted_batch_dim
 
 
+def _batch_invariant_token_align(token_count: int) -> int:
+    """Round a CUDA-graph bucket token count UP to a 64-multiple (min 64).
+
+    Under batch-invariant mode every graphed step must execute norms/GEMMs in
+    the same M-alignment class as eager steps: TE rmsnorm (and other
+    M-sensitive kernels) switch reduction codepaths at M % 32, and the eager
+    path already pads token counts to TOKEN_ROUNDER (64) multiples.  Without
+    this alignment, auto-sizing injects 1- and 2-token decode buckets whose
+    graphed norms execute in a different bit-class, breaking cross-batch
+    bit-equality.  Request counts are untouched (mirrors eager semantics).
+    """
+    rounded_up = math.ceil(token_count / TOKEN_ROUNDER) * TOKEN_ROUNDER
+    return max(TOKEN_ROUNDER, rounded_up)
+
+
+def _batch_invariant_mode_enabled() -> bool:
+    # Lazy import to avoid a circular dependency at module import time.
+    from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+        is_batch_invariant_mode_enabled,
+    )
+
+    return is_batch_invariant_mode_enabled()
+
+
 class CUDAGraphBatchDimensionBuilder:
     """Builder for creating and managing CUDA graph batch dimensions.
 
@@ -278,6 +310,11 @@ def _calculate_cuda_graph_token_counts(
         ), f"cuda_graph_max_tokens must be > 0, got {cuda_graph_max_tokens}"
 
         rounder = CUDAGraphBatchDimensionBuilder.CUDA_GRAPH_ROUNDER
+        if _batch_invariant_mode_enabled():
+            # Batch-invariant mode: TOKEN_ROUNDER-multiple token ladder (see
+            # _batch_invariant_token_align).
+            rounder = TOKEN_ROUNDER
+            cuda_graph_max_tokens = _batch_invariant_token_align(cuda_graph_max_tokens)
 
         # Cuda graph step size.
         cuda_graph_step_size = cuda_graph_max_tokens / num_cuda_graphs
@@ -311,7 +348,8 @@ def _calculate_cuda_graph_token_counts(
 
         # Always include the endpoints: cuda_graph_max_tokens (largest) and tp_size (smallest).
         sizes.add(cuda_graph_max_tokens)
-        sizes.add(tp_size)
+        # Batch-invariant mode: smallest bucket is TOKEN_ROUNDER, never tp_size.
+        sizes.add(TOKEN_ROUNDER if _batch_invariant_mode_enabled() else tp_size)
 
         cuda_graph_token_counts = sorted(sizes, reverse=True)
 
@@ -341,6 +379,10 @@ def _calculate_token_counts_linear(
             sizes = (
                 [1, 2, 4] + list(range(8, 256, 8)) + list(range(256, cuda_graph_max_tokens + 1, 16))
             )
+            if _batch_invariant_mode_enabled():
+                # Batch-invariant mode: align every bucket up to a 64-multiple
+                # (see _batch_invariant_token_align) and dedupe collisions.
+                sizes = [_batch_invariant_token_align(s) for s in sizes]
             # TP-align and dedupe in order; preserve original ordering for parity.
             sizes = list(dict.fromkeys(round_up_to_nearest_multiple(s, tp_size) for s in sizes))
             sizes = [s for s in sizes if s <= cuda_graph_max_tokens]
@@ -431,7 +473,21 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int
             """Helper to create and append batch dimension to list only if it's valid."""
             batch_dim = InferenceBatchDimensions(token_count, prefill_req_count, decode_req_count)
             if batch_dim.is_valid(max_requests, max_sequence_length, num_speculative_tokens):
-                cuda_graph_batch_dimensions_list.append(batch_dim)
+                if _batch_invariant_mode_enabled():
+                    # Batch-invariant mode: align the bucket's token count up
+                    # to a 64-multiple (see _batch_invariant_token_align). The
+                    # alignment is PADDING, mirroring the eager path's
+                    # TOKEN_ROUNDER (which already yields token counts above
+                    # what the requests produce), so validity is judged on the
+                    # unpadded dims; request counts are untouched. Aligning can
+                    # collide previously-distinct buckets, so skip duplicates.
+                    batch_dim = InferenceBatchDimensions(
+                        _batch_invariant_token_align(token_count),
+                        prefill_req_count,
+                        decode_req_count,
+                    )
+                if batch_dim not in cuda_graph_batch_dimensions_list:
+                    cuda_graph_batch_dimensions_list.append(batch_dim)
 
         # Cuda graph token-counts
         # (i.e., token counts used by cuda-graph steps, both decode and non-decode).

megatron/core/inference/contexts/dynamic_context.py

@@ -12,6 +12,7 @@
 from torch import Tensor  # type: ignore
 
 from megatron.core import parallel_state
+from megatron.core.inference.batch_dimensions_utils import TOKEN_ROUNDER as _TOKEN_ROUNDER
 from megatron.core.inference.batch_dimensions_utils import (
     CUDAGraphBatchDimensionBuilder,
     InferenceBatchDimensions,
@@ -315,7 +316,7 @@ class DynamicInferenceContext(BaseInferenceContext):
     """
 
     DEFAULT_MAX_TOKENS = 16384
-    TOKEN_ROUNDER = 64
+    TOKEN_ROUNDER = _TOKEN_ROUNDER
     REQUEST_ROUNDER = 4
     TMS_TAG = "inference_context"
 

megatron/core/inference/moe/batch_invariant.py

@@ -21,6 +21,7 @@
 try:
     import triton
     import triton.language as tl
+    from triton.language.extra import libdevice
 
     HAVE_TRITON = True
 except ImportError:
@@ -86,6 +87,143 @@ def _squared_relu_with_probs_kernel(
                     tl.store(output_ptr + row * hidden_size + cols, value, mask=mask)
 
 
+@triton.jit
+def _swiglu_with_probs_kernel(
+    input_ptr,
+    output_ptr,
+    permutation_map_ptr,
+    n_used_ptr,
+    probs_ptr,
+    ffn_size,  # output width; input row width is 2*ffn_size (gate | up)
+    max_rows,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+):
+    """Apply gated SiLU (SwiGLU) and router probabilities in training order.
+
+    Matches the training fused weighted-swiglu rounding: SiLU(gate)*up*prob is
+    computed in FP32 with a single BF16 round at the end. Input row width is
+    2*ffn_size: gate = first half, up = second half (megatron chunk
+    convention). Fixed NUM_BLOCKS CTAs iterating rows -> CUDA-graph safe.
+    """
+    pid = tl.program_id(0)
+    n_used = tl.load(n_used_ptr)
+    if pid >= n_used:
+        return
+    two_n = 2 * ffn_size
+
+    for row in tl.range(pid, max_rows, NUM_BLOCKS):
+        if row < n_used:
+            if tl.load(permutation_map_ptr + row) >= 0:
+                prob = tl.load(probs_ptr + row)
+                for offset in tl.range(0, ffn_size, BLOCK_SIZE):
+                    cols = offset + tl.arange(0, BLOCK_SIZE)
+                    mask = cols < ffn_size
+                    gate = tl.load(input_ptr + row * two_n + cols, mask=mask).to(tl.float32)
+                    up = tl.load(input_ptr + row * two_n + ffn_size + cols, mask=mask).to(
+                        tl.float32
+                    )
+                    value = gate * tl.sigmoid(gate) * up * prob
+                    tl.store(output_ptr + row * ffn_size + cols, value.to(tl.bfloat16), mask=mask)
+
+
+def swiglu_with_probs(
+    x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor, probs: torch.Tensor
+) -> torch.Tensor:
+    """Gated-SiLU counterpart of squared_relu_with_probs (SwiGLU models)."""
+    num_rows, two_ffn = x.shape
+    ffn_size = two_ffn // 2
+    out = torch.empty(num_rows, ffn_size, dtype=x.dtype, device=x.device)
+    block_size = min(triton.next_power_of_2(ffn_size), 1024)
+    num_blocks = min(num_rows, 512)
+    _swiglu_with_probs_kernel[(num_blocks,)](
+        x,
+        out,
+        permutation_map,
+        n_used,
+        probs,
+        ffn_size,
+        num_rows,
+        BLOCK_SIZE=block_size,
+        NUM_BLOCKS=num_blocks,
+    )
+    return out
+
+
+@triton.jit
+def _weighted_silu_mul_bounded_kernel(
+    in_ptr0, in_ptr1, out_ptr0, bound_ptr, xnumel, HALF_N: tl.constexpr, XBLOCK: tl.constexpr
+):
+    """Device-bounded weighted SwiGLU with training-parity rounding.
+
+    The per-element instruction sequence is copied VERBATIM from Inductor's
+    emitted Triton for the training fused weighted-swiglu
+    (bf16 -> fp32 silu(gate) * up * prob -> bf16, single final rounding), so a
+    token's activation bits match the training forward exactly. Elementwise
+    kernels have no cross-element reduction, so only the per-element sequence
+    determines bits; the schedule below is a persistent 1D grid (static launch,
+    CUDA-graph-safe) striding while xoffset < a DEVICE element bound
+    (= valid_tokens * topk * HALF_N — the live prefix of the flat token-major
+    layout). Rows beyond the bound are neither read nor written.
+    """
+    xbound = tl.load(bound_ptr)
+    num_progs = tl.num_programs(0)
+    xoffset = tl.program_id(0) * XBLOCK
+    while xoffset < xbound:
+        xindex = xoffset + tl.arange(0, XBLOCK)[:]
+        xmask = (xindex < xbound) & (xindex < xnumel)
+        x0 = xindex % HALF_N
+        x1 = xindex // HALF_N
+        tmp0 = tl.load(in_ptr0 + (x0 + 2 * HALF_N * x1), xmask).to(tl.float32)
+        tmp8 = tl.load(in_ptr0 + (HALF_N + x0 + 2 * HALF_N * x1), xmask).to(tl.float32)
+        tmp11 = tl.load(in_ptr1 + (x1), xmask, eviction_policy='evict_last')
+        tmp1 = tmp0.to(tl.float32)
+        tmp2 = -tmp1
+        tmp3 = libdevice.exp(tmp2)
+        tmp4 = tl.full([1], 1.0, tl.float32)
+        tmp5 = tmp3 + tmp4
+        tmp6 = tmp1 / tmp5
+        tmp7 = tmp6.to(tl.float32)
+        tmp9 = tmp7 * tmp8
+        tmp10 = tmp9.to(tl.float32)
+        tmp12 = tmp10 * tmp11
+        tmp13 = tmp12.to(tl.float32)
+        tl.store(out_ptr0 + xindex, tmp13, xmask)
+        xoffset += num_progs * XBLOCK
+
+
+def weighted_silu_mul_bounded(
+    y: torch.Tensor,
+    weights_flat: torch.Tensor,
+    bound_elems: torch.Tensor,
+    num_programs: Optional[int] = None,
+    xblock: int = 1024,
+) -> torch.Tensor:
+    """SwiGLU with routing weights applied at the activation (training parity).
+
+    y: [rows, 2*half_n] bf16 (gate | up); weights_flat: [rows] fp32 routing
+    probabilities; bound_elems: device scalar = live_rows * half_n.
+    Returns [rows, half_n] bf16; rows beyond the live bound are untouched.
+
+    num_programs defaults to SMs * 8 waves (Inductor's persistent-grid sizing;
+    1184 on the B200 this was captured from). Grid size cannot affect bits:
+    the kernel is elementwise with each program owning a disjoint strided
+    index range, so it is an occupancy knob only.
+    """
+    rows, two_half_n = y.shape
+    half_n = two_half_n // 2
+    if num_programs is None:
+        # Lazy import: permute.py imports this module at its top level.
+        from megatron.core.inference.moe.permute import _get_num_sms
+
+        num_programs = _get_num_sms(y.device) * 8
+    out = torch.empty(rows, half_n, dtype=y.dtype, device=y.device)
+    _weighted_silu_mul_bounded_kernel[(num_programs,)](
+        y, weights_flat, out, bound_elems, rows * half_n, HALF_N=half_n, XBLOCK=xblock
+    )
+    return out
+
+
 def squared_relu_with_probs(
     x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor, probs: torch.Tensor
 ) -> torch.Tensor:

megatron/core/inference/moe/fused_moe.py

@@ -210,9 +210,14 @@ def mcore_fused_moe(
     n_used = offs[-1:]
     if batch_invariant_mode:
         # Match training: BF16 activation, FP32 probability multiply, then BF16 before FC2.
-        activation_out = batch_invariant.squared_relu_with_probs(
-            fc1_output, permutation_map, n_used, permuted_probs
-        )
+        if activation_type == ActivationType.SWIGLU:
+            activation_out = batch_invariant.swiglu_with_probs(
+                fc1_output, permutation_map, n_used, permuted_probs
+            )
+        else:
+            activation_out = batch_invariant.squared_relu_with_probs(
+                fc1_output, permutation_map, n_used, permuted_probs
+            )
     else:
         activation_out = activation_func(fc1_output, permutation_map, n_used)
     # Fused activation+quant returns MXFP8Tensor; otherwise quantize separately.