Skip to the content.

Native device mesh and collective seam

This is the design contract for issue #25: a shared device-mesh and collective-comms substrate under compute.Backend, so native tensor parallelism (TP) and MoE expert parallelism (EP) do not grow two incompatible communication layers.

Scope: design only. No NCCL/RCCL binding, no multi-device CUDA change, and no MoE expert sharding lands here. This is R3+ scope-gated. Until the native communicator exists, fak rides external TP/EP through Track A (vLLM/SGLang/Dynamo workers) and does not chase raw single-GPU throughput parity with vLLM.

Ground truth

Claim Status Pointer
Base compute.Backend whole-op surface is still the forward-loop target (MatMul, BatchedMatMul, RMSNorm, RoPE, SwiGLU, Attention, Argmax). [SHIPPED] internal/compute/compute.go:318-344
The tensor collective seam exists as the optional compute.CollectiveBackend interface with AllReduceSum, AllGather, ReduceScatter, and AllToAll. [SEAM-ONLY] internal/compute/compute.go:346-408
CPU reference collectives are single-box, rank-order exact, fail closed on malformed parts, and the single-rank case is identity. [SHIPPED] internal/compute/collective.go:97-179
Cross-process host-f32 collectives exist through model.DistComm, but they are not a device/NCCL communicator. [PARTIAL] internal/model/dist_collective.go, docs/serving/multi-node-compute.md
CUDA is single-device by construction today. [GAP] internal/compute/cuda_kernels.cu:52-57 (cudaSetDevice(0))
MoE routing and selected experts run in-process: router top-k, then per-expert SwiGLU, then weighted accumulation. [PARTIAL] internal/model/moe.go:258-289, internal/model/moe.go:342-362
Native TP explicitly refuses MoE/GLM-MoE-DSA instead of mis-serving those decompositions. [GAP] internal/model/tensor_parallel_forward.go:79-82

Line numbers drift; re-anchor with:

rg -n 'type Backend interface|type CollectiveBackend interface|func \(c \*cpuBackend\) AllToAll' internal/compute
rg -n 'cudaSetDevice|func route\(|type moeFFN|ForwardTP does not yet shard MoE' internal

Collective seam

compute.CollectiveBackend is the additive method surface under compute.Backend. Backends that do not implement it remain valid single-device backends. Backends that do implement it are bound to one communicator/process group, so tensors reduced by that backend are rank-local shards in the same group.

The process-group layer owns:

A future concrete shape can be a runtime object rather than fields on base Backend:

type CollectiveGroup struct {
	Name      string
	Rank      int
	WorldSize int
	LocalRank int
	DeviceID  int
	Coord     MeshCoord
	Backend   compute.CollectiveBackend
}

The four primitives stay exactly the HAL primitives already named:

The CPU reference group is the degenerate group: rank=0, world_size=1, local_rank=0, device_id=0. All four collectives are identity for one rank, preserving the bit-exact single-device path. A real NCCL/RCCL backend is correct only if its single-rank behavior and rank-order conformance match the CPU reference before any throughput claim is made.

Device mesh

The native mesh coordinate is:

rank := (dp, pp, tp, ep)
world_size := DP * PP * TP * EP

The planner derives process groups from fixed coordinates:

Group Varies Holds fixed Communication
DP dp pp,tp,ep Inference replicas and placement/routing; no per-token collective by default.
PP pp dp,tp,ep Point-to-point stage handoff through StageTransport; deliberately not a collective.
TP tp dp,pp,ep AllGather, AllReduceSum, and later ReduceScatter inside dense/attention blocks.
EP ep dp,pp,tp AllToAll dispatch to expert owners and AllToAll/local reduce combine back to token order.

Topology is part of placement, not an afterthought:

Dependency chain

  1. Collective seam / device mesh - multi-engineer-month. Land communicator binding, rank/world-size, mesh placement, and a device CollectiveBackend over NCCL/RCCL, witnessed against CPU-ref. This issue (#25) is the design gate.
  2. Native TP - multi-engineer-month. Consume the seam for Megatron-style column/row sharding over real device tensors. The host and CPU-ref decomposition exists; the missing rung is the real communicator, per-rank device binding, and measured multi-GPU run.
  3. EP-for-MoE - multi-engineer-month after TP. Add expert ownership and load-aware routing on top of the same mesh. The EP delta over TP is AllToAll token dispatch to expert owners, per-owner expert execution, and dispatch-combine back to token order. It does not mint a second collective substrate.

Coordination notes

The migrated #25 body cites old internal tracker IDs #274 and #492. On live GitHub those numbers resolve to unrelated issues. The corrected live coordination targets are:

Do not post coordination comments to live GitHub #274 or #492 for this topic; doing so would attach TP/EP design notes to unrelated tickets. If those internal tracker IDs are ever imported with their original meanings, this doc should be cited there.

Non-goals