Infrastructure Prerequisites: Reading Kimi K3 as a Systems Paper
This chapter supplies the systems vocabulary that the Kimi K3 infrastructure section assumes. Read it before or alongside Course 4 — Infrastructure. It is deliberately a prerequisite chapter, not a claim that this repository reproduces a frontier distributed-training or serving system.
Evidence contract
We use three labels throughout:
- Paper-reported — a mechanism attributed to Kimi K3 Section 5 (and, where relevant, Section 7). The paper reports the design; this repository does not measure its throughput, scale, or reliability.
- Background synthesis — standard systems vocabulary and explanatory diagrams. These make the paper readable, but are not claims that K3 invented the technique or used every possible implementation variant.
- Local miniature — a small repository artifact. It can illustrate an API or invariant, never cluster-scale performance.
The right question is not “can I recite the acronym?” It is: what state is partitioned, what communication is induced, what lifetime is shortened, and what correctness boundary must remain exact?
1. One training step, five different ways to divide it
Let a model have parameters \(\theta\), a microbatch \(B\), and loss \(\ell_B(\theta)\). A distributed system can partition different axes at the same time. The following techniques are complementary rather than competing.
| Axis | Technique | Divide | Must communicate |
|---|---|---|---|
| Batch | Data parallelism (DP) | examples | gradients or gradient shards |
| Depth | Pipeline parallelism (PP) | consecutive layers | activations and activation gradients |
| Experts | Expert parallelism (EP) | MoE experts | routed token activations |
| Sequence | Context parallelism (CP) | token positions | attention/KDA summaries |
| State | ZeRO | optimizer and/or gradient state | shards when an owner needs them |
Source: Author-created distributed-axis synthesis for the methods used in Kimi K3 §5; not a paper table.
The important consequence is that a worker can be, for example, one pipeline stage, in one expert group, for one context shard, with a data-parallel replica group around it. “GPU 17” has no useful meaning until those memberships are specified.
1.1 Data parallelism and ZeRO-1
Background synthesis. In ordinary data parallelism, every DP rank holds a model replica. Rank \(r\) receives \(B_r\), computes a local gradient
Source: Standard data-parallel gradient notation used to explain Kimi K3 §5.2.1; background source: ZeRO [Rajbhandari et al., 2020].
and an all-reduce produces the global gradient
Source: Standard data-parallel reduction used to explain Kimi K3 §5.2.1; background source: ZeRO [Rajbhandari et al., 2020].
Every replica then takes the same optimizer step, so the replicas remain equal. DP buys batch throughput; it does not make one model replica fit in less memory.
For Adam-like optimizers, each parameter often has parameter, gradient, and two optimizer-moment tensors. ZeRO-1 shards the optimizer state across the DP group: rank \(r\) owns optimizer state for only one subset \(\theta_r\). The model parameters and gradients are still replicated. Before an update, each owner updates its shard; the updated parameter pieces are made available to the other replicas. The memory win is roughly on optimizer state, not on every resident tensor.
DP without ZeRO-1 DP with ZeRO-1
rank 0: θ, g, m, v rank 0: θ, g, m[0], v[0]
rank 1: θ, g, m, v rank 1: θ, g, m[1], v[1]
... ...
Do not confuse the label with an implementation detail: different frameworks overlap collectives and updates differently. The invariant is ownership — each optimizer-state element has one DP owner rather than \(D\) copies.
1.2 Pipeline parallelism (PP), bubbles, virtual stages, and 1F1B
Background synthesis. PP splits a depth-ordered network into stages. If stage 0 owns early layers and stage 1 owns later layers, stage 0 sends the boundary activation forward and stage 1 sends its gradient backward. A whole batch would leave most stages idle, so the batch is divided into microbatches.
time →
stage 0: F(µ0) F(µ1) F(µ2) ... B(µ0) B(µ1)
stage 1: F(µ0) F(µ1) ... B(µ0) B(µ1)
stage 2: F(µ0) ... B(µ0) B(µ1)
The empty diagonal regions during fill and drain are pipeline bubbles. More microbatches reduce their fraction, but increase the number of activations that must survive until their corresponding backwards pass.
1F1B means one forward, one backward in steady state. After warm-up, a stage alternates forward work for a newer microbatch with backward work for an older one. This bounds activation residency better than doing all forwards first and all backwards later, while retaining a pipeline schedule.
A virtual pipeline stage (VP) divides the layers held by one physical rank
into multiple smaller, interleaved chunks. A physical GPU may therefore act as
stage 0a → stage 1a → ... → stage 0b → stage 1b, rather than owning one
continuous block. Interleaving can reduce bubbles, but introduces a more
irregular activation-lifetime pattern. VP is a scheduling abstraction; it is
not a second GPU.
Paper-reported connection. K3 Section 5.2 uses interleaved PP while describing memory pressure and hiding vision work in pipeline bubbles. The paper-specific claim is not that 1F1B was invented there; it is that the infrastructure exploits the resulting uneven lifetime distribution.
1.3 Why PP memory is uneven
An activation created by an early PP rank waits while its microbatch travels through later stages and then waits for the backward wave to return. A later rank creates its activation closer to its backwards pass. Hence “the same number of layers per rank” does not imply the same peak activation memory.
early rank: create ─────────────────────────────── consume in backward
late rank: create ─── consume in backward
Paper-reported connection. Section 5.2 reports remote activation offload according to lifetime and PP-rank imbalance. Treat this as a global memory budget: a rank with spare memory may temporarily hold a tensor whose owner is an early, overloaded rank. Correctness requires a known tensor identity, completion event, and return path before that tensor’s backward use.
1.4 Expert parallelism (EP)
Background synthesis. In a mixture-of-experts layer, a router chooses a small set of experts for each token. EP places different expert weights on different ranks. Routing therefore turns a local tensor into a transport problem:
tokens → router → pack by destination expert/rank → exchange → expert GEMMs
← unpack and combine ← exchange expert outputs ←
The mathematical sparse choice may be balanced on average while a particular step is badly skewed. One overloaded EP rank then determines latency and peak memory for everyone.
Paper-reported: MoonEP, Section 5.2. K3 reports a GPU planning phase that places bounded redundant expert replicas to make the received token workload per EP rank fixed. The reported bound is at most \(E/R\) redundant expert replicas per rank for \(E\) experts and \(R\) EP ranks. Once every destination has a fixed receive count, offsets and GEMM shapes can be preplanned.
This explains three terms that otherwise sound like marketing:
- Perfect balance: each rank receives the scheduled count, not merely a good expected count.
- Zero-copy dispatch: data is written directly to its preplanned, expert-grouped destination slot rather than arriving in an intermediate receive buffer and being repacked. It does not mean that network hardware transports bytes without DMA.
- Expert-GEMM overlap: while one statically shaped expert group computes, later groups can be transferred or prefetched. The dependency is between a group and its inputs, not between all communication and all computation.
The local StaticZeroCopyBuffer is intentionally only a contrast: a local
preallocated copy_ cannot demonstrate network zero-copy, balance, or overlap.
1.5 Context parallelism (CP) and KDA Context Parallelism (KCP)
Background synthesis. CP partitions a long sequence by positions. Unlike DP, the pieces represent one example and must jointly produce the same causal answer as the unpartitioned sequence.
For conventional attention, the required exchange often concerns distributed key/value blocks. KDA has a recurrent state, so a segment cannot simply be run from zero and summed with other segment outputs. Its segment transition can be written as
Source: Kimi K3 §5.1.2, KDA Context Parallelism affine state summary, arXiv:2607.24653v2.
Two segments compose associatively:
Source: Derived here from the associative KDA state composition used in Kimi K3 §5.1.2.
Paper-reported: KCP, Section 5.1.2. Each rank first computes its local pair \((M,\widetilde S)\); a prefix scan composes the pairs and gives every rank its correct incoming state. The communication object is fixed-size with respect to sequence length, although its tensor dimensions depend on the KDA state. This is the key distinction:
naïve sequential recurrence: rank 0 → rank 1 → rank 2 → ...
KCP: local affine summaries + associative prefix scan → all ranks receive state
1.5.1 Chunkwise KDA kernels and intra-device CP
Paper-reported motivation in Section 5.1; execution explanation below is background synthesis. “Linear” attention does not automatically mean every token can run independently. KDA still has a causal state transition. A useful kernel decomposition chooses chunks \(C_0,C_1,\ldots\) and separates two jobs:
- compute the data-parallel work within each chunk — projections, gates, local recurrence contributions, and outputs that can be prepared without an earlier chunk’s final state;
- communicate or scan only the compact state-transition summary between chunks, then complete the state-dependent pieces.
tokens: [ C0 ][ C1 ][ C2 ][ C3 ]
local work: └──┘ └──┘ └──┘ └──┘ wide GPU work
state summary: ───► ───► ───► ───► causal composition
The point is overlap, not a claim that causality disappeared. A kernel that first materializes every intermediate in high-bandwidth memory can lose the benefit of the mathematics. Fusing nearby projections, recurrence work, and state handoff reduces traffic and launch overhead, provided the numerical order of operations stays within the intended precision contract.
For long prefill, a single GPU can also have too little independent work when tensor parallelism has left it with only a few heads. Intra-device context parallelism partitions the long sequence across SMs on the same GPU. Each SM or cooperative group computes a segment transition summary, then an on-device scan composes summaries before state-dependent output work completes. It is CP without an inter-GPU collective:
one GPU
SM group 0: segment 0 → (M0, S~0)
SM group 1: segment 1 → (M1, S~1) } on-device prefix composition
SM group 2: segment 2 → (M2, S~2)
K3 reports FlashKDA-style overlap for training and prefill, plus intra-device context planning for long prefixes. This chapter does not claim a particular tile size, CUTLASS schedule, or measured speedup; those are kernel- and device- specific facts, not consequences of the recurrence alone.
1.6 Pipeline ZeRO-2
Background synthesis. ZeRO-2 normally shards both optimizer state and gradients across the DP group. A rank keeps only gradient shards it owns once the reduction is complete. “Pipeline ZeRO-2” names the integration problem: the pipeline schedule determines when a gradient becomes ready, while the ZeRO owner/shard schedule determines where it may reside and when it may be offloaded.
The safe order is conceptual rather than API-specific:
- a microbatch’s local gradient contribution becomes available;
- it is reduced to its owner shard;
- that shard is retained, staged, or offloaded according to the next consumer;
- the optimizer owner updates it at the synchronization boundary.
Paper-reported connection. K3 Section 5.2 reports sharding and staging gradient or optimizer state, with pipeline communication overlapped with the work that consumes it. The exact collective order and buffer layout are not reproduced in this repository; this explanation is background synthesis.
2. Memory is a schedule, not a pile of tensors
Peak memory is determined by tensors whose live intervals overlap:
Source: Author-created liveness abstraction for Kimi K3 §5.2.2; not a paper equation.
Reducing one tensor’s size helps only if it lies on the peak. The useful questions are: who creates it, when is it next read, can it be recomputed, and where can it wait safely?
2.1 Unified activation manager
Paper-reported, Section 5.2; background formulation below. A unified activation manager treats every saved-for-backward tensor as a record with metadata and a storage policy, instead of hard-coding one checkpoint rule into each model layer.
ActivationRecord
id, producer, consumer, bytes, ready_event
policy ∈ {keep, quantize, recompute, local-offload, remote-offload}
location, restoration_recipe, release_event
This is powerful because decisions can be tensor-specific:
keepis appropriate when recomputation costs more than storage;recomputestores a recipe and reruns cheap elementwise work;quantizetrades a controlled conversion for less residency;local-offloadmoves a tensor to a slower local tier;remote-offloaduses spare memory on another rank.
The abstraction does not make memory free. It makes the cost explicit: each policy chooses a point on a storage–bandwidth–recomputation frontier. A correct manager also has to respect streams/events; freeing a tensor after “enqueue” is not safe if a GPU kernel has not finished reading it.
2.2 Memory-efficient MoE backward
Paper-reported mechanism, Section 5.2; no private derivation is claimed. Some MoE backward expressions require a forward output in addition to router inputs and upstream gradient. Retaining that forward output extends its live interval across the entire forward/backward gap.
K3 reports algebraically transforming the MoE gradient so that this retained output dependency can be removed, at the cost of lightweight elementwise work. The systems lesson is broader than the unpublished symbolic identity:
before: router input → forward output ───────────┐
backward needs both
after: router input + upstream gradient → transformed backward
(forward output may be released earlier)
Never replace a saved tensor with a “clever identity” casually. Verify the identity under the actual router normalization, top-k behavior, precision, and masked-token semantics.
2.3 Memory-efficient AttnRes
Paper-reported connection, Section 5.2. Block Attention Residuals retrieve over block-level depth summaries. Storing every intermediate depth source over the full sequence would make memory grow with both sequence length and source count. K3 reports checkpointing AttnRes at block boundaries.
The idea is analogous to a travel itinerary: retain an exact checkpoint at the start of a block, recompute cheap or local work inside the block when backward needs it, and discard transient interiors. The checkpoint boundary must include all state necessary to reproduce the block; an incomplete boundary silently changes gradients.
2.4 Why PP rank imbalance changes the offload decision
Remote offload is not “send tensors away whenever VRAM is full.” A useful policy considers at least:
Source: Author-created offload heuristic for Kimi K3 §5.2.3; not the undisclosed production policy.
An early PP rank tends to have high wait-until-use; a late rank can have a temporarily empty memory window. This makes an early-to-late transfer useful only when it fits before the later rank’s own peak. Interleaved 1F1B changes those windows every microbatch, which is why fixed “offload layer 12” rules are usually too crude.
3. Muon ownership and point-to-point orthogonalization
Background synthesis with paper-reported connection in Section 5.2. Muon uses a Newton–Schulz-style polynomial iteration to turn a matrix update into a better-conditioned, approximately orthogonal direction. For a matrix \(X\), the exact polynomial coefficients and normalization matter, but the systems fact is simple: this transformation must see a coherent matrix shard.
With sharded parameters, a wasteful design first gathers every shard to every rank, lets every rank run the same orthogonalization, and then discards most results. P2P-based Muon orthogonalization instead assigns an owner for each parameter shard:
gradient shard → owning rank → normalize / Newton–Schulz iterations
→ updated shard or required result → dependent rank
Point-to-point (P2P) means only ranks that need a shard exchange it. Chunks can be pipelined: while the owner orthogonalizes chunk \(i\), the next chunk is in flight and a previous result is returning. This avoids an unnecessary global all-gather, but does not remove the need for precise ownership, ordering, and numerical validation.
The repository’s Muon notebook demonstrates local mathematical behavior. It does not demonstrate P2P transport or distributed overlap.
4. Decoupling the vision encoder from the language critical path
4.1 Decoupled Encoder Process (DEP)
Background synthesis; paper-reported motivation in Section 5.2.3. A Decoupled Encoder Process is a scheduling pattern: run a variable-cost vision encoder as a separately managed producer of embeddings, then hand those embeddings to the language-model pipeline at a declared boundary. Its purpose is not to change the vision representation; it prevents image patch count from making every language stage wait.
image → DEP: patchify/encode ──┐
├─ join at multimodal boundary → language PP
text tokens ───────────────────┘
The join needs a contract: batch/sample identity, encoder version, shape, precision, readiness event, and failure/cancellation semantics. Otherwise a fast text path can consume embeddings for the wrong sample or block on a late one without observability.
K3 reports two complementary controls: dynamic CP for large patch sequences, and scheduling encoder work into PP bubbles. A DEP is the conceptual lens that makes those controls legible; this repository does not claim a specific production process topology beyond the paper’s description.
5. Million-token RL: state should outlive a GPU allocation
5.1 External KV/KDA state pools
Paper-reported, Section 5.3. A long rollout alternates among decoding, tool execution, queueing, partial-rollout pauses, and later optimization. GPU memory cannot be the sole home for a trajectory that is idle most of the time.
An external state-pool record needs more than “the KV cache”:
trajectory_id, model/version, token boundary,
MLA KV pages, KDA state checkpoint, sampler/RNG state,
environment snapshot reference, integrity/version metadata
The exact boundary is crucial. Resuming token \(b+1\) requires KV data through \(b\) and KDA state after \(b\), from the same model version and decoding configuration. A state pool is a lifecycle system, not simply CPU RAM.
Write-back versus write-through. A generic write-through design copies each new block out immediately. A write-back design keeps a hot block on GPU and copies it only when eviction makes that necessary. The latter avoids copies for blocks that remain hot, but demands durable bookkeeping: a failed eviction must not leave the only valid copy in neither tier.
5.2 Rollout auto-throttling
Paper-reported: adaptive throttling, Section 5.3; control law below is background synthesis. Admission should react to state pressure, not merely number of active requests. One conceptual controller is
Source: Author-created control-law model of Kimi K3 §5.3.1 rollout auto-throttling; not a paper or production equation.
where \(c\) is rollout concurrency, \(h\) is cache headroom, \(q\) is an eviction or queue-pressure signal, and \(e\) is restore/backlog pressure. This is a teaching controller, not a published K3 formula.
Good throttling includes hysteresis and cooldowns. Without them, a controller can oscillate: admit too many rollouts, evict state, observe pressure, stop all work, recover, then immediately repeat.
5.3 Gradient-buffer reuse for non-policy forwarding
Paper-reported connection, Section 5.3. Reference or teacher models may be needed for scoring/regularization but do not receive a training gradient in that forward pass. K3 reports streamed reference-model forwarding and reuse of the policy model’s FP32 gradient-buffer storage for non-policy parameters.
This is safe only when lifetime intervals do not overlap:
reference-weight chunk uses buffer ────────┐
policy gradient will write that buffer ────────┐
no overlap required
The allocator must fence the reuse, prove that the reference chunk is no longer read, and restore/stream the next material before use. “The reference does not backpropagate” is necessary, not sufficient; asynchronous kernels can still be reading its storage.
5.4 OverlayBD and AgentENV density
Paper-reported term in the AgentENV discussion; background explanation. An overlay block-device format such as OverlayBD can present a shared, immutable base image plus a small writable overlay per sandbox. Copy-on-write means a page is physically copied only when a sandbox modifies it.
shared base OS image (read-only)
├─ sandbox A writable overlay
├─ sandbox B writable overlay
└─ sandbox C writable overlay
This supports efficient fork/snapshot semantics: a child begins by referencing
the same pages and pays storage only for divergence. It is not a substitute for
isolation. AgentENV still needs process/VM boundaries, resource accounting,
network policy, cleanup, and a durable identity for snapshots. The local
ResumableMicroVMSandbox is an API sketch, not evidence of OverlayBD,
microVM isolation, or sandbox density.
6. Hybrid KDA–MLA prefix caching: an exact boundary, not a best effort
6.1 The mismatch
MLA retains token-addressable KV entries. KDA carries a recurrent state:
Source: Author-created notation for the fixed KDA prefix state described in Kimi K3 §5.4.1.
The two objects have different convenient granularities. KV can be reused at fine token-block boundaries. A KDA checkpoint is larger and should be stored less frequently. A correct cache therefore has three layers:
logical hash blocks: identify matching token prefixes finely
physical pages: allocate / transfer storage coarsely
KDA checkpoints: restore recurrence at selected causal boundaries
6.2 A concrete logical layout
Paper-reported design principle, Section 5.4; layout below is background synthesis. For every prefix boundary \(b\), a cache index should be able to answer:
PrefixKey(model_revision, tokenizer/config, tokens[0:b])
→ KV page references covering [0, b]
→ latest KDA checkpoint (c, S_c), where c ≤ b
→ replay recipe for tokens (c, b] if c < b
→ integrity metadata: layer layout, dtype, RoPE/NoPE and sampler compatibility
The serving system may store logical hash blocks smaller than physical pages. That lets a request match a fine-grained prefix without forcing every small match to be independently allocated or transferred. A KDA checkpoint is either at the matched boundary or at an earlier valid boundary followed by replay.
6.3 Restore protocol
For a prefix hit at \(b\):
- Validate the cache key. A token match from a different model revision, layer layout, or incompatible precision policy is a miss.
- Pin the KV pages covering the reusable prefix through \(b\).
- Retrieve the newest valid checkpoint \((c,S_c)\) with \(c\le b\).
- If \(c<b\), replay the required token projections through KDA to construct \(S_b\). Do not pretend \(S_c\) is the state after \(b\).
- Install MLA KV data through \(b\) and KDA state \(S_b\) as one logical session boundary. Decode begins only after both are visible.
The non-negotiable invariant is
Source: Author-created invariant for the unified KDA–MLA cache in Kimi K3 §5.4.1; not a numbered paper equation.
Combining KV through \(b\) with \(S_{b-\Delta}\) produces an output that may look plausible while being causally wrong. This is why “state-aware” matters.
6.4 Eviction, reuse, and scheduling
Eviction has to respect dependency lineage. A KDA checkpoint cannot be kept as a useful continuation object if the replay material or compatible KV pages it requires are destroyed, unless its boundary itself is enough for the next request. Likewise, a logical-prefix entry must not point to a physical page that has been recycled.
Paper-reported connection. K3 describes fine logical hashing, coarser physical pages, sparse KDA checkpoints, cache-aware affinity scheduling, and budget-based admission. These are one system:
- affinity prefers the replica already holding the valid lineage;
- admission refuses work whose expected cache/replay cost violates a class budget;
- eviction preserves objects that enable likely reuse, not merely the most recently touched bytes.
The local StateAwarePrefixCache stores one supplied tensor under a key. It
is suitable for testing an API boundary only; it implements none of the layout,
replay, atomicity, or fleet policy above.
7. Attention and kernel vocabulary: DSA, KDA, Triton, MiniTriton, cuBLAS
7.1 DeepSeek Sparse Attention (DSA) versus KDA
Background synthesis; K3’s Section 7 case reports runtime results, not an architecture ablation. Sparse attention and recurrent linear attention solve long-context cost in different ways:
| Question | Sparse attention family (e.g. DSA) | KDA |
|---|---|---|
| What is reduced? | number of token-to-token interactions evaluated | state carried from the past |
| Long-history representation | selected/retrieved KV interactions | fixed-size recurrent state plus hybrid MLA layers |
| Typical systems pressure | selection, gather/layout, sparse kernel efficiency | state propagation, checkpointing, associative scan |
| Key correctness question | did the sparse policy select sufficient context? | is the recurrent state restored at the exact boundary? |
Source: Author-created comparison of the DSA and KDA families discussed in Kimi K3 §§2.1.1 and 7; not a paper table.
Neither column makes the other “obsolete”; they make different accuracy, layout, and parallelism trade-offs. The K3 paper’s kernel-optimization case reports runtime reductions for DSA and KDA operators. That evidence does not show that one attention family is universally faster, better, or interchangeable with the other.
7.2 cuBLAS, Triton, and MiniTriton occupy different layers
Background synthesis; MiniTriton is paper-reported as a Section 7 artifact.
Model/operator intent
├─ call a tuned library primitive ───────→ cuBLAS (vendor GEMM library)
├─ write a GPU kernel DSL ───────────────→ Triton (language/compiler/runtime)
└─ build a compiler stack yourself ──────→ MiniTriton (the reported project)
- cuBLAS is NVIDIA’s optimized dense-linear-algebra library. Calling it is often the right implementation choice for a standard GEMM; it does not express an arbitrary fused kernel.
- Triton is a GPU programming language and compiler ecosystem for writing specialized kernels. It may lower a custom tiling/fusion strategy instead of calling a library routine unchanged.
- MiniTriton is the K3-reported compiler-development artifact: a DSL, MLIR lowering, PTX generation, tensor operations, autograd, distributed execution components, roofline analysis, and GPT training-parity work.
The learning point is compositional depth. A compiler project is not proven by emitting one kernel, and it does not follow that MiniTriton replaces cuBLAS for every dense operation. A serious stack must specify its ABI, layouts, numerical semantics, runtime scheduling, collective behavior, and end-to-end validation.
Retrieval checks
- A PP rank has spare compute but no spare memory. Why might adding microbatches improve utilization yet make the job fail?
- Explain why MoonEP’s fixed receive counts enable both zero-copy placement and static expert GEMM shapes.
- At a prefix hit at token \(b\), what two objects must agree before K3 can generate token \(b+1\)?
- Why can a reference-model weight safely reuse a gradient buffer only during a carefully fenced interval?
- Contrast a DSA selection problem with a KDA state-restoration problem.
If those answers are clear, return to Course 4 — Infrastructure and read each reported system as an answer to a specific state-lifecycle problem rather than a list of acronyms.