# Kimi K3 — Training and Infrastructure Notes > Pre-training, post-training, RL environments, and the systems work behind **Kimi K3** (Kimi Team, > Moonshot AI, July 2026). Companion to [`k3_architecture_notes.md`](k3_architecture_notes.md). > **[report]** = stated in the technical report; **[analysis]** = our inference. --- ## 1. Pre-training ### 1.1 Data A curated corpus spanning four text domains — **Web Text, Code, Mathematics, Knowledge** — plus a large-scale vision corpus. **Text.** Each domain is filtered by a combination of rule-based heuristics, classifier-based quality scoring, and deduplication, with **domain-specific sampling rates determined by ablation studies on smaller models**. Following K2's rephrasing recipe, knowledge and mathematics corpora are rephrased with style- and perspective-diverse prompting, chunk-wise autoregressive generation, and **fidelity verification against the source documents**. **Vision.** Open-source collections combined with in-house pipelines for filtering, synthesis, and deduplication. Two details worth extracting: - **Coordinate supervision is provided in both absolute and normalized `[0,1]` formats**, enabling localization that is both precise and resolution-robust. - Beyond classical text-captioned images, they **substantially scale up programmatic multimodal data, pairing code snippets with their rendered visuals** across domain-specific formats including SVG, 3D assets, Webpage, **Game**, and CAD schematics. > **[analysis]** The programmatic pairing idea generalizes cheaply: anywhere you can *render* from a > known state, you get ground-truth-annotated visual supervision for free, with none of the data > governance burden that attaches to captured or user-contributed imagery. ### 1.2 Scaling law — a methodology worth copying Report Fig. 7 gives fitted scaling curves for K2 and K3; at equal FLOPs, K3's held-out validation loss corresponds to a **2.5× compute-equivalent gain**. The more transferable content is **how they compared learning-rate schedules**: > Their scaling-law study consistently favors cosine decay over Warmup-Stable-Decay (WSD). **But** — > even under the same model size and training-token budget, the two schedules' optimal peak learning > rates and batch sizes **differ substantially**. As a result, **comparing the two schedules using a > shared set of hyperparameters may unfairly favor one simply because those hyperparameters are better > aligned with it.** To ensure a fair comparison, they conducted an **independent scaling-law search for > each schedule**. Under their respective optimal settings, cosine decay consistently achieves a lower > final loss than WSD. > **[analysis]** This discipline is more valuable than the conclusion, and it applies far below > frontier scale: **any A/B comparison between two training configurations is invalid if their optimal > hyperparameters do not sit in the same neighborhood.** Choice of adapter rank, teacher model, > objective weighting — all of these move the optimum. The cheap fix is a small independent LR/batch > sweep per arm before declaring a winner; the honest fallback is to label the result > "hyperparameters not aligned, indicative only". ### 1.3 Recipe and the context curriculum - **Native multimodal training**: language and vision are **jointly optimized from the start of training**, not grafted together via a post-hoc alignment stage. Visual and textual tokens are interleaved within a single next-token prediction objective. - Per-Head Muon plus the weight-clipping mechanism introduced in K2; QB for MoE load balancing; **cosine LR schedule with 1% linear warmup**; **weight decay 0.1 throughout**. - **Four-stage context curriculum**: **8K → 64K during pre-training**, then **256K → 1M during cooldown**. Concentrating the costly long-sequence computation within a small fraction of the overall training budget keeps the curriculum economical while still letting the model adapt gradually to increasingly long-range dependencies. - **Long-context data cleaning.** Long documents and videos from natural sources contain a substantial amount of low-quality content — near-duplicates, binary blobs, truncated files, video clips, invalid machine-generated logs. The pipeline combines exact and fuzzy deduplication, **perceptual hashing over frames for video**, heuristic and classifier-based quality filtering, and structural validation. Because genuinely long and coherent documents are scarce relative to short text, they are **upsampled** so the long-context distribution is not overwhelmed by short sequences during cooldown. - **Length alone does not confer long-range capability.** They additionally **synthesize** long-context data by carefully permuting and concatenating multimodal documents and sub-tasks, **so that the embedded tasks can be solved only by attending to information scattered across the full 1M-token context**. This trains the attention mechanism at the intended scale and prevents it from degenerating into local patterns. --- ## 2. Post-training ``` SFT (cold start) → RL (3 domains × 3 effort levels = 9 experts) → MOPD (consolidate into one) ``` ### 2.1 Supervised fine-tuning Data trajectories are synthesized using **domain-specialized models from the prior Kimi series**, followed by multi-stage verification and human-in-the-loop annotation. All data is serialized with the **XTML** chat template (§2.5). **Quantization-aware training is applied from the SFT stage onward.** ### 2.2 Reinforcement learning: 3 domains × 3 effort levels Rather than training specialized RL models per task, RL is scaled across three broad domains, each spanning a wide spectrum of sub-tasks, with **one expert trained per domain at every reasoning-effort level**: | Domain | Coverage | |---|---| | **General tasks** | General experience, vision, reasoning, faithfulness, search, knowledge work | | **General agents** | Long-horizon assistant tasks, deep research, paragraph-level writing | | **Coding agents** | Software engineering, coding experience, kernel tasks, web development | Crossed with `{low, high, max}` → **nine expert models**. Report Fig. 8 shows that **as RL FLOPs scale, tool-call steps scale up consistently**, accompanied by comprehensive improvement in overall capability. **Partial rollout.** Long-horizon tasks have severe long-tail latency. For each of `N` prompts they sample `K` completions, maintaining `N × K` active trajectories. **Rather than waiting for all rollouts to terminate**, the generation phase pauses as soon as a fraction `λ ∈ (0,1)` completes (i.e. `λNK`), allowing policy optimization to proceed without execution stragglers. Paused rollouts are enqueued and prioritized for resumption at the start of the next iteration — powered by the sandbox infrastructure (§4.4). Once all `K` responses for a prompt complete, they are immediately dispatched for policy optimization. **The cost and the mitigation.** An individual long-horizon trajectory naturally spans multiple iterations, **introducing data staleness that threatens training stability**. Their policy optimization algorithm inherently tolerates an extreme off-policy regime through a **per-token regularization**: by constraining policy updates within a localized neighborhood, it robustly handles highly stale data and sustains training stability. **Reasoning Effort RL.** A **per-problem token budget control mechanism**: each problem `x` gets an initial token budget `b₀(x)` estimated from the cold-start model, and the task reward is **overridden with −1** for trajectories whose total budget `T(y)` exceeds a scaled threshold `τ · b₀(x)`. For general tasks `T(y)` measures thinking tokens; for agentic tasks it accounts for cumulative output tokens including both reasoning traces and tool-call arguments. Training follows a **stage-wise curriculum over the budget multiplier τ**: first a *max-budget* variant with relatively large τ (while still capping the maximum to suppress excessive overthinking), then **τ is annealed to smaller values** to obtain the *high*- and *low*-effort experts. τ is configured per domain under human-in-the-loop guidance. **Agentic Generative Reward Model.** For non-verifiable general tasks, a tournament-style group reward with binary comparisons. Beyond generic agentic capabilities, the agentic judge must follow a **mandatory protocol**: (1) read the outcome, product, or text output; (2) generate a rubric; (3) score each candidate against the rubric; (4) record the rubric-assigned scores in a scorepad. To mitigate reward hacking toward increasingly verbose outputs, a **budget-based verbosity control** mirrors the reasoning-effort control: given an initial verbosity `ℓ₀` estimated from the cold-start model and a multiplier `σ`, a candidate whose output length exceeds `σ · ℓ₀` **automatically loses the binary comparison**. ### 2.3 Multi-Teacher On-Policy Distillation (MOPD) Consolidating the nine experts into one model. For a given domain `d` and sampled effort level `e ∈ {low, high, max}`, optimization is guided by the corresponding teacher `π_teacher^{(d,e)}`. The per-token OPD reward on `y_t`: ``` r_opd(y_t | e, x, y_ 0` clips extreme advantage signals, stabilizing RL training. **This dense reward signal integrates seamlessly into their RL framework** — meaning infrastructure-level optimizations such as partial rollout naturally apply to distillation training for long-horizon tasks. They also experimented with more fine-grained top-`k` distillation objectives and **observed no clear advantage in either convergence speed or final performance**. ### 2.4 Deployment-aware post-training **MXFP4 quantization-aware training.** The **MoE expert weights — which dominate the model's parameter memory — are quantized to MXFP4**, with activations computed in MXFP8, while **all non-expert components (attention projections, latent MoE projections, shared experts, MoE routers) remain in higher precision**. QAT runs throughout the entire post-training stage, covering both SFT and RL, so the model adapts to quantization-induced precision loss. **During RL, rollout and training share the same quantization scheme — eliminating the train–inference mismatch.** **Draft model fine-tuning.** K3 is pre-trained with a multi-token-prediction layer that mirrors the structure of a backbone block. That MTP layer is fine-tuned into an **EAGLE-3-style draft model** — target model frozen, only the draft layer and its feature-fusion projection trained. The draft is **unrolled for seven steps during training**; beyond the first step, it consumes its own outputs from earlier steps, mirroring the recurrent drafting procedure at inference. Two details worth extracting: - **Feature fusion with an identity-preserving initialization.** The draft input fuses low-, mid- and high-level features of the target model, taken from the outputs of the **1st, 4th, and final AttnRes blocks**. These are concatenated and projected to hidden size by a bias-free matrix `W_E3`, **initialized as `[0 0 I]`** so that at initialization the fused representation **coincides exactly with the high-level input `h_h` on which the MTP layer was pre-trained** — and gradually learns to incorporate the low- and mid-level features during fine-tuning. - **Optimize acceptance rate directly, not a KL surrogate.** Speculative decoding speedup is governed by the per-token acceptance rate `Σ_x min(p(x), q(x))` under lossless speculative sampling. **Minimizing the conventional KL-divergence surrogate does not guarantee maximizing this rate** for a capacity-limited draft model. So they directly optimize the likelihood-based **LK loss** — the negative logarithm of the acceptance rate itself: `L_LK = −log Σ_x min(p(x), q(x))`, with `p` and `q` evaluated at temperature 1 and **no auxiliary ground-truth cross-entropy term**. ### 2.5 The XTML chat template (Appendix F) Three design goals: **extensibility** (new capabilities via backward-compatible message formats rather than template revisions), **a low alignment tax** (the format should be learnable with minimal supervised data, supporting a pipeline where a lightly fine-tuned pre-trained model proceeds directly to RL), and **decoding friendliness** (simple encoders, streaming parsers, grammar-constrained enforcers). - **Syntax.** Angle-bracket syntax is replaced by three reserved special tokens — `[open]`, `[sep]`, `[close]` — plus `[end_of_msg]` as the generation stop marker. An element `[open]tag attr="value"[sep] ... [close]tag[sep]` is isomorphic to its XML counterpart, but **every structural boundary is an explicit special token**, which removes tokenization ambiguity at element boundaries and simplifies constrained decoding. - **Messages and zones.** *Input messages* cover system/user/tool/assistant. *Option messages* are placed by scope: **global options** (`tool-declare`, `thinking-effort`) appear **before** all input messages — they govern the whole session, rarely change, and modifying them invalidates the KV cache anyway. **One-shot options** (`tool_choice`, `response_format`) are appended **after** the input messages, **so per-request changes leave the history KV cache intact**. A third kind, the *input option message*, is interleaved with input messages to supplement or override a global option mid-session — which is what supports **dynamically loaded tools**: tools retrieved during a conversation are announced through an additional `tool-declare` message, expanding the toolset **without rebuilding the preceding context**. - **Channels** (inspired by OpenAI's Harmony format). An assistant message body is organized into `think` (reasoning trace), `response` (user-visible answer), and `tools` (tool calls). The two generation modes are selected **purely through the generation prefix** — `[open]think[sep]` for thinking mode, `[open]response[sep]` for instruct mode — **rather than through separate templates**. K3 supports only **preserved thinking**: in thinking mode the think channel is **always retained in history, kept even when its content is empty**, so the model observes a consistent message structure across turns. In instruct mode historical messages contain only the response and tools channels. - **Tool calling.** Each call carries `tool` and `index` attributes; index numbers parallel calls within a message, and each tool-result message repeats the same `tool`/`index` pair and follows the order of its call, **so results are unambiguously associated with calls**. Arguments are **typed**: string arguments appear as raw text, other JSON types are compactly serialized. **Free-form text such as code is therefore a first-class citizen rather than an escaped JSON string.** A pure-JSON fallback block covers inputs whose arguments cannot be decomposed into typed argument blocks; it occurs only in input tokens, never in model outputs, and **its loss is masked during training**. - **Reasoning effort** is exposed as a global option message of type `thinking-effort`, inserted after the tool declaration and before the input messages. The schema reserves four levels (`low`, `medium`, `high`, `max`), of which K3 supports a subset. This **decouples the effort interface from the template syntax** and aligns directly with the effort-conditioned training of §2.2. --- ## 3. RL task synthesis and agentic environments ### 3.1 Unified white-box RL environment Training with a single fixed agent harness can cause a model to **overfit to a particular tool schema, system prompt, context management mechanism, or interaction protocol**. K3's environment represents an agent harness as **a collection of configurable, composable modules** — tool interfaces, system prompts, context management strategies, skills, memories, subagents, and other components. Composing these through configuration, the environment can instantiate mainstream harnesses such as **Kimi Code, Claude Code, Codex, OpenClaw, and Hermes**, as well as entirely new ones. During RL training they **dynamically construct different harness configurations for different task groups**, exposing K3 to diverse combinations of these modules rather than the conventions of any single harness. ### 3.2 Knowledge-graph-guided task synthesis The quality and diversity of post-training tasks are largely determined by their source materials. Retrieval guided by fine-grained concepts surfaces specialized and underrepresented knowledge, while sampling across diverse concepts broadens domain coverage. K3 builds a **self-evolving, hierarchically organized knowledge graph** that agents continuously expand through web-scale exploration. **Construction.** A directed acyclic graph built by recursive, agent-driven expansion. Starting from predefined coarse-grained seed nodes, an agent instance is assigned to each node and performs multiple web searches to investigate that concept. **Before adding new nodes, the agent explores the existing graph to identify equivalent or related concepts, reuse existing nodes where appropriate, and minimize duplication.** Edges are always directed from the coarser concept to the finer one, **regardless of which endpoint the agent discovers first**. Newly added nodes are subsequently assigned to agents for further exploration. **A branch stops expanding when the assigned agent determines that the current concept is sufficiently atomic.** **Material retrieval and task synthesis.** To target a desired distribution across domains and task types, the system samples nodes at varying levels of granularity, either individually or in related combinations. Keywords derived from the sampled nodes are combined with **contextual information from their ancestors in the knowledge graph** to formulate web queries. The retrieved real-world materials are assembled so that a synthesis agent produces training tasks of various types. ### 3.3 Verifiable problems in agentic environments | Task family | Content and verification | |---|---| | **Multi-step complex information searching** | The model plans its research, gathers evidence from the web step by step, and produces a verifiable answer | | **Real day-to-day professional work** | Investment banking, data analysis, legal practice. The model decomposes a complex request, operates domain tools in a sandbox, and completes a deliverable **over dozens to hundreds of steps** | | **Multi-step verifiable visual reasoning** | STEM problems, visual puzzles, chart understanding. Each trajectory is generated in an agent environment equipped with a **Python interpreter in an isolated sandbox**: the model iteratively writes and executes code to crop, zoom, or transform the input image, perform precise computation, or verify intermediate results, and **receives the execution outputs — including generated images — as new observations** over multiple interaction steps | | **GPU kernel optimization** | Single-operator kernels to fused mega-kernels, sourced from high-quality GitHub repositories such as Flash Linear Attention. Covers CUDA, Triton, CuTe DSL, Gluon, ThunderKittens, TileLang, and BF16/FP8/FP4. **Rewards evaluate both correctness and performance**: each kernel provides a PyTorch reference implementation, and **solutions exceeding a predefined numerical error threshold receive zero reward**. Performance is scored against an expert implementation — **matching yields 0.5, approaching the hardware roofline increases the reward toward 1**. A **hacking-detection system** penalizes reward-hacking strategies such as CUDA graph replay, input caching, and precision reduction, and is **continuously extended with new safeguards as new hacking strategies are observed** | | **Long-horizon personal assistant** | **Realistic mock implementations** of widely used applications (Gmail, Notion, Slack, Canvas) that preserve core semantics while enabling reproducible, large-scale interaction **without external APIs or rate limits**. The agent operates in a **persistent, evolving environment over multiple simulated days**, encountering dozens of interdependent events distributed across applications. **A single rollout may involve up to thousands of tool calls and millions of context tokens.** Each event carries its own evaluation criterion, assessed by deterministic rules or LLM-based evaluators | | **Autonomous Execution Tasks (AET)** | Each task specifies an initial state, a constrained goal, a tool-based action space, execution budgets, and **an independent verifier**. Agents see **only** the objective, context, constraints, and verification interfaces — **no reference trajectories or predefined procedures** — and must autonomously perform task decomposition, tool selection, planning, error recovery, and termination. **Rewards are grounded in the verifier's evaluation of the final environment state rather than the agent's self-reported completion.** Environments include black-box system replication, quantitative factor discovery, and tax auditing. Reward hacking is mitigated by **isolating agents from verifiers**, **pairing public verifiers that offer diagnostic feedback with hidden verifiers that evaluate held-out scenarios**, and applying **penalty-based rewards under limited submission budgets** | | **Web development** | Expert-curated; inputs range from one-line scene descriptions to multi-paragraph specifications; artifacts span websites, interactive games, 3D/WebGL scenes, data visualizations, SVGs, and full-stack applications. **Every task runs in a containerized sandbox and is rolled out under diverse agent scaffolds rather than a single fixed harness**, to promote cross-scaffold generalization. Rewards combine **deterministic checks** (functional tests, plus structural and pixel-level similarity for tasks replicating a reference) with **model judging** (source-code inspection, and looking at and interacting with the output artifact). **The reward is zeroed when a project fails to build, runs with errors, or fakes rather than implements the artifact** | > **[analysis]** Two patterns here transfer to evaluation design at any scale. First, **grade the > environment state, not the agent's report** — self-reported completion is the single easiest thing > for a policy to learn to fake. Second, **split verifiers into a public one that gives diagnostic > feedback and a hidden one that scores held-out scenarios**; this is the agentic analogue of a > train/test split, and without it a public metric will be optimized instead of the capability. --- ## 4. Infrastructure K3 combines three system challenges rarely encountered in a single model: **hybrid KDA attention, 3T-class sparse MoE, and million-token agentic RL workloads.** ### 4.1 KDA algorithm–system co-design KDA replaces a growing KV cache with a fixed-size recurrent state. **The serial dependence poses challenges in parallel execution, in exchange for a fixed-size state that is cheap to transfer and reuse** — exploited at two levels. **FlashKDA (open source, MIT).** The chunkwise form is parallel within each chunk but serial across chunks; executed naively the two phases alternate, leaving SMs idle during serial propagation. FlashKDA is a **CUTLASS-based chunkwise kernel that overlaps intra-chunk computation with cross-chunk state propagation**, decomposing the work into token-parallel stages and a head-parallel recurrence, each scheduled and tuned independently. It substantially outperforms the Triton reference, **serves both training and inference prefill**, and is auto-dispatched as a backend of flash-linear-attention. **Intra-device context parallelism for long-context prefill.** Tensor parallelism partitions heads across devices but **never shortens the recurrence**, so under pure TP, prefilling an ultra-long sequence leaves most SMs idle when each rank holds only a few heads. The key observation: **the state transition of each segment can be evaluated independently of the incoming state and composed exactly afterward.** An automatic SM-level context-parallel planner partitions the sequence across the SMs of a single rank, evaluates the segment transitions in parallel, and merges them to recover each segment's exact state. **Entirely intra-device; no cross-device communication.** **KDA Context Parallelism (KCP), cross-device.** For softmax attention, context parallelism requires exchanging KV blocks whose size grows with sequence length. Linear attention instead carries only a fixed-size state — but **direct summation is insufficient for KDA**: the delta rule applies a token-dependent matrix `M_t = (I − β_t k_t k_tᵀ) Diag(α_t)` to the **incoming** state before adding the current write, so the effect of a local segment depends on the state entering that segment and **cannot be determined from the state computed with `S = 0` alone**. KCP decomposes each segment's effect into two locally computable quantities — a **cumulative transition `M` acting on the incoming state**, and a **state `S̃` generated locally from zero**: ``` S^t_{[i+1]} = S̃^t_{[i+1]} + M^{t←1}_{[i+1]} · Σ_{j=1}^{i} ( ∏_{l←j+1} M^{T_l←1}_{[l]} ) S̃^{T_j}_{[j]} ``` Every state is composed purely from locally computed fragments. These rank-level updates **compose associatively**, so the incoming state of each rank can be recovered by a **prefix scan**. Each rank first computes `M^{T+1←1}` and `S̃^{T}` locally, then **exchanges both tensors with a single `all-gather`**. → **KCP requires only a fixed-size all-gather for recurrent-state synchronization and achieves linear compute scaling.** Implementation is upstream in [FLA PR #691](https://github.com/fla-org/flash-linear-attention/pull/691) (merged 2026-01-20). ### 4.2 Infrastructure for 3T-class pre-training Parallelism: **PP with virtual stages + EP + ZeRO-1 DP + Pipeline ZeRO-2 gradient sharding + CP**. MoE layers replicate shared experts across EP ranks, and the all-to-all for expert dispatch and combine is overlapped with computation. **(a) MoonEP — perfectly balanced expert-parallel MoE training (open source, MIT)** In conventional EP schemes token loads are imbalanced across ranks; the resulting computational imbalance degrades throughput, and the dynamically varying shapes of routed-expert activations cause substantial memory fragmentation. MoonEP achieves **perfect load balance with dynamic redundant experts**: in the forward pass it plans redundant experts from the router outputs of the current micro-batch and layer and **prefetches them before the routed-expert computation**; in the backward pass it **stages their gradients in a local reduce buffer** and reduces them back to their home ranks once computation completes. - **Perfect balance with bounded redundant experts.** MoonEP requires every rank to receive exactly `S × K` tokens. The report **proves that a balanced plan always exists with at most `E/R` redundant experts per rank, and that this bound is essentially tight** (Appendix E, with a tightness construction). Reserving `E/R` redundant-expert slots per rank therefore **guarantees a feasible solution always exists, so training is never interrupted**. In contrast, prior work such as ECHO and UltraEP presets the number of redundant experts or imposes a per-rank token cap — **training is then forced to stop whenever no feasible plan exists within the cap**, the cap itself requires manual tuning, and residual imbalance remains. - **Online planning.** Computing the exact optimum at every step is prohibitively expensive. Exact solutions are computed offline with integer linear programming as references, and a **GPU planning kernel** is designed that is near-optimal, incurs negligible overhead, and **always respects the `E/R` upper bound**. - **Zero-copy communication.** A fused permute/unpermute operator in which **the planning kernel precomputes the destination of every token**, so tokens are sent directly to their expert-grouped positions on remote ranks and views of the communication buffer are returned directly to the computation, **eliminating intermediate copies**. Under worst-case imbalance the zero-copy data path in DeepEP requires a buffer of size `S × K × R`, whereas **MoonEP requires only a fixed `S × K`** owing to the perfect balance. - **Sync-free execution with static shapes.** In conventional MoE implementations the per-expert token counts vary across steps and layers, and the host must synchronize with the device at every layer to obtain the actual computation shapes, **stalling the pipeline between layers**. With perfect balance, **the computation shapes of all layers are statically known** — eliminating the per-layer MoE host synchronization and alleviating host-side kernel-launch overhead. - **Expert-GEMM scheduling and overlap.** Even with perfectly balanced aggregate load, per-expert token counts within each rank remain skewed, and a fixed-order, workload-oblivious schedule turns this skew into an imbalanced makespan across SM workers. A **workload-aware scheduler** adapts its parameters to the current token distribution before launch and keeps them fixed during execution; a lightweight heuristic selects these using an **analytical cost model of hardware metrics**, with key coefficients calibrated through offline autotuning. Shared-expert GEMMs are dispatched to a **separate stream** so they overlap with other kernels. **(b) Memory-efficient training** - **Unified activation manager.** Every tensor saved for the backward pass is associated with a **pluggable storage backend**; recomputation, quantization, and offload/remote-offload are **merely storage policies under this abstraction**, freely composable at tensor granularity, declared via lightweight annotations on tensors and **fully decoupled from the model code**. Recomputation is performed at function granularity (supporting cross-layer recomputation). All GPU memory is allocated on the main compute stream and managed within a single memory pool, **avoiding multi-stream fragmentation and host-bound recomputation**; activations are prefetched back at layer granularity and overlapped with computation. In K3, most activations use **block-wise FP8 quantization** combined with offload/remote-offload, and element-wise operators are configured with recomputation. - **Memory-efficient MoE.** In the native implementation the gradient computation of permuted probs depends on the forward output `output`. Inspired by SonicMoE, this gradient is **rewritten through a mathematical transformation into a form depending only on the intermediate activation `act_output` and the upstream gradient `doutput`**, eliminating the backward dependency on `output` at the cost of an additional lightweight element-wise computation. Furthermore, in the forward pass of the group GEMM they **save only the input of the dispatch operation**; during the backward pass, the input of the group GEMM is **recovered by recomputing dispatch**, and the communication introduced by this recomputation can be **overlapped with part of the group-GEMM backward**, eliminating this portion of activation storage at negligible cost. - **Memory-efficient Attention Residual.** The block representation is **generated once at the boundary layer and shared by all subsequent layers, residing directly on the GPU**. The AttnRes computation is **entirely wrapped with checkpointing**, so the activation saved for the backward pass at each layer is **identical to that of the standard residual architecture**. For pipeline parallelism they adopt **cache-based pipeline communication**, in which only newly generated blocks are incrementally transferred between stages and released as soon as the micro-batch finishes, **reaching the theoretical lower bound on memory footprint**. - **Balancing activations across PP ranks.** Under interleaved 1F1B, activations are unevenly distributed across PP ranks due to pipeline warmup, and the number of resident activations decreases as the PP rank increases. To avoid OOM, activations are **remotely offloaded to the memory of other PP ranks** using the Mooncake Transfer Engine. - **Pipeline ZeRO-2 gradient sharding and offloading.** Gradients are sharded across DP ranks, and the **sharded gradients are stored in CPU memory** to reduce GPU usage while keeping the double grad buffer on the GPU. After gradients are reduced across DP ranks into the double grad buffer, they are accumulated into the CPU shards. - **P2P-based Muon orthogonalization.** The distributed optimizer shards parameters evenly across DP ranks, whereas Newton–Schulz orthogonalization requires the **full parameter matrix**, necessitating a communication step to gather complete parameters before each update. The naive approach performs an all-gather over the entire parameter buffer on every rank, incurring a substantial memory footprint on top of making communication the primary bottleneck at scale. Instead, **each rank retrieves only the shards of its locally owned parameters via peer-to-peer communication with the corresponding owner ranks**, eliminating the full-parameter buffer and reducing both memory usage and communication volume. Communication and computation are further pipelined at the granularity of model-chunk buffers. **(c) Multimodal encoder optimization** - **Dynamic CP in the multimodal encoder.** In long-context multimodal training, large images and long videos substantially increase vision-encoder computation time and cause significant load imbalance across devices. Context parallelism is extended to such large samples: **a single large image is partitioned along the patch dimension across multiple devices**, and attention is computed by gathering key–value pairs across CP ranks. In addition, each CP group is divided into several **sub-CP groups** and large images are distributed among them in a load-balanced manner, **preventing the communication fraction from growing with scale**. This reduces both the encoder latency of large visual samples and the cross-device load imbalance, **allowing the remaining encoder computation to be hidden in pipeline bubbles**. - **Encoder computation in PP bubbles.** K2.5 introduced the Decoupled Encoder Process, splitting ViT and text training into separate stages. K3 observes that **under the interleaved 1F1B schedule, the text forward passes of the first PP micro-batches are all scheduled at the very beginning, while the text backward passes of the last PP micro-batches finish only at the very end.** They therefore further decompose the ViT computation: **the ViT forward passes of the first PP micro-batches are executed synchronously upfront, the remaining forward passes are scheduled into pipeline bubbles**, and the backward passes are handled analogously. As a result, **most of the ViT computation is hidden within pipeline bubbles, largely eliminating the effective overhead of the vision encoder.** ### 4.3 Infrastructure for 1M agentic RL - **Co-located RL training** keeps each 1M-context RL experiment within a few hundred GPUs, and partial rollouts reduce tail latency. This achieves good hardware utilization but **introduces memory contention between the rollout KV cache that must be persisted for the next iteration and the memory needed for training** — more severe in long-context RL. - **External KV cache pool.** At 1M-context multi-step rollout, a prefix KV-cache miss is extremely expensive. Partial rollout exacerbates this (many unfinished long prefill requests from the previous iteration arrive at the same time), and speculative decoding further accelerates request turnover within relatively fixed tool-call intervals, **increasing prefix-block churn** and lowering the cache hit rate. The fix decouples prefix retention from GPU residency with a **write-back design**: active decoding blocks remain in GPU KV cache, while **reusable prefixes are written back to an external KV cache pool in CPU DRAM only when evicted from GPU**, and prefetched back before the next reuse. **KDA states are offloaded and prefetched together with the corresponding MLA KV cache blocks, keeping their lifecycles aligned.** Compared with a write-through strategy, this incurs CPU DRAM usage and transfer bandwidth only for prefixes that leave the active decode path, avoiding redundant CPU copies of blocks that are still resident and active on GPU. Sufficient DRAM is available because **training states (model weights and optimizer states) are offloaded to NVMe after a training iteration finishes**; the pool is released after a rollout iteration to avoid contention with training workloads. - **Rollout auto-throttling scheduler.** In multi-step rollout, contexts grow progressively as the trajectory advances, making fixed concurrency based on the full-trajectory average length both hard to estimate and overly conservative early on. Conversely, setting concurrency too high creates KV cache pressure in later stages and can trigger preemption. An **auto-throttling mechanism at the LLM request scheduling layer** uses runtime signals such as active request count, queued request count, and KV cache utilization to dynamically control how many requests are sent to the inference engine — keeping early rollout well utilized while reducing concurrency as KV cache pressure rises, **avoiding both under-saturation and overload without manual tuning**. - **Gradient-buffer reuse for non-policy model forwarding.** RL loss computation often requires forward-only non-policy models, such as reference models, whose weights are too large to keep resident on GPU. These weights are **kept in CPU memory and materialized only when needed, backing their parameter tensors by the policy model's FP32 gradient-buffer storage.** This reuses existing GPU memory without extra allocation or fragmentation, and **remains safe because the buffers are overwritten when real gradients are later computed**. With ZeRO-2 gradient sharding and offloading, each GPU retains gradient buffers for only two VPP chunks; reference weights are **streamed into these slots chunk by chunk — one slot for the current forward computation while the other prefetches the next chunk** — hiding copy overhead without increasing GPU memory. ### 4.4 AgentENV — the microVM sandbox (open source, MIT) K3 employs multiple sandbox runtimes: a traditional container-based runtime, a GPU sandbox runtime, and most notably **AgentENV**, a microVM-based sandbox developed in collaboration with partners. Three core design goals: 1. **High-fidelity isolated sandbox runtime.** As agents become more capable and tasks more difficult, they explore more aggressively and may even attempt reward hacking. In early experiments with traditional container-based sandbox runtimes they **observed several kernel panics and deadlocks caused by unintended agent operations**. On the other hand, they want to permit as much exploration as possible: complex tasks require a sandbox close to a real-world environment — agents should be able to mount disks, run containers, or even launch virtual machines at will. **Firecracker microVMs** provide a level of isolation and fidelity that container-based runtimes cannot match. 2. **Flexible sandbox life-cycles for agentic RL.** At the low level, AgentENV supports **incremental checkpointing and resuming**, where only memory pages dirtied since the last checkpoint are saved — achieving **checkpoint and resume latencies as low as 133 ms and 49 ms** respectively. On top of this, three high-level operations: - **Pause and Resume** — a paused sandbox consumes no memory or CPU, and a sandbox can therefore be paused while the agent is waiting for the model's inference result, **which can account for as much as 98% of the sandbox lifetime**. - **Fork** — creates a new sandbox from the exact state of the original while keeping the original running, **useful for reward judging without side effects**. - **Snapshot** — snapshots at regular intervals for error recovery. 3. **High efficiency and high density.** Tens of thousands of sandboxes, each with a unique set of images, may need to be created within seconds. They adopt **OverlayBD** as the image format, together with a custom ublk driver implementation, storage-layer sharing, and P2P transport, **achieving sub-second launch latency at large scale**. Copy-on-write memory and page-cache optimizations further reduce memory usage, achieving a **memory overcommit ratio of up to 6.5× in real workloads**. > **Scale.** Throughout K3's training and evaluation, a total of **51,219,741 sandboxes across > 1,505,678 images** were created. **[report]** ### 4.5 Inference and online serving **KDA-aware prefix cache management.** The hybrid architecture complicates prefix caching: the KDA recurrent state and the MLA KV cache **differ fundamentally in size and lifetime**, yet a cached prefix is reusable only when both can be restored together at the same boundary. - **Unified cache layout.** KDA states are **packed into the same paged block pool as MLA KV**, with pages unified to the same byte size so both page types share one implementation of allocation, eviction, and transfer. Within a page, the states of all heads are **stored contiguously head by head**, so each head's byte stream is self-contained and serves as the minimal unit of cross-node transfer. Under prefill/decode disaggregation, when prefill and decode nodes adopt different TP degrees, **re-layout is performed on the transfer path with zero GPU-side reshuffling**. This asymmetry proved useful during development: **any type-confused access yields garbage rather than plausible data — a zero-overhead sanity check on the pooled layout.** - **The granularity problem.** Block-hash-based prefix caching reuses the KV cache at the granularity of one physical block; only complete blocks are hashed. This coupling breaks down in K3. Block-hash matching requires one block size shared by all layers, and a prefix hit is reusable only if the KDA state at the hit boundary has been persisted. A KDA layer maintains **a single large recurrent state per sequence** rather than per-token entries, so state snapshots are affordable only at sparse boundaries — and the shared block size is therefore **forced to 1024–6144 tokens**, with the hash granularity tied to the storage block. At such a coarse granularity **caching is nearly useless**: requests shorter than one block can never be reused, and chunked prefill exports no cacheable prefix until it crosses a full block boundary. - **Decouple the two granularities.** Prefix hashing runs on fine **hash blocks** (e.g. 512 tokens) inside MLA pages, while the physical block remains the coarse allocation unit. Alignment runs the other way for KDA: **checkpoints of the recurrent state are saved only at (a sparse subset of) MLA's hash endpoints** — the only positions a lookup can ever reference. - During prefill, a partially filled MLA page is **registered in the prefix-cache index under the chained hash of its last complete hash block**, where each hash covers all preceding hash blocks so that matching an endpoint certifies the whole prefix up to it; the registered endpoint advances as the page fills. Meanwhile, the KDA kernel persists the recurrent state **at the last hash-aligned position processed on each forward pass**. Checkpoints are large, so intermediate checkpoints are superseded as the request advances and recycled, while **those at conversation-turn boundaries are retained** for cross-request reuse. Cached checkpoints are **read-only snapshots**: a hit restores the state by copying it into the request's private running state before the next forward pass, and new checkpoints are written to fresh slots, **so a checkpoint visible to other requests is never mutated in place**. - Lookup proceeds in two stages. The MLA stage matches whole physical blocks by chained hash and, **at the first missing block, falls back to the hash endpoints inside it**, so partially filled pages remain hittable. The KDA stage then **requires a checkpoint at the candidate boundary in every KDA cache group**. The hit is the longest boundary satisfying both stages — **always a multiple of the hash block, and never required to be a multiple of the physical block.** - **Consistency under concurrent scheduling**, dictated by three concrete failure modes: ① all cache groups draw blocks from one shared free list, so allocating a private copy for one group could evict a block that another group has just hit → **every hit block is pinned across all groups before anything is allocated**; ② the copy into the private block executes on the GPU immediately before the forward pass, so a block allocated or registered within the current scheduling step would still hand the previous owner's bytes to a reader → **such blocks are excluded from matching until their copies land**; ③ a checkpoint can restore a request only if it exists in every KDA group, so **evicting one group's checkpoint atomically invalidates its siblings** — a checkpoint is either hittable in every group or in none. - **Result:** prefix caching for hybrid KDA–MLA models **reaches the same generality as for full-attention models — any shared prefix is reusable at any 512-token boundary, independently of request length, chunking, or scheduling interleaving.** **High-performance kernels.** - **KDA decoding.** The bottleneck shifts from exploiting parallelism to efficiently managing the evolving recurrent state, which is updated in place at every decoding step. This in-place update becomes problematic in MTP-based speculative decoding: **if verification rejects a subset of the drafted tokens, the state has already advanced beyond the last accepted token and cannot be trivially rolled back.** Maintaining a state snapshot per draft position would enable rollback but would multiply state traffic — a cost that dominates at the large batch sizes typical of online serving. The observation: **the state after any accepted draft prefix is fully determined by the projected inputs of the draft tokens, which are far smaller than the state itself.** So they **cache only these projected inputs**, rebuild the states of accepted tokens on-chip, and write back the states of the verified and bonus tokens — a design independently proposed in the concurrent work ReplaySSM. The replayed tokens, the bonus token, and the next draft window **share one recurrent loop inside a single fused kernel** covering short convolution, input normalization, gating, the KDA recurrence, and output normalization. Verification latency grows **sub-linearly** with the number of tokens verified and remains below that of state-caching baselines. **Because the projection caches never leave the decode stage, prefix caching and prefill–decode disaggregation operate on the same payload as in non-speculative serving.** - **Block AttnRes.** A two-stage schedule: a batched inter-block pass reads the cached block representations once per block, after which each layer folds in the intra-block partial sum through an **online-softmax merge**. Memory access accounts for a substantial fraction of the cost of these kernels in both prefill and decoding, so both stages focus primarily on memory efficiency. **For prefill**, materializing the block representations on every tensor-parallel rank would incur substantial redundant memory consumption; they therefore adopt **sequence parallelism for activations** — the TP all-reduce of the AttnRes output is decomposed into a reduce-scatter and an all-gather, with the intra-block kernel inserted between the two collectives, operating on the sequence-sharded hidden states so that the block representations of each token are **materialized on exactly one rank**. **For decoding**, the inter-block kernel is launched on a side stream so that it overlaps with independent computation on the main stream; the intra-block kernel is instead streamlined through **fusion** — the merging of the AttnRes output with its partial-sum update, together with the subsequent RMSNorm, is fused into the preceding TP all-reduce, eliminating a dedicated kernel for the intra-block phase and reducing its memory traffic. - **Stable LatentMoE.** Three optimizations mitigate the overhead of the latent GEMMs: ① **fuse the latent down-projection with the MoE router into a single GEMM**; ② shard latent weight matrices across ranks and **fuse the output all-gather into the GEMM epilogue using multimem store instructions**; ③ overlap the resulting communication with other operators, such as the shared-expert computation. **For routed experts at small batch sizes**, the group GEMMs reduce to memory-bound streaming of weight matrices — a regime for which conventional tile-centric kernels are poorly suited due to their compute-oriented design and preprocessing overheads. They instead build the MoE decoding kernel upon the **token-centric design of WarpDecode**, in which **each warp is responsible for one output neuron and streams the associated weights directly from memory**; each warp is further subdivided into finer-grained lane teams, each processing a disjoint subset of experts, followed by a warp-wide reduction of the partial results. The weight layout is **permuted offline at a one-time preprocessing cost**, substantially reducing the runtime dequantization overhead. **Fleet-level scheduling.** - **Cache-aware affinity scheduling.** At 1M context, a typical coding input carries a prefix of 400K tokens but requires a prefill increment of only ~4K tokens, so **a prefix-cache hit avoids re-prefilling the entire prefix and is orders of magnitude cheaper than a miss**. Each request is therefore routed to the cluster that holds its prefix cache, as moving the cache to another cluster would require transferring it over inter-cluster links far slower than the intra-cluster fabric. This affinity, however, binds each session to a single cluster, whose failure would interrupt all sessions bound to it. **Consistent hashing therefore pins each session to two clusters**, a primary that serves its traffic and a pre-assigned secondary that takes over when the primary fails. The secondary holds none of the session's prefix cache and must re-prefill it upon failover; but since consistent hashing **distributes the secondary assignments of different sessions uniformly across the fleet**, this re-prefill load is divided among many clusters rather than concentrated on one. Cache locality is preserved in the common case while the impact of any single cluster failure remains bounded. - **Budget-based admission control.** Production traffic mixes short requests under 2K tokens with ultra-long requests up to 1M tokens, so **the per-request cost spans roughly three orders of magnitude** and the total load imposed by any fixed number of requests is highly unpredictable. Capacity planning, queueing models, and rate-limiting quotas based on the "average request" all break down under this variance. In a typical failure mode, a burst of long-context requests saturates the available compute, and short requests arriving afterwards cannot be scheduled promptly, **degrading time to first token across all traffic**. They therefore allocate **separate resource budgets to different request classes**, so that bursty long-context traffic consumes at most its own share of the capacity and **cannot degrade the system-wide SLOs experienced by other classes**.