dockyard_rl.models.dtensor.moe.dispatch

Token dispatch / combine for MoE expert routing.

Reorders tokens by expert assignment so each expert’s tokens are contiguous for the grouped matmul, then scatter-adds the expert outputs back to token positions. Distilled from torchtitan’s LocalTokenDispatcher (pytorch/torchtitan, BSD-3-Clause; Copyright (c) Meta Platforms, Inc. and affiliates) into a plain dockyard module — no torchtitan Configurable / ops framework.

LocalTokenDispatcher is the EP=1 (single expert-parallel rank) path: pure local reordering, no all-to-all. AllToAllTokenDispatcher (EP>1) adds the cross-rank token shuffle (the DeepEP-equivalent all_to_all_single): the all-to-all collectives are device-bound (NCCL, HV-10), but the reorder / permute index math and the EP-disabled fallback are pure tensor ops and are unit-tested on CPU.

Shape suffixes: T = num tokens (BL), D = model dim, K = top-k, E = num (local) experts, N = TK routed slots, R = routed tokens (= N for EP=1).

Module Contents

Classes

LocalDispatchMetadata	State carried from dispatch() to combine().
LocalTokenDispatcher	Local token reordering for expert compute (EP=1; base for EP variants).
AllToAllDispatchMetadata	State carried from AllToAllTokenDispatcher.dispatch() to combine().
AllToAllTokenDispatcher	EP>1 token dispatcher: local reorder + cross-rank all-to-all (HV-10).

API

class dockyard_rl.models.dtensor.moe.dispatch.LocalDispatchMetadata

State carried from dispatch() to combine().

token_indices_experts_sorted_N: torch.Tensor

None

topk_scores_experts_sorted_N: torch.Tensor

None

class dockyard_rl.models.dtensor.moe.dispatch.LocalTokenDispatcher(num_experts: int, top_k: int, score_before_experts: bool = True)

Local token reordering for expert compute (EP=1; base for EP variants).

Not an nn.Module — holds no learnable state.

Args: num_experts: Number of (local) experts. top_k: Experts selected per token. score_before_experts: If True, routing scores are applied to inputs before the expert FFN; otherwise to the outputs in combine().

Initialization

dispatch(x_TD: torch.Tensor, topk_scores_TK: torch.Tensor, topk_expert_ids_TK: torch.Tensor, num_local_tokens_per_expert_E: torch.Tensor) → tuple[torch.Tensor, torch.Tensor, dockyard_rl.models.dtensor.moe.dispatch.LocalDispatchMetadata]

Reorder tokens by expert; returns (routed_input_RD, counts_E, metadata).

combine(routed_output_RD: torch.Tensor, metadata: dockyard_rl.models.dtensor.moe.dispatch.LocalDispatchMetadata, x_TD: torch.Tensor) → torch.Tensor

Scatter-add expert outputs back to their original token rows.

Each token’s top-k expert outputs accumulate into one row, weighted by the routing score (applied here when score_before_experts is False).

class dockyard_rl.models.dtensor.moe.dispatch.AllToAllDispatchMetadata

Bases: dockyard_rl.models.dtensor.moe.dispatch.LocalDispatchMetadata

State carried from AllToAllTokenDispatcher.dispatch() to combine().

Extends the local metadata with the EP all-to-all bookkeeping needed to reverse the shuffle: the pre-permute shape and the rank-major->expert-major permutation, plus the per-rank token splits (reused, swapped, by the reverse all-to-all in combine()).

input_shape: torch.Size

None

permuted_indices: torch.Tensor

None

input_splits: list[int]

None

output_splits: list[int]

None

class dockyard_rl.models.dtensor.moe.dispatch.AllToAllTokenDispatcher(num_experts: int, top_k: int, score_before_experts: bool = True)

Bases: dockyard_rl.models.dtensor.moe.dispatch.LocalTokenDispatcher

EP>1 token dispatcher: local reorder + cross-rank all-to-all (HV-10).

Distributes routed tokens across the expert-parallel ranks so each rank’s grouped-GEMM sees only the tokens destined for its LOCAL experts, then reverses the shuffle in combine. Distilled from torchtitan’s AllToAllTokenDispatcher (BSD-3-Clause).

The ep_mesh is wired AFTER construction (wire_ep_mesh) because the device mesh does not exist when the surgery/config path builds the dispatcher. Until wired (ep_mesh is None) dispatch/combine fall back to the local EP=1 path, so a misconfigured or pre-parallelize call is correct (local) rather than a crash.

Sequence parallelism is NOT supported on the MoE path (see parallelize.py::_parallelize_qwen3_moe), so there is no SP shard offset — combine scatters into a full (T, D) buffer directly.

The all_to_all_single collectives are device-bound (NCCL); the _permute/_unpermute index math and the EP-disabled fallback are pure and CPU-tested.

Initialization

wire_ep_mesh(ep_mesh: Optional[torch.distributed.device_mesh.DeviceMesh]) → None

Install the 1-D EP submesh used by dispatch/combine (None disables EP).

dispatch(x_TD: torch.Tensor, topk_scores_TK: torch.Tensor, topk_expert_ids_TK: torch.Tensor, num_local_tokens_per_expert_E: torch.Tensor) → tuple[torch.Tensor, torch.Tensor, Union[dockyard_rl.models.dtensor.moe.dispatch.LocalDispatchMetadata, dockyard_rl.models.dtensor.moe.dispatch.AllToAllDispatchMetadata]]

Local reorder, then all-to-all tokens to their experts’ EP ranks.

num_local_tokens_per_expert_E is this rank’s token count over the GLOBAL experts (E,); experts are assigned contiguously per EP rank (see mesh.experts_for_rank), so view(ep_size, -1) rows are the per-destination-rank groups. Returns the expert-major routed inputs for this rank’s LOCAL experts and their global token counts.

combine(routed_output_RD: torch.Tensor, metadata: dockyard_rl.models.dtensor.moe.dispatch.AllToAllDispatchMetadata, x_TD: torch.Tensor) → torch.Tensor

Reverse the dispatch: unpermute, all-to-all back, score, scatter-add.