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

BLOCK_SIZE: triton.language.constexpr,

NUM_BLOCKS: triton.language.constexpr,

)

Apply squared ReLU and router probabilities in training order.

core.inference.moe.batch_invariant.squared_relu_with_probs(

x: torch.Tensor,

permutation_map: torch.Tensor,

n_used: torch.Tensor,

probs: torch.Tensor,

) → torch.Tensor

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

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.