Skip to content

Infrastructure for Frontier-Scale Intelligence

Section 5 of the Kimi K3 paper is not an implementation appendix. It explains how the architecture changes the physical system required to train, post-train, and serve it. This chapter separates the paper-reported system from the repository's much smaller local miniatures.

If DP, ZeRO-1/2, PP/VP/1F1B, EP, CP, P2P Muon, DEP, OverlayBD, DSA, MiniTriton, or cuBLAS are not already familiar, begin with Infrastructure Prerequisites. It defines each foundation, then returns here to show how K3 composes them.

The machine is part of the model

Four regimes stress four different boundaries:

  1. KDA execution: a recurrent state must move forward while GPUs prefer wide parallel work.
  2. 3T-class pretraining: sparse expert demand, activation memory, and variable visual computation must remain bounded per rank.
  3. Million-token agentic RL: model cache and environment state live much longer than one policy update.
  4. Online serving: hybrid KDA state and MLA KV entries must be cached, reconstructed, scheduled, and evicted consistently.

The unifying systems question is: where does state live, when may work proceed, and which boundary must remain exact?

Executable companion: notebooks/09_infrastructure_mechanics.ipynb verifies five small mechanics used below. Its retained outputs are local arithmetic evidence—not distributed-system or performance evidence.

5.1 Algorithm–system co-design for KDA

KDA kernels across three execution regimes

Training, prefill, and decoding execute the same recurrence under different shapes:

Regime Available parallelism Dominant pressure Reported systems response
Training Many tokens and chunks Preserve dependencies while filling tensor cores FlashKDA overlaps intra-chunk computation and state propagation
Long prefill One very long prefix Too few heads under tensor parallelism can leave device resources idle Intra-device context planning partitions sequence work across SMs
Decode One or few new tokens Launch and state-update latency Specialized recurrent update path and speculative-state handling

Source: Author-created execution-regime comparison based on Kimi K3 §5.1.1, arXiv:2607.24653v2; not a paper table.

The important idea is overlap, not removal of causality. Work inside a chunk can be reorganized, but later state still depends on earlier state.

KDA Context Parallelism as affine prefix scan

For a sequence segment, write the state transition abstractly as

\[S_{\mathrm{out}} = M S_{\mathrm{in}} + \widetilde S.\]

Source: Kimi K3 §5.1.2, KDA Context Parallelism affine state summary, arXiv:2607.24653v2.

Two adjacent segments compose as

\[ (M_2,\widetilde S_2) \circ (M_1,\widetilde S_1) = (M_2M_1,\; M_2\widetilde S_1 + \widetilde S_2). \]

Source: Derived here from the associative composition used by Kimi K3 §5.1.2; not a numbered paper equation.

This operator is associative. Each rank can therefore compute its local pair before knowing its incoming state, exchange fixed-size summaries, and use a prefix scan to obtain the composed transform before every rank. Applying the resulting prefix transform reveals each rank's correct incoming state.

This is fundamentally different from merely passing a tensor to the next rank. The repository's KDAContextParallelism reports the next ring rank and state shape; it does not implement the affine maps, scan, distributed communication, or performance measurement.

5.2 Infrastructure for 3T-class pretraining

Three independent bottlenecks

At this scale, the global average hides the failure mode. A training step is limited by the rank with the most routed tokens, the largest live-memory set, or the longest vision-encoder delay. Section 5.2 addresses these separately:

  • expert-parallel load must be bounded and statically schedulable;
  • activation, gradient, optimizer, and communication lifetimes must fit together;
  • variable image workloads must leave the language-model critical path.

MoonEP: balance before dispatch

Let \(E\) be the number of experts and \(R\) the number of expert-parallel ranks. The paper reports that at most \(E/R\) redundant expert replicas per rank are sufficient for its balancing construction. A GPU planning phase identifies overload and assigns redundant copies so each rank receives a fixed token workload.

Perfectly balanced counts unlock the rest of the design:

  1. destination offsets can be planned before dispatch;
  2. tokens can land directly in expert-grouped remote positions;
  3. receive and GEMM shapes become static;
  4. the host no longer waits for per-layer token counts;
  5. communication, prefetch, and expert GEMMs can be overlapped.

The local StaticZeroCopyBuffer is deliberately not evidence of this system. It preallocates a local tensor and calls copy_; it performs neither network zero-copy dispatch nor expert balancing.

Memory is a lifetime schedule

The reported memory system combines several policies rather than looking for one universal checkpoint rule:

  • recompute inexpensive element-wise results instead of retaining them;
  • keep reusable allocations in managed pools;
  • checkpoint AttnRes at block boundaries;
  • transform MoE gradients so communication buffers can be reused;
  • offload activations according to lifetime and pipeline-rank imbalance;
  • shard and stage gradient or optimizer state;
  • pipeline communication with the compute that consumes it.

Under interleaved pipeline parallelism, early ranks can hold more live activations than late ranks. Remote offload turns otherwise idle memory on later ranks into part of the global activation budget. The local PipelineExpertParallelism only returns a stage label and tensor shape.

Remove the vision encoder from the critical path

Large images create variable patch counts and therefore variable encoder latency. K3 reports two complementary controls:

  • dynamic context parallelism partitions large patch sequences across an appropriately sized group, while multiple images can occupy subgroups;
  • pipeline decomposition schedules vision work into pipeline bubbles so most encoder computation is not on the step's critical path.

The target is not only lower isolated encoder latency. It is lower effective latency for the joint multimodal training step.

5.3 Infrastructure for one-million-token agentic RL

Externalize rollout state without losing the trajectory

A long agent trajectory spans policy inference, tool waits, environment execution, partial rollout boundaries, and later optimization. Keeping every KV block and KDA state resident on the active GPU would bind trajectory lifetime to expensive accelerator memory.

The paper reports external KV/KDA-state pools, adaptive throttling, partial rollout preservation, and streamed reference-model forwarding. Together they allow a rollout worker to release active compute while preserving the exact state required to continue. Streaming reference weights or shards avoids requiring a full reference-model replica to remain resident throughout every comparison.

The repository does not contain this distributed runtime. Its post-training queue illustrates pause/resume scheduling only.

AgentENV: world state has a lifecycle

AgentENV treats the sandbox as durable external state rather than a disposable subprocess:

  • pause/resume releases resources during model inference or queueing;
  • fork creates branches from a shared past;
  • snapshot/checkpoint supports recovery and replay;
  • copy-on-write avoids duplicating unchanged disk or memory pages;
  • isolation prevents one trajectory from corrupting another.

The local ResumableMicroVMSandbox stores caller-supplied Python objects in a dictionary. It is useful only as a lifecycle API sketch—not as microVM, isolation, copy-on-write, or density evidence.

5.4 Inference and online serving

Joint KDA-state and MLA-KV prefix caching

MLA KV cache and KDA recurrent state have different natural granularities. MLA entries can be hashed at fine token boundaries. A KDA checkpoint is a larger state object and is economical only at sparser boundaries. The serving system therefore separates:

  • a fine logical hash block used to identify reusable prefixes;
  • a coarser physical page used for storage and transfer;
  • sparse KDA checkpoints aligned with valid restoration boundaries.

If a prefix match ends at token \(b\), the system must restore MLA entries through \(b\) and a KDA state representing exactly the same causal boundary. Mixing objects from different boundaries silently changes the computation.

The local StateAwarePrefixCache clones one supplied tensor under a string key. It does not implement joint state, hashing, pages, admission, or consistency.

Specialized hybrid serving kernels

The serving path contains several unlike operators:

  • KDA performs recurrent state updates;
  • Block AttnRes retrieves over depth summaries;
  • LatentMoE dispatches sparse expert work;
  • speculative decoding proposes, verifies, accepts, and reconstructs state.

The paper reports operator-specific kernels and state reconstruction rather than forcing these paths through one generic attention implementation. No production kernel, device profile, or serving speedup is reproduced in this repository.

Cache-aware fleet scheduling

Routing a request to the replica that already owns its prefix can avoid transfer or recomputation. That preference competes with memory pressure, queue delay, fairness, and the request's latency budget. A useful admission score is therefore multi-objective, conceptually:

\[\text{score}(r,n)=w_a\,\text{affinity}(r,n)-w_q\,\text{queue}(n)-w_m\,\text{pressure}(n)-w_b\,\text{budgetRisk}(r,n).\]

Source: Author-created scheduler abstraction for the mechanisms in Kimi K3 §5.4.3; not a paper equation or disclosed production score.

This equation is a teaching abstraction, not a published K3 scheduling formula. Its purpose is to show why consistent hashing alone is insufficient: locality is valuable only while the selected replica remains feasible.

Synthesis: architecture becomes an operating system

FlashKDA and KCP control recurrent dependencies. MoonEP controls sparse work placement. Memory and vision scheduling control the training critical path. External state pools and AgentENV control long-lived trajectories. Hybrid caching and fleet scheduling control state identity across serving replicas.

The common language is lifecycle control: summarize, place, overlap, pause, restore, and admit—without changing the mathematical boundary the model expects.

Retrieval check

Choose one mechanism from training, RL, and serving. For each, identify the state being moved or retained, the bottleneck being controlled, and the full-scale result this repository does not reproduce.