`core.inference.moe.batch_invariant`#

Batch-invariant inference MoE helpers.

Module Contents#

Functions#

`enabled`	Return whether global batch-invariant mode is active.
`grouped_mm`	Batch-invariant BF16 grouped GEMM used by inference fused MoE.
`grouped_mm_alignment`	Per-expert row alignment required by the batch-invariant grouped GEMM.
`_squared_relu_with_probs_kernel`	Apply squared ReLU and router probabilities in training order.
`_swiglu_with_probs_kernel`	Apply gated SiLU (SwiGLU) and router probabilities in training order.
`swiglu_with_probs`	Gated-SiLU counterpart of squared_relu_with_probs (SwiGLU models).
`_weighted_silu_mul_bounded_kernel`	Device-bounded weighted SwiGLU with training-parity rounding.
`weighted_silu_mul_bounded`	SwiGLU with routing weights applied at the activation (training parity).
`squared_relu_with_probs`	Match training’s BF16 squared-ReLU rounding before the FP32 probability multiply.
`_ordered_reduce_scatter_v_kernel`	Reduce peer rows with an explicit rank-order FP32 sum.
`ordered_reduce_scatter_v`	Reduce-scatter variable token rows with a fixed FP32 rank order.
`_unpermute_tokens_in_expert_order_kernel`	Token-local batch-invariant unpermute.
`unpermute_tokens_in_expert_order`	Reduce local expert contributions token-by-token in fixed expert order.

API#

core.inference.moe.batch_invariant.enabled() → bool#: Return whether global batch-invariant mode is active.

core.inference.moe.batch_invariant.grouped_mm( x_bf16: torch.Tensor, weight: torch.Tensor, offs: torch.Tensor, ) → torch.Tensor#: Batch-invariant BF16 grouped GEMM used by inference fused MoE.

core.inference.moe.batch_invariant.grouped_mm_alignment() → int#: Per-expert row alignment required by the batch-invariant grouped GEMM.

core.inference.moe.batch_invariant._squared_relu_with_probs_kernel( input_ptr, output_ptr, permutation_map_ptr, n_used_ptr, probs_ptr, hidden_size, max_rows, clamp_scale, CLAMP: triton.language.constexpr, ZERO_PADDING: triton.language.constexpr, BLOCK_SIZE: triton.language.constexpr, NUM_BLOCKS: triton.language.constexpr, )#

Apply squared ReLU and router probabilities in training order.

With CLAMP set this reproduces training’s fused weighted_clamped_squared_relu bit for bit: the soft-clamped pre-activation and the square both stay in FP32, and the only BF16 round is the final one after the FP32 routing probability is applied.

Without CLAMP the square is materialized in BF16 first, matching the unclamped weighted_squared_relu, which squares a BF16 ReLU output.

core.inference.moe.batch_invariant._swiglu_with_probs_kernel( input_ptr, output_ptr, permutation_map_ptr, n_used_ptr, probs_ptr, ffn_size, max_rows, ZERO_PADDING: triton.language.constexpr, BLOCK_SIZE: triton.language.constexpr, NUM_BLOCKS: triton.language.constexpr, )#

Apply gated SiLU (SwiGLU) and router probabilities in training order.

Matches the training fused weighted-swiglu rounding: SiLU(gate)upprob 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.

core.inference.moe.batch_invariant.swiglu_with_probs( x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor, probs: torch.Tensor, zero_padding: bool = False, ) → torch.Tensor#

Gated-SiLU counterpart of squared_relu_with_probs (SwiGLU models).

zero_padding initializes aligned dummy rows because MXFP8 quantization and grouped GEMM have no permutation map with which to skip them.

core.inference.moe.batch_invariant._weighted_silu_mul_bounded_kernel( in_ptr0, in_ptr1, out_ptr0, bound_ptr, xnumel, HALF_N: triton.language.constexpr, XBLOCK: triton.language.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.

core.inference.moe.batch_invariant.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.

core.inference.moe.batch_invariant.squared_relu_with_probs( x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor, probs: torch.Tensor, clamp_scale: Optional[float] = None, zero_padding: bool = False, ) → torch.Tensor#

Match training’s BF16 squared-ReLU rounding before the FP32 probability multiply.

Parameters:

clamp_scale – config.activation_func_tanh_clamp_scale. If set, precondition the input with the tanh soft clamp s * tanh(x / s).
zero_padding – initialize aligned dummy rows because MXFP8 quantization and grouped GEMM have no permutation map with which to skip them.

core.inference.moe.batch_invariant._ordered_reduce_scatter_v_kernel( local_ptr, buffer_ptrs_dev, signal_pad_ptrs, local_tokens, rank_token_offset_ptr, ep_max_tokens_ptr, input_byte_offset, HIDDEN_SIZE: triton.language.constexpr, BLOCK_SIZE: triton.language.constexpr, RANK: triton.language.constexpr, WORLD_SIZE: triton.language.constexpr, )#: Reduce peer rows with an explicit rank-order FP32 sum.

core.inference.moe.batch_invariant.ordered_reduce_scatter_v(

output_tensor: torch.Tensor,

input_tensor: torch.Tensor,

symm_mem_hdl: torch._C._distributed_c10d._SymmetricMemory,

rank_token_offset: torch.Tensor,

ep_max_tokens: torch.Tensor,

per_rank_max_tokens: int,

input_byte_offset: int = 0,

**kwargs,

) → torch.Tensor#: Reduce-scatter variable token rows with a fixed FP32 rank order.

core.inference.moe.batch_invariant._unpermute_tokens_in_expert_order_kernel( expert_out_ptr, inverse_map_ptr, valid_tokens_ptr, output_ptr, hidden_dim, num_local_experts: triton.language.constexpr, BLOCK_H: triton.language.constexpr, )#

Token-local batch-invariant unpermute.

Each program owns one output token and one hidden tile. Contributions are accumulated in fp32 by increasing local expert id, avoiding atomic-add order.

core.inference.moe.batch_invariant.unpermute_tokens_in_expert_order( expert_output: torch.Tensor, inverse_map: torch.Tensor, valid_tokens: torch.Tensor, out: Optional[torch.Tensor], ) → torch.Tensor#: Reduce local expert contributions token-by-token in fixed expert order.

core.inference.moe.batch_invariant#

Module Contents#

Functions#

API#

`core.inference.moe.batch_invariant`#