core.transformer.moe.batch_invariant

Batch-invariant MoE permutation helpers.

Module Contents

Functions

build_inverse_permutation_map	Build token/top-k -> permuted-row and expert-id map for batch-invariant unpermute.
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.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.