From 04bf0c3289ea5ce945f0fa30575c7b3ee9905b8b Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:45:55 -0700 Subject: [PATCH 01/26] Batch-invariant mode: 64-multiple floor for CUDA-graph token buckets Under batch_invariant_mode, graphed steps must execute norms/GEMMs in the same M-alignment class as eager steps (TE rmsnorm switches reduction codepaths at M % 32; the eager path pads token counts to TOKEN_ROUNDER=64). The bucket auto-sizing paths (exponential ladder endpoint tp_size; linear -1 ladder [1, 2, 4, 8, ...]) inject 1- and 2-token decode buckets whose graphed norms execute in a different bit-class, breaking cross-batch bit-equality. Floor every generated bucket token count to a 64-multiple (min 64) when batch-invariant mode is enabled: in both sizing distributions and centrally in add_if_valid (with duplicate suppression after collisions). Request counts are untouched, mirroring eager TOKEN_ROUNDER semantics. Verified on the NeMo-RL true-on-policy determinism campaign: without the floor, graphed decode diverges bitwise from eager scoring at M in {1,2} buckets; with it, full-CUDA-graph decode is bitwise-exact vs the TE scoring path over 20-step GRPO runs (gen_kl == 0.0). Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../core/inference/batch_dimensions_utils.py | 43 ++++++++++++++++++- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index 1229d333d0a..3aee65eb862 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/megatron/core/inference/batch_dimensions_utils.py @@ -205,6 +205,29 @@ def adjust_batch_dims_for_expert_parallelism( return adjusted_batch_dim +def _batch_invariant_token_floor(token_count: int) -> int: + """Floor a CUDA-graph bucket token count 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 floor, 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). + """ + return max(64, ((token_count + 63) // 64) * 64) + + +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 +301,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: 64-multiple token ladder (see + # _batch_invariant_token_floor). + rounder = 64 + cuda_graph_max_tokens = _batch_invariant_token_floor(cuda_graph_max_tokens) # Cuda graph step size. cuda_graph_step_size = cuda_graph_max_tokens / num_cuda_graphs @@ -311,7 +339,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 64, never tp_size. + sizes.add(64 if _batch_invariant_mode_enabled() else tp_size) cuda_graph_token_counts = sorted(sizes, reverse=True) @@ -341,6 +370,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: floor every bucket to a 64-multiple + # (see _batch_invariant_token_floor) and dedupe collisions. + sizes = [_batch_invariant_token_floor(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] @@ -429,9 +462,15 @@ def generate_cuda_graph_batch_dimensions_list( def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int) -> None: """Helper to create and append batch dimension to list only if it's valid.""" + if _batch_invariant_mode_enabled(): + # Batch-invariant mode: floor EVERY bucket's token count to a + # 64-multiple (see _batch_invariant_token_floor); the flooring + # can collide previously-distinct buckets, so skip duplicates. + token_count = _batch_invariant_token_floor(token_count) 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_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). From b6479a77b5444dd3f4fd99deed44e6b35b66eaa5 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:45:56 -0700 Subject: [PATCH 02/26] Batch-invariant MoE: support gated SwiGLU activations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit batch_invariant_mode hard-wired the MoE activation to squared_relu_with_probs, so gated-unit models (e.g. Qwen3 MoE, enabled for inference_optimized by #5700) fail with a grouped-GEMM K mismatch (FC1 output is 2*ffn wide and never halved). Add swiglu_with_probs — the gated-SiLU counterpart with the same graph-safe fixed-CTA structure — and select it from mcore_fused_moe when activation_type is SWIGLU. SiLU(gate)*up*prob is computed in FP32 with a single BF16 round, matching the training fused weighted-swiglu rounding. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../core/inference/moe/batch_invariant.py | 63 +++++++++++++++++++ megatron/core/inference/moe/fused_moe.py | 11 +++- 2 files changed, 71 insertions(+), 3 deletions(-) diff --git a/megatron/core/inference/moe/batch_invariant.py b/megatron/core/inference/moe/batch_invariant.py index 06dfddc2869..53721da5066 100644 --- a/megatron/core/inference/moe/batch_invariant.py +++ b/megatron/core/inference/moe/batch_invariant.py @@ -86,6 +86,69 @@ 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 + + def squared_relu_with_probs( x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor, probs: torch.Tensor ) -> torch.Tensor: diff --git a/megatron/core/inference/moe/fused_moe.py b/megatron/core/inference/moe/fused_moe.py index c131255b224..4f917726b36 100644 --- a/megatron/core/inference/moe/fused_moe.py +++ b/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. From 993cdec2af35c688e005125b2122697ae671b72e Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:45:56 -0700 Subject: [PATCH 03/26] Batch-invariant mode: support the vLLM Triton fused-MoE backend batch_invariant_mode previously forced inference_grouped_gemm_backend='torch' (DeepGEMM grouped GEMM), which costs ~6x generation throughput vs the vLLM Triton fused-MoE backend on the same engine (measured on Qwen3-30B-A3B, EP8/TP1, 8xB200: ~4.1k vs ~25.4k tok/s at BS256/OSL1024). This change makes the vLLM backend batch-invariant and allows it under the flag; the equivalent deterministic configuration reaches ~24k tok/s (0.99x of the engine's non-invariant throughput) and is certified bitwise-identical to the Megatron training forward over full-learning-rate GRPO runs (gen_kl == 0.0) in NeMo-RL true on-policy training. Three changes, active only under batch_invariant_mode: - vllm_fused_moe: pin the kernel launch config. _get_default_config selects tile shapes from the token-count hint, so different co-batch sizes change the fp32 accumulation grouping (batch-variant bits). Grid sizing may still use the hint: the kernel strides, so grid size never changes per-tile math. - SwiGLU activation: apply routing probabilities at the activation (before FC2) with the training kernel's exact rounding sequence, via a new device-bounded weighted SiLU-mul kernel (persistent grid, bounded by the valid_tokens*topk device scalar; CUDA-graph safe). Matches the convention already used by the torch-backend batch-invariant activation. - _moe_sum: unit weights when probs were applied at the activation, and fp64 accumulation so the topk reduction is order-independent by precision. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../core/inference/moe/batch_invariant.py | 65 +++++++++++++++++++ megatron/core/inference/moe/vllm_fused_moe.py | 60 ++++++++++++++--- .../core/transformer/transformer_config.py | 8 ++- 3 files changed, 122 insertions(+), 11 deletions(-) diff --git a/megatron/core/inference/moe/batch_invariant.py b/megatron/core/inference/moe/batch_invariant.py index 53721da5066..0d2a277c95a 100644 --- a/megatron/core/inference/moe/batch_invariant.py +++ b/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: @@ -149,6 +150,70 @@ def swiglu_with_probs( 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: int = 1184, + 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. + """ + rows, two_half_n = y.shape + half_n = two_half_n // 2 + 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: diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index f6087ebbfe9..8a77327e56e 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/megatron/core/inference/moe/vllm_fused_moe.py @@ -30,6 +30,7 @@ tl = MagicMock() from megatron.core.inference.moe.activations import bounded_silu_mul +from megatron.core.inference.moe import batch_invariant from megatron.core.inference.moe.fused_moe import ActivationType from megatron.core.inference.moe.permute import ( _get_num_sms, @@ -452,6 +453,8 @@ def _moe_sum_kernel( BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr, NUM_K_BLOCKS: tl.constexpr, + APPLY_WEIGHTS: tl.constexpr = True, + ACC_FP64: tl.constexpr = False, ): """Reduce topk dimension with routing weight application. @@ -483,16 +486,23 @@ def _moe_sum_kernel( offs_k = k_idx * BLOCK_K + tl.arange(0, BLOCK_K) k_mask = offs_k < K - acc = tl.zeros([BLOCK_K], dtype=tl.float32) + acc = tl.zeros([BLOCK_K], dtype=tl.float64 if ACC_FP64 else tl.float32) for t in range(topk): eid = tl.load(routing_map_ptr + token_id * topk + t) lid = eid - local_expert_start if lid >= 0 and lid < num_local_experts: v = tl.load(input_ptr + base + t * K + offs_k, mask=k_mask, other=0.0) - w = tl.load(topk_weights_ptr + token_id * topk + t) - acc += v.to(tl.float32) * w + if ACC_FP64: + if APPLY_WEIGHTS: + w = tl.load(topk_weights_ptr + token_id * topk + t) + acc += v.to(tl.float64) * w.to(tl.float64) + else: + acc += v.to(tl.float64) + else: + w = tl.load(topk_weights_ptr + token_id * topk + t) + acc += v.to(tl.float32) * w - tl.store(output_ptr + token_id_i64 * K + offs_k, acc, mask=k_mask) + tl.store(output_ptr + token_id_i64 * K + offs_k, acc.to(tl.float32), mask=k_mask) def _moe_sum( @@ -506,6 +516,8 @@ def _moe_sum( local_expert_start: int, num_local_experts: int, out: Optional[torch.Tensor] = None, + apply_weights: bool = True, + acc_fp64: bool = False, ) -> torch.Tensor: """Fused topk reduction: [max_tokens*topk, K] bf16 → [max_tokens, K]. @@ -536,6 +548,8 @@ def _moe_sum( BLOCK_M=BLOCK_M, BLOCK_K=BLOCK_K, NUM_K_BLOCKS=NUM_K_BLOCKS, + APPLY_WEIGHTS=apply_weights, + ACC_FP64=acc_fp64, ) return out @@ -596,7 +610,17 @@ def vllm_fused_moe( # Mirror upstream vLLM: pick the full launch config (tile sizes, warps, # stages) host-side from the token-count hint, not from the worst-case # buffer size. Same config is used for both FC1 and FC2 (matches vLLM). - config = _get_default_config(M=effective_tokens, E=num_local_experts, top_k=topk) + batch_invariant_mode = batch_invariant.enabled() + if batch_invariant_mode: + # Batch-invariant mode: the config must not depend on the token count. + # _get_default_config is a step function of M, so different co-batch + # sizes would otherwise select different tile shapes and change the + # fp32 accumulation grouping (batch-variant bits). Pin the large-M + # config for every launch; grid SIZING below may still use the hint + # (the kernel strides, so grid size never changes per-tile math). + config = _get_default_config(M=1 << 30, E=num_local_experts, top_k=topk) + else: + config = _get_default_config(M=effective_tokens, E=num_local_experts, top_k=topk) sorted_token_ids, expert_ids, num_post_padded = _moe_align_block_size_cuda_graphable( routing_map, config['BLOCK_SIZE_M'], num_local_experts, local_expert_start, valid_tokens @@ -645,10 +669,23 @@ def vllm_fused_moe( fuse_squared_relu=not is_swiglu, ) if is_swiglu: - # intermediate1 is [num_valid, 2N] (gate | up); reduce to [num_valid, N] via - # SiLU(gate) * up over the valid_tokens*topk live rows only. - n_rows = (valid_tokens * topk).to(torch.int32) - intermediate1 = bounded_silu_mul(intermediate1, n_rows) + if batch_invariant_mode: + # Match training: routing probabilities multiply at the activation + # (before FC2), with the training kernel's exact rounding sequence + # (single bf16 round of fp32 silu(gate)*up*prob). The reduction + # below then sums with unit weights. Device-bounded to the live + # valid_tokens*topk prefix; CUDA-graph safe. + bound_elems = valid_tokens.to(torch.int64) * ( + topk * (intermediate1.shape[1] // 2) + ) + intermediate1 = batch_invariant.weighted_silu_mul_bounded( + intermediate1, topk_weights_flat, bound_elems + ) + else: + # intermediate1 is [num_valid, 2N] (gate | up); reduce to [num_valid, N] via + # SiLU(gate) * up over the valid_tokens*topk live rows only. + n_rows = (valid_tokens * topk).to(torch.int32) + intermediate1 = bounded_silu_mul(intermediate1, n_rows) # FC2: [max_tokens*topk, N] → [max_tokens*topk, K], without routing weights. # Routing weights are applied in the reduction kernel to avoid an extra @@ -687,4 +724,9 @@ def vllm_fused_moe( local_expert_start, num_local_experts, out=out, + # Batch-invariant mode: probs were already applied at the activation + # for SwiGLU (training parity), so the reduction uses unit weights; + # fp64 accumulation makes the topk sum order-independent by precision. + apply_weights=not (batch_invariant_mode and is_swiglu), + acc_fp64=batch_invariant_mode, ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e783a017056..0babb76cc69 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1630,9 +1630,13 @@ def __post_init__(self): ) if self.batch_invariant_mode: - if self.inference_grouped_gemm_backend != InferenceGroupedGemmBackend.TORCH: + if self.inference_grouped_gemm_backend not in ( + InferenceGroupedGemmBackend.TORCH, + InferenceGroupedGemmBackend.VLLM, + ): raise ValueError( - "batch_invariant_mode requires " "inference_grouped_gemm_backend='torch'." + "batch_invariant_mode requires inference_grouped_gemm_backend " + "'torch' or 'vllm'." ) if ( self.expert_model_parallel_size > 1 From 3b62e8db42986c1106ad280e55622a5fb73ddf00 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:45:57 -0700 Subject: [PATCH 04/26] Batch-invariant mode: add 'te_native' GEMM backend (workspace starvation) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a third batch-invariant GEMM backend that keeps the NATIVE cuBLASLt kernels for every dense GEMM (aten and TE) and obtains batch invariance by starving the cuBLASLt workspace (~1KB): split-K reduction variants require workspace, so starving it disqualifies them and pins every M to the same serial-K reduction recipe. M/N tile selection may still vary with M, but tiling does not affect bits — bf16 products are exact in fp32, so only the K-reduction order matters. This is the lowest-overhead invariant dense path: no kernel substitution, native speed. It is also the configuration certified bitwise-identical to the Megatron/TE training forward in NeMo-RL true on-policy GRPO (gen_kl == 0.0 over full-learning-rate runs, Qwen3-30B-A3B EP8/TP1 8xB200). Notes: - TE (<= 2.15 verified) hardcodes a 32MiB workspace in get_cublas_workspace_size_bytes() and ignores CUBLASLT_WORKSPACE_SIZE, so the env pin alone never engages for TE-launched GEMMs; the backend patches the size fn and clears its lru_cache. - Under te_native, aten::mm/addmm are left unpatched and TE general_gemm is not substituted (skip_gemm); the non-GEMM batch-invariant patches (log_softmax, mean, RMSNorm, attention gate, Mamba autotuner pins) still apply. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../custom_layers/batch_invariant_kernels.py | 115 +++++++++++++----- 1 file changed, 85 insertions(+), 30 deletions(-) diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index 5b8c6e678e6..508a6e739e4 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -526,7 +526,7 @@ def mean_dim( # Production uses DeepGEMM for bf16 and the deterministic Triton kernel for # intentional higher-precision operations such as the fp32 MoE router. -_BATCH_INVARIANT_BACKENDS = ("deepgemm", "triton") +_BATCH_INVARIANT_BACKENDS = ("deepgemm", "triton", "te_native") _BATCH_INVARIANT_BACKEND: str = "deepgemm" @@ -625,47 +625,53 @@ def _import_module_if_available(name: str): return importlib.import_module(name) -def _te_patch_for_batch_invariant(): +def _te_patch_for_batch_invariant(skip_gemm: bool = False): """Patch Transformer Engine modules to use batch-invariant GEMM and RMSNorm. This monkey-patches TE's GEMM and RMSNorm entry points to dispatch to the batch-invariant implementations when batch-invariant mode is enabled. Safe no-op if TE is unavailable. + + Args: + skip_gemm: leave TE's native general_gemm in place (used by the + "te_native" backend, where GEMM batch invariance comes from + cuBLASLt workspace starvation rather than kernel substitution). """ global _TE_GENERAL_GEMM_ORIG, _TE_RMSNORM_ORIG_FWD, _MEG_TE_GENERAL_GEMM_ORIG import transformer_engine.pytorch as te import transformer_engine.pytorch.cpp_extensions as te_cpp - # Patch general_gemm once - if _TE_GENERAL_GEMM_ORIG is None and hasattr(te_cpp, "general_gemm"): - _TE_GENERAL_GEMM_ORIG = te_cpp.general_gemm - te_cpp.general_gemm = _te_general_gemm_patched + if not skip_gemm: + # Patch general_gemm once + if _TE_GENERAL_GEMM_ORIG is None and hasattr(te_cpp, "general_gemm"): + _TE_GENERAL_GEMM_ORIG = te_cpp.general_gemm + te_cpp.general_gemm = _te_general_gemm_patched - # Also patch the symbol imported inside TE's module.linear - # (from ..cpp_extensions import general_gemm) - import transformer_engine.pytorch.module.linear as te_linear_mod + # Also patch the symbol imported inside TE's module.linear + # (from ..cpp_extensions import general_gemm) + import transformer_engine.pytorch.module.linear as te_linear_mod - if hasattr(te_linear_mod, "general_gemm"): - if "module.linear.general_gemm" not in _TE_GEMM_FUNC_ORIGS: - _TE_GEMM_FUNC_ORIGS["module.linear.general_gemm"] = te_linear_mod.general_gemm - te_linear_mod.general_gemm = _te_general_gemm_patched + if hasattr(te_linear_mod, "general_gemm"): + if "module.linear.general_gemm" not in _TE_GEMM_FUNC_ORIGS: + _TE_GEMM_FUNC_ORIGS["module.linear.general_gemm"] = te_linear_mod.general_gemm + te_linear_mod.general_gemm = _te_general_gemm_patched - # Also patch the symbol imported inside TE's module.layernorm_linear - import transformer_engine.pytorch.module.layernorm_linear as te_layernorm_linear_mod + # Also patch the symbol imported inside TE's module.layernorm_linear + import transformer_engine.pytorch.module.layernorm_linear as te_layernorm_linear_mod - if hasattr(te_layernorm_linear_mod, "general_gemm"): - if "module.layernorm_linear.general_gemm" not in _TE_GEMM_FUNC_ORIGS: - _TE_GEMM_FUNC_ORIGS["module.layernorm_linear.general_gemm"] = ( - te_layernorm_linear_mod.general_gemm - ) - te_layernorm_linear_mod.general_gemm = _te_general_gemm_patched + if hasattr(te_layernorm_linear_mod, "general_gemm"): + if "module.layernorm_linear.general_gemm" not in _TE_GEMM_FUNC_ORIGS: + _TE_GEMM_FUNC_ORIGS["module.layernorm_linear.general_gemm"] = ( + te_layernorm_linear_mod.general_gemm + ) + te_layernorm_linear_mod.general_gemm = _te_general_gemm_patched - # Also patch the symbol imported into Megatron's TE wrapper module - import megatron.core.extensions.transformer_engine as meg_te + # Also patch the symbol imported into Megatron's TE wrapper module + import megatron.core.extensions.transformer_engine as meg_te - if _MEG_TE_GENERAL_GEMM_ORIG is None and hasattr(meg_te, "general_gemm"): - _MEG_TE_GENERAL_GEMM_ORIG = meg_te.general_gemm - meg_te.general_gemm = _te_general_gemm_patched + if _MEG_TE_GENERAL_GEMM_ORIG is None and hasattr(meg_te, "general_gemm"): + _MEG_TE_GENERAL_GEMM_ORIG = meg_te.general_gemm + meg_te.general_gemm = _te_general_gemm_patched # Patch RMSNorm.forward once (class may be on te or te.pytorch) rms_cls = getattr(te, "RMSNorm", None) @@ -1636,6 +1642,42 @@ def is_batch_invariant_mode_enabled(): return _batch_invariant_MODE +_TE_NATIVE_WORKSPACE_BYTES = 1024 + + +def _enable_te_native_workspace_starvation(workspace_bytes: int = _TE_NATIVE_WORKSPACE_BYTES): + """Make the NATIVE cuBLASLt GEMM kernels batch-invariant via workspace starvation. + + cuBLASLt selects split-K reduction variants per M (the co-batch token count), + which changes the fp32 accumulation order across batch compositions. Split-K + variants require workspace; starving the workspace to ~1KB disqualifies them, + pinning every M to the same serial-K reduction recipe — batch-invariant at + native kernel speed, with no kernel substitution. (Tile shape selection may + still vary with M, but M/N tiling does not affect bits: bf16 products are + exact in fp32, so only the K-reduction order matters.) + + Two knobs must both land: + - cuBLASLt env pins for kernels that honor them (set before first use). + - Transformer Engine's own workspace: TE (<= 2.15 verified) hardcodes 32MiB in + get_cublas_workspace_size_bytes() and ignores CUBLASLT_WORKSPACE_SIZE, so the + env pin alone never engages for TE-launched GEMMs — patch the size fn and + clear its lru_cache. + """ + import os + + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":0:0") + os.environ.setdefault("CUBLASLT_WORKSPACE_SIZE", "0") + try: + import transformer_engine.pytorch.cpp_extensions.gemm as _te_gemm_mod + + if hasattr(_te_gemm_mod, "get_cublas_workspace_size_bytes"): + _te_gemm_mod.get_cublas_workspace_size_bytes = lambda: workspace_bytes + if hasattr(getattr(_te_gemm_mod, "get_cublas_workspace", None), "cache_clear"): + _te_gemm_mod.get_cublas_workspace.cache_clear() + except ImportError: + pass + + def enable_batch_invariant_mode(backend: str = "deepgemm"): """Enable global batch-invariant mode and patch Aten/TE kernels. @@ -1663,12 +1705,25 @@ def enable_batch_invariant_mode(backend: str = "deepgemm"): dispatch_key = getattr(torch.accelerator.current_accelerator(), "type", "cpu").upper() _batch_invariant_MODE = True _batch_invariant_LIB = torch.library.Library("aten", "IMPL") - _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, dispatch_key) - _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, dispatch_key) + if backend == "te_native": + # Keep the NATIVE cuBLASLt kernels for every dense GEMM (aten and TE); + # batch invariance comes from workspace starvation (split-K + # disqualified => fixed K-reduction recipe at every M) instead of + # kernel substitution — native speed, verified bitwise-identical to + # the training forward in NeMo-RL true on-policy GRPO (gen_kl == 0.0). + _enable_te_native_workspace_starvation() + else: + _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, dispatch_key) + _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, dispatch_key) _batch_invariant_LIB.impl("aten::_log_softmax", _log_softmax_batch_invariant, dispatch_key) _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, dispatch_key) - # Also patch Transformer Engine kernels when available - _te_patch_for_batch_invariant() + # Also patch Transformer Engine kernels when available (te_native keeps + # TE's native GEMMs — invariance comes from the starved workspace — but + # still applies the non-GEMM TE patches, e.g. the attention gate). + if backend == "te_native": + _te_patch_for_batch_invariant(skip_gemm=True) + else: + _te_patch_for_batch_invariant() # Pin the Mamba autotuners so rollout and training processes can't end # up on different tile configs (and therefore different fp32 reduction # orders) through autotune timing noise. From 6b9e1d6076538648ffcf9a4558a192d07263564d Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:45:58 -0700 Subject: [PATCH 05/26] Batch-invariant mode: make the GEMM backend selectable via config Add TransformerConfig.batch_invariant_backend ('deepgemm' | 'triton' | 'te_native') and plumb it through training initialization, so the backend added in the previous commit is reachable from the CLI (--batch-invariant-backend) instead of being hardcoded to the default. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- megatron/core/transformer/transformer_config.py | 7 +++++++ megatron/training/initialize.py | 5 +++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 0babb76cc69..15f002d0efd 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1167,6 +1167,13 @@ class TransformerConfig(ModelParallelConfig): training and inference as the kernels are not full optimized. Defaults to False.""" + batch_invariant_backend: str = "deepgemm" + """Which batch-invariant GEMM backend to use when batch_invariant_mode is + enabled: "deepgemm" (DeepGEMM bf16 kernels), "triton" (persistent Triton + matmul; any dtype), or "te_native" (keep the native cuBLASLt kernels and + obtain invariance via workspace starvation — lowest overhead, and the + configuration verified bitwise-identical to the TE training forward).""" + use_te_activation_func: bool = False """Whether to use ffn activation functions implemented by TransformerEngine""" diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 61374a4ae08..e21f9ab716c 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -100,8 +100,9 @@ def state_restore_func(state_dict): ) if args.batch_invariant_mode: - print_rank_0("Enabling batch invariant mode globally") - enable_batch_invariant_mode() + backend = getattr(args, "batch_invariant_backend", "deepgemm") + print_rank_0(f"Enabling batch invariant mode globally (backend={backend})") + enable_batch_invariant_mode(backend) # torch.distributed initialization def finish_mpu_init(): From da6ad41d37fc010a03acdcdd91875ca01fbf339f Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:45:59 -0700 Subject: [PATCH 06/26] Batch-invariant te_native backend: keep native TE RMSNorm TE RMSNorm switches its reduction codepath (and therefore bit pattern on rare-value rows) at M % 32. The default batch-invariant backends handle this by substituting a batch-invariant RMSNorm kernel; the te_native backend can instead keep the NATIVE kernel, because the 64-multiple alignment discipline (CUDA-graph bucket floor, eager TOKEN_ROUNDER, scoring-side sequence-length rounding) holds every launch in the same M%32 bit-class. Native RMSNorm is faster and is the configuration certified bitwise-identical to the training forward. Applies to both the inference-optimized RMSNorm call site and the TE class/module-level patches (skip_rmsnorm under te_native). Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../core/tensor_parallel/inference_layers.py | 6 +++++- .../custom_layers/batch_invariant_kernels.py | 16 ++++++++++++++-- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/megatron/core/tensor_parallel/inference_layers.py b/megatron/core/tensor_parallel/inference_layers.py index 102cdbcb4b9..521aaef0000 100644 --- a/megatron/core/tensor_parallel/inference_layers.py +++ b/megatron/core/tensor_parallel/inference_layers.py @@ -28,6 +28,7 @@ from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( is_batch_invariant_mode_enabled, rmsnorm_batch_invariant, + get_batch_invariant_backend, ) from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import get_tensor_model_parallel_group_if_none @@ -47,7 +48,10 @@ def _te_rms_norm_kernel(x: torch.Tensor, weight: torch.Tensor, eps: float): # Use the same RMSNorm kernel as the training recompute. - if is_batch_invariant_mode_enabled(): + if is_batch_invariant_mode_enabled() and get_batch_invariant_backend() != "te_native": + # te_native keeps the native TE RMSNorm: the 64-multiple alignment + # discipline holds its M%32 reduction bit-class constant, so kernel + # substitution is unnecessary (and native is faster). return rmsnorm_batch_invariant(x, weight, eps).to(x.dtype) x_shape = x.shape x = x.view(-1, x.size(-1)) diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index 508a6e739e4..75946f97726 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -625,7 +625,7 @@ def _import_module_if_available(name: str): return importlib.import_module(name) -def _te_patch_for_batch_invariant(skip_gemm: bool = False): +def _te_patch_for_batch_invariant(skip_gemm: bool = False, skip_rmsnorm: bool = False): """Patch Transformer Engine modules to use batch-invariant GEMM and RMSNorm. This monkey-patches TE's GEMM and RMSNorm entry points to dispatch to the @@ -673,6 +673,9 @@ def _te_patch_for_batch_invariant(skip_gemm: bool = False): _MEG_TE_GENERAL_GEMM_ORIG = meg_te.general_gemm meg_te.general_gemm = _te_general_gemm_patched + if skip_rmsnorm: + return + # Patch RMSNorm.forward once (class may be on te or te.pytorch) rms_cls = getattr(te, "RMSNorm", None) if rms_cls is None: @@ -1678,6 +1681,11 @@ def _enable_te_native_workspace_starvation(workspace_bytes: int = _TE_NATIVE_WOR pass +def get_batch_invariant_backend() -> str: + """Return the active batch-invariant GEMM backend name.""" + return _BATCH_INVARIANT_BACKEND + + def enable_batch_invariant_mode(backend: str = "deepgemm"): """Enable global batch-invariant mode and patch Aten/TE kernels. @@ -1721,7 +1729,11 @@ def enable_batch_invariant_mode(backend: str = "deepgemm"): # TE's native GEMMs — invariance comes from the starved workspace — but # still applies the non-GEMM TE patches, e.g. the attention gate). if backend == "te_native": - _te_patch_for_batch_invariant(skip_gemm=True) + # te_native also keeps TE's NATIVE RMSNorm: its M%32 reduction + # bit-class is held constant by the 64-multiple alignment discipline + # (CUDA-graph bucket floor + eager TOKEN_ROUNDER + scoring-side + # sequence-length rounding), so no kernel substitution is needed. + _te_patch_for_batch_invariant(skip_gemm=True, skip_rmsnorm=True) else: _te_patch_for_batch_invariant() # Pin the Mamba autotuners so rollout and training processes can't end From c4166a592ef8c69670ab8052a73baaa62812bda3 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:45:59 -0700 Subject: [PATCH 07/26] Batch-invariant vLLM MoE backend: pin only the K-reduction recipe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous change pinned the entire launch config to the large-M tuning, which penalizes decode (small-M) steps. The fused-MoE kernel accumulates in fp32 with no split-K, so its bit pattern depends only on the K-loop grouping: pin BLOCK_SIZE_K and keep the M/N tile shapes, tile grouping and pipeline depth hint-adaptive. bf16 products are exact in fp32, so tile shape changes reorder nothing in the accumulation — only the K addition order matters. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- megatron/core/inference/moe/vllm_fused_moe.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index 8a77327e56e..3039c99e0ee 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/megatron/core/inference/moe/vllm_fused_moe.py @@ -611,16 +611,17 @@ def vllm_fused_moe( # stages) host-side from the token-count hint, not from the worst-case # buffer size. Same config is used for both FC1 and FC2 (matches vLLM). batch_invariant_mode = batch_invariant.enabled() + config = _get_default_config(M=effective_tokens, E=num_local_experts, top_k=topk) if batch_invariant_mode: - # Batch-invariant mode: the config must not depend on the token count. - # _get_default_config is a step function of M, so different co-batch - # sizes would otherwise select different tile shapes and change the - # fp32 accumulation grouping (batch-variant bits). Pin the large-M - # config for every launch; grid SIZING below may still use the hint - # (the kernel strides, so grid size never changes per-tile math). - config = _get_default_config(M=1 << 30, E=num_local_experts, top_k=topk) - else: - config = _get_default_config(M=effective_tokens, E=num_local_experts, top_k=topk) + # Batch-invariant mode: pin only the K-reduction recipe. The kernel + # accumulates in fp32 with no split-K, so bits depend solely on the + # K-loop grouping (BLOCK_SIZE_K); M/N tile shapes, tile grouping and + # pipeline depth reorder nothing in the accumulation (bf16 products + # are exact in fp32 — only the addition order matters). Keeping the + # M/N tiling adaptive preserves the decode-tuned configs; pinning + # BLOCK_SIZE_K removes the one field _get_default_config varies that + # could change the summation order across co-batch sizes. + config['BLOCK_SIZE_K'] = 64 sorted_token_ids, expert_ids, num_post_padded = _moe_align_block_size_cuda_graphable( routing_map, config['BLOCK_SIZE_M'], num_local_experts, local_expert_start, valid_tokens From 3e85a90c6ebb647db57f2b1aa117805fd355b7bf Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:46:00 -0700 Subject: [PATCH 08/26] Tests: batch-invariant vLLM fused-MoE backend Unit tests for the batch-invariant vLLM-backend work: - CUDA-graph bucket token counts are 64-multiples under batch-invariant mode (all sizing distributions incl. num_cuda_graphs=-1 auto), with a guard test documenting that the non-BI auto ladder still injects 1/2-token buckets. - swiglu_with_probs / weighted_silu_mul_bounded: value correctness (tolerance-based vs a torch reference), bitwise repeat-determinism, row-locality across co-batch sizes, and device-bound soundness under NaN-poisoned tails. - _moe_sum apply_weights/acc_fp64 options: unit-weight fp64 reduction is bitwise-exact vs an fp64 reference; the default weighted fp32 path is unchanged. - End-to-end batch invariance: the same tokens produce bitwise-identical outputs across co-batch sizes/hint classes (exercises the pinned K-reduction recipe with adaptive M/N tiling and warp counts). - te_native backend registration and enable/disable round-trip (native aten::mm stays unpatched). Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../test_vllm_fused_moe_batch_invariant.py | 251 ++++++++++++++++++ 1 file changed, 251 insertions(+) create mode 100644 tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py diff --git a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py new file mode 100644 index 00000000000..775f2a762ed --- /dev/null +++ b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py @@ -0,0 +1,251 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for batch-invariant mode on the vLLM Triton fused-MoE backend. + +Covers: +- CUDA-graph bucket token counts floored to 64-multiples under batch-invariant mode +- swiglu_with_probs / weighted_silu_mul_bounded: training-parity rounding vs reference +- _moe_sum apply_weights / acc_fp64 options +- vllm_fused_moe end-to-end batch invariance for gated (SwiGLU) models +- te_native backend registration +""" + +import os +import tempfile + +os.environ.setdefault("TRITON_CACHE_DIR", os.path.join(tempfile.gettempdir(), "triton_test_cache")) + +import pytest +import torch + +from megatron.core.inference.batch_dimensions_utils import CUDAGraphBatchDimensionBuilder +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + _BATCH_INVARIANT_BACKENDS, + set_batch_invariant_mode, +) + + +def _vt(n): + return torch.tensor(n, dtype=torch.int32, device="cuda") + + +# --------------------------------------------------------------------------- +# CUDA-graph bucket 64-multiple floor +# --------------------------------------------------------------------------- + + +class TestCudaGraphBucket64Floor: + + @pytest.mark.parametrize("num_cuda_graphs", [-1, 8, 16]) + def test_bucket_token_counts_are_64_multiples(self, num_cuda_graphs): + with set_batch_invariant_mode(True, backend="triton"): + dims, token_counts = ( + CUDAGraphBatchDimensionBuilder.generate_cuda_graph_batch_dimensions_list( + tp_size=1, + num_cuda_graphs=num_cuda_graphs, + cuda_graph_max_tokens=2048, + cuda_graph_mixed_prefill_request_count=None, + max_requests=512, + max_tokens=2048, + max_sequence_length=4096, + use_cuda_graphs_for_non_decode_steps=False, + ) + ) + for bd in dims: + assert bd.token_count % 64 == 0 and bd.token_count >= 64, ( + f"bucket token_count {bd.token_count} violates the 64-multiple floor " + f"(num_cuda_graphs={num_cuda_graphs})" + ) + + def test_auto_sizing_injects_small_buckets_without_bi(self): + # Guard for the non-BI behavior this floor exists to counteract: the + # auto (-1) ladder includes 1/2-token buckets when BI mode is off. + dims, _ = CUDAGraphBatchDimensionBuilder.generate_cuda_graph_batch_dimensions_list( + tp_size=1, + num_cuda_graphs=-1, + cuda_graph_max_tokens=2048, + cuda_graph_mixed_prefill_request_count=None, + max_requests=512, + max_tokens=2048, + max_sequence_length=4096, + use_cuda_graphs_for_non_decode_steps=False, + ) + assert any(bd.token_count < 64 for bd in dims) + + +# --------------------------------------------------------------------------- +# Training-parity weighted SwiGLU kernels +# --------------------------------------------------------------------------- + + +def _weighted_swiglu_reference(y, probs_flat): + """bf16(fp32 silu(gate) * up * prob) with a single final rounding.""" + half = y.shape[1] // 2 + gate = y[:, :half].float() + up = y[:, half:].float() + return (torch.nn.functional.silu(gate) * up * probs_flat[:, None]).to(y.dtype) + + +class TestWeightedSwigluKernels: + + def test_swiglu_with_probs_value_deterministic_and_row_local(self): + from megatron.core.inference.moe import batch_invariant + + torch.manual_seed(7) + rows, ffn = 512, 256 + y = (torch.randn(rows, 2 * ffn, device="cuda") * 2.0).bfloat16() + probs = torch.rand(rows, device="cuda", dtype=torch.float32) + perm_map = torch.arange(rows, device="cuda", dtype=torch.int32) + n_used = _vt(rows) + out = batch_invariant.swiglu_with_probs(y, perm_map, n_used, probs) + # value correctness (tolerance-based: sigmoid instruction sequences may + # legitimately differ from the torch reference by 1 ulp on rare values) + # compare at bf16 (bf16 tolerances): sigmoid instruction sequences may + # legitimately differ from the torch reference by 1 bf16 ulp + torch.testing.assert_close(out, _weighted_swiglu_reference(y, probs)) + # bitwise repeat-determinism + for _ in range(5): + assert torch.equal( + batch_invariant.swiglu_with_probs(y, perm_map, n_used, probs), out + ) + # row-locality: a row's bits do not depend on co-batch size + half_out = batch_invariant.swiglu_with_probs( + y[:128].contiguous(), perm_map[:128], _vt(128), probs[:128] + ) + assert torch.equal(half_out, out[:128]) + + def test_weighted_silu_mul_bounded_bound_and_invariance(self): + from megatron.core.inference.moe import batch_invariant + + torch.manual_seed(8) + rows, ffn, live = 512, 256, 300 + y = (torch.randn(rows, 2 * ffn, device="cuda") * 2.0).bfloat16() + probs = torch.rand(rows, device="cuda", dtype=torch.float32) + bound = torch.tensor(live * ffn, dtype=torch.int64, device="cuda") + out = batch_invariant.weighted_silu_mul_bounded(y, probs, bound) + torch.testing.assert_close( + out[:live], _weighted_swiglu_reference(y[:live], probs[:live]) + ) + # rows beyond the device bound are neither read nor written: NaN-poison + # the tail and require the live rows to stay BITWISE identical + y2 = y.clone() + y2[live:] = float("nan") + out2 = batch_invariant.weighted_silu_mul_bounded(y2, probs, bound) + assert torch.equal(out2[:live], out[:live]) + + +# --------------------------------------------------------------------------- +# _moe_sum options +# --------------------------------------------------------------------------- + + +class TestMoeSumOptions: + + def _setup(self): + torch.manual_seed(9) + max_tokens, topk, K, E = 64, 4, 128, 8 + inp = (torch.randn(max_tokens * topk, K, device="cuda")).bfloat16() + probs = torch.rand(max_tokens, topk, device="cuda", dtype=torch.float32) + routing = torch.randint(0, E, (max_tokens, topk), device="cuda", dtype=torch.int64) + return inp, probs, routing, max_tokens, topk, K, E + + def test_unit_weights_fp64_matches_fp64_reference(self): + from megatron.core.inference.moe.vllm_fused_moe import _moe_sum + + inp, probs, routing, max_tokens, topk, K, E = self._setup() + out = _moe_sum( + inp, probs, max_tokens, topk, K, _vt(max_tokens), routing, 0, E, + apply_weights=False, acc_fp64=True, + ) + ref = ( + inp.view(max_tokens, topk, K).to(torch.float64).sum(dim=1).to(torch.float32) + ) + assert torch.equal(out, ref) + + def test_default_weighted_fp32_deterministic_and_correct(self): + from megatron.core.inference.moe.vllm_fused_moe import _moe_sum + + inp, probs, routing, max_tokens, topk, K, E = self._setup() + out = _moe_sum(inp, probs, max_tokens, topk, K, _vt(max_tokens), routing, 0, E) + # value correctness (tolerance-based: the kernel's acc += v*w compiles + # to FMA, which a mul-then-add torch reference cannot match bitwise) + ref = torch.zeros(max_tokens, K, device="cuda", dtype=torch.float32) + for t in range(topk): + ref += inp.view(max_tokens, topk, K)[:, t].float() * probs[:, t : t + 1] + torch.testing.assert_close(out, ref) + # bitwise repeat-determinism of the default path + for _ in range(5): + assert torch.equal( + _moe_sum(inp, probs, max_tokens, topk, K, _vt(max_tokens), routing, 0, E), + out, + ) + + +# --------------------------------------------------------------------------- +# End-to-end batch invariance (gated / SwiGLU) +# --------------------------------------------------------------------------- + + +class TestVllmFusedMoeBatchInvariance: + + def test_same_tokens_bitwise_across_cobatch_sizes(self): + from megatron.core.inference.moe import ActivationType + from megatron.core.inference.moe.vllm_fused_moe import vllm_fused_moe + + torch.manual_seed(11) + max_tokens, K, ffn, E, topk = 256, 128, 64, 8, 4 + hidden = (torch.randn(max_tokens, K, device="cuda") * 0.05).bfloat16() + fc1 = (torch.randn(E, 2 * ffn, K, device="cuda") * 0.02).bfloat16() + fc2 = (torch.randn(E, K, ffn, device="cuda") * 0.02).bfloat16() + routing = torch.stack( + [torch.randperm(E, device="cuda")[:topk] for _ in range(max_tokens)] + ).long() + probs = torch.rand(max_tokens, topk, device="cuda", dtype=torch.float32) + probs = probs / probs.sum(-1, keepdim=True) + + def run(valid, hint): + with set_batch_invariant_mode(True, backend="triton"): + return vllm_fused_moe( + hidden, probs, fc1, fc2, + activation_type=ActivationType.SWIGLU, + num_local_experts=E, local_expert_start=0, + valid_tokens=_vt(valid), routing_map=routing, + num_tokens_hint=hint, + ) + + # the first 64 tokens, computed in co-batches of different sizes and + # hint classes, must be bitwise identical + small = run(valid=64, hint=64)[:64] + large = run(valid=256, hint=256)[:64] + assert torch.equal(small, large) + + +# --------------------------------------------------------------------------- +# te_native backend registration +# --------------------------------------------------------------------------- + + +class TestTeNativeBackend: + + def test_backend_registered(self): + assert "te_native" in _BATCH_INVARIANT_BACKENDS + + def test_enable_disable_roundtrip(self): + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + disable_batch_invariant_mode, + enable_batch_invariant_mode, + get_batch_invariant_backend, + is_batch_invariant_mode_enabled, + ) + + try: + enable_batch_invariant_mode("te_native") + assert is_batch_invariant_mode_enabled() + assert get_batch_invariant_backend() == "te_native" + # te_native must NOT reroute aten::mm — native kernels stay + a = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + b = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + torch.mm(a, b) # should not raise / not require DeepGEMM + finally: + disable_batch_invariant_mode() + assert not is_batch_invariant_mode_enabled() From 1ed83b98fdfda1ddf785ea59c7b88cca4a16c662 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:46:00 -0700 Subject: [PATCH 09/26] Docstrings: document te_native backend and _moe_sum options Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- megatron/core/inference/moe/vllm_fused_moe.py | 7 +++++++ .../transformer/custom_layers/batch_invariant_kernels.py | 5 ++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index 3039c99e0ee..67774ab298a 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/megatron/core/inference/moe/vllm_fused_moe.py @@ -529,6 +529,13 @@ def _moe_sum( (downstream RSV reads only the valid range). Only accumulates contributions from local experts; non-local topk slots are skipped (their values in `input` are undefined). + + Args: + apply_weights: multiply each slot by its routing probability (default). + Pass False when the probabilities were already applied upstream + (e.g. at the activation in batch-invariant mode). + acc_fp64: accumulate the topk sum in fp64, making the reduction + order-independent by precision. """ if out is None: out = torch.empty(max_tokens, K, dtype=torch.float32, device=input.device) diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index 75946f97726..b87570d68c4 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -1694,7 +1694,10 @@ def enable_batch_invariant_mode(backend: str = "deepgemm"): "deepgemm" (default) routes bf16 CUDA inputs through DeepGEMM `bf16_gemm_nn`. "triton" routes through the batch-invariant Triton `matmul_persistent` kernel (works for bf16/fp16/fp32 and - on any CUDA device). Grouped GEMM always uses DeepGEMM regardless. + on any CUDA device). "te_native" keeps the native cuBLASLt + kernels (aten and TE) and obtains invariance via workspace + starvation instead of kernel substitution. Grouped GEMM uses + DeepGEMM for "deepgemm"/"triton"; "te_native" leaves it native. """ global _batch_invariant_MODE, _batch_invariant_LIB, _BATCH_INVARIANT_BACKEND if _batch_invariant_MODE: From d5ece8416c7c6e2de79c96987405b4f3c275a7bf Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 11:46:01 -0700 Subject: [PATCH 10/26] Apply tools/autoformat.sh (black/isort/ruff) Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../core/inference/batch_dimensions_utils.py | 2 +- megatron/core/inference/moe/vllm_fused_moe.py | 6 +-- .../core/tensor_parallel/inference_layers.py | 2 +- .../test_vllm_fused_moe_batch_invariant.py | 39 +++++++++++-------- 4 files changed, 27 insertions(+), 22 deletions(-) diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index 3aee65eb862..376a05a080a 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/megatron/core/inference/batch_dimensions_utils.py @@ -175,7 +175,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: diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index 67774ab298a..6c2b7ec40dd 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/megatron/core/inference/moe/vllm_fused_moe.py @@ -29,8 +29,8 @@ triton.jit = null_decorator tl = MagicMock() -from megatron.core.inference.moe.activations import bounded_silu_mul from megatron.core.inference.moe import batch_invariant +from megatron.core.inference.moe.activations import bounded_silu_mul from megatron.core.inference.moe.fused_moe import ActivationType from megatron.core.inference.moe.permute import ( _get_num_sms, @@ -683,9 +683,7 @@ def vllm_fused_moe( # (single bf16 round of fp32 silu(gate)*up*prob). The reduction # below then sums with unit weights. Device-bounded to the live # valid_tokens*topk prefix; CUDA-graph safe. - bound_elems = valid_tokens.to(torch.int64) * ( - topk * (intermediate1.shape[1] // 2) - ) + bound_elems = valid_tokens.to(torch.int64) * (topk * (intermediate1.shape[1] // 2)) intermediate1 = batch_invariant.weighted_silu_mul_bounded( intermediate1, topk_weights_flat, bound_elems ) diff --git a/megatron/core/tensor_parallel/inference_layers.py b/megatron/core/tensor_parallel/inference_layers.py index 521aaef0000..2da50cc66bf 100644 --- a/megatron/core/tensor_parallel/inference_layers.py +++ b/megatron/core/tensor_parallel/inference_layers.py @@ -26,9 +26,9 @@ reduce_scatter_to_sequence_parallel_region, ) from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + get_batch_invariant_backend, is_batch_invariant_mode_enabled, rmsnorm_batch_invariant, - get_batch_invariant_backend, ) from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import get_tensor_model_parallel_group_if_none diff --git a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py index 775f2a762ed..4a690290a41 100644 --- a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py +++ b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py @@ -105,9 +105,7 @@ def test_swiglu_with_probs_value_deterministic_and_row_local(self): torch.testing.assert_close(out, _weighted_swiglu_reference(y, probs)) # bitwise repeat-determinism for _ in range(5): - assert torch.equal( - batch_invariant.swiglu_with_probs(y, perm_map, n_used, probs), out - ) + assert torch.equal(batch_invariant.swiglu_with_probs(y, perm_map, n_used, probs), out) # row-locality: a row's bits do not depend on co-batch size half_out = batch_invariant.swiglu_with_probs( y[:128].contiguous(), perm_map[:128], _vt(128), probs[:128] @@ -123,9 +121,7 @@ def test_weighted_silu_mul_bounded_bound_and_invariance(self): probs = torch.rand(rows, device="cuda", dtype=torch.float32) bound = torch.tensor(live * ffn, dtype=torch.int64, device="cuda") out = batch_invariant.weighted_silu_mul_bounded(y, probs, bound) - torch.testing.assert_close( - out[:live], _weighted_swiglu_reference(y[:live], probs[:live]) - ) + torch.testing.assert_close(out[:live], _weighted_swiglu_reference(y[:live], probs[:live])) # rows beyond the device bound are neither read nor written: NaN-poison # the tail and require the live rows to stay BITWISE identical y2 = y.clone() @@ -154,12 +150,19 @@ def test_unit_weights_fp64_matches_fp64_reference(self): inp, probs, routing, max_tokens, topk, K, E = self._setup() out = _moe_sum( - inp, probs, max_tokens, topk, K, _vt(max_tokens), routing, 0, E, - apply_weights=False, acc_fp64=True, - ) - ref = ( - inp.view(max_tokens, topk, K).to(torch.float64).sum(dim=1).to(torch.float32) + inp, + probs, + max_tokens, + topk, + K, + _vt(max_tokens), + routing, + 0, + E, + apply_weights=False, + acc_fp64=True, ) + ref = inp.view(max_tokens, topk, K).to(torch.float64).sum(dim=1).to(torch.float32) assert torch.equal(out, ref) def test_default_weighted_fp32_deterministic_and_correct(self): @@ -176,8 +179,7 @@ def test_default_weighted_fp32_deterministic_and_correct(self): # bitwise repeat-determinism of the default path for _ in range(5): assert torch.equal( - _moe_sum(inp, probs, max_tokens, topk, K, _vt(max_tokens), routing, 0, E), - out, + _moe_sum(inp, probs, max_tokens, topk, K, _vt(max_tokens), routing, 0, E), out ) @@ -206,10 +208,15 @@ def test_same_tokens_bitwise_across_cobatch_sizes(self): def run(valid, hint): with set_batch_invariant_mode(True, backend="triton"): return vllm_fused_moe( - hidden, probs, fc1, fc2, + hidden, + probs, + fc1, + fc2, activation_type=ActivationType.SWIGLU, - num_local_experts=E, local_expert_start=0, - valid_tokens=_vt(valid), routing_map=routing, + num_local_experts=E, + local_expert_start=0, + valid_tokens=_vt(valid), + routing_map=routing, num_tokens_hint=hint, ) From ad1c023d1a4a23064c94f11334f1fe482aed1791 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 14:41:58 -0700 Subject: [PATCH 11/26] Address review: clearer rounding expression in _batch_invariant_token_floor Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- megatron/core/inference/batch_dimensions_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index 376a05a080a..9f169193819 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/megatron/core/inference/batch_dimensions_utils.py @@ -216,7 +216,8 @@ def _batch_invariant_token_floor(token_count: int) -> int: graphed norms execute in a different bit-class, breaking cross-batch bit-equality. Request counts are untouched (mirrors eager semantics). """ - return max(64, ((token_count + 63) // 64) * 64) + rounded_up = math.ceil(token_count / 64) * 64 + return max(64, rounded_up) def _batch_invariant_mode_enabled() -> bool: From cc584955938c8f9e6126d110c44a94c2d4ca3daf Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 14:47:50 -0700 Subject: [PATCH 12/26] Address review: make te_native the default batch-invariant backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit te_native keeps the native cuBLASLt kernels (invariance via workspace starvation), has the lowest overhead, needs no extra dependencies, and is the configuration verified bitwise-identical to the TE training forward. Also scope the DeepGEMM MoE requirement to the backend combinations that actually use it (deepgemm/triton backends, or the torch inference grouped-GEMM path); te_native with the vLLM inference backend — and the training path, where TE grouped GEMM stays native — no longer require DeepGEMM to be installed. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../custom_layers/batch_invariant_kernels.py | 7 +++--- .../core/transformer/transformer_config.py | 25 +++++++++++++------ megatron/training/initialize.py | 2 +- 3 files changed, 23 insertions(+), 11 deletions(-) diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index b87570d68c4..457d63cd3eb 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -1686,13 +1686,14 @@ def get_batch_invariant_backend() -> str: return _BATCH_INVARIANT_BACKEND -def enable_batch_invariant_mode(backend: str = "deepgemm"): +def enable_batch_invariant_mode(backend: str = "te_native"): """Enable global batch-invariant mode and patch Aten/TE kernels. Args: backend: which kernel to dispatch `aten::mm`/`aten::addmm` through. - "deepgemm" (default) routes bf16 CUDA inputs through DeepGEMM - `bf16_gemm_nn`. "triton" routes through the batch-invariant + "te_native" (default) keeps the native cuBLASLt kernels and obtains + invariance via workspace starvation. "deepgemm" routes bf16 CUDA + inputs through DeepGEMM `bf16_gemm_nn`. "triton" routes through the batch-invariant Triton `matmul_persistent` kernel (works for bf16/fp16/fp32 and on any CUDA device). "te_native" keeps the native cuBLASLt kernels (aten and TE) and obtains invariance via workspace diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 15f002d0efd..7660221d5d6 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1167,12 +1167,13 @@ class TransformerConfig(ModelParallelConfig): training and inference as the kernels are not full optimized. Defaults to False.""" - batch_invariant_backend: str = "deepgemm" + batch_invariant_backend: str = "te_native" """Which batch-invariant GEMM backend to use when batch_invariant_mode is - enabled: "deepgemm" (DeepGEMM bf16 kernels), "triton" (persistent Triton - matmul; any dtype), or "te_native" (keep the native cuBLASLt kernels and - obtain invariance via workspace starvation — lowest overhead, and the - configuration verified bitwise-identical to the TE training forward).""" + enabled: "te_native" (default: keep the native cuBLASLt kernels and obtain + invariance via workspace starvation — lowest overhead, no extra + dependencies, and the configuration verified bitwise-identical to the TE + training forward), "deepgemm" (DeepGEMM bf16 kernels), or "triton" + (persistent Triton matmul; any dtype).""" use_te_activation_func: bool = False """Whether to use ffn activation functions implemented by TransformerEngine""" @@ -3191,9 +3192,19 @@ def _scope_to_str(s): "Batch-invariant MoE training requires " "moe_token_dispatcher_type='alltoall'." ) - assert HAVE_DEEPGEMM_BF16, ( + # DeepGEMM is used by the "deepgemm"/"triton" backends, and by + # the torch inference grouped-GEMM path under any backend. The + # "te_native" backend with the vLLM inference backend (or the + # training path, where TE grouped GEMM stays native) does not + # need it. + needs_deepgemm = self.batch_invariant_backend in ("deepgemm", "triton") or ( + self.transformer_impl == "inference_optimized" + and self.inference_grouped_gemm_backend == InferenceGroupedGemmBackend.TORCH + ) + assert not needs_deepgemm or HAVE_DEEPGEMM_BF16, ( "batch_invariant_mode=True with MoE requires DeepGEMM with bf16 " - "grouped-GEMM bindings (m_grouped_bf16_gemm_nt_contiguous). " + "grouped-GEMM bindings (m_grouped_bf16_gemm_nt_contiguous) for " + "this backend combination. " "Install via `uv pip install -e .[batch_invariant]`." ) assert not ( diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index e21f9ab716c..2ddcaa3c5a5 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -100,7 +100,7 @@ def state_restore_func(state_dict): ) if args.batch_invariant_mode: - backend = getattr(args, "batch_invariant_backend", "deepgemm") + backend = getattr(args, "batch_invariant_backend", "te_native") print_rank_0(f"Enabling batch invariant mode globally (backend={backend})") enable_batch_invariant_mode(backend) From 41f8dd3deb714f1a0395a37a1a42fd3dbbfc6ba4 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 14:51:18 -0700 Subject: [PATCH 13/26] Address review: make _moe_sum apply_weights/acc_fp64 truly orthogonal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit APPLY_WEIGHTS was nested inside the ACC_FP64 branch, so apply_weights=False was silently ignored on the fp32 path (weights applied anyway). Latent today — the single call site couples the flags — but the docstring advertises them as independent. Hoist APPLY_WEIGHTS out and add the missing (apply_weights=False, acc_fp64=False) test case. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- megatron/core/inference/moe/vllm_fused_moe.py | 14 ++++++----- .../test_vllm_fused_moe_batch_invariant.py | 24 +++++++++++++++++++ 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index 6c2b7ec40dd..5105cedb3f4 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/megatron/core/inference/moe/vllm_fused_moe.py @@ -492,15 +492,17 @@ def _moe_sum_kernel( lid = eid - local_expert_start if lid >= 0 and lid < num_local_experts: v = tl.load(input_ptr + base + t * K + offs_k, mask=k_mask, other=0.0) - if ACC_FP64: - if APPLY_WEIGHTS: - w = tl.load(topk_weights_ptr + token_id * topk + t) + if APPLY_WEIGHTS: + w = tl.load(topk_weights_ptr + token_id * topk + t) + if ACC_FP64: acc += v.to(tl.float64) * w.to(tl.float64) else: - acc += v.to(tl.float64) + acc += v.to(tl.float32) * w else: - w = tl.load(topk_weights_ptr + token_id * topk + t) - acc += v.to(tl.float32) * w + if ACC_FP64: + acc += v.to(tl.float64) + else: + acc += v.to(tl.float32) tl.store(output_ptr + token_id_i64 * K + offs_k, acc.to(tl.float32), mask=k_mask) diff --git a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py index 4a690290a41..ed6b6a201cb 100644 --- a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py +++ b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py @@ -165,6 +165,30 @@ def test_unit_weights_fp64_matches_fp64_reference(self): ref = inp.view(max_tokens, topk, K).to(torch.float64).sum(dim=1).to(torch.float32) assert torch.equal(out, ref) + def test_unit_weights_fp32_matches_sequential_reference(self): + from megatron.core.inference.moe.vllm_fused_moe import _moe_sum + + inp, probs, routing, max_tokens, topk, K, E = self._setup() + out = _moe_sum( + inp, + probs, + max_tokens, + topk, + K, + _vt(max_tokens), + routing, + 0, + E, + apply_weights=False, + acc_fp64=False, + ) + # unit weights => pure sequential fp32 adds (no FMA), so a same-order + # torch reference is bitwise-reproducible + ref = torch.zeros(max_tokens, K, device="cuda", dtype=torch.float32) + for t in range(topk): + ref += inp.view(max_tokens, topk, K)[:, t].float() + assert torch.equal(out, ref) + def test_default_weighted_fp32_deterministic_and_correct(self): from megatron.core.inference.moe.vllm_fused_moe import _moe_sum From e5bc58b263748051c152f0c665f7ad9fd997e538 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 15:12:34 -0700 Subject: [PATCH 14/26] Address review: pad buckets to alignment instead of invalidating them Validate batch dimensions BEFORE the 64-multiple alignment: the floor 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 and request budgets are untouched. Previously a non-64-multiple max_requests (e.g. 100) lost its largest decode bucket to the is_valid token-sufficiency check. Regression test with max_requests=100 added. Also rename _batch_invariant_token_floor -> _batch_invariant_token_align (it rounds up). Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../core/inference/batch_dimensions_utils.py | 30 +++++++++----- .../test_vllm_fused_moe_batch_invariant.py | 41 +++++++++++++++++++ 2 files changed, 60 insertions(+), 11 deletions(-) diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index 9f169193819..d2c26fc196c 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/megatron/core/inference/batch_dimensions_utils.py @@ -205,8 +205,8 @@ def adjust_batch_dims_for_expert_parallelism( return adjusted_batch_dim -def _batch_invariant_token_floor(token_count: int) -> int: - """Floor a CUDA-graph bucket token count to a 64-multiple (min 64). +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 @@ -304,9 +304,9 @@ def _calculate_cuda_graph_token_counts( rounder = CUDAGraphBatchDimensionBuilder.CUDA_GRAPH_ROUNDER if _batch_invariant_mode_enabled(): # Batch-invariant mode: 64-multiple token ladder (see - # _batch_invariant_token_floor). + # _batch_invariant_token_align). rounder = 64 - cuda_graph_max_tokens = _batch_invariant_token_floor(cuda_graph_max_tokens) + 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 @@ -373,8 +373,8 @@ def _calculate_token_counts_linear( ) if _batch_invariant_mode_enabled(): # Batch-invariant mode: floor every bucket to a 64-multiple - # (see _batch_invariant_token_floor) and dedupe collisions. - sizes = [_batch_invariant_token_floor(s) for s in sizes] + # (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] @@ -463,13 +463,21 @@ def generate_cuda_graph_batch_dimensions_list( def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int) -> None: """Helper to create and append batch dimension to list only if it's valid.""" - if _batch_invariant_mode_enabled(): - # Batch-invariant mode: floor EVERY bucket's token count to a - # 64-multiple (see _batch_invariant_token_floor); the flooring - # can collide previously-distinct buckets, so skip duplicates. - token_count = _batch_invariant_token_floor(token_count) batch_dim = InferenceBatchDimensions(token_count, prefill_req_count, decode_req_count) if batch_dim.is_valid(max_requests, max_sequence_length, num_speculative_tokens): + if _batch_invariant_mode_enabled(): + # Batch-invariant mode: floor the bucket's token count to a + # 64-multiple (see _batch_invariant_token_align). The floor + # is alignment 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. Flooring 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) diff --git a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py index ed6b6a201cb..6c2f89f6670 100644 --- a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py +++ b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py @@ -57,6 +57,31 @@ def test_bucket_token_counts_are_64_multiples(self, num_cuda_graphs): f"(num_cuda_graphs={num_cuda_graphs})" ) + def test_non_multiple_max_requests_keeps_largest_decode_bucket(self): + # Regression: the floor must PAD buckets, not invalidate them. With + # max_requests=100 (not a 64-multiple), the largest decode bucket must + # survive with its request budget intact and an aligned token count. + with set_batch_invariant_mode(True, backend="triton"): + dims, _ = CUDAGraphBatchDimensionBuilder.generate_cuda_graph_batch_dimensions_list( + tp_size=1, + num_cuda_graphs=16, + cuda_graph_max_tokens=2048, + cuda_graph_mixed_prefill_request_count=None, + max_requests=100, + max_tokens=2048, + max_sequence_length=4096, + use_cuda_graphs_for_non_decode_steps=False, + ) + decode_dims = [d for d in dims if d.prefill_req_count == 0 and d.decode_req_count > 0] + assert decode_dims, "no decode buckets survived" + largest = max(decode_dims, key=lambda d: d.decode_req_count) + assert ( + largest.decode_req_count == 100 + ), f"largest decode bucket lost its request budget: {largest}" + assert ( + largest.token_count % 64 == 0 and largest.token_count >= 128 + ), f"largest decode bucket not aligned/padded: {largest}" + def test_auto_sizing_injects_small_buckets_without_bi(self): # Guard for the non-BI behavior this floor exists to counteract: the # auto (-1) ladder includes 1/2-token buckets when BI mode is off. @@ -269,10 +294,21 @@ def test_enable_disable_roundtrip(self): is_batch_invariant_mode_enabled, ) + try: + import transformer_engine.pytorch.cpp_extensions.gemm as te_gemm_mod + + ws_fn_before = te_gemm_mod.get_cublas_workspace_size_bytes + have_te = True + except ImportError: + have_te = False + env_before = os.environ.get("CUBLASLT_WORKSPACE_SIZE") try: enable_batch_invariant_mode("te_native") assert is_batch_invariant_mode_enabled() assert get_batch_invariant_backend() == "te_native" + assert os.environ.get("CUBLASLT_WORKSPACE_SIZE") == "0" + if have_te: + assert te_gemm_mod.get_cublas_workspace_size_bytes() == 1024 # te_native must NOT reroute aten::mm — native kernels stay a = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) b = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) @@ -280,3 +316,8 @@ def test_enable_disable_roundtrip(self): finally: disable_batch_invariant_mode() assert not is_batch_invariant_mode_enabled() + # the workspace patch and env pin must be fully restored (no leak + # into subsequent non-BI work in the same process) + if have_te: + assert te_gemm_mod.get_cublas_workspace_size_bytes is ws_fn_before + assert os.environ.get("CUBLASLT_WORKSPACE_SIZE") == env_before From 93e933c9ab31949a94576da9b01a8a4f92767676 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 15:12:34 -0700 Subject: [PATCH 15/26] Address review: restore te_native workspace state on disable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Save and restore TE's get_cublas_workspace_size_bytes across enable/disable so the 1KB workspace cannot leak into non-batch-invariant work in the same process, and force-pin CUBLASLT_WORKSPACE_SIZE (with save/restore, a warning when overriding a preset value, and a warning when enabling after CUDA initialization — earlier cuBLASLt handles may retain their original workspace). Round-trip test asserts full restoration. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- .../custom_layers/batch_invariant_kernels.py | 61 ++++++++++++++++++- 1 file changed, 58 insertions(+), 3 deletions(-) diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index 457d63cd3eb..f0061e505f9 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -1647,6 +1647,10 @@ def is_batch_invariant_mode_enabled(): _TE_NATIVE_WORKSPACE_BYTES = 1024 +# Originals saved by _enable_te_native_workspace_starvation for restoration. +_TE_WORKSPACE_SIZE_FN_ORIG = None +_TE_NATIVE_ENV_ORIG: dict = {} + def _enable_te_native_workspace_starvation(workspace_bytes: int = _TE_NATIVE_WORKSPACE_BYTES): """Make the NATIVE cuBLASLt GEMM kernels batch-invariant via workspace starvation. @@ -1666,14 +1670,41 @@ def _enable_te_native_workspace_starvation(workspace_bytes: int = _TE_NATIVE_WOR env pin alone never engages for TE-launched GEMMs — patch the size fn and clear its lru_cache. """ + import logging import os - os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":0:0") - os.environ.setdefault("CUBLASLT_WORKSPACE_SIZE", "0") + global _TE_WORKSPACE_SIZE_FN_ORIG + logger = logging.getLogger(__name__) + + # CUBLASLT_WORKSPACE_SIZE must be pinned (not setdefault): a preset value + # (e.g. from a determinism launcher) large enough for split-K would make + # te_native silently non-invariant on the aten GEMM path. + for var, pinned in (("CUBLASLT_WORKSPACE_SIZE", "0"),): + prev = os.environ.get(var) + if prev is not None and prev != pinned: + logger.warning( + "te_native batch-invariant backend overriding %s=%s with %s " + "(split-K disqualification requires a starved workspace).", + var, + prev, + pinned, + ) + _TE_NATIVE_ENV_ORIG[var] = prev + os.environ[var] = pinned + if torch.cuda.is_initialized(): + logger.warning( + "te_native backend enabled after CUDA initialization; cuBLASLt " + "handles created earlier may retain their original workspace and " + "remain batch-variant. Enable batch-invariant mode before the " + "first GEMM." + ) try: import transformer_engine.pytorch.cpp_extensions.gemm as _te_gemm_mod - if hasattr(_te_gemm_mod, "get_cublas_workspace_size_bytes"): + if _TE_WORKSPACE_SIZE_FN_ORIG is None and hasattr( + _te_gemm_mod, "get_cublas_workspace_size_bytes" + ): + _TE_WORKSPACE_SIZE_FN_ORIG = _te_gemm_mod.get_cublas_workspace_size_bytes _te_gemm_mod.get_cublas_workspace_size_bytes = lambda: workspace_bytes if hasattr(getattr(_te_gemm_mod, "get_cublas_workspace", None), "cache_clear"): _te_gemm_mod.get_cublas_workspace.cache_clear() @@ -1681,6 +1712,29 @@ def _enable_te_native_workspace_starvation(workspace_bytes: int = _TE_NATIVE_WOR pass +def _disable_te_native_workspace_starvation(): + """Restore the TE workspace function and env pinned by the te_native backend.""" + import os + + global _TE_WORKSPACE_SIZE_FN_ORIG + if _TE_WORKSPACE_SIZE_FN_ORIG is not None: + try: + import transformer_engine.pytorch.cpp_extensions.gemm as _te_gemm_mod + + _te_gemm_mod.get_cublas_workspace_size_bytes = _TE_WORKSPACE_SIZE_FN_ORIG + if hasattr(getattr(_te_gemm_mod, "get_cublas_workspace", None), "cache_clear"): + _te_gemm_mod.get_cublas_workspace.cache_clear() + except ImportError: + pass + _TE_WORKSPACE_SIZE_FN_ORIG = None + for var, prev in _TE_NATIVE_ENV_ORIG.items(): + if prev is None: + os.environ.pop(var, None) + else: + os.environ[var] = prev + _TE_NATIVE_ENV_ORIG.clear() + + def get_batch_invariant_backend() -> str: """Return the active batch-invariant GEMM backend name.""" return _BATCH_INVARIANT_BACKEND @@ -1755,6 +1809,7 @@ def disable_batch_invariant_mode(): _batch_invariant_LIB = None # Restore Transformer Engine kernels if previously patched _te_unpatch_for_batch_invariant() + _disable_te_native_workspace_starvation() _unpin_mamba_autotuners() From ed496e849f66bd44cb5bc958eaf6aefa9aad00b2 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 15:12:34 -0700 Subject: [PATCH 16/26] Address review: scope acc_fp64 to SwiGLU; type the backend field Squared-ReLU's weighted fp32 reduction is already order-deterministic, so fp64 accumulation is only enabled for the SwiGLU (unit-weight) path. batch_invariant_backend becomes Literal[...] for argparse choices and config validation. Co-Authored-By: Claude Fable 5 Signed-off-by: Utkarsh Utkarsh --- megatron/core/inference/moe/vllm_fused_moe.py | 2 +- megatron/core/transformer/transformer_config.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index 5105cedb3f4..7d183d50c13 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/megatron/core/inference/moe/vllm_fused_moe.py @@ -736,5 +736,5 @@ def vllm_fused_moe( # for SwiGLU (training parity), so the reduction uses unit weights; # fp64 accumulation makes the topk sum order-independent by precision. apply_weights=not (batch_invariant_mode and is_swiglu), - acc_fp64=batch_invariant_mode, + acc_fp64=batch_invariant_mode and is_swiglu, ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 7660221d5d6..b4b4dcd0b9c 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1167,7 +1167,7 @@ class TransformerConfig(ModelParallelConfig): training and inference as the kernels are not full optimized. Defaults to False.""" - batch_invariant_backend: str = "te_native" + batch_invariant_backend: Literal["te_native", "deepgemm", "triton"] = "te_native" """Which batch-invariant GEMM backend to use when batch_invariant_mode is enabled: "te_native" (default: keep the native cuBLASLt kernels and obtain invariance via workspace starvation — lowest overhead, no extra From ed183b7fd3669c225c6a40a2498763119b15e73a Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 13 Aug 2026 18:10:27 -0700 Subject: [PATCH 17/26] Address review: derive weighted-swiglu grid size from SM count MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit num_programs=1184 was the Inductor launch config captured on B200 (148 SMs x 8 waves) — device-specific. Derive it as SMs * 8 with the explicit argument kept as an override, and document why grid size is bit-inert (persistent elementwise loop, disjoint per-program ranges). Add a test asserting bitwise-identical output across grid sizes (1, 148, 1184, 4096). GPU-validated: 14/14 (job 2390801). Signed-off-by: Utkarsh Utkarsh --- megatron/core/inference/moe/batch_invariant.py | 12 +++++++++++- .../test_vllm_fused_moe_batch_invariant.py | 16 ++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/moe/batch_invariant.py b/megatron/core/inference/moe/batch_invariant.py index 0d2a277c95a..e8e5fdc03ed 100644 --- a/megatron/core/inference/moe/batch_invariant.py +++ b/megatron/core/inference/moe/batch_invariant.py @@ -196,7 +196,7 @@ def weighted_silu_mul_bounded( y: torch.Tensor, weights_flat: torch.Tensor, bound_elems: torch.Tensor, - num_programs: int = 1184, + num_programs: Optional[int] = None, xblock: int = 1024, ) -> torch.Tensor: """SwiGLU with routing weights applied at the activation (training parity). @@ -204,9 +204,19 @@ def weighted_silu_mul_bounded( 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 diff --git a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py index 6c2f89f6670..89a0f7da23f 100644 --- a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py +++ b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py @@ -154,6 +154,22 @@ def test_weighted_silu_mul_bounded_bound_and_invariance(self): out2 = batch_invariant.weighted_silu_mul_bounded(y2, probs, bound) assert torch.equal(out2[:live], out[:live]) + def test_weighted_silu_mul_bounded_grid_size_is_bit_inert(self): + # num_programs is an occupancy knob only: the kernel is elementwise + # with disjoint per-program index ranges, so any grid size must give + # bitwise-identical output (incl. 1184 = the B200 Inductor capture). + from megatron.core.inference.moe import batch_invariant + + torch.manual_seed(11) + rows, ffn, live = 512, 256, 300 + y = (torch.randn(rows, 2 * ffn, device="cuda") * 2.0).bfloat16() + probs = torch.rand(rows, device="cuda", dtype=torch.float32) + bound = torch.tensor(live * ffn, dtype=torch.int64, device="cuda") + ref = batch_invariant.weighted_silu_mul_bounded(y, probs, bound) # derived default + for np_ in (1, 148, 1184, 4096): + out = batch_invariant.weighted_silu_mul_bounded(y, probs, bound, num_programs=np_) + assert torch.equal(out[:live], ref[:live]), f"bits changed at num_programs={np_}" + # --------------------------------------------------------------------------- # _moe_sum options From da0407b4c126a8fc0b27c57df7c8d8299e73e6bf Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Fri, 14 Aug 2026 12:28:05 -0700 Subject: [PATCH 18/26] Address review: validate batch_invariant_backend in __post_init__ The Literal annotation covers argparse, but programmatic TransformerConfig construction bypassed it and only failed inside enable_batch_invariant_mode() after model init. Guard at construction time against _BATCH_INVARIANT_BACKENDS (single source of truth). GPU-validated: 24/24 (job 2391673). Signed-off-by: Utkarsh Utkarsh --- megatron/core/transformer/transformer_config.py | 11 +++++++++++ .../transformer/test_transformer_config.py | 13 +++++++++++++ 2 files changed, 24 insertions(+) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index b4b4dcd0b9c..2fad8084dbd 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -3158,6 +3158,17 @@ def _scope_to_str(s): ) if self.batch_invariant_mode: + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + _BATCH_INVARIANT_BACKENDS, + ) + + # argparse validates via the Literal annotation; guard here too so + # programmatic TransformerConfig construction fails at build time + # rather than inside enable_batch_invariant_mode() after model init. + assert self.batch_invariant_backend in _BATCH_INVARIANT_BACKENDS, ( + f"Unknown batch_invariant_backend {self.batch_invariant_backend!r}; " + f"expected one of {_BATCH_INVARIANT_BACKENDS}." + ) assert self.params_dtype == torch.bfloat16, ( "Batch invariant mode supports BF16 model parameters only; " f"got {self.params_dtype}." diff --git a/tests/unit_tests/transformer/test_transformer_config.py b/tests/unit_tests/transformer/test_transformer_config.py index 24339c12b5a..43ce5cca0e4 100644 --- a/tests/unit_tests/transformer/test_transformer_config.py +++ b/tests/unit_tests/transformer/test_transformer_config.py @@ -32,6 +32,19 @@ def test_ep_a2a_overlap_rejects_unsupported_mtp_layer_counts(mtp_num_layers: int _make_overlap_config(mtp_num_layers) +def test_batch_invariant_backend_rejects_unknown_value_at_construction(): + # Programmatic construction bypasses argparse's Literal choices, so + # __post_init__ must catch typos before model init. + with pytest.raises(AssertionError, match="Unknown batch_invariant_backend"): + TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + batch_invariant_mode=True, + batch_invariant_backend="te-native", + ) + + def test_gdp_num_householder_defaults_to_three(): config = TransformerConfig(num_layers=1, hidden_size=128, num_attention_heads=4) From 371772da6ac573ad034afa1d9c324fbda42e9c46 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Fri, 14 Aug 2026 12:36:21 -0700 Subject: [PATCH 19/26] Address review: finish align-up wording; tie rounder to TOKEN_ROUNDER MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remaining 'floor' comment wording said the value could only shrink — the exact misreading behind the dropped-bucket bug. Reworded to 'align up'. The 64 literal now derives from DynamicInferenceContext.TOKEN_ROUNDER (lazy import; dynamic_context imports this module at top level) so the graph-bucket alignment and the eager path's padding multiple cannot drift apart. Behavior-identical. GPU-validated: 24/24 (job 2391676). Signed-off-by: Utkarsh Utkarsh --- .../core/inference/batch_dimensions_utils.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index d2c26fc196c..00742bd4b7d 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/megatron/core/inference/batch_dimensions_utils.py @@ -212,12 +212,18 @@ def _batch_invariant_token_align(token_count: int) -> int: 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 floor, auto-sizing injects 1- and 2-token decode buckets whose + 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 / 64) * 64 - return max(64, rounded_up) + # Lazy import (dynamic_context imports this module at its top level). + # Tying to TOKEN_ROUNDER keeps the graph-bucket alignment and the eager + # path's padding multiple from drifting apart. + from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext + + rounder = DynamicInferenceContext.TOKEN_ROUNDER + rounded_up = math.ceil(token_count / rounder) * rounder + return max(rounder, rounded_up) def _batch_invariant_mode_enabled() -> bool: @@ -372,7 +378,7 @@ def _calculate_token_counts_linear( [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: floor every bucket to a 64-multiple + # 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. @@ -466,12 +472,12 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int batch_dim = InferenceBatchDimensions(token_count, prefill_req_count, decode_req_count) if batch_dim.is_valid(max_requests, max_sequence_length, num_speculative_tokens): if _batch_invariant_mode_enabled(): - # Batch-invariant mode: floor the bucket's token count to a - # 64-multiple (see _batch_invariant_token_align). The floor - # is alignment PADDING, mirroring the eager path's + # 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. Flooring can + # unpadded dims; request counts are untouched. Aligning can # collide previously-distinct buckets, so skip duplicates. batch_dim = InferenceBatchDimensions( _batch_invariant_token_align(token_count), From e69fe65b3dd751f0c39a1be544cec5bf7ff757f0 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Mon, 17 Aug 2026 11:30:43 -0700 Subject: [PATCH 20/26] Address review: scope skip_rmsnorm to normalization patches only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The early return also skipped _te_patch_general_grouped_gemm (a GEMM concern); moved it under skip_gemm where it belongs — under te_native grouped GEMMs stay native, covered by the same workspace starvation. Documented skip_rmsnorm in the docstring and rewrote the wrong call-site comment (te_native substitutes no TE kernels; the attention gate is a standalone assert, not part of this patch function). Both existing call sites behave identically. GPU-validated: 24/24 (job 2395658). Signed-off-by: Utkarsh Utkarsh --- .../custom_layers/batch_invariant_kernels.py | 32 +++++++++++++------ 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py index f0061e505f9..9ab0fae8058 100644 --- a/megatron/core/transformer/custom_layers/batch_invariant_kernels.py +++ b/megatron/core/transformer/custom_layers/batch_invariant_kernels.py @@ -633,9 +633,15 @@ def _te_patch_for_batch_invariant(skip_gemm: bool = False, skip_rmsnorm: bool = Safe no-op if TE is unavailable. Args: - skip_gemm: leave TE's native general_gemm in place (used by the - "te_native" backend, where GEMM batch invariance comes from - cuBLASLt workspace starvation rather than kernel substitution). + skip_gemm: leave TE's native GEMMs (general_gemm and grouped GEMM) in + place (used by the "te_native" backend, where GEMM batch + invariance comes from cuBLASLt workspace starvation rather than + kernel substitution). + skip_rmsnorm: leave TE's native normalization in place (RMSNorm + class/module functions and the fused-module apply_normalization + entry). Used by "te_native", where the norm's M%32 reduction + bit-class is held constant by the 64-multiple alignment + discipline instead of kernel substitution. """ global _TE_GENERAL_GEMM_ORIG, _TE_RMSNORM_ORIG_FWD, _MEG_TE_GENERAL_GEMM_ORIG import transformer_engine.pytorch as te @@ -673,7 +679,15 @@ def _te_patch_for_batch_invariant(skip_gemm: bool = False, skip_rmsnorm: bool = _MEG_TE_GENERAL_GEMM_ORIG = meg_te.general_gemm meg_te.general_gemm = _te_general_gemm_patched + # Patch TE.general_grouped_gemm at every known import site so that + # TEGroupedMLP (forward + dgrad + wgrad) goes through DeepGEMM in bf16. + # Grouped GEMMs are a GEMM concern: under skip_gemm ("te_native") they + # stay native, covered by the same cuBLASLt workspace starvation. + _te_patch_general_grouped_gemm() + if skip_rmsnorm: + # Everything below patches normalization only (RMSNorm class/module + # functions and the fused-module apply_normalization entry). return # Patch RMSNorm.forward once (class may be on te or te.pytorch) @@ -717,10 +731,6 @@ def _patched(*args, **kwargs): _TE_RMSNORM_FUNC_ORIGS[name] = orig setattr(te_layernorm_mod, name, _make_rmsnorm_patched(orig)) - # Patch TE.general_grouped_gemm at every known import site so that - # TEGroupedMLP (forward + dgrad + wgrad) goes through DeepGEMM in bf16. - _te_patch_general_grouped_gemm() - # Patch the fused-module normalization entry (`apply_normalization`). TE's # fused LayerNormLinear / LayerNormMLP call this instead of RMSNorm.forward, # so without this patch their internal RMSNorm runs TE's tex kernel, whose @@ -1783,9 +1793,11 @@ def enable_batch_invariant_mode(backend: str = "te_native"): _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, dispatch_key) _batch_invariant_LIB.impl("aten::_log_softmax", _log_softmax_batch_invariant, dispatch_key) _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, dispatch_key) - # Also patch Transformer Engine kernels when available (te_native keeps - # TE's native GEMMs — invariance comes from the starved workspace — but - # still applies the non-GEMM TE patches, e.g. the attention gate). + # Also patch Transformer Engine kernels when available. Under te_native + # BOTH skips are set, so no TE kernel is substituted at all: GEMMs (dense + # and grouped) stay native under the starved workspace, and norms stay + # native under the 64-multiple alignment discipline. (The TE attention + # version gate is a separate standalone assert, not part of this patch.) if backend == "te_native": # te_native also keeps TE's NATIVE RMSNorm: its M%32 reduction # bit-class is held constant by the 64-multiple alignment discipline From 2cdc4b6863dd0005a1d912c4c4e4b4d872a96dc0 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Mon, 17 Aug 2026 11:30:43 -0700 Subject: [PATCH 21/26] Address review: skip BI kernel tests without CUDA/Triton Module-level pytestmark so CPU-only or Triton-less environments (e.g. the lts lane) report a legible skip instead of hard-failing at cuda allocation. Reuses the library's own HAVE_TRITON flag. GPU-validated: 24/24 (job 2395658). Signed-off-by: Utkarsh Utkarsh --- .../inference/test_vllm_fused_moe_batch_invariant.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py index 89a0f7da23f..fabf03a5c3d 100644 --- a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py +++ b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py @@ -23,6 +23,12 @@ _BATCH_INVARIANT_BACKENDS, set_batch_invariant_mode, ) +from megatron.core.inference.moe.batch_invariant import HAVE_TRITON + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not HAVE_TRITON, + reason="batch-invariant MoE kernels require CUDA and Triton", +) def _vt(n): From bf46b6e74f1bfd6305e0d032f7a98b40d2cb6835 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Mon, 17 Aug 2026 11:30:43 -0700 Subject: [PATCH 22/26] Address review: drop getattr default for batch_invariant_backend The field is auto-registered from TransformerConfig so argparse always populates it; a silent fallback would mask a real wiring bug (and its hardcoded default had already drifted once). GPU-validated: 24/24 (job 2395658). Signed-off-by: Utkarsh Utkarsh --- megatron/training/initialize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 2ddcaa3c5a5..c0b46db1d4d 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -100,7 +100,7 @@ def state_restore_func(state_dict): ) if args.batch_invariant_mode: - backend = getattr(args, "batch_invariant_backend", "te_native") + backend = args.batch_invariant_backend print_rank_0(f"Enabling batch invariant mode globally (backend={backend})") enable_batch_invariant_mode(backend) From e3aa4385b622d93d450b73257a85b7d7de84744d Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Mon, 17 Aug 2026 12:07:29 -0700 Subject: [PATCH 23/26] =?UTF-8?q?Address=20review:=20=5Fmoe=5Fsum=20docstr?= =?UTF-8?q?ing=20=E2=80=94=20fp32=20or=20fp64=20accumulation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 'Accumulates in fp32' was stale after acc_fp64 was added; it is the sentence a caller reads to judge reduction precision. Docstring-only. Signed-off-by: Utkarsh Utkarsh --- megatron/core/inference/moe/vllm_fused_moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index 7d183d50c13..a14a40261f3 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/megatron/core/inference/moe/vllm_fused_moe.py @@ -524,7 +524,8 @@ def _moe_sum( """Fused topk reduction: [max_tokens*topk, K] bf16 → [max_tokens, K]. Applies routing weights and reduces over topk in a single kernel. - Accumulates in fp32. When `out` is None, allocates and returns an fp32 + Accumulates in fp32, or fp64 when `acc_fp64` is set (the output buffer + stays fp32 either way). When `out` is None, allocates and returns an fp32 buffer. When `out` is provided (e.g. the RSV symmetric memory tensor), writes directly into it — tl.store handles the cast to the buffer's dtype. Only writes the first valid_tokens rows; rows beyond are left untouched From 90517517a6e5b4e696e41ff8ff8ce791ad57b718 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Wed, 19 Aug 2026 18:51:30 -0700 Subject: [PATCH 24/26] Address review: single module-level TOKEN_ROUNDER constant Factor the batch-invariant alignment multiple into a module-level TOKEN_ROUNDER in batch_dimensions_utils, used by all three batch-invariant sites (align function, token ladder, smallest bucket). DynamicInferenceContext.TOKEN_ROUNDER now references the same constant (that import direction already exists), so the eager rounder and the batch-invariant alignment share one definition and the earlier lazy import is removed. Behavior-identical. GPU-validated: 24/24 (job 2403997). Signed-off-by: Utkarsh Utkarsh --- .../core/inference/batch_dimensions_utils.py | 26 ++++++++++--------- .../inference/contexts/dynamic_context.py | 3 ++- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index 00742bd4b7d..71572a86037 100644 --- a/megatron/core/inference/batch_dimensions_utils.py +++ b/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: @@ -216,14 +224,8 @@ def _batch_invariant_token_align(token_count: int) -> int: graphed norms execute in a different bit-class, breaking cross-batch bit-equality. Request counts are untouched (mirrors eager semantics). """ - # Lazy import (dynamic_context imports this module at its top level). - # Tying to TOKEN_ROUNDER keeps the graph-bucket alignment and the eager - # path's padding multiple from drifting apart. - from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext - - rounder = DynamicInferenceContext.TOKEN_ROUNDER - rounded_up = math.ceil(token_count / rounder) * rounder - return max(rounder, rounded_up) + rounded_up = math.ceil(token_count / TOKEN_ROUNDER) * TOKEN_ROUNDER + return max(TOKEN_ROUNDER, rounded_up) def _batch_invariant_mode_enabled() -> bool: @@ -309,9 +311,9 @@ def _calculate_cuda_graph_token_counts( rounder = CUDAGraphBatchDimensionBuilder.CUDA_GRAPH_ROUNDER if _batch_invariant_mode_enabled(): - # Batch-invariant mode: 64-multiple token ladder (see + # Batch-invariant mode: TOKEN_ROUNDER-multiple token ladder (see # _batch_invariant_token_align). - rounder = 64 + rounder = TOKEN_ROUNDER cuda_graph_max_tokens = _batch_invariant_token_align(cuda_graph_max_tokens) # Cuda graph step size. @@ -346,8 +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) - # Batch-invariant mode: smallest bucket is 64, never tp_size. - sizes.add(64 if _batch_invariant_mode_enabled() else 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) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 4de2c64f06f..ab8c185b1e5 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/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" From 8a256893b9f86ff44a6dd7e7bea792ac2e3924a8 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Thu, 20 Aug 2026 09:36:53 -0700 Subject: [PATCH 25/26] Apply tools/autoformat.sh (isort import order in BI test) Signed-off-by: Utkarsh Utkarsh --- .../unit_tests/inference/test_vllm_fused_moe_batch_invariant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py index fabf03a5c3d..0836e2e61fa 100644 --- a/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py +++ b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py @@ -19,11 +19,11 @@ import torch from megatron.core.inference.batch_dimensions_utils import CUDAGraphBatchDimensionBuilder +from megatron.core.inference.moe.batch_invariant import HAVE_TRITON from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( _BATCH_INVARIANT_BACKENDS, set_batch_invariant_mode, ) -from megatron.core.inference.moe.batch_invariant import HAVE_TRITON pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or not HAVE_TRITON, From a5aa08867a6a5cf374fe1ed5d4ca829a943e3285 Mon Sep 17 00:00:00 2001 From: Utkarsh Utkarsh Date: Fri, 21 Aug 2026 10:51:02 -0700 Subject: [PATCH 26/26] Add batch_invariant_backend to Mamba MoE golden config The config-drift guard in test_hybrid_moe_model.py snapshots every TransformerConfig field; register the new batch_invariant_backend field with its default value (te_native), following the test's ADDED ARGS guidance. The field only takes effect when batch_invariant_mode is enabled, so downstream model configs are unaffected. Signed-off-by: Utkarsh Utkarsh --- tests/unit_tests/models/test_hybrid_moe_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 037dbae3ee4..dabe6b57103 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -49,6 +49,7 @@ "attention_softmax_in_fp32": False, "autocast_dtype": "torch.bfloat16", "barrier_with_L1_time": True, + "batch_invariant_backend": "te_native", "batch_invariant_mode": False, "batch_p2p_comm": True, "batch_p2p_sync": True,