`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.

Matches the training fused weighted-squared-relu rounding, clamped (CLAMP) and unclamped alike: the square and the routing probability are applied in FP32 with a single BF16 round at the end. Rounding the square first, as the BF16 input dtype invites, costs a routed element its last bit.

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`#