diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py index 1229d333d0a..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: @@ -175,7 +183,7 @@ def adjust_batch_dims_for_expert_parallelism( if ep_zmq_communicator is not None: # CPU-only sync via ZMQ: avoids a NCCL AllReduce kernel on the # compute stream plus the H2D/D2H pair that sandwiches it. - (max_token_count, max_is_non_decode) = ep_zmq_communicator.sync_all_reduce_max( + max_token_count, max_is_non_decode = ep_zmq_communicator.sync_all_reduce_max( local_batch_dims.token_count, int(is_non_decode) ) else: @@ -205,6 +213,30 @@ def adjust_batch_dims_for_expert_parallelism( return adjusted_batch_dim +def _batch_invariant_token_align(token_count: int) -> int: + """Round a CUDA-graph bucket token count UP to a 64-multiple (min 64). + + Under batch-invariant mode every graphed step must execute norms/GEMMs in + the same M-alignment class as eager steps: TE rmsnorm (and other + M-sensitive kernels) switch reduction codepaths at M % 32, and the eager + path already pads token counts to TOKEN_ROUNDER (64) multiples. Without + this alignment, auto-sizing injects 1- and 2-token decode buckets whose + graphed norms execute in a different bit-class, breaking cross-batch + bit-equality. Request counts are untouched (mirrors eager semantics). + """ + rounded_up = math.ceil(token_count / TOKEN_ROUNDER) * TOKEN_ROUNDER + return max(TOKEN_ROUNDER, rounded_up) + + +def _batch_invariant_mode_enabled() -> bool: + # Lazy import to avoid a circular dependency at module import time. + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + is_batch_invariant_mode_enabled, + ) + + return is_batch_invariant_mode_enabled() + + class CUDAGraphBatchDimensionBuilder: """Builder for creating and managing CUDA graph batch dimensions. @@ -278,6 +310,11 @@ def _calculate_cuda_graph_token_counts( ), f"cuda_graph_max_tokens must be > 0, got {cuda_graph_max_tokens}" rounder = CUDAGraphBatchDimensionBuilder.CUDA_GRAPH_ROUNDER + if _batch_invariant_mode_enabled(): + # Batch-invariant mode: TOKEN_ROUNDER-multiple token ladder (see + # _batch_invariant_token_align). + rounder = TOKEN_ROUNDER + cuda_graph_max_tokens = _batch_invariant_token_align(cuda_graph_max_tokens) # Cuda graph step size. cuda_graph_step_size = cuda_graph_max_tokens / num_cuda_graphs @@ -311,7 +348,8 @@ def _calculate_cuda_graph_token_counts( # Always include the endpoints: cuda_graph_max_tokens (largest) and tp_size (smallest). sizes.add(cuda_graph_max_tokens) - sizes.add(tp_size) + # Batch-invariant mode: smallest bucket is TOKEN_ROUNDER, never tp_size. + sizes.add(TOKEN_ROUNDER if _batch_invariant_mode_enabled() else tp_size) cuda_graph_token_counts = sorted(sizes, reverse=True) @@ -341,6 +379,10 @@ def _calculate_token_counts_linear( sizes = ( [1, 2, 4] + list(range(8, 256, 8)) + list(range(256, cuda_graph_max_tokens + 1, 16)) ) + if _batch_invariant_mode_enabled(): + # Batch-invariant mode: align every bucket up to a 64-multiple + # (see _batch_invariant_token_align) and dedupe collisions. + sizes = [_batch_invariant_token_align(s) for s in sizes] # TP-align and dedupe in order; preserve original ordering for parity. sizes = list(dict.fromkeys(round_up_to_nearest_multiple(s, tp_size) for s in sizes)) sizes = [s for s in sizes if s <= cuda_graph_max_tokens] @@ -431,7 +473,21 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int """Helper to create and append batch dimension to list only if it's valid.""" batch_dim = InferenceBatchDimensions(token_count, prefill_req_count, decode_req_count) if batch_dim.is_valid(max_requests, max_sequence_length, num_speculative_tokens): - cuda_graph_batch_dimensions_list.append(batch_dim) + if _batch_invariant_mode_enabled(): + # Batch-invariant mode: align the bucket's token count up + # to a 64-multiple (see _batch_invariant_token_align). The + # alignment is PADDING, mirroring the eager path's + # TOKEN_ROUNDER (which already yields token counts above + # what the requests produce), so validity is judged on the + # unpadded dims; request counts are untouched. Aligning can + # collide previously-distinct buckets, so skip duplicates. + batch_dim = InferenceBatchDimensions( + _batch_invariant_token_align(token_count), + prefill_req_count, + decode_req_count, + ) + if batch_dim not in cuda_graph_batch_dimensions_list: + cuda_graph_batch_dimensions_list.append(batch_dim) # Cuda graph token-counts # (i.e., token counts used by cuda-graph steps, both decode and non-decode). 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" diff --git a/megatron/core/inference/moe/batch_invariant.py b/megatron/core/inference/moe/batch_invariant.py index 06dfddc2869..e8e5fdc03ed 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: @@ -86,6 +87,143 @@ def _squared_relu_with_probs_kernel( tl.store(output_ptr + row * hidden_size + cols, value, mask=mask) +@triton.jit +def _swiglu_with_probs_kernel( + input_ptr, + output_ptr, + permutation_map_ptr, + n_used_ptr, + probs_ptr, + ffn_size, # output width; input row width is 2*ffn_size (gate | up) + max_rows, + BLOCK_SIZE: tl.constexpr, + NUM_BLOCKS: tl.constexpr, +): + """Apply gated SiLU (SwiGLU) and router probabilities in training order. + + Matches the training fused weighted-swiglu rounding: SiLU(gate)*up*prob is + computed in FP32 with a single BF16 round at the end. Input row width is + 2*ffn_size: gate = first half, up = second half (megatron chunk + convention). Fixed NUM_BLOCKS CTAs iterating rows -> CUDA-graph safe. + """ + pid = tl.program_id(0) + n_used = tl.load(n_used_ptr) + if pid >= n_used: + return + two_n = 2 * ffn_size + + for row in tl.range(pid, max_rows, NUM_BLOCKS): + if row < n_used: + if tl.load(permutation_map_ptr + row) >= 0: + prob = tl.load(probs_ptr + row) + for offset in tl.range(0, ffn_size, BLOCK_SIZE): + cols = offset + tl.arange(0, BLOCK_SIZE) + mask = cols < ffn_size + gate = tl.load(input_ptr + row * two_n + cols, mask=mask).to(tl.float32) + up = tl.load(input_ptr + row * two_n + ffn_size + cols, mask=mask).to( + tl.float32 + ) + value = gate * tl.sigmoid(gate) * up * prob + tl.store(output_ptr + row * ffn_size + cols, value.to(tl.bfloat16), mask=mask) + + +def swiglu_with_probs( + x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor, probs: torch.Tensor +) -> torch.Tensor: + """Gated-SiLU counterpart of squared_relu_with_probs (SwiGLU models).""" + num_rows, two_ffn = x.shape + ffn_size = two_ffn // 2 + out = torch.empty(num_rows, ffn_size, dtype=x.dtype, device=x.device) + block_size = min(triton.next_power_of_2(ffn_size), 1024) + num_blocks = min(num_rows, 512) + _swiglu_with_probs_kernel[(num_blocks,)]( + x, + out, + permutation_map, + n_used, + probs, + ffn_size, + num_rows, + BLOCK_SIZE=block_size, + NUM_BLOCKS=num_blocks, + ) + return out + + +@triton.jit +def _weighted_silu_mul_bounded_kernel( + in_ptr0, in_ptr1, out_ptr0, bound_ptr, xnumel, HALF_N: tl.constexpr, XBLOCK: tl.constexpr +): + """Device-bounded weighted SwiGLU with training-parity rounding. + + The per-element instruction sequence is copied VERBATIM from Inductor's + emitted Triton for the training fused weighted-swiglu + (bf16 -> fp32 silu(gate) * up * prob -> bf16, single final rounding), so a + token's activation bits match the training forward exactly. Elementwise + kernels have no cross-element reduction, so only the per-element sequence + determines bits; the schedule below is a persistent 1D grid (static launch, + CUDA-graph-safe) striding while xoffset < a DEVICE element bound + (= valid_tokens * topk * HALF_N — the live prefix of the flat token-major + layout). Rows beyond the bound are neither read nor written. + """ + xbound = tl.load(bound_ptr) + num_progs = tl.num_programs(0) + xoffset = tl.program_id(0) * XBLOCK + while xoffset < xbound: + xindex = xoffset + tl.arange(0, XBLOCK)[:] + xmask = (xindex < xbound) & (xindex < xnumel) + x0 = xindex % HALF_N + x1 = xindex // HALF_N + tmp0 = tl.load(in_ptr0 + (x0 + 2 * HALF_N * x1), xmask).to(tl.float32) + tmp8 = tl.load(in_ptr0 + (HALF_N + x0 + 2 * HALF_N * x1), xmask).to(tl.float32) + tmp11 = tl.load(in_ptr1 + (x1), xmask, eviction_policy='evict_last') + tmp1 = tmp0.to(tl.float32) + tmp2 = -tmp1 + tmp3 = libdevice.exp(tmp2) + tmp4 = tl.full([1], 1.0, tl.float32) + tmp5 = tmp3 + tmp4 + tmp6 = tmp1 / tmp5 + tmp7 = tmp6.to(tl.float32) + tmp9 = tmp7 * tmp8 + tmp10 = tmp9.to(tl.float32) + tmp12 = tmp10 * tmp11 + tmp13 = tmp12.to(tl.float32) + tl.store(out_ptr0 + xindex, tmp13, xmask) + xoffset += num_progs * XBLOCK + + +def weighted_silu_mul_bounded( + y: torch.Tensor, + weights_flat: torch.Tensor, + bound_elems: torch.Tensor, + num_programs: Optional[int] = None, + xblock: int = 1024, +) -> torch.Tensor: + """SwiGLU with routing weights applied at the activation (training parity). + + y: [rows, 2*half_n] bf16 (gate | up); weights_flat: [rows] fp32 routing + probabilities; bound_elems: device scalar = live_rows * half_n. + Returns [rows, half_n] bf16; rows beyond the live bound are untouched. + + num_programs defaults to SMs * 8 waves (Inductor's persistent-grid sizing; + 1184 on the B200 this was captured from). Grid size cannot affect bits: + the kernel is elementwise with each program owning a disjoint strided + index range, so it is an occupancy knob only. + """ + rows, two_half_n = y.shape + half_n = two_half_n // 2 + if num_programs is None: + # Lazy import: permute.py imports this module at its top level. + from megatron.core.inference.moe.permute import _get_num_sms + + num_programs = _get_num_sms(y.device) * 8 + out = torch.empty(rows, half_n, dtype=y.dtype, device=y.device) + _weighted_silu_mul_bounded_kernel[(num_programs,)]( + y, weights_flat, out, bound_elems, rows * half_n, HALF_N=half_n, XBLOCK=xblock + ) + return out + + def squared_relu_with_probs( x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor, probs: torch.Tensor ) -> torch.Tensor: 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. diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py index f6087ebbfe9..a14a40261f3 100644 --- a/megatron/core/inference/moe/vllm_fused_moe.py +++ b/megatron/core/inference/moe/vllm_fused_moe.py @@ -29,6 +29,7 @@ triton.jit = null_decorator tl = MagicMock() +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 ( @@ -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,25 @@ 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 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.float32) * w + else: + 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, mask=k_mask) + tl.store(output_ptr + token_id_i64 * K + offs_k, acc.to(tl.float32), mask=k_mask) def _moe_sum( @@ -506,17 +518,27 @@ 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]. 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 (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) @@ -536,6 +558,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 +620,18 @@ 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). + 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: 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 @@ -645,10 +680,21 @@ 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 +733,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 and is_swiglu, ) diff --git a/megatron/core/tensor_parallel/inference_layers.py b/megatron/core/tensor_parallel/inference_layers.py index 102cdbcb4b9..2da50cc66bf 100644 --- a/megatron/core/tensor_parallel/inference_layers.py +++ b/megatron/core/tensor_parallel/inference_layers.py @@ -26,6 +26,7 @@ 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, ) @@ -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 5b8c6e678e6..9ab0fae8058 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,70 @@ 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, 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 batch-invariant implementations when batch-invariant mode is enabled. Safe no-op if TE is unavailable. + + Args: + 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 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 - - # 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 - - # 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 - - # 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 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 + + 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 + + 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 + + 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 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) rms_cls = getattr(te, "RMSNorm", None) @@ -708,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 @@ -1636,15 +1655,114 @@ def is_batch_invariant_mode_enabled(): return _batch_invariant_MODE -def enable_batch_invariant_mode(backend: str = "deepgemm"): +_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. + + 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 logging + import os + + 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 _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() + except ImportError: + 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 + + +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). 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: @@ -1663,12 +1781,31 @@ 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. 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 + # (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 # up on different tile configs (and therefore different fp32 reduction # orders) through autotune timing noise. @@ -1684,6 +1821,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() diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e783a017056..2fad8084dbd 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1167,6 +1167,14 @@ class TransformerConfig(ModelParallelConfig): training and inference as the kernels are not full optimized. Defaults to False.""" + 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 + 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""" @@ -1630,9 +1638,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 @@ -3146,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}." @@ -3180,9 +3203,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 61374a4ae08..c0b46db1d4d 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 = args.batch_invariant_backend + print_rank_0(f"Enabling batch invariant mode globally (backend={backend})") + enable_batch_invariant_mode(backend) # torch.distributed initialization def finish_mpu_init(): 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..0836e2e61fa --- /dev/null +++ b/tests/unit_tests/inference/test_vllm_fused_moe_batch_invariant.py @@ -0,0 +1,345 @@ +# 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.inference.moe.batch_invariant import HAVE_TRITON +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + _BATCH_INVARIANT_BACKENDS, + set_batch_invariant_mode, +) + +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): + 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_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. + 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]) + + 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 +# --------------------------------------------------------------------------- + + +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_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 + + 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: + 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) + torch.mm(a, b) # should not raise / not require DeepGEMM + 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 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, 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)