core.transformer.moe.batch_invariant

Batch-invariant MoE permutation helpers.

Module Contents

Classes

_RankTreeUnpermute

Fused unpermute for probs=None: same arithmetic tree as the loop, one launch.

Functions

build_inverse_permutation_map	Build token/top-k -> permuted-row and expert-id map for batch-invariant unpermute.
_rank_tree_unpermute_kernel	One program per (token, hidden block): the rank-ordered combine tree of unpermute.
_rank_tree_unpermute_bwd_kernel	One program per (token, hidden block): scatter the token’s gradient to its rows.
unpermute_with_bf16_rank_partials	Forward-only unpermute with the FlashInfer routed kernels’ per-rank finalize.
_unpermute_loop	Reference form of unpermute: a loop over EP ranks and top-k slots.
unpermute	Batch-invariant MoE unpermute.

API

core.transformer.moe.batch_invariant.build_inverse_permutation_map(

routing_map: torch.Tensor,

flat_sorted: torch.Tensor,

sorted_indices: torch.Tensor,

num_out_tokens: int,

) → torch.Tensor

Build token/top-k -> permuted-row and expert-id map for batch-invariant unpermute.

The regular permutation map is row -> token. Batch-invariant unpermute needs the inverse ownership model so each output token can read its routed rows and add them in a fixed order.

core.transformer.moe.batch_invariant._rank_tree_unpermute_kernel(

rows_ptr,

inverse_rows_ptr,

inverse_experts_ptr,

probs_ptr,

out_ptr,

hidden,

EP_SIZE: triton.language.constexpr,

EXPERTS_PER_RANK: triton.language.constexpr,

TOPK: triton.language.constexpr,

CROSS_RANK_FP64: triton.language.constexpr,

BF16_RANK_PARTIALS: triton.language.constexpr,

BLOCK_H: triton.language.constexpr,

)

One program per (token, hidden block): the rank-ordered combine tree of unpermute.

With BF16_RANK_PARTIALS, each rank’s partial is the FlashInfer routed kernels’ finalize instead: FP32 fused multiply-adds of probability and row in slot order, from zero, quantized to BF16.

core.transformer.moe.batch_invariant._rank_tree_unpermute_bwd_kernel(

grad_out_ptr,

grad_out_stride_token,

grad_out_stride_hidden,

inverse_rows_ptr,

valid_ptr,

grad_rows_ptr,

hidden,

TOPK: triton.language.constexpr,

BLOCK_H: triton.language.constexpr,

)

One program per (token, hidden block): scatter the token’s gradient to its rows.

Each permuted row belongs to one (token, slot), so the stores never collide. The grid depends only on tensor shapes, which keeps the backward CUDA-graph capturable.

class core.transformer.moe.batch_invariant._RankTreeUnpermute

Bases: torch.autograd.Function

Fused unpermute for probs=None: same arithmetic tree as the loop, one launch.

static forward(

ctx,

permuted_tokens,

inverse_rows,

inverse_experts,

ep_size,

num_experts,

)

static backward(ctx, grad_output)

core.transformer.moe.batch_invariant.unpermute_with_bf16_rank_partials(

permuted_tokens: torch.Tensor,

route_rows: torch.Tensor,

route_experts: torch.Tensor,

probs: torch.Tensor,

num_experts: int,

ep_size: int,

) → torch.Tensor

Forward-only unpermute with the FlashInfer routed kernels’ per-rank finalize.

Each EP rank’s partial adds the token’s routes on that rank in the order of route_rows, as FP32 fused multiply-adds of probs and row, from zero, and quantizes the sum to BF16. The ranks are then added as unpermute adds them. route_rows, route_experts and probs (FP32) are [tokens, topk].

core.transformer.moe.batch_invariant._unpermute_loop(

permuted_tokens: torch.Tensor,

restore_shape: torch.Size,

probs: Optional[torch.Tensor],

num_experts: int,

inverse_map: torch.Tensor,

ep_size: int,

) → torch.Tensor

Reference form of unpermute: a loop over EP ranks and top-k slots.

core.transformer.moe.batch_invariant.unpermute(

permuted_tokens: torch.Tensor,

restore_shape: torch.Size,

*,

probs: Optional[torch.Tensor],

num_experts: int,

inverse_map: torch.Tensor,

) → torch.Tensor

Batch-invariant MoE unpermute.

Accumulation is token-owned. The AllToAll inverse map avoids data-dependent shapes and adds contributions by EP rank then top-k slot, matching the inference NVLS rank-ordered combine. Without probs the tree runs as one Triton kernel; with probs it runs as a loop over EP ranks and top-k slots.