[kalebbroo]

Summary

Multi-rank sharding for the Kimi K3.1 decoder, verified on CPU with virtual ranks. Nothing here runs on CUDA yet, and the engine does not route Kimi through this path.

• Collectives. ICollectiveComm.AllToAll (variable-size, counts matrix). HostStagedComm implements it by copy. NcclComm implements it as grouped ncclSend/ncclRecv per ordered rank pair, skipping zero counts. Both use one layout helper, AllToAllLayout.
• Expert parallel. MoeFeedForward can hold a contiguous expert range (ShardExperts). The router and latent projections stay whole. New internal entry points: RouteRows, ApplyLocalExperts, LatentEnter, LatentExit. Default behaviour is unchanged: the whole-bank Forward path refuses a shard.
• ExpertParallelMoe (new, generic). Each rank sources a contiguous token slice, routes it, sends latent rows to the owning rank, runs experts there, returns outputs, and combines in latent space. The plan comes from the existing ExpertParallelDispatch.Plan. The routed output is zero outside the source slice, so the caller adds other partials and runs one AllReduce.
• Kimi decoder. KimiK31Model takes a TpPlacement, or one backend for the one-rank case, which runs the same code path with no collectives. Heads (KDA and MLA), the dense and shared FFN intermediates, and the vocabulary split across ranks. The DSA indexer, latent projections, router and norms are replicated. Each sublayer ends in one AllReduce; logits use AllGather. KimiK31State is per rank.
• Refusals. KimiK31Sharding.Require names the dimension that does not divide the rank count: heads, vocab, experts, dense intermediate, shared intermediate.
• Engine fixes.
  • The attention residual mix includes the running partial only inside a block (NeMo _forward_attn_residual). The MLP and output mixes always include it. The zero row this removes was a common scale factor before RMSNorm, so logits move only by RMSNorm eps.
  • A_log is read one value per head and folded into each channel of its head. The loader accepts a vector with trailing padding and reads the leading Heads values.

Tests run (CPU lane)

• HartsyInference.LLM.Tests, CPU filter: 192 passed, 0 failed. Includes the Kimi folder, the sharded equivalence tests, ExpertParallelMoeTests, ExpertParallelDispatchTests, and the Ernie and MoE tests.
• HartsyInference.Core.Tests, CPU filter: 134 passed, 0 failed. Includes HostStagedAllToAllTests (2).
• HartsyInference.Engine.Tests, filter Ernie5|MiniMax|Moe: 33 passed, 0 failed.
• Sharded equivalence (KimiK31ShardedForwardTests.ShardedForward_MatchesSingleRank) at 2 and 4 virtual ranks against one rank, same weights, tolerance 1e-4 over six tokens.
• Attention-residual source rule (KimiK31AttnResSourceTests) and per-head A_log fold (KimiK31ALogTests).

Not run

• NcclAllToAllTests (tests/HartsyInference.Core.Tests/Distributed/NcclAllToAllTests.cs), category GpuIntegration. 2-GPU filter: FullyQualifiedName~NcclAllToAllTests.TwoGpus. 8-GPU filter: FullyQualifiedName~NcclAllToAllTests.EightGpus. Both compile; neither was run (no GPU runs on the shared card). The grouped send/recv self-pair (i == j) is assumed valid in NCCL and is unverified.
• No CUDA or Vulkan Kimi test, and no real-checkpoint run.
• The A_log tensor length in the published shard header could not be read from here. The loader accepts [Heads] or longer with the head values leading. Confirm against a real shard before relying on it.

Memory per GPU, computed from shapes (bf16 layout, 8 ranks, excluding KV cache and activations)

• Routed experts: 112 experts × 92 layers × 17.5 MB (MXFP4 with scales) ≈ 180 GB.
• Sharded dense (KDA/MLA heads, dense layer 0, shared experts, embed and lm_head): ≈ 51.0 B params → 102 GB total → 12.8 GB.
• Replicated (latent projections 4.73 B, router 0.59 B, MLA latent down 0.36 B, indexer 0.18 B): ≈ 5.9 B params → 11.7 GB.
• Total ≈ 205 GB of 288. The current host loader stores dense as F32, which doubles the dense share to about 230 GB total.
• For comparison, fully replicated dense (114 GB) plus experts is 294 GB, which is why the dense path shards.

Known limits

• Sharding copies host F32 tensors per rank. A production loader must slice at load time.
• Latent projections and router are replicated; a later change can shard the latent projections.
• Engine wiring (TextService placement for Kimi) is out of scope; the decoder is reachable only through KimiK31Model(TpPlacement, ...).
• A single-token decode has one source rank, so multiple all-to-all sources are exercised only by ExpertParallelMoeTests.