Kimi K3: the paper and proofs side by side

Robert Joseph George

Kimi K3: the paper and proofs side by side

Robert Joseph George · TorchLean Verified Examples

Kimi K3 is an open-weight model from Moonshot AI that works with text, images, and video. Its technical report describes 2.78 trillion learned parameters and a context window of up to one million tokens. It combines attention with recurrent memory and selects a subset of its experts for each token. Kimi K3 technical report ↗

Mathematicians have spent years formalizing research papers in Lean. I think it's time we did the same for ML papers. We now have TorchLean for tensor programs and model definitions, and FloatLib for floating-point arithmetic and its proofs. I wanted to use them to work through K3's equations and understand how the model fits together.

Read by chapter

Taking autoformalization into ML

Mathematicians have been putting research papers into Lean for years. The Liquid Tensor Experiment and the Polynomial Freiman–Ruzsa project showed how people could work through substantial mathematics together, writing the definitions and filling in the arguments until Lean could check the proofs. I want us to do that for machine-learning papers too.

Now AI is helping with the translation itself. In September 2026, Anthropic reported a complete Lean formalization of Fermat's Last Theorem, building in part on Kevin Buzzard's community project. This is what people mean by autoformalization: taking mathematics written for people and producing definitions and proofs a proof assistant can check. Seeing this happen for research mathematics makes me want to push further with ML :)

There is also recent work around the Riemann hypothesis. OpenAI's Lean formalization of the quasi-Riemann hypothesis describes a zero-free region $\operatorname{Re}(s)>7/8$ for the Riemann zeta function. The Riemann hypothesis itself would put every nontrivial zero on $\operatorname{Re}(s)=1/2$; the zero-free-region result is a different statement. What interests me here is the practice of publishing the mathematical argument together with a formal statement that readers can inspect.

A model report contains equations we can translate into tensor functions: how a recurrent state changes, how attention reads a cache, how a router chooses experts, and how their outputs are combined. We can then ask whether an implementation computes those equations and prove properties that matter when we use it. Those definitions can use different weights and dimensions, so the proofs can cover a whole family of models. I'd call this model autoformalization. AI could help translate the equations and find proofs, while Lean checks the arguments. Here we'll work through the calculations ourselves, starting with small examples before the full tensor equations.

We can write the tensor computations and their mathematical properties together. FloatLib supplies arithmetic definitions and rounding theory for the next question: what changes when we calculate with finite precision? That matters in ML, where the choice of precision and accumulation order can change a model's output. I want those choices to be part of the formalization, close to the equations that use them.

K3 is a good place to start because several of its architectural choices interact. Kimi Delta Attention (KDA) updates a recurrent memory, while Multi-head Latent Attention (MLA) looks back at individual cached tokens. The model also retrieves summaries from earlier layers and sends each token through selected groups of weights called experts. An image goes through a vision encoder before joining the language model. To formalize the model, we have to understand how those calculations share their inputs and state. Figure 2 and §2, pp. 3–4 ↗

That leads to concrete questions. Does splitting a sequence into two chunks preserve the output when we carry the state between them? Does the graph of tensor operations compute the same function as the paper equation? How large can a capped activation become? We can answer these with theorems about the definitions, and later ask how rounding changes the calculations.

In the Kimi team's introduction to K3, the architecture serves a larger ambition: a model that can keep working through long coding and research tasks. I like that they explain the memory and routing choices behind it. Those choices also give us mathematical questions we can settle. If a serving system pauses a calculation, resumes from a cache, or changes how it groups tokens, we want to know which changes preserve the function and which change its numerical behavior.

We'll start with the transformer calculations that K3 builds on, then work through its memory, attention, experts, and vision model. I'll keep the report beside the explanation and use small examples before the full tensor equations. The training and evaluation sections give us another part of the story: how the authors used those mechanisms and what they measured.

Start with a transformer

At each generation step, a language model assigns probabilities to the next token using the preceding context. K3 does this with 93 layers, 2.78 trillion parameters, and a million-token context. A transformer gives us a starting point for understanding the calculations that K3 changes.

1
Tokenize. The text gets split into vocabulary pieces called tokens — roughly word fragments — and each gets an integer ID. "unhappy" might become two tokens: "un" and "happy." Everything downstream works with these integer IDs.
2
Embed. Each token ID selects a learned vector from a lookup table. A vector has ordered numerical coordinates; K3 uses 7,168 of them per token. Training adjusts those numbers so the later layers can use them when predicting the next token. A concept need not live in one particular coordinate.
3
Compute match scores. Three learned matrices project each token's vector into a query, a key, and a value. A dot product compares the current query with each key. The resulting scores determine how much each value contributes to the attention output.
4
Mask the future. When predicting the next token, a position can only look at itself and earlier tokens, never later ones. This causal mask excludes future positions from the weighted sum. The usual attention notation writes their scores as negative infinity.
5
Weight the answers. Softmax converts all the match scores for one token into nonnegative weights that sum to one in the real-valued equation. For example, weights of 0.4 and 0.3 multiply their corresponding value vectors by those amounts. Adding the weighted values gives the attention output.
6
Update and continue. A residual connection adds the old vector back to the attention output. Then a feed-forward network processes that updated vector, and another residual addition carries the result to the next layer. The residual path lets a block add an update to the existing vector rather than replace it entirely.

The standard attention equation compresses steps three through five into one line. Capital letters stand for tables with one row per token: holds the queries, the keys, and the values.

takes every allowed query–key dot product. A dot product multiplies matching coordinates and sums them, giving one match score per position pair.
Dividing by controls the scale of those scores. If the coordinate products are independent and each has variance one, their sum has variance ; this division brings it back to one. Large score differences can otherwise make softmax concentrate heavily on one position. The variance calculation explains the usual scaling choice, not a law that every wider learned key must produce a larger score.
is zero at allowed positions and is written as negative infinity at forbidden positions, excluding them from the normalization.
The final multiplication by forms a weighted combination of the value vectors from all allowed positions.

Negative infinity here is shorthand for excluding a position. A finite negative score still has a positive exponential over the real numbers, however small it is. The real-valued hard-mask definitions therefore give forbidden positions exactly zero weight. A finite sentinel in a numerical implementation has to be checked against the masking behavior we intend.

Conventional multi-head attention runs this calculation across many attention heads in parallel. Each head uses its own separate query, key, and value projection matrices, so different heads can learn to pay attention to different things — one might track syntactic structure, another semantic role, another pronoun reference. K3 uses 96 heads. Their outputs get concatenated and projected back to the model width. A conventional cache with separate keys and values for every head multiplies its payload by the head count. K3's MLA cache instead retains a shared latent representation; we'll count that smaller representation separately.

Dot products alone do not encode where a token occurs in the sequence. Many transformers add positional information; rotary positional embeddings (RoPE), for example, rotate query and key coordinates according to position. K3 uses NoPE MLA, whose shared-key coordinates remain unrotated. Its KDA layers carry order through sequential state updates: changing the order of inputs changes which state each update receives.

Every block computes , and the representation going into the next block is . If is small, the original vector survives nearly unchanged. The residual stream accumulates the blocks’ contributions, so a later block receives information written at several earlier depths in one sum. K3’s AttnRes keeps earlier block sums available for separate retrieval.

The feed-forward half of each block is what holds most of the model's parameters. It takes the current hidden vector, projects it up to a wider hidden space, applies a nonlinearity, and projects back. In K3 this becomes a mixture-of-experts layer rather than one dense network — each token only activates a small fraction of the available capacity. But the basic pattern is the same.

Training processes many token positions in parallel. Inference is different: the model generates one token, appends it, and runs again. Keys and values from previous steps don't change, so implementations save them in a KV cache. The cache grows by one entry per token per attention layer. This is the trade-off the original transformer design made: instead of compressing history into a fixed recurrent state (like an LSTM would), it keeps the full history directly accessible. Storing that history costs more as the context grows. K3's hybrid attention is one way to manage the growing cache.

The Attention Is All You Need paper from 2017 made this pattern work at scale. Later decoder models have changed normalization, positional encodings, activations, and attention layout. But the alternating rhythm — attention to mix between positions, then feed-forward to transform within one position, then residual addition, then the next layer — stayed recognizable all the way through to K3.

The cost of a longer, larger model

We can make a transformer larger by adding layers, widening its vectors, or using longer sequences. Longer sequences enlarge the attention cache; more experts increase weight storage even when each token uses only a few. Let's look at how K3 handles these costs.

Attention and a growing cache

For a full causal sequence of length , one attention layer has permitted query–key pairs. Efficient kernels can avoid storing the full score matrix, but they do not turn quadratic attention into a constant-size recurrence. During generation, the KV cache grows by one entry per new token in every full-attention layer.

Models reduce this cost in two broad ways: store less information for each previous token, or replace the growing table with a fixed state that is updated after every token. K3 uses both. Multi-head Latent Attention (MLA) keeps one compressed entry for each earlier token, so a layer can still look back at individual positions. Kimi Delta Attention (KDA) keeps a fixed-size recurrent matrix instead, so its memory does not grow with sequence length.

More parameters, more matrix work

The feed-forward network often owns most of a transformer's parameters. Sparse mixture-of-experts models keep many feed-forward networks but route each token through only a few. The Switch Transformer demonstrated this at trillion-parameter scale. Only the selected experts compute an output for a token, but all the expert weights still need storage. Routing also adds load-balancing and communication costs.

One residual stream across many layers

A conventional residual stack carries a single running sum. By layer ninety, that vector contains contributions from every earlier block, but the current layer cannot directly choose an old representation. It receives the sum. Attention Residuals changes the depth path: block summaries remain visible and the current layer retrieves a weighted mixture. K3 stores summaries of groups of layers, rather than a separate summary for every layer.

Bringing images into the sequence

A vision transformer turns an image or video into patch tokens. A multimodal language model must decide where those tokens enter, how their feature width matches the language backbone, and whether training treats visual and text examples as one sequence from the beginning. We need to specify the projection to the language backbone's width and where the visual tokens enter the sequence.

Kimi's earlier Kimi Linear work developed Kimi Delta Attention as a gated delta-rule recurrence and paired mostly recurrent layers with occasional full attention. K3 brings that hybrid into a much larger system, adds depth retrieval and a latent expert space, and trains the result as a native multimodal model.

K3's architectural choices

KDA and MLA split the sequence problem: KDA carries a fixed-size summary forward, while periodic MLA layers can still inspect individual cached tokens. AttnRes retrieves earlier summaries across depth. Stable LatentMoE expands capacity across width without running every expert. MoonViT-V2 turns images and video into vectors the language backbone can consume.

The resulting model has 2.78 trillion parameters, activates 104 billion of them for each token, accepts text, images, and video, and supports a context window of 1,048,576 tokens. Moonshot released the weights. This gives the formalization a concrete object to aim at, rather than an architecture reconstructed from an API description.

2.8Ttotal parameters

→

104Bactive per token

1Mmaximum context

sequence
KDA compresses the earlier tokens into a fixed-size matrix that changes after each new token.
global
Three KDA layers are followed by one MLA layer to look back at individual cached tokens; layer 93 is also forced to MLA.
depth
AttnRes keeps block summaries so a later layer can choose among earlier depth states instead of receiving only one running sum.
width
LatentMoE sends each token through 16 of 896 experts. An expert is a separate feed-forward network; sparse routing avoids running all 896.
vision
MoonViT-V2 pools images and video into visual tokens with the same 7,168-wide representation expected by the language model.

Together they let K3 grow in context, depth, and parameter count while accepting both text and visual inputs. Figure 2 and §2, pp. 3–4 ↗

Across the architecture, data, and training recipe together, the Kimi team reports about 2.5× better scaling efficiency than Kimi K2. The number comes from a fitted loss-versus-compute curve on a validation set deliberately different from the training data. It is an empirical result, not a serving-speed measurement or a Lean theorem. The report does not split the gain among KDA, AttnRes, LatentMoE, data, and optimizer changes. §3.2 and Figure 7, pp. 10–11 ↗

The report also gives results for coding, search, tool use, and web development. I return later to the scores and the evaluation conditions behind them. Abstract and Figure 1, pp. 1–2 ↗

The explicit equations are enough to get started :) We can take one expert, write its matrix products and activation in Lean, and prove that the graph computes that equation. Then we can use the expert inside a routed layer. Once that calculation is clear, I want to ask what changes when we store the weights and run the operations at lower precision. FloatLib gives us arithmetic definitions and rounding bounds for that next question.

Reading the paper with Lean beside it

Okay... imma try something a little different here :) I'd like to read the paper on one side and follow the Lean definitions and proofs alongside it. This is how I imagine working through an ML formalization: we read an equation, try a small tensor example, and see exactly what we've proved about it before moving to the next part of the model.

Click “Open PDF” on the left and follow along. I'll start with one expert, then work through K3's recurrent memory, attention, and complete decoder. You can change the numerical examples as we go and open the corresponding Lean source. I'd love future model reports to be shared this way, with the mathematics and its formalization close enough to read together.

Model architecture

KDA, Gated MLA, AttnRes, Stable LatentMoE, MoonViT-V2, and Per-Head Muon.

reference + selected graph proofs Read the annotations ↓

Pre-training

Data, scaling curves, joint image-and-text training, and the long-context curriculum.

selected loss formula Read the explanation ↓

Post-training

Supervised fine-tuning, nine reinforcement-learning specialists, teacher–student distillation, low-precision numbers, and agent tasks.

selected equations Read the explanation ↓

Infrastructure

KDA kernels, MoonEP, activation memory, AgentENV, prefix caches, and serving.

reported systems Read the explanation ↓

Evaluations

Reasoning, coding, agents, vision, cyber security, third-party results, and cost.

reported results Read the evidence ↓

Case studies

Kernel work, research, office workflows, web development, and long trajectories.

worked examples Read the examples ↓

Report Figure 2, p. 3

Recurrent memory, depth attention, and experts

Figure 2 is easier to follow as three separate directions of information flow. Along the token sequence, KDA remembers and MLA looks back globally. Along model depth, AttnRes chooses earlier summaries. Across model width, LatentMoE chooses a small set of experts.

The backbone remains one 93-layer stack. KDA and MLA mix information across tokens, AttnRes retrieves across depth, and LatentMoE selects transformations across width.

In Lean

The same separation appears in the types: KDA and MLA are sequence mixers, AttnRes owns the depth state, and LatentMoE is a channel mixer.

KDA and MLA need different state

Choosing KDA selects a recurrent matrix state; choosing MLA selects a cache of latent entries. The step function matches on that choice and applies its transition. forward repeats the transition across the input tokens, and forward_length checks that each token produces one output.

/-- Advance either KDA or Gated MLA by one token while retaining its causal state. -/
def step (mixer : SequenceMixer α cfg decayRank) (state : mixer.CausalState)
    (token : Tensor α [cfg.hiddenDim]) :
    mixer.CausalState × Tensor α [cfg.hiddenDim] :=
  match mixer with
  | .kda layer logFloor => layer.scanStep logFloor state token
  | .mla layer => layer.step state token

/-- Run one sequence-mixing layer from an empty sequence cache/state. -/
def forward (mixer : SequenceMixer α cfg decayRank)
    (tokens : List (Tensor α [cfg.hiddenDim])) :
    List (Tensor α [cfg.hiddenDim]) :=
  (CausalScan.run mixer.step mixer.initialState tokens).2

/-- Both sequence mixers return one representation per input token. -/
theorem forward_length (mixer : SequenceMixer α cfg decayRank)
    (tokens : List (Tensor α [cfg.hiddenDim])) :
    (mixer.forward tokens).length = tokens.length := by
  exact CausalScan.outputs_length mixer.step mixer.initialState tokens

Full source

Both transitions take a cache and return an updated cache and output tensor. Their cache types differ: KDA stores a recurrent matrix, while MLA stores earlier token entries.

Model.lean, dependent sequence transition ↗

Writing one expert in Lean

An “expert” sounds like a whole model, but here it's a feed-forward network with three learned matrices. Two matrices make the gate and value branches; SiTU combines their coordinates; the third matrix produces the output. Routing decides which experts to call. We can understand this one function before worrying about choosing among 896 of them.

Ok lets start with the sums and shapes. We'll write the tensor function, build the same calculation from graph operations, and use the existing equality theorem on a tiny example. The proof uses real-valued tensors. Rounding the arithmetic is a further question we can tackle once those two descriptions agree.

Lean keeps the program and the proof together

The part I find useful about Lean is that it's both a programming language and a theorem prover. We can define a layer as a function, run a supported numerical version of it, and state a theorem about the definition in the same language. The tensor types also record the input and output shapes.

For example, K3 caps the two branches of its SiTU activation before multiplying them. We can write that as a function of two real numbers and ask how large its result can get. The paper uses caps of 4 and 25; the theorem below bounds the scalar output by their product, 100, for every pair of real inputs. We don't need to choose particular weights to prove this property of the equation.

/-- Smoothly cap a scalar at magnitude `cap` using `cap * tanh (x / cap)`. -/
def softCap {α : Type} [Context α] (cap x : α) : α :=
  cap * MathFunctions.tanh (x / cap)

/-- Kimi K3's caps imply the concrete coordinate bound `100`. -/
theorem paper_caps_bound (gate up : ℝ) : |scalar 4 25 gate up| ≤ 100 := by
  have h := abs_scalar_le_mul_caps (gateCap := (4 : ℝ)) (upCap := (25 : ℝ))
    (by positivity) (by positivity) gate up
  norm_num at h ⊢
  exact h

Full source

Let's evaluate the cap at 2 and then the full SiTU response at gate 2, up 6. These commands use Float so we can run the calculation. Hover over softCap or the theorem name to inspect the types and assumptions. The theorem below concerns real numbers; the two decimal outputs are rounded evaluations.

import KimiK3.FeedForward.Activation

open KimiK3.SiTU

-- Float gives us a runnable value; the real-valued theorem quantifies over every input.
#eval softCap (4 : Float) 2
#eval scalar (4 : Float) 25 2 6
#check softCap
#check paper_caps_bound

Lean output

1.848469

9.585417

KimiK3.SiTU.softCap {α : Type} [Context α] (cap x : α) : α

KimiK3.SiTU.paper_caps_bound (gate up : ℝ) : |scalar 4 25 gate up| ≤ 100

Dependent types catch a different class of bug. A KDA state carries its two dimensions in the type Tensor α [keyDim, valueDim]. A matrix with those axes reversed has a different type when those dimensions differ. When the graph connects an output wire to the next layer's input, Lean checks that the shapes agree. This catches incompatible shapes before execution; it does not detect every possible mix-up between axes of the same size.

I also build the calculation as a graph. The GraphSpec files in this example connect typed tensor operations such as multiplication and tanh. Evaluating those nodes over real numbers should return exactly the same tensor as the reference function. A denotation theorem proves that equality, so we can check the graph's wiring against the equation.

We can lower that graph to a program for execution. Connecting a fused, low-precision GPU implementation to the real-valued graph needs a further proof: how much can its answer differ under the chosen arithmetic? The graph equality and that numerical bound answer different questions. We'll keep both calculations visible as we work through the model.

We'll use this approach for the recurrent update and the sequence of decoder steps too. Their statements are larger, but we can still follow the function, its assumptions, and the proof.

Writing the model in Lean

We write a layer as a function between tensors. A recurrent layer also takes a state and returns an updated one. Here is the KDA step: it updates the memory, then reads from that new memory to produce the current token's output.

A recurrent step and its return type

The colon in a Lean definition separates a name and its arguments from its type. In the block below, step takes a state and an input and returns a pair: the updated state and an output tensor. The return type appears after the final colon. The state type is the same before and after the call, so the update must preserve the matrix dimensions. The symbol := means "is defined to be," and everything after it is the actual calculation.

/-- One recurrence step returns the updated state and its current-token output. -/
def step (state : State α keyDim valueDim) (input : KDA.Input α keyDim valueDim) :
    State α keyDim valueDim × Tensor α [valueDim] :=
  let next := update state input
  (next, read next input.query)

Full source

We can try this with a zero memory, a key and query of [1, 0], and a value of [3]. The update writes into the first row; reading that row returns 3. Hover over memory, token, or step to inspect their types.

import KimiK3.Sequence.KDA
import NN.Tensor

open TorchLean

-- One head with two key coordinates and one value coordinate.
def memory : KimiK3.KDA.State Float32 2 1 := [[0], [0]]
def token : KimiK3.KDA.Input Float32 2 1 :=
  { query := [1, 0]
    key := [1, 0]
    value := [3]
    retention := [1, 1]
    writeStrength := 1 }

-- Read the first output coordinate.
#eval (KimiK3.KDA.step memory token).2.getScalar ⟨0, by decide⟩
#check KimiK3.KDA.step

Lean output

3.000000

KimiK3.KDA.step {α : Type} [Storage α] [Context α] {keyDim valueDim : ℕ} (state : KimiK3.KDA.State α keyDim valueDim)
  (input : KimiK3.KDA.Input α keyDim valueDim) : KimiK3.KDA.State α keyDim valueDim × Tensor α [valueDim]

Scalar types and tensor shapes

The Greek letter α (alpha) names the scalar type. Most proofs here use real numbers; an executable calculation can use a finite format such as a 16-bit float. Definitions that leave α open specify the operations they need, so we can reuse their tensor layout with different arithmetic. A theorem about real arithmetic still needs a refinement argument before it describes rounded execution.

A type like Tensor α [4, 2] means a tensor of α-values with four rows and two columns. The shape [4, 2] is right there in the type, not hidden in a runtime assertion. Lean checks these dimensions when it typechecks the program. Connecting incompatible shapes produces a type error before the calculation runs.

There is one small shape convention worth trying before the larger equations. Tensor α [] holds one scalar: it has no axes, rather than no values. Tensor α [0] is an empty vector, because its one axis has length zero. The number of entries is the product of the axis lengths; the product of no lengths is one. This is how a reduction can return a scalar tensor without introducing a different container type.

Indices, tensors, and growing caches

Fin n pairs a natural number with evidence that it is less than n. An expert ID of type Fin 896 therefore ranges from 0 to 895. Passing 1000 would require a proof of the false inequality 1000 < 896, so Lean rejects it.

The public sequence API takes tensors and returns tensors. Some internal caches use a List because they grow as tokens arrive, while the graph takes packed tensors with fixed shapes for one step. To relate the two, we need to show that packing the cache preserves the update and the output. Fixing the shape alone would leave the order of the cached entries unchecked.

Definitions and theorems

A structure bundles named fields. A model record can hold weight matrices; a cache record can hold state; a probability record can pair a number with evidence that it is positive. To see how those fields affect the model, follow the functions that consume them.

A def defines a function or value. It consumes structures and produces new ones, or produces tensors. This is where the actual computation happens. A theorem makes a claim about what a def returns. After theorem, everything up to := by is the statement being proved. The lines after by are tactic commands that Lean uses to build a proof term that its kernel verifies. If you see simp or ring or linarith, those tactics simplify expressions, normalize polynomial identities, or solve linear inequalities. Lean checks the proof terms they produce.

Describing tensor operations with Term

For the graph version, we use Term Γ shape. This is a typed expression in a computation graph. Its result is a tensor of shape shape, and Γ lists the shapes of its available inputs. You build expressions by composing Term values with Term.op, naming each primitive operation explicitly. Term.eval gives the graph a mathematical meaning by plugging in actual tensors. specFwd evaluates a whole graph model at once.

We'll often prove an equality of this form: model.specFwd params inputs = referenceFunction inputs. That equation says the graph and the reference function produce the same result. We can inspect the function first, then check how the graph implements each operation.

From an expert equation to its graph

Let's put an equation next to each part of Expert.forward. I'll use a row vector for the input, matching the [input, output] matrix layout in our source. That convention matters when we write the products.

Three matrix products and SiTU

Write $d$ for the input width, $h$ for the hidden width, and $o$ for the output width. The expert has gate and up matrices $G,U$, and a down matrix $D$:

Input and the two projection matrices.

Output projection and hidden branches.

A projection makes each hidden coordinate by weighting and adding the input coordinates. For $0\le j<h$, the two sums are:

Gate projection.

Up projection.

Now take positive caps $c_g,c_u$. The sigmoid $\sigma(t)=1/(1+e^{-t})$ lies between zero and one. SiTU uses it to gate two smoothly capped branches, giving a hidden vector $z$:

The SiTU hidden coordinate.

The final output projection.

Notice that the sigmoid sees the original $g_j$, before capping. Applying it to $c_g\tanh(g_j/c_g)$ instead would give a different activation, even though every tensor shape would still match. That is one of the choices our equation and graph need to agree on.

/-- Evaluate one SiTU-GLU expert. -/
def forward (expert : Expert α inputDim hiddenDim outputDim) (gateCap upCap : α)
    (x : Tensor α [inputDim]) : Tensor α [outputDim] :=
  let gate := vecMatMulSpec x expert.gateWeight
  let up := vecMatMulSpec x expert.upWeight
  vecMatMulSpec (SiTU.vector gateCap upCap gate up) expert.downWeight

Full source

This is the same calculation in Lean. gate and up both have shape [hiddenDim], so SiTU can combine matching coordinates. The final multiplication returns [outputDim].

Each weight matrix has one axis for the incoming coordinates and another for the outgoing coordinates. For an input of width inputDim, both gateWeight and upWeight have shape [inputDim, hiddenDim]. The two products therefore meet at the same hidden width, where SiTU combines matching coordinates. downWeight has shape [hiddenDim, outputDim]. Reading those shapes tells us which matrix acts at each stage before we inspect any values.

The same calculation as a graph

/-- Build the typed DAG term for one SiTU expert from explicit input and parameter terms.

Use this form to embed an expert in a larger graph such as Stable LatentMoE. Both input
projections and the final output projection remain explicit graph operations.
-/
def term {Γ : List Shape} (inputDim hiddenDim outputDim : Nat)
    (input : Term Γ [inputDim])
    (gateWeight upWeight : Term Γ [inputDim, hiddenDim])
    (downWeight : Term Γ [hiddenDim, outputDim])
    (gateCap upCap : Term Γ []) : Term Γ [outputDim] :=
  let gate := Term.op (PrimOp.broadcastVecMat [] [] [] inputDim hiddenDim
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
    (.cons input (.cons gateWeight .nil))
  let up := Term.op (PrimOp.broadcastVecMat [] [] [] inputDim hiddenDim
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
    (.cons input (.cons upWeight .nil))
  let cappedGate := GraphSpec.softCapTerm [hiddenDim] gateCap gate
  let sigmoidGate := Term.op (NN.GraphSpec.DAG.PrimOp.sigmoid [hiddenDim])
    (.cons gate .nil)
  let gated := Term.op (NN.GraphSpec.DAG.PrimOp.mul [hiddenDim])
    (.cons cappedGate (.cons sigmoidGate .nil))
  let cappedUp := GraphSpec.softCapTerm [hiddenDim] upCap up
  let hidden := Term.op (NN.GraphSpec.DAG.PrimOp.mul [hiddenDim])
    (.cons gated (.cons cappedUp .nil))
  Term.op (PrimOp.broadcastVecMat [] [] [] hiddenDim outputDim
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
    (.cons hidden (.cons downWeight .nil))

Full source

The graph keeps those operations visible. Its two matrix products make $g$ and $u$; softCapTerm makes each capped branch; the sigmoid node reads the uncapped gate. Multiplication combines them before the down projection. Lean checks the shape of every connection.

I'm checking two things here. The types tell us that the dimensions fit. The equality theorem tells us that this particular wiring computes the intended function. Exchanging two same-shaped matrices would pass the first check, but would generally change the equation we have to prove.

A tiny expert we can work through

For a hand calculation, choose two input coordinates, two hidden coordinates, and one output. These toy matrices make the first two products easy:

Toy input and gate weights.

Toy up and down weights.

Then $g=(1,2)$ and $u=(2,6)$. With the paper's caps of 4 and 25, we get these values to six decimal places:

Illustrative hidden activation.

Illustrative output.

Here's how I instantiate the graph theorem for that example. The square brackets construct tensors because the weight fields and input have tensor types; Lean checks the rows and columns while elaborating them. We use $\mathbb R$ to prove a statement about real-valued tensors; the decimals above help us work through the calculation.

import KimiK3.GraphSpec.Expert

open TorchLean

namespace ExpertExample

-- Two input coordinates, two hidden coordinates, and one output coordinate.
noncomputable def expert : KimiK3.Expert ℝ 2 2 1 :=
  { gateWeight := [[1, 0], [0, 1]]
    upWeight := [[2, 0], [0, 3]]
    downWeight := [[1], [-1]] }

noncomputable def input : Tensor ℝ [2] := [1, 2]

-- Instantiate the existing theorem with these weights and the paper's caps.
example :
    (KimiK3.GraphSpec.Expert.model 2 2 1).specFwd
        (KimiK3.GraphSpec.Expert.parameters expert)
        (KimiK3.GraphSpec.Expert.inputs input 4 25) =
      expert.forward 4 25 input := by
  exact KimiK3.GraphSpec.Expert.specFwd_eq_forward expert input 4 25

end ExpertExample

We reuse the existing expert and its proof; we don't need a new model wrapper for these dimensions. The theorem below leaves the weights and widths general.

Graph equivalence for arbitrary weights

/-- The typed executable graph denotes the original SiTU expert equation. -/
theorem specFwd_eq_forward {inputDim hiddenDim outputDim : Nat}
    (expert : KimiK3.Expert ℝ inputDim hiddenDim outputDim)
    (input : Tensor ℝ [inputDim]) (gateCap upCap : ℝ) :
    (model inputDim hiddenDim outputDim).specFwd
        (parameters expert) (inputs input gateCap upCap) =
      expert.forward gateCap upCap input := by
  simp [model, parameters, inputs, NN.GraphSpec.DAG.Model.specFwd,
    TorchLean.TensorPack.append, Args.vars, Args.weakenLeft, Args.rename,
    Term.rename, Term.eval, Env.tget,
    KimiK3.Expert.forward, Params, Inputs]

Full source

The source proof expands the model's parameter and input packaging, then uses the evaluation rules for the graph operations. Those rules leave the three matrix products and SiTU expression from Expert.forward. That's why the result works for every choice of weights and dimensions, not only our tiny example.

We can then lower the graph with toProgram. To connect a GPU answer to this real-valued theorem, we still need a rounding argument for the operations that actually execute. The graph equality gives us the equation that argument must use.

A bound on the expert's output

For the hidden activation, we can get a bound without inspecting any learned weights. Since $|\tanh(t)|\le1$ and $0<\sigma(t)<1$, the paper's caps $c_g=4$ and $c_u=25$ give:

The paper's hidden-coordinate bound.

I want to be careful about where that 100 applies. The down matrix can amplify the hidden vector. For an output coordinate, the triangle inequality gives:

Bound the absolute value of the output sum.

Use the SiTU bound for each hidden coordinate.

In the toy example, the absolute down weights sum to two, so this gives $|y_0|\le200$. It's a loose bound, but it holds for every real input to those weights. The example above gives one output; the inequality tells us what can happen across all inputs. That's a useful first result to have before putting the expert inside a larger routed layer :)

The output bound in Lean

We reuse the SiTU theorem for each hidden coordinate, expand the vector–matrix product, and apply mathlib's bound on the absolute value of a finite sum.

import KimiK3.GraphSpec.Expert
import NN.Proofs.Tensor.Algebra

open TorchLean
open scoped BigOperators

example {d h o : Nat} (expert : KimiK3.Expert ℝ d h o)
    (input : Tensor ℝ [d]) (k : Fin o) :
    |(expert.forward 4 25 input).getScalar k| ≤
      100 * ∑ j : Fin h, |Spec.get2 expert.downWeight j k| := by
  let z := KimiK3.SiTU.vector 4 25
    (Spec.vecMatMulSpec input expert.gateWeight)
    (Spec.vecMatMulSpec input expert.upWeight)
  have hz (j : Fin h) : |z.getScalar j| ≤ 100 := by
    simpa [z, KimiK3.SiTU.vector, Tensor.stackLeading] using
      KimiK3.SiTU.paper_caps_bound
        ((Spec.vecMatMulSpec input expert.gateWeight).getScalar j)
        ((Spec.vecMatMulSpec input expert.upWeight).getScalar j)
  change |(Spec.vecMatMulSpec z expert.downWeight).getScalar k| ≤ _
  rw [Proofs.TensorAlgebra.getScalar_vec_mat_mul_spec]
  calc
    |∑ j : Fin h, z.getScalar j * Spec.get2 expert.downWeight j k| ≤
        ∑ j : Fin h, |z.getScalar j * Spec.get2 expert.downWeight j k| :=
      Finset.abs_sum_le_sum_abs _ _
    _ ≤ ∑ j : Fin h, 100 * |Spec.get2 expert.downWeight j k| := by
      apply Finset.sum_le_sum
      intro j _
      rw [abs_mul]
      exact mul_le_mul_of_nonneg_right (hz j) (abs_nonneg _)
    _ = 100 * ∑ j : Fin h, |Spec.get2 expert.downWeight j k| := by
      rw [Finset.mul_sum]

GraphSpec/Expert.lean, complete worked pattern ↗

Figure 2, p. 3 · Table 1, p. 11

Checking K3's published dimensions

The same Lean definitions work for tiny examples and the published 93-layer model. All published dimensions live in one checked configuration; the 2.78-trillion-parameter architecture is one value of those general definitions.

7,168language width

93backbone layers

96attention heads

896routed experts

16active experts

1,048,576maximum context

Explore the 93-layer schedule

KDA fills three positions, MLA fills the fourth, and layer 93 is forced to MLA so the backbone ends with global attention.

69 KDA24 MLASelect a layer

Dimension checks and the released configuration

The shapes impose arithmetic conditions. Dimensions must be positive; 16 selected experts must fit inside a pool of 896; the relevant widths must divide into 32-value quantization blocks. The vision width must also divide among its heads, and its output width must match the language input. Config.WF collects these conditions; paperConfig_wf checks them for the published configuration.

/-- Kimi K3's `3 KDA : 1 MLA` schedule, with a final global-attention layer.

Layers are represented by zero-based `Fin` indices.  The report numbers layers from one, hence the
`layer.val + 1` conversion below.
-/
def attentionKindAt (cfg : TextConfig) (layer : Fin cfg.numLayers) : AttentionKind :=
  let paperIndex := layer.val + 1
  if paperIndex = cfg.numLayers || paperIndex % 4 = 0 then .mla else .kda

/-- The paper schedule contains 69 recurrent KDA layers. -/
theorem paperConfig_kda_layer_count : paperConfig.text.countAttentionKind .kda = 69 := by
  decide

/-- The paper schedule contains 24 global Gated MLA layers, including the final layer. -/
theorem paperConfig_mla_layer_count : paperConfig.text.countAttentionKind .mla = 24 := by
  decide

/-- At the published dimensions, MLA stores 576 scalars per token instead of 30,720. The integer
identity records the exact compression ratio `160 / 3 = 53 1/3` without introducing division. -/
theorem paperConfig_mla_cache_compression :
    paperConfig.text.mlaUncompressedScalarsPerToken * 3 =
      paperConfig.text.mlaCompressedScalarsPerToken * 160 := by
  decide

Full source

Two different consistency checks run. paperConfig_wf proves the arithmetic invariants used by the Lean model. validateFile reads a released config.json and rejects it when an architecture field disagrees with the hand-written specification. Agreement cannot by itself tell us which source is right; that is why the page also links the exact report and released configuration.

/-- Read and validate a released Kimi K3 Hugging Face configuration file. -/
def validateFile (path : System.FilePath) : IO Unit := do
  let json ← TorchLean.Json.readFile path
  match validateJson json with
  | .ok () => pure ()
  | .error message => throw <| IO.userError s!"{path}: {message}"

/-- The dimensions transcribed from the paper and released config are internally consistent. -/
theorem paperConfig_wf : paperConfig.WF := by
  constructor <;> decide

Full source

IO Unit means this function may read a file and either finish without a result or raise an error.
validateJson checks the dimensions, layer lists, quantization group size, and vision-to-text width against the released-checkpoint specification.
The theorem is different: it proves the in-memory paperConfig satisfies every field of Config.WF.

Checked here: the concrete schedule and cache-width arithmetic evaluate to the report's values.

Not checked here: the record does not allocate the checkpoint or derive 2.78T parameters from loaded weight tensors.

Config.lean, attention schedule ↗ Config.lean, released-config check ↗

§2.1.1, pp. 4–5 · Equation 1

KDA's fixed-size recurrent memory

For one head, KDA carries a matrix $S_t$ from one token to the next. Its rows correspond to key channels and its columns to value channels. That matrix is the memory. It has a fixed size even when the sequence keeps growing.

Equation 1 packs three actions into one line. Reading left to right matters: fade the old memory by α, erase the part aligned with the current key, then write the current value back along that same key. Only after the update does the query read the state.

(1)

S_t = (I - beta k k^T) Diag(alpha) S + beta k v^T

o-tilde_t = S_t^T q_t

01
Decay. Diag(alpha)S scales each key row. Small retention forgets that row quickly.
02
Correct. k^T Diag(alpha)S asks what value the old state already associates with this key.
03
Write and read. The delta rule adjusts that association toward v_t, with strength β_t; the query q_t then reads the updated matrix.

Let's reduce this to one key coordinate and one value coordinate. The memory is now one number s. Set the key and query to one. Decay first changes the old memory to αs; correction and writing then give:

The new memory adds a scaled prediction error to the decayed memory.

Suppose the old memory is 6, retention is 0.8, the current value is 10, and write strength is 0.5. The old prediction after decay is 4.8. Its error relative to the current value is 5.2, so the update adds half that error and leaves memory 7.4. A write strength of one would replace 4.8 with 10; a strength of zero would leave only the decay. These are small real-number examples of the equation, not measurements from the released model.

Token order still matters. Starting from zero with retention and write strength both 0.5, values 2 then 10 leave memory 5.25. Values 10 then 2 leave memory 2.25. Parallel implementations can regroup compositions of state updates, but cannot permute the tokens and expect the same result. The matrix version stores many key-to-value associations at once; the same correction principle applies along the current key.

Run a 2 × 2 delta updatesmall dimensions, same equation

retention α0.80 write β0.55

incoming S

1.000.20−0.300.70

→

decay

→

correct + write

Writing the KDA update as a tensor function

/-- Recurrent KDA state `S ∈ R^(d_k × d_v)`. -/
abbrev State (α : Type) [Storage α] (keyDim valueDim : Nat) :=
  Tensor α [keyDim, valueDim]

/--
One single-head KDA recurrence (Eq. 1):

`S_t = (I - β_t k_t k_tᵀ) Diag(α_t) S_(t-1) + β_t k_t v_tᵀ`.

The implementation follows the factorized form used by an efficient kernel: decay the rows of the
state, read the decayed state with the key, subtract the key-aligned correction, and finally add
the rank-one write. Every operation is an ordinary TorchLean tensor operation.
-/
def update (state : State α keyDim valueDim) (input : KDA.Input α keyDim valueDim) :
    State α keyDim valueDim :=
  let decayed := Tensor.stackLeading fun row => Tensor.stackLeading fun column => Tensor.scalar <|
    Tensor.getScalar input.retention row * get2 state row column
  let correction := vecMatMulSpec input.key decayed
  Tensor.addSpec
    (Tensor.subSpec decayed
      (Tensor.mulSpec
        (Tensor.full [keyDim, valueDim] input.writeStrength)
        (outerProductSpec input.key correction)))
    (Tensor.mulSpec
      (Tensor.full [keyDim, valueDim] input.writeStrength)
      (outerProductSpec input.key input.value))

Full source

/-- Coordinate form of the factorized KDA update.

The correction coordinate is the key's read from the already decayed state. This statement makes
the order of decay and delta correction explicit without duplicating the implementation.
-/
theorem update_get2 (state : State ℝ keyDim valueDim)
    (input : KDA.Input ℝ keyDim valueDim) (row : Fin keyDim) (column : Fin valueDim) :
  let decayed := Tensor.stackLeading fun source => Tensor.stackLeading fun target => Tensor.scalar
    (Tensor.getScalar input.retention source * get2 state source target)
  get2 (update state input) row column =
      Tensor.getScalar input.retention row * get2 state row column -
        input.writeStrength * Tensor.getScalar input.key row *
          Tensor.getScalar (vecMatMulSpec input.key decayed) column +
        input.writeStrength * Tensor.getScalar input.key row *
          Tensor.getScalar input.value column := by
  simp [update, Tensor.stackLeading, Tensor.addSpec, Tensor.subSpec, Tensor.mulSpec]
  ring

Full source

update_get2 proves that the function computes the paper's update at each row and column. To follow its proof, expand the tensor additions and products at that coordinate. The remaining scalar expression can be rearranged by ring. The row and column are arbitrary, so the result covers the whole matrix.

No numerical approximation appears here.

The recurrence is interpreted over an abstract scalar with the operations it needs; the coordinate theorem specializes to real numbers. Finite-precision kernel behavior is a separate claim.

Sequence/KDA.lean, state and update ↗

§2.1.1, pp. 4–5 · Equations 2–6 · Figure 3

Preparing the queries, keys, values, and decay

Before KDA updates memory, the token is turned into query, key, and value vectors. A short convolution mixes each vector with a few earlier tokens; a smooth activation is applied, and the queries and keys receive regularized L2 normalization. A separate smaller projection chooses how quickly each row of memory decays.

The regularizer is part of the equation. For a projected vector z, we divide each coordinate by . The result has norm strictly below one over the real numbers, rather than exactly one; a zero vector stays zero. The regularizer keeps its denominator positive even in that case. KDA uses 10⁻⁶ here, while its RMSNorm operations use a different stabilizer.

(2)

q, k projections

(5)

bounded decay

(6)

output gate

What the slider is showing

The sigmoid keeps the log-decay between g_min and zero. Exponentiating turns that into a retention strictly between exp(g_min) and one. Move the logit right to forget faster; move it left to keep more of the row.

Lower-bounded retention$g_{\min}=-5$

decay logit z 0.00

log-decay g−2.5000

retention α = exp(g)0.0821

proved intervale⁻⁵ < α < 1

The lower bound keeps the accumulated log-decay for a 16-token chunk above −80. This is why the report says the chunk remains in a numerically safe range. Lean proves the real-number bound, not the rounding behavior of the low-precision GPU kernel.

Proving the KDA graph matches the update

/-- Lower-bounded real retention from Eq. 5. -/
noncomputable def retentionReal (logFloor logScale logit : ℝ) : ℝ :=
  Real.exp (logFloor * Real.sigmoid (Real.exp logScale * logit))

/-- For Kimi K3, every one-step retention lies strictly between `exp(-5)` and `1`. -/
theorem paper_retention_bounds (logScale logit : ℝ) :
    Real.exp (-5) < retentionReal (-5) logScale logit ∧
      retentionReal (-5) logScale logit < 1 := by
  exact retentionReal_bounds (by norm_num) logScale logit

Full source

/-- Packed recurrent update used by the second half of a KDA layer.

This term contains only the state transition from Eq. 1 of the K3 report. Projection,
convolution, and output readout remain separate graph terms, so each equation has an independently
checkable refinement theorem. -/
def updateTerm {Γ : List Shape} (heads keyDim valueDim : Nat)
    (state : Term Γ [heads, keyDim, valueDim])
    (key : Term Γ [heads, keyDim])
    (value : Term Γ [heads, valueDim])
    (retention : Term Γ [heads, keyDim])
    (beta : Term Γ [heads]) :
    Term Γ [heads, keyDim, valueDim] :=
  let decayed := Term.op (PrimOp.mul [heads, keyDim, valueDim])
    (.cons (expandTerm [heads, keyDim] [valueDim] retention) (.cons state .nil))
  let correction := batchedVecMatTerm heads keyDim valueDim key decayed
  let correctionOuter := outerTerm [heads] keyDim valueDim key correction
  let writeOuter := outerTerm [heads] keyDim valueDim key value
  let scaledCorrection := Term.op (PrimOp.mul [heads, keyDim, valueDim])
    (.cons (expandTerm [heads] [keyDim, valueDim] beta) (.cons correctionOuter .nil))
  let scaledWrite := Term.op (PrimOp.mul [heads, keyDim, valueDim])
    (.cons (expandTerm [heads] [keyDim, valueDim] beta) (.cons writeOuter .nil))
  let corrected := Term.op
    (NN.GraphSpec.DAG.PrimOp.sub [heads, keyDim, valueDim])
    (.cons decayed (.cons scaledCorrection .nil))
  Term.op
    (NN.GraphSpec.DAG.PrimOp.add [heads, keyDim, valueDim])
    (.cons corrected (.cons scaledWrite .nil))

/-- `updateTerm` denotes one independent KDA recurrence update for every packed head. -/
theorem eval_updateTerm {Γ : List Shape} (env : TorchLean.TensorPack ℝ Γ)
    (heads keyDim valueDim : Nat)
    (state : Term Γ [heads, keyDim, valueDim])
    (query : Term Γ [heads, keyDim])
    (key : Term Γ [heads, keyDim])
    (value : Term Γ [heads, valueDim])
    (retention : Term Γ [heads, keyDim])
    (beta : Term Γ [heads]) :
    Term.eval env (updateTerm heads keyDim valueDim state key value retention beta) =
      Tensor.stackLeading fun head =>
        KDA.update (Spec.get (Term.eval env state) head)
          { query := Spec.get (Term.eval env query) head
            key := Spec.get (Term.eval env key) head
            value := Spec.get (Term.eval env value) head
            retention := Spec.get (Term.eval env retention) head
            writeStrength := (Spec.get (Term.eval env beta) head).item }

Read the proof

by
  have decayed :
      Tensor.mulSpec
        (Tensor.mapLeading [heads, keyDim] (fun x => Tensor.full [valueDim] x.item)
          (Term.eval env retention)) (Term.eval env state) =
        Tensor.stackLeading (fun head => Tensor.stackLeading (fun row => Tensor.stackLeading (fun column =>
          Tensor.scalar (((Term.eval env retention).unstack head).getScalar row *
            Spec.get2 ((Term.eval env state).unstack head) row column)))) := by
    apply TorchLean.Tensor.Internal.Rep.ext
    rintro ⟨head, row, column, ⟨⟩⟩
    simp [Tensor.mapLeading, Tensor.mulSpec, Tensor.map2Spec, Tensor.dim, Tensor.full,
      Tensor.getScalar, Tensor.item, Tensor.unstack, Spec.get2, Spec.get]
  simp only [updateTerm, Term.eval_op, Term.evalArgs,
    eval_expandTerm, PrimOp.mul_specFwd,
    eval_batchedVecMatTerm,
    eval_outerTerm,
    NN.GraphSpec.DAG.PrimOp.sub_specFwd, NN.GraphSpec.DAG.PrimOp.add_specFwd]
  rw [decayed]
  simp [KDA.update, Tensor.addSpec, Tensor.subSpec, Tensor.mulSpec,
    Tensor.zipEach, Tensor.mapLeading, Tensor.map2Spec, Tensor.dim, Tensor.full,
    Tensor.getScalar, Spec.get2, Spec.get, Tensor.item, Tensor.unstack]
  apply TorchLean.Tensor.Internal.Rep.ext
  rintro ⟨head, row, column, ⟨⟩⟩
  simp

Full source

/-- The complete packed KDA graph denotes the fixed-window mathematical layer step. -/
theorem layerModel_specFwd_eq_stepWindow
    {modelDim heads keyDim valueDim convWidth decayRank : Nat}
    (hModel : 0 < modelDim) (hHeads : 0 < heads) (hKey : 0 < keyDim)
    (hValue : 0 < valueDim) (hWidth : 0 < convWidth)
    (layer : KDA.Layer ℝ modelDim heads keyDim valueDim convWidth decayRank)
    (window : Tensor ℝ [convWidth, modelDim])
    (state : KDA.Layer.State ℝ heads keyDim valueDim) (logFloor : ℝ) :
    (layerModel modelDim heads keyDim valueDim convWidth decayRank
      hModel hHeads hKey hValue hWidth).specFwd (layerParameters layer)
        (layerInputs window state logFloor Normalize.l2Epsilon) =
      let result := layer.stepWindow hWidth logFloor window state
      .cons result.1 (.cons result.2 .nil)

Read the proof

by
  simp only [NN.GraphSpec.DAG.MultiModel.specFwd, layerModel]
  rw [Block.eval_andThen]
  rw [eval_prepareBlock_eq_preparedTensors hHeads hKey hValue hWidth layer window state logFloor]
  simp only [preparedTensors]
  rw [eval_updateReadoutBlock_eq_stepPrepared]
  simp [KDA.Layer.stepWindow]

Full source

The real-valued bound and the graph theorem answer different questions. The first pins down the paper's decay interval. The second checks that projection, window preparation, recurrence, and readout are wired into the same fixed-window token step.

reference stepKDA.Layer.stepWindow

proved equal

typed graph denotationKDA.layerModel.specFwd

separate refinement needed

kernelfused finite-precision KDA

GraphSpec/KDA.lean, graph denotation ↗

§2.1.2, pp. 5–6 · Equation 7

MLA compresses the attention cache

DeepSeek-V2 introduced MLA. At an MLA layer, a token still gets global attention. The saving comes from what survives in the cache. Instead of storing a separate key and value for all 96 heads, the model keeps one 512-wide key/value latent and one shared 64-wide key.

When a head needs the old token, it reconstructs that head's 128-wide content key and 128-wide value from the shared latent. The reconstruction costs computation, but the cache grows by 576 scalars per token. A fully expanded per-head layout would store 30,720. Drag the token count below and watch the difference accumulate linearly.

The smaller number is 512 latent coordinates plus 64 shared-key coordinates. All heads can read that stored entry through their learned reconstruction maps. The larger baseline stores a content key, shared-key coordinates, and value for each of 96 heads: 96 × (128 + 64 + 128) = 30,720. This comparison counts stored scalar coordinates; converting it to bytes also requires the representation used for each cache.

That expanded baseline repeats the shared-key coordinates for every head. A layout that retained those 64 coordinates only once would use 96 × (128 + 128) + 64 = 24,640 scalars instead. The 53⅓ ratio below compares the specific fully expanded layout with the 576-coordinate latent cache; it is not a claim about every possible cache layout or an inference-speed ratio.

MLA still keeps an entry for every token. If a prompt doubles in length, this cache doubles too. KDA's recurrent matrix stays the same size, but no longer provides a separate stored entry that a query can inspect for every old token. K3 retains both mechanisms so the sequence calculation can use recurrent memory at most layers and global attention at selected layers.

Fully expanded per-head layout96 heads × (128 key + 64 shared + 128 value)30,720 scalars / token

÷ 53⅓

Gated MLA cache512 KV latent + 64 shared key576 scalars / token

Cache growthone compressed entry per token per MLA layer

tokens 128

uncompressed3,932,160scalars

MLA73,728scalars

saved3,858,432scalars

From causal MLA to a packed-cache graph

/-- One causal MLA token step.  The current latent is inserted before attention is evaluated. -/
def step
    (layer :
      GatedMLA α modelDim heads queryLatentDim kvLatentDim contentKeyDim sharedKeyDim valueDim)
    (cache : Cache α kvLatentDim sharedKeyDim) (x : Tensor α [modelDim]) :
    Cache α kvLatentDim sharedKeyDim × Tensor α [modelDim] :=
  let queryLatent := RMSNorm.scale (vecMatMulSpec x layer.queryDown) layer.queryNormScale
  let currentKV := RMSNorm.scale (vecMatMulSpec x layer.kvDown) layer.kvNormScale
  let current : Entry α kvLatentDim sharedKeyDim :=
    { latent := currentKV, sharedKey := vecMatMulSpec x layer.sharedKeyDown }
  let cache' := cache ++ [current]
  let headOutput := fun h => layer.attend h queryLatent cache'
  let gate := fun h channel => Activation.Math.sigmoidSpec <|
    (List.finRange modelDim).foldl (fun total input =>
      total + Tensor.getScalar x input * get2 (Spec.get layer.gateWeight input) h channel) 0
  let output := Tensor.stackLeading (fun feature => Tensor.scalar <|
    (List.finRange heads).foldl (fun acrossHeads h =>
      (List.finRange valueDim).foldl (fun total channel =>
        total + gate h channel * Tensor.getScalar (headOutput h) channel *
          get2 (Spec.get layer.outputWeight h) channel feature) acrossHeads) 0)
  (cache', output)

/-- One MLA step appends exactly one compressed latent to the streaming cache. -/
theorem step_cache_length
    (layer :
      GatedMLA α modelDim heads queryLatentDim kvLatentDim contentKeyDim sharedKeyDim valueDim)
    (cache : Cache α kvLatentDim sharedKeyDim) (x : Tensor α [modelDim]) :
    (layer.step cache x).1.length = cache.length + 1 := by
  simp [step]

Full source

The graph version appends the same two current-token representations with typed concatAxis nodes, then exposes score addition, scaling, softmax, and value retrieval as ordinary operations.

Exact excerpt from the body of stepModel; the surrounding input terms and shape abbreviations are defined immediately above it in the source.

/-- The complete fixed-context Gated MLA token step as a typed, multi-output DAG. -/
def stepModel (pastTokens modelDim heads queryLatentDim kvLatentDim contentKeyDim
    sharedKeyDim valueDim : Nat) (hQueryLatent : 0 < queryLatentDim)
    (hKVLatent : 0 < kvLatentDim) :
    NN.GraphSpec.DAG.MultiModel
      (LayerParams modelDim heads queryLatentDim kvLatentDim contentKeyDim sharedKeyDim valueDim)
      (StepInputs pastTokens modelDim kvLatentDim sharedKeyDim)
      (StepOutputs pastTokens modelDim kvLatentDim sharedKeyDim) :=
  let params :=
    LayerParams modelDim heads queryLatentDim kvLatentDim contentKeyDim sharedKeyDim valueDim
  let inputs := StepInputs pastTokens modelDim kvLatentDim sharedKeyDim
  let Γ := params ++ inputs
  let parameterTerms : Args Γ params :=
    Args.rename (Var.inLeft inputs) (Args.vars params)
  let inputTerms : Args Γ inputs :=
    Args.rename (Var.inRight params) (Args.vars inputs)
  let queryDown := Args.get parameterTerms .head
  let queryNormScale := Args.get parameterTerms (.tail .head)
  let kvDown := Args.get parameterTerms (.tail (.tail .head))
  let kvNormScale := Args.get parameterTerms (.tail (.tail (.tail .head)))
  let sharedKeyDown := Args.get parameterTerms (.tail (.tail (.tail (.tail .head))))
  let queryContentUp := Args.get parameterTerms
    (.tail (.tail (.tail (.tail (.tail .head)))))
  let querySharedUp := Args.get parameterTerms
    (.tail (.tail (.tail (.tail (.tail (.tail .head))))))
  let keyUp := Args.get parameterTerms
    (.tail (.tail (.tail (.tail (.tail (.tail (.tail .head)))))))
  let valueUp := Args.get parameterTerms
    (.tail (.tail (.tail (.tail (.tail (.tail (.tail (.tail .head))))))))
  let gateWeight := Args.get parameterTerms
    (.tail (.tail (.tail (.tail (.tail (.tail (.tail (.tail (.tail .head)))))))))
  let outputWeight := Args.get parameterTerms
    (.tail (.tail (.tail (.tail (.tail (.tail (.tail (.tail (.tail (.tail .head))))))))))
  let pastLatentCache := Args.get inputTerms .head
  let pastSharedKeyCache := Args.get inputTerms (.tail .head)
  let x := Args.get inputTerms (.tail (.tail .head))
  let scoreScale := Args.get inputTerms (.tail (.tail (.tail .head)))
  let queryProjected : Term Γ [queryLatentDim] :=
    Term.op (PrimOp.broadcastVecMat [] [] [] modelDim queryLatentDim
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
      (.cons x (.cons queryDown .nil))
  let queryLatent : Term Γ [queryLatentDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.rmsNorm [] queryLatentDim hQueryLatent)
      (.cons queryProjected (.cons queryNormScale .nil))
  let kvProjected : Term Γ [kvLatentDim] :=
    Term.op (PrimOp.broadcastVecMat [] [] [] modelDim kvLatentDim
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
      (.cons x (.cons kvDown .nil))
  let currentKV : Term Γ [kvLatentDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.rmsNorm [] kvLatentDim hKVLatent)
      (.cons kvProjected (.cons kvNormScale .nil))
  let currentShared : Term Γ [sharedKeyDim] :=
    Term.op (PrimOp.broadcastVecMat [] [] [] modelDim sharedKeyDim
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
      (.cons x (.cons sharedKeyDown .nil))
  let currentKVRow : Term Γ [1, kvLatentDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons currentKV .nil)
  let currentSharedRow : Term Γ [1, sharedKeyDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons currentShared .nil)
  let latentCache : Term Γ [pastTokens + 1, kvLatentDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.concatAxis
      [pastTokens, kvLatentDim] 0 pastTokens 1)
      (.cons pastLatentCache (.cons currentKVRow .nil))
  let sharedKeyCache : Term Γ [pastTokens + 1, sharedKeyDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.concatAxis
      [pastTokens, sharedKeyDim] 0 pastTokens 1)
      (.cons pastSharedKeyCache (.cons currentSharedRow .nil))
  let queryBatchSource : Term Γ [1, 1, queryLatentDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons queryLatent .nil)
  let queryBatch : Term Γ [heads, 1, queryLatentDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.broadcast
      (Shape.CanBroadcastTo.dim_1_to_n
        (Shape.CanBroadcastTo.dim_eq
          (Shape.CanBroadcastTo.dim_eq
            (Shape.CanBroadcastTo.scalarTo []))))) (.cons queryBatchSource .nil)
  let latentBatchSource :
      Term Γ [1, pastTokens + 1, kvLatentDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons latentCache .nil)
  let latentBatch :
      Term Γ [heads, pastTokens + 1, kvLatentDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.broadcast
      (Shape.CanBroadcastTo.dim_1_to_n
        (Shape.CanBroadcastTo.dim_eq
          (Shape.CanBroadcastTo.dim_eq
            (Shape.CanBroadcastTo.scalarTo []))))) (.cons latentBatchSource .nil)
  let sharedBatchSource :
      Term Γ [1, pastTokens + 1, sharedKeyDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons sharedKeyCache .nil)
  let sharedBatch :
      Term Γ [heads, pastTokens + 1, sharedKeyDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.broadcast
      (Shape.CanBroadcastTo.dim_1_to_n
        (Shape.CanBroadcastTo.dim_eq
          (Shape.CanBroadcastTo.dim_eq
            (Shape.CanBroadcastTo.scalarTo []))))) (.cons sharedBatchSource .nil)
  let queryContent : Term Γ [heads, 1, contentKeyDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.matmul [heads] [heads] [heads]
      1 queryLatentDim contentKeyDim (.refl [heads]) (.refl [heads]))
      (.cons queryBatch (.cons queryContentUp .nil))
  let queryShared : Term Γ [heads, 1, sharedKeyDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.matmul [heads] [heads] [heads]
      1 queryLatentDim sharedKeyDim (.refl [heads]) (.refl [heads]))
      (.cons queryBatch (.cons querySharedUp .nil))
  let keys :
      Term Γ [heads, pastTokens + 1, contentKeyDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.matmul [heads] [heads] [heads]
      (pastTokens + 1) kvLatentDim contentKeyDim (.refl [heads]) (.refl [heads]))
      (.cons latentBatch (.cons keyUp .nil))
  let values : Term Γ [heads, pastTokens + 1, valueDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.matmul [heads] [heads] [heads]
      (pastTokens + 1) kvLatentDim valueDim (.refl [heads]) (.refl [heads]))
      (.cons latentBatch (.cons valueUp .nil))
  let keysTranspose :
      Term Γ [heads, contentKeyDim, pastTokens + 1] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.swapAdjacentAtDepth
      [heads, pastTokens + 1, contentKeyDim] 1)
      (.cons keys .nil)
  let sharedTranspose :
      Term Γ [heads, sharedKeyDim, pastTokens + 1] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.swapAdjacentAtDepth
      [heads, pastTokens + 1, sharedKeyDim] 1)
      (.cons sharedBatch .nil)
  let contentScores : Term Γ [heads, 1, pastTokens + 1] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.matmul [heads] [heads] [heads]
      1 contentKeyDim (pastTokens + 1) (.refl [heads]) (.refl [heads]))
      (.cons queryContent (.cons keysTranspose .nil))
  let sharedScores : Term Γ [heads, 1, pastTokens + 1] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.matmul [heads] [heads] [heads]
      1 sharedKeyDim (pastTokens + 1) (.refl [heads]) (.refl [heads]))
      (.cons queryShared (.cons sharedTranspose .nil))
  let scoreShape : Shape := [heads, 1, pastTokens + 1]
  let scores : Term Γ scoreShape :=
    Term.op (NN.GraphSpec.DAG.PrimOp.add scoreShape)
      (.cons contentScores (.cons sharedScores .nil))
  let scaledScores : Term Γ scoreShape :=
    scaleTerm scoreShape scoreScale scores
  let weights : Term Γ scoreShape :=
    Term.op (NN.GraphSpec.DAG.PrimOp.softmax scoreShape 2) (.cons scaledScores .nil)
  let headOutput3 : Term Γ [heads, 1, valueDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.matmul [heads] [heads] [heads]
      1 (pastTokens + 1) valueDim (.refl [heads]) (.refl [heads]))
      (.cons weights (.cons values .nil))
  let headOutput : Term Γ [heads, valueDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons headOutput3 .nil)
  let gateMatrix : Term Γ [modelDim, heads * valueDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons gateWeight .nil)
  let gateFlat : Term Γ [heads * valueDim] :=
    Term.op (PrimOp.broadcastVecMat [] [] [] modelDim (heads * valueDim)
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
      (.cons x (.cons gateMatrix .nil))
  let gateUnactivated : Term Γ [heads, valueDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons gateFlat .nil)
  let gate : Term Γ [heads, valueDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.sigmoid [heads, valueDim])
      (.cons gateUnactivated .nil)
  let gatedHeads : Term Γ [heads, valueDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.mul [heads, valueDim])
      (.cons gate (.cons headOutput .nil))
  let gatedFlat : Term Γ [heads * valueDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size]))
      (.cons gatedHeads .nil)
  let outputMatrix : Term Γ [heads * valueDim, modelDim] :=
    Term.op (NN.GraphSpec.DAG.PrimOp.reshape _ _ (by simp [Shape.size, Nat.mul_assoc]))
      (.cons outputWeight .nil)
  let output : Term Γ [modelDim] :=
    Term.op (PrimOp.broadcastVecMat [] [] [] (heads * valueDim) modelDim
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
      (.cons gatedFlat (.cons outputMatrix .nil))
  { initParams :=
      initialLayerParams modelDim heads queryLatentDim kvLatentDim contentKeyDim sharedKeyDim
        valueDim
    body := Block.ret (.cons latentCache (.cons sharedKeyCache (.cons output .nil))) }

Full source

/-- The graph returns the same extended latent cache, extended shared-key cache, and output
as the reference token step. Both paths append the current token before attending. -/
theorem stepModel_specFwd_eq_stepFixed
    {pastTokens modelDim heads queryLatentDim kvLatentDim contentKeyDim sharedKeyDim valueDim : Nat}
    (hQueryLatent : 0 < queryLatentDim) (hKVLatent : 0 < kvLatentDim)
    (layer :
      GatedMLA ℝ modelDim heads queryLatentDim kvLatentDim contentKeyDim sharedKeyDim valueDim)
    (pastLatentCache : Tensor ℝ [pastTokens, kvLatentDim])
    (pastSharedKeyCache : Tensor ℝ [pastTokens, sharedKeyDim])
    (x : Tensor ℝ [modelDim]) (scoreScale : ℝ) :
    stepGraphOutputs pastTokens modelDim heads queryLatentDim kvLatentDim contentKeyDim
      sharedKeyDim valueDim hQueryLatent hKVLatent layer pastLatentCache pastSharedKeyCache x
        scoreScale =
      stepFixedOutputs pastTokens modelDim heads queryLatentDim kvLatentDim contentKeyDim
        sharedKeyDim valueDim layer pastLatentCache pastSharedKeyCache x scoreScale

Read the proof

by
  -- Interpret the packed graph operations as the projections, cache appends, and attention
  -- in the reference step. The three output tensors must agree separately.
  simp only [stepGraphOutputs, stepFixedOutputs, NN.GraphSpec.DAG.MultiModel.specFwd,
    stepModel, layerParameters, stepInputs,
    Args.vars, Args.weakenLeft, Args.rename,
    Term.rename, Block.eval, Term.evalArgs, Term.eval,
    TorchLean.TensorPack.append, NN.GraphSpec.DAG.PrimOp.concatAxis,
    NN.GraphSpec.DAG.PrimOp.concatAxisSpec, Shape.replaceAxis,
    NN.GraphSpec.DAG.PrimOp.reshape_specFwd,
    NN.GraphSpec.DAG.PrimOp.rmsNorm_specFwd, NN.GraphSpec.DAG.PrimOp.rmsNormSpec,
    NN.GraphSpec.DAG.PrimOp.matmul_specFwd,
    NN.GraphSpec.DAG.PrimOp.swapAdjacentAtDepth_specFwd,
    NN.GraphSpec.DAG.PrimOp.add_specFwd, eval_scaleTerm,
    NN.GraphSpec.DAG.PrimOp.sigmoid_specFwd, NN.GraphSpec.DAG.PrimOp.mul_specFwd,
    PrimOp.broadcastVecMat, NN.GraphSpec.DAG.PrimOp.broadcast,
    NN.GraphSpec.DAG.PrimOp.softmax, GatedMLA.stepFixed,
    GraphSpec.rmsNormVectorSpec_eq_scale]
  simp [Tensor.permuteByAdjacentSwaps]
  constructor
  · with_unfolding_all rfl
  · constructor
    · with_unfolding_all rfl
    · with_unfolding_all rfl

Full source

The list cache in GatedMLA.step is convenient for causal recursion. The two leading-axis tensors in stepFixed give GraphSpec a static input and output shape. The theorem above proves the typed graph equals stepFixed.

The representation bridge is still open.

There is not yet a theorem proving that packing the list cache and calling stepFixed agrees with GatedMLA.step. The graph theorem also stops before the report's 32-bit output tile and fused training kernel.

GraphSpec/MLA.lean, typed cache step ↗ GraphSpec/MLA.lean, denotation theorem ↗

§2.2, p. 6 · Equations 8–10

AttnRes: attention across depth

A normal residual stream gives layer ell one running sum of everything below it. AttnRes keeps several depth sources visible. The current layer scores those sources, normalizes the scores, and retrieves a weighted combination.

K3 does not retain all 93 layer outputs separately. Every 12 layers are folded into a completed block sum, while the unfinished block remains separate. The embedding is always source zero. This is why the memory grows by blocks rather than by layers.

(9)

depth attention weights

depth retrieval

Try layer 12, then layer 13

Processing the twelfth member commits the completed sum immediately and resets the partial size to zero. Layer thirteen starts the next partial block. At every stored state, the partial size stays strictly below twelve.

Walk through the block stateblock size 12

layers completed 17

After all 93 layers, the final retrieval can see nine depth sources when the original embedding is counted. Seven complete groups of twelve account for 84 layers. Nine layers remain in the partial group. Together with the embedding, that gives seven completed sums, one partial sum, and one embedding source. Completing a group changes which tensor is stored; it must not duplicate the completed layers in both the full and partial lists.

The query chooses how to combine these sources. If all source scores are equal, softmax gives them equal weight. A different query can favor one source over another. Normalization is used to compute the scores, while the weighted sum uses the original source vectors. Normalizing those values as well would change the retrieved vector, which is why the distinction appears in both the function and the graph proof.

The AttnRes graph

The first definition is the mathematical reference. It computes tensors. Read its local names in order: normalize the stored sources, score them, exponentiate the scores, divide by their sum, and use the resulting weights to mix the original sources.

/-- Fixed-shape form of depth attention used at graph and backend boundaries.

The first axis enumerates the visible depth representations. Rows are RMS-normalized, scored
against the query, exponentiated, and normalized by their total mass. The resulting weights form a
weighted vector--matrix product with the unnormalized sources, exactly as in Eq. 9--10.
-/
def attendPacked {sources : Nat} (hModel : 0 < modelDim)
    (query : Tensor α [modelDim])
    (values : Tensor α [sources, modelDim]) :
    Tensor α [modelDim] :=
  let unitScale := Tensor.ones [modelDim]
  let normalized := Tensor.stackLeading fun row =>
    RMSNorm.scalePositive hModel (Spec.get values row) unitScale
  let scores := vecMatMulSpec query (Tensor.swapAdjacentAxes normalized 0)
  let kernels := Tensor.expSpec scores
  let denominator := Tensor.sumSpec kernels
  let weights := Tensor.map (fun weight => weight / denominator) kernels
  vecMatMulSpec weights values

Full source

The next definition builds that calculation from graph terms. Its nodes compute the normalization, exponentials, sum, and weighted values separately, using the Term interface we introduced earlier.

/-- AttnRes retrieval as a reusable term in an arbitrary surrounding graph context. -/
def term {Γ : List Shape} (sources modelDim : Nat)
    (hModel : 0 < modelDim) (query : Term Γ [modelDim])
    (values : Term Γ [sources, modelDim]) :
    Term Γ [modelDim] :=
  let unitScale := Term.op (NN.GraphSpec.DAG.PrimOp.one [modelDim]) .nil
  let normalized := Term.op
    (NN.GraphSpec.DAG.PrimOp.rmsNorm [sources] modelDim hModel)
    (.cons values (.cons unitScale .nil))
  let transposed := Term.op (NN.GraphSpec.DAG.PrimOp.swapAdjacentAtDepth
      [sources, modelDim] 0)
    (.cons normalized .nil)
  let scores := Term.op (PrimOp.broadcastVecMat [] [] [] modelDim sources
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
    (.cons query (.cons transposed .nil))
  let kernels := Term.op (NN.GraphSpec.DAG.PrimOp.exp [sources])
    (.cons scores .nil)
  -- Bind the exponentials once: the denominator and weighted values use the same scores.
  Term.let1 kernels <|
    let boundKernels : Term (Γ ++ [[sources]]) [sources] :=
      Term.var (Var.last Γ)
    let denominator := Term.op (NN.GraphSpec.DAG.PrimOp.sum [sources])
      (.cons boundKernels .nil)
    let inverse := Term.op (NN.GraphSpec.DAG.PrimOp.inv []) (.cons denominator .nil)
    let weights := scaleTerm [sources] inverse boundKernels
    Term.op (PrimOp.broadcastVecMat [] [] [] sources modelDim
      (Shape.CanBroadcastTo.refl []) (Shape.CanBroadcastTo.refl []))
      (.cons weights (.cons (Term.weakenRight values) .nil))

Full source

/-- The reusable graph term has exactly the fixed-shape AttnRes semantics. -/
theorem eval_term {Γ : List Shape} (env : TorchLean.TensorPack ℝ Γ)
    (sources modelDim : Nat) (hModel : 0 < modelDim)
    (query : Term Γ [modelDim])
    (values : Term Γ [sources, modelDim]) :
    Term.eval env (term sources modelDim hModel query values) =
      KimiK3.AttnRes.attendPacked hModel (Term.eval env query) (Term.eval env values)

Read the proof

by
  -- The let-bound exponentials extend the environment; the original values are weakened into it.
  unfold term
  rw [eval_let1]
  simp only [Term.eval_op, Term.evalArgs, eval_scaleTerm,
    Term.eval_var_last_append, Term.eval_weakenRight]
  simp only [NN.GraphSpec.DAG.PrimOp.rmsNorm_specFwd,
    NN.GraphSpec.DAG.PrimOp.swapAdjacentAtDepth_specFwd,
    NN.GraphSpec.DAG.PrimOp.inv_specFwd,
    PrimOp.broadcastVecMat, NN.GraphSpec.DAG.PrimOp.exp,
    NN.GraphSpec.DAG.PrimOp.sum, NN.GraphSpec.DAG.PrimOp.one]
  -- Match normalized scoring and the weighted sum of the original, unnormalized values.
  simp [KimiK3.AttnRes.attendPacked, NN.GraphSpec.DAG.PrimOp.rmsNormSpec,
    NN.GraphSpec.DAG.PrimOp.Internal.rmsNormVectorSpec, RMSNorm.scalePositive,
    PrimOp.Internal.mapFinalAxis, Tensor.mapLeading, Tensor.stackLeading, Tensor.ones,
    Tensor.invSpec, div_eq_mul_inv, mul_comm, Spec.get]

Full source

Term.eval env replaces each external wire with the tensor supplied by env, then gives every graph node its mathematical meaning.
The equality says the evaluated graph returns the same tensor as attendPacked, for every real input of the stated shapes.
Separately, finishLayer_wf proves that committing a layer either increments the partial depth block or resets it at the block boundary.

The packed/list bridge is not proved yet.

The graph theorem targets attendPacked. The public BlockState.retrieve uses list-valued sources, and no theorem currently identifies those two representations.

GraphSpec/AttnRes.lean, retrieval term ↗ GraphSpec/AttnRes.lean, evaluation theorem ↗

§2.3, pp. 6–8 · Equations 11–13

Each token selects 16 of 896 routed experts

Shared experts operate at the full 7,168-wide model dimension. Routed experts operate at the 3,584-wide latent dimension. A token is projected down once, processed by its 16 selected experts, normalized, and projected back up.

tokenx ∈ ℝ⁷¹⁶⁸

shared pathrouted path

2 shared expertsfull model width

W↓ → 16 of 896 → RMSNorm → W↑latent width 3,584

sumy ∈ ℝ⁷¹⁶⁸

(11)

routed latent sum

shared and routed output

A
The shared experts see the original 7,168 coordinates and run for every token.
B
The routed branch projects once to 3,584 coordinates, then sends the same latent vector to 16 distinct experts.
C
The selected outputs are weighted, then RMSNorm rescales their combined vector by its root-mean-square size. A final projection restores 7,168 coordinates and adds the shared path.

The router first passes each raw score through sigmoid, a smooth function whose output lies between zero and one. Top-k then means choosing the 16 largest adjusted scores. The selection bias can change which experts enter that set. After selection, the weighted sum uses the original sigmoid scores, so the bias does not directly change how strongly a chosen expert contributes.

For a smaller example, imagine four experts with sigmoid scores 0.8, 0.7, 0.6, and 0.2, and select two. Without a bias, the first two win. Add a selection bias of 0.3 to the third expert, and its adjusted score becomes 0.9. The selected pair is now the third and first experts. Their mixture weights still use 0.6 and 0.8, normalized by their sum: 3/7 and 4/7. Using 0.9 in that mixture would implement a different rule.

The shared experts have no top-k decision: they process every token at model width. The routed experts all receive the same down-projected latent input, then their weighted outputs are combined before normalization and the up-projection. Sparse selection saves those unselected expert computations; it does not remove their weights from the checkpoint.

Routing scores and the selected experts

/--
Stable LatentMoE forward pass (Eq. 11): shared experts operate at model width, while the selected
routed experts operate in latent width and are normalized before the up-projection.
-/
def forward
    (moe : StableLatentMoE α modelDim latentDim sharedHidden routedHidden numShared numRouted
      activeExperts)
    (route : Route numRouted activeExperts) (gateCap upCap : α)
    (x : Tensor α [modelDim]) : Tensor α [modelDim] :=
  let shared := moe.sharedOutput gateCap upCap x
  let routed := moe.routedAggregate route gateCap upCap x
  shared + vecMatMulSpec (RMSNorm.scale routed moe.routedNormScale) moe.upProject

/-- Route a real-valued token using Eq. 13's bias-adjusted scores. -/
noncomputable def route
    (moe : StableLatentMoE ℝ modelDim latentDim sharedHidden routedHidden numShared numRouted
      activeExperts)
    (hActive : activeExperts ≤ numRouted)
    (x : Tensor ℝ [modelDim]) : Route numRouted activeExperts :=
  Route.chooseTopK (fun expert => Tensor.getScalar (moe.adjustedRouterScores x) expert) hActive

/-- The route computed by a Stable LatentMoE layer is top-k for its adjusted router scores. -/
theorem route_isTopK
    (moe : StableLatentMoE ℝ modelDim latentDim sharedHidden routedHidden numShared numRouted
      activeExperts)
    (hActive : activeExperts ≤ numRouted)
    (x : Tensor ℝ [modelDim]) :
    (moe.route hActive x).IsTopK
      (fun expert => Tensor.getScalar (moe.adjustedRouterScores x) expert) :=
  Route.chooseTopK_isTopK _ hActive

Full source

/--
Specialize the sparse graph to the route selected by this model on this token.

The resulting DAG still contains the ordinary projections, activations, expert evaluations, and
mixture arithmetic. Only the discrete expert indices are fixed while the graph is assembled. This
is the appropriate GraphSpec boundary until the DAG language carries integer-valued intermediate
results: it does not conceal top-k inside an opaque floating-point operation, and it does not permit
an unrelated route to stand in for K3 routing.

Because routing depends on both `moe` and `input`, this graph must be rebuilt if either changes
enough to change the selected experts. Runtime implementations may instead compute and check a
route before executing `modelGivenRoute`; the theorem below identifies the exact route that such a
checker must produce.
-/
noncomputable def modelAtInput
    {modelDim latentDim sharedHidden routedHidden numShared numRouted activeExperts : Nat}
    (hActive : activeExperts ≤ numRouted) (hLatent : 0 < latentDim)
    (moe : KimiK3.StableLatentMoE ℝ modelDim latentDim sharedHidden routedHidden numShared
      numRouted activeExperts)
    (input : Tensor ℝ [modelDim]) :
    NN.GraphSpec.DAG.Model
      (Params modelDim latentDim sharedHidden routedHidden numShared numRouted)
      (Inputs modelDim) [modelDim] :=
  modelGivenRoute modelDim latentDim sharedHidden routedHidden numShared numRouted activeExperts
    (moe.route hActive input) hLatent

/--
The route-specialized graph denotes K3's real-valued Stable LatentMoE semantics, including the
deterministic score-based routing decision. This strengthens `modelGivenRoute_specFwd_eq_forward`:
the right-hand side is `forwardReal`, whose route is derived from the adjusted router scores rather
than supplied by the caller.
-/
theorem modelAtInput_specFwd_eq_forwardReal
    {modelDim latentDim sharedHidden routedHidden numShared numRouted activeExperts : Nat}
    (hActive : activeExperts ≤ numRouted) (hLatent : 0 < latentDim)
    (hActivePos : 0 < activeExperts)
    (moe : KimiK3.StableLatentMoE ℝ modelDim latentDim sharedHidden routedHidden numShared
      numRouted activeExperts)
    (input : Tensor ℝ [modelDim]) (gateCap upCap : ℝ) :
    (modelAtInput hActive hLatent moe input).specFwd
        (parameters moe) (inputs input gateCap upCap) =
      moe.forwardReal hActive gateCap upCap input := by
  exact modelGivenRoute_specFwd_eq_forward hLatent hActivePos moe
    (moe.route hActive input) input gateCap upCap

Full source

Fin activeExperts fixes the number of slots, Fin numExperts keeps every ID in range, and injectivity forbids duplicates. The route-specialized graph then carries the projections, selected expert calls, score normalization, RMSNorm, and final sum.

Top-k is not a graph node.

modelAtInput chooses the top experts mathematically while constructing a graph for one concrete real input. A runtime implementation can compute and check that route before executing modelGivenRoute; if the selected experts change, the staged graph changes.

GraphSpec/MoE.lean, route-specialized graph ↗ FeedForward/MoE.lean, computed route is top-k ↗

§2.3.2, p. 7 · Appendix B, p. 43 · Equations 18–19

SiTU caps both linear branches

SwiGLU is the gated activation inside a feed-forward expert: one transformed branch decides how much of another transformed branch passes through. Two of its factors can grow with the input, so its output has no fixed ceiling. SiTU smoothly caps both factors while keeping the sigmoid gate. Near zero it still behaves like SwiGLU; at large magnitude it stays bounded. The two squashing functions make the bound possible: sigmoid stays between zero and one, while tanh stays between minus one and one.

(19)

SiTU definition

SiTU bound

Drag past 25. The purple SwiGLU curve keeps climbing; the blue SiTU curve bends toward its ceiling. The formal theorem below proves that ceiling for every real pair (g,u), including points away from the plotted diagonal.

Scalar responsedrag the preactivation

x 12.0

SiTU-GLUSwiGLU

Both branches use the same input, g = u = x. SwiGLU continues above the plotted range; SiTU approaches 100.

Proving the SiTU bound for every real input

/-- Smoothly cap a scalar at magnitude `cap` using `cap * tanh (x / cap)`. -/
def softCap {α : Type} [Context α] (cap x : α) : α :=
  cap * MathFunctions.tanh (x / cap)

/--
The SiTU-GLU scalar response.  The sigmoid is applied to the uncapped gate preactivation, while
the linear factors in both GLU branches are smoothly capped.
-/
def scalar {α : Type} [Context α] (gateCap upCap gate up : α) : α :=
  (softCap gateCap gate * Activation.Math.sigmoidSpec gate) * softCap upCap up

/--
SiTU-GLU has the paper's advertised uniform bound.  Unlike a numerical test, this statement holds
for every pair of real preactivations.
-/
theorem abs_scalar_le_mul_caps {gateCap upCap : ℝ}
    (hGate : 0 ≤ gateCap) (hUp : 0 ≤ upCap) (gate up : ℝ) :
    |scalar gateCap upCap gate up| ≤ gateCap * upCap

Read the proof

by
  have hTanhGate : |Real.tanh (gate / gateCap)| ≤ 1 := (Real.abs_tanh_lt_one _).le
  have hTanhUp : |Real.tanh (up / upCap)| ≤ 1 := (Real.abs_tanh_lt_one _).le
  have hSigmoid : |Real.sigmoid gate| ≤ 1 := by
    rw [abs_of_nonneg (Real.sigmoid_nonneg gate)]
    exact Real.sigmoid_le_one gate
  have hSoftGate : |softCap gateCap gate| ≤ gateCap := by
    rw [softCap, tanh_real, abs_mul, abs_of_nonneg hGate]
    exact mul_le_of_le_one_right hGate hTanhGate
  have hSoftUp : |softCap upCap up| ≤ upCap := by
    rw [softCap, tanh_real, abs_mul, abs_of_nonneg hUp]
    exact mul_le_of_le_one_right hUp hTanhUp
  rw [scalar, sigmoidSpec_real, abs_mul, abs_mul]
  calc
    |softCap gateCap gate| * |Real.sigmoid gate| * |softCap upCap up|
        ≤ gateCap * 1 * upCap := by gcongr
    _ = gateCap * upCap := by ring

/-- Kimi K3's caps imply the concrete coordinate bound `100`. -/
theorem paper_caps_bound (gate up : ℝ) : |scalar 4 25 gate up| ≤ 100 := by
  have h := abs_scalar_le_mul_caps (gateCap := (4 : ℝ)) (upCap := (25 : ℝ))
    (by positivity) (by positivity) gate up
  norm_num at h ⊢
  exact h

Full source

/-- Smoothly cap every coordinate of a tensor by a scalar graph input.

The returned term contains reciprocal, scalar-broadcast, multiplication, and `tanh` nodes computing
`cap * tanh(x / cap)`. Keeping the cap as an input records the numerical convention in the graph
ABI without hiding the calculation behind a K3-specific operation.
-/
def softCapTerm {Γ : List Shape} (shape : Shape)
    (cap : Term Γ []) (input : Term Γ shape) : Term Γ shape :=
  let inverse := Term.op (NN.GraphSpec.DAG.PrimOp.inv []) (.cons cap .nil)
  let inverseTensor := Term.op
    (NN.GraphSpec.DAG.PrimOp.broadcast (Shape.CanBroadcastTo.scalarTo shape))
    (.cons inverse .nil)
  let scaled := Term.op (NN.GraphSpec.DAG.PrimOp.mul shape)
    (.cons input (.cons inverseTensor .nil))
  let cappedUnit := Term.op (NN.GraphSpec.DAG.PrimOp.tanh shape) (.cons scaled .nil)
  let capTensor := Term.op
    (NN.GraphSpec.DAG.PrimOp.broadcast (Shape.CanBroadcastTo.scalarTo shape))
    (.cons cap .nil)
  Term.op (NN.GraphSpec.DAG.PrimOp.mul shape)
    (.cons cappedUnit (.cons capTensor .nil))

/-- Evaluating `softCapTerm` gives the coordinatewise cap `cap * tanh (input / cap)`. -/
@[simp] theorem eval_softCapTerm {α : Type} [Storage α] [Context α] {Γ : List Shape}
    (env : TorchLean.TensorPack α Γ) (shape : Shape)
    (cap : Term Γ []) (input : Term Γ shape) :
    Term.eval env (softCapTerm shape cap input) =
      Tensor.mulSpec
        (Activation.tanhSpec
          (Tensor.mulSpec (Term.eval env input)
            (Tensor.full shape (1 / (Term.eval env cap).item))))
        (Tensor.full shape (Term.eval env cap).item) := by
  simp [softCapTerm, Term.eval, Term.evalArgs,
    NN.GraphSpec.DAG.PrimOp.inv, NN.GraphSpec.DAG.PrimOp.broadcast,
    NN.GraphSpec.DAG.PrimOp.mul, NN.GraphSpec.DAG.PrimOp.tanh]

Full source

The proof combines |tanh x| ≤ 1, 0 ≤ sigmoid x ≤ 1, and nonnegative caps. The graph uses reciprocal, broadcast, multiply, tanh, sigmoid, and vector–matrix nodes; SiTU is not hidden inside a custom opaque primitive.

Coordinate-wise bound.

This bounds each scalar coordinate. A bound on a whole network activation norm would require the surrounding linear maps and their norms.

FeedForward/Activation.lean, SiTU bound ↗ GraphSpec, transparent SiTU graph ↗

§2.3.3, p. 8 · Appendices C–D, pp. 43–44

Balancing expert loads with score quantiles

K3 chooses the 16 experts with the highest adjusted scores for each token. For one expert, subtract each token’s 16th-highest cutoff from its raw score. That difference is the expert’s margin: adding a bias can move it above or below the selection threshold. A quantile is a cutoff in an ordered set of scores, eg the point above which a chosen number of margins lie. Quantile balancing chooses the bias to aim for a target load. We can check the resulting count when the cutoff is exact, characterize thresholds that minimize the coordinate objective, and bound histogram interpolation error within a known bin. The production histogram configuration and an exact-cutoff guarantee remain open.

routing margin

selection rule

quantile bias

target count

The two bars for each expert answer different questions. The raw-score bar shows the sigmoid score used in the mixture. The biased-score bar adds the balancing bias and decides whether the expert crosses the top-k cutoff.

Bias changes selection, not mixture weightfour experts, top-2

Counting the tokens above a routing threshold

/-- An exact target quantile has precisely `targetLoad` strict exceedances.

Values equal to the threshold are deliberately excluded, matching the routing convention in
Appendix C.
-/
def IsExactTargetQuantile {tokens : Nat} (targetLoad : Nat)
    (margins : Fin tokens → ℝ) (threshold : ℝ) : Prop :=
  exceedanceCount margins threshold = targetLoad

/-- Number of tokens routed to an expert after applying a candidate bias. -/
noncomputable def routedTokenCount {tokens : Nat}
    (rawScore cutoff : Fin tokens → ℝ) (bias : ℝ) : Nat :=
  (Finset.univ.filter fun token => cutoff token < rawScore token + bias).card

/-- Negating an exact margin quantile gives the expert exactly its target load. -/
theorem routedTokenCount_neg_eq {tokens targetLoad : Nat}
    (rawScore cutoff : Fin tokens → ℝ) (threshold : ℝ)
    (hQuantile : IsExactTargetQuantile targetLoad
      (fun token => rawScore token - cutoff token) threshold) :
    routedTokenCount rawScore cutoff (-threshold) = targetLoad := by
  rw [← hQuantile]
  apply congrArg Finset.card
  ext token
  simp only [Finset.mem_filter, Finset.mem_univ, true_and]
  constructor <;> intro h <;> linarith

Full source

Subtracting the threshold from a margin selects exactly the margins strictly above that threshold. IsExactTargetQuantile assumes there are targetLoad of them, and routedTokenCount_neg_eq carries that count through the bias update. The same threshold minimizes the finite coordinate dual objective by coordinateObjective_le_of_exact_quantile.

For a histogram approximation, shard counts pool exactly. The interpolation error is at most one bin width if the true quantile lies in the chosen bin. The missing step is showing that bin selection puts it there. Mean-centering has separate results, centerBias_score_ge_iff and centerBias_isTopK_iff.

There is a small but important tie case here. If two scores are equal and we want one selection, no strict threshold can select exactly one. The tied cutoff can still minimize the objective. Appendix C sets ties aside in practice, but we can prove minimization without excluding them: count the margins strictly above the threshold, then count those at or above it. If the target lies between those counts, coordinateObjective_le_of_counts proves that the threshold minimizes the objective. This does not make a strict threshold select only part of a tie.

The wording in Appendix C, Equation 24 is worth looking at directly. The authors say the objective is:

minimized exactly when k margins lie above α

That equivalence misses an endpoint. With distinct margins 1 and 0 and target load 1, every threshold from 0 to 1 minimizes the objective. At threshold 1 the strict count is zero. We checked this counterexample in Lean. The report's chosen threshold 0 still minimizes the objective. This is an error in the stated equivalence, not evidence that the model or balancing update fails. coordinateObjective_minimum_iff proves the corrected equivalence: a threshold minimizes the objective exactly when the strict and weak counts bracket the target. It covers ties and fractional target loads; it does not construct a threshold.

We can work through the endpoint example directly. For margins 1 and 0 and target load 1, the coordinate objective at threshold θ is

The two-margin coordinate objective.

Between zero and one, this reduces to θ + (1 − θ) = 1. Below zero it is 1 − θ, and above one it is θ, both larger than one. So the whole closed interval minimizes the objective. At its upper endpoint, neither margin is strictly above the threshold, but one is at least the threshold. The corrected strict-and-weak count condition includes this minimizer.

A tied example makes the same distinction unavoidable. If both margins are zero and the target is one, a strict threshold selects either both or neither. Nevertheless, zero minimizes the objective: the strict count is zero and the weak count is two, with the target between them. An actual router needs a tie-breaking rule to select one expert. Minimizing the threshold objective alone does not choose that rule.

Ties and one-sided slopes

We can understand the corrected condition by looking at the objective on either side of a threshold. For target load k and margins mᵢ, its general form is

The balancing objective is a sum of hinge terms and a linear term.

Between margin values, every active hinge has slope −1 and every inactive hinge has slope zero. The total slope is the target load minus the number of margins strictly above the threshold. At a margin value, the slope can jump: moving left activates the equal margins, while moving right leaves them inactive.

The two one-sided slopes count ties differently.

A convex objective is minimized where the left slope is at most zero and the right slope is at least zero. That gives strict count ≤ target ≤ weak count. This is the intuition behind the corrected theorem. The Lean proof works directly with finite hinge inequalities, so it doesn't assume an ordinary derivative exists at a tied margin.

Checking one chosen threshold would miss the endpoint issue: the report's choice still minimizes the objective in our example. We found the error by asking which thresholds minimize it, including the endpoints. The corrected condition gives us all of them. We still need an algorithm to find a threshold and a rule for selecting within a tie.

/-- The complete minimizer condition for the QB hinge objective, with ties and real target loads.

The objective has right slope target minus the strict count and left slope target minus the
weak count. A minimum requires the first to be nonnegative and the second nonpositive. The
finite-gap argument below tests these slopes without a differentiability assumption at ties.
-/
theorem coordinateObjective_minimum_iff {tokens : ℕ} (targetLoad : ℝ)
    (margins : Fin tokens → ℝ) (threshold : ℝ) :
    (∀ candidate, coordinateObjective targetLoad margins threshold ≤
      coordinateObjective targetLoad margins candidate) ↔
      (exceedanceCount margins threshold : ℝ) ≤ targetLoad ∧
        targetLoad ≤ (atLeastCount margins threshold : ℝ)

Read the proof

by
  constructor
  · intro hMin
    exact ⟨strict_count_le_of_minimum margins threshold hMin,
      target_le_weak_count_of_minimum margins threshold hMin⟩
  · rintro ⟨hStrict, hWeak⟩ candidate
    exact coordinateObjective_le_of_counts margins threshold hStrict hWeak candidate

Full source

FeedForward/Quantile.lean, coordinate minimum ↗ FeedForward/Histogram.lean, pooled estimator ↗

§2.4, pp. 9–10

MoonViT turns images and video into visual embeddings

MoonViT-V2 alternates spatial attention within each frame and temporal attention across frames at a fixed patch position. Temporal pooling collapses the frame axis. A 2 × 2 spatial merge trades four neighboring tokens for four times as many channels before the projector maps into the language width.

spatial merge preserves the number of scalars

four neighboring tokens become four channel groups

Switch from image to video and follow the third stage. With one frame the temporal attention axis has one position, but the block still applies its normalization, projections, output map, and residual. With several frames the same code attends along that axis at each spatial location.

The merge preserves the number of scalars by grouping four spatial tokens into one wider row. The language model sees fewer visual positions, each carrying more channels before projection.

Pooling frames and projecting visual tokens

mergeAndProject takes a tensor with frame, row, column, and feature axes. It averages over frames, groups neighboring patches into channels, and projects each group to the language model’s width. Its output has one visual-token axis and one feature axis.

/--
Temporal pooling, pixel-shuffle merge, and projection into the language width. The output sequence
has one token for each merged spatial location; the frame axis has been pooled away.
-/
def mergeAndProject
    (model : Model α cfg frames (rows * cfg.mergeHeight)
      (columns * cfg.mergeWidth) patchFeatures)
    (grid : Grid α frames (rows * cfg.mergeHeight) (columns * cfg.mergeWidth) cfg.hiddenDim)
    (hFrames : 0 < frames)
    (hText : 0 < cfg.textHiddenDim) :
    Tensor α [rows * columns, cfg.textHiddenDim] :=
  let pooled : Tensor α [rows * cfg.mergeHeight, columns * cfg.mergeWidth, cfg.hiddenDim] :=
    Tensor.reduceMean 0 grid (Shape.hasNonemptyAxisZeroOfPos hFrames).proof
  let interleavedShape : Shape :=
    [rows, cfg.mergeHeight, columns, cfg.mergeWidth, cfg.hiddenDim]
  let groupedShape : Shape :=
    [rows, columns, cfg.mergeHeight, cfg.mergeWidth, cfg.hiddenDim]
  let mergedDim := cfg.mergeHeight * cfg.mergeWidth * cfg.hiddenDim
  let mergedShape : Shape := [rows * columns, mergedDim]
  have hInterleaved :
      Shape.size [rows * cfg.mergeHeight, columns * cfg.mergeWidth, cfg.hiddenDim] =
        Shape.size interleavedShape := by
    simp [interleavedShape, Shape.size, Nat.mul_assoc]
  have hMerged : Shape.size groupedShape = Shape.size mergedShape := by
    simp [groupedShape, mergedShape, mergedDim, Shape.size, Nat.mul_assoc]
  let interleaved := Tensor.reshapeSpec pooled hInterleaved
  have hGrouped : interleavedShape.swapAdjacentAtDepth 1 = groupedShape := by
    simp [interleavedShape, groupedShape, Shape.swapAdjacentAtDepth]
  let grouped : Tensor α groupedShape :=
    hGrouped ▸ Tensor.swapAdjacentAxes interleaved 1
  let merged := Tensor.reshapeSpec grouped hMerged
  let hidden := Activation.geluSpec (Tensor.matmulSpec (Shape.CanBroadcastTo.refl [])
    (Shape.CanBroadcastTo.refl []) merged model.projector.firstWeight)
  let projected := Tensor.matmulSpec (Shape.CanBroadcastTo.refl [])
    (Shape.CanBroadcastTo.refl []) hidden model.projector.secondWeight
  RMSNorm.rows hText projected model.projector.outputNormScale

Full source

reduceMean 0 averages the frame axis, so image and video both become one spatial grid.
The two reshape proofs establish equal element counts. The axis swap changes row, merge-row, column, merge-column into groups of neighboring patches.
Two learned matrix multiplications with GELU between them project each merged patch into the text model's width; RMSNorm finishes the visual token.

The result shape checks pooling, merging, and projection. The axes are positional rather than semantically tagged, however: two axes of the same length are not distinguished as “frame” and “row” by separate Lean types.

Component graphs are checked.

GraphSpec/Vision.lean builds patch embedding, one divided-attention block, and merge/project entirely from typed graph primitives. Three denotation theorems identify their graph evaluations with the reference equations. A checkpoint-sized graph still requires the concrete list of 27 released block parameter tensors.

Report claim outside Lean.

The gradient-norm stability comparison in Figure 6 depends on measurements. This file specifies the architecture; it does not turn that plot into a theorem.

Vision.lean, merge and projection ↗ Vision.lean, complete reference path ↗ GraphSpec/Vision.lean, denotation proofs ↗

Figure 2, p. 3 · §§2.1–2.4

How the K3 language model fits together

A backbone layer follows the same order every time. It makes two depth queries. The first retrieves input for KDA or MLA. Its output joins the partial depth block. The second query retrieves input for the dense or sparse channel mixer. Only after both submodules run does the block counter advance.

1
Retrieve a sequence-mixer input from the current depth sources.
2
Normalize and mix sequence with KDA or Gated MLA.
3
Retrieve again with the sequence output included.
4
Mix channels with the dense FFN or routed LatentMoE.
5
Commit once to the AttnRes block state.

Text and visual positions remain distinct until embedding. A visual slot indexes MoonViT's projected sequence; a text slot indexes the vocabulary table. Both become 7,168-wide tensors.

Keeping the sequence cache and depth memory in one layer

/-- Advance one backbone layer at a single autoregressive token while preserving the sequence
mixer's causal state. This is the primitive used by the EAGLE draft and streaming decoders. -/
noncomputable def forwardToken (layer : BackboneLayer ℝ cfg decayRank)
    (hActive : cfg.activeExperts ≤ cfg.numRoutedExperts)
    (sequenceState : layer.sequence.CausalState)
    (depthState : AttnRes.BlockState ℝ cfg.hiddenDim) :
    layer.sequence.CausalState × AttnRes.BlockState ℝ cfg.hiddenDim :=
  let sequenceInput := RMSNorm.scale (depthState.retrieve layer.sequenceQuery)
    layer.sequenceNormScale
  let sequenceResult := layer.sequence.step sequenceState sequenceInput
  let channelInput := depthState.retrieveAfterSequence sequenceResult.2 layer.channelQuery
  let channelOutput := layer.channel.forward hActive
    (RMSNorm.scale channelInput layer.channelNormScale)
  (sequenceResult.1,
    depthState.finishLayer cfg.attnResBlockSize sequenceResult.2 channelOutput)

Full source

/-- A one-token backbone step preserves the AttnRes block-position invariant. -/
theorem forwardToken_depth_wf (layer : BackboneLayer ℝ cfg decayRank)
    (hActive : cfg.activeExperts ≤ cfg.numRoutedExperts)
    (sequenceState : layer.sequence.CausalState)
    (depthState : AttnRes.BlockState ℝ cfg.hiddenDim)
    (hDepth : depthState.WF cfg.attnResBlockSize) :
    (layer.forwardToken hActive sequenceState depthState).2.WF cfg.attnResBlockSize := by
  simp only [forwardToken]
  exact depthState.finishLayer_wf cfg.attnResBlockSize _ _ hDepth

Full source

The first retrieval produces the input to the sequence mixer. Depending on the layer, sequence.step is a KDA recurrence or an MLA cache update.
The second retrieval can see the new sequence output before the expert network runs.
finishLayer commits both outputs once. The theorem proves this step keeps the AttnRes partial-block counter inside its legal range.

The return type depends on the concrete mixer, so KDA and MLA keep their own causal-state types. The packed GraphSpec file separately proves the four KDA/MLA × dense/sparse layer cases equal their corresponding packed specStep.

Giving text and image tokens the same width

/-- Resolve a mixed text/visual sequence into language-width embeddings. The slot list is the
formal counterpart of expanding a serialized media placeholder into the visual tokens produced by
MoonViT. -/
def embedMultimodal {sequenceLength visualTokens : Nat}
    (model : LanguageModel α cfg decayRank)
    (visual : Tensor α [visualTokens, cfg.hiddenDim])
    (slots : Fin sequenceLength → InputSlot cfg.vocabSize visualTokens) :
    Tensor α [sequenceLength, cfg.hiddenDim] :=
  Tensor.stackLeading fun position =>
    match slots position with
    | .text token => Spec.get model.tokenEmbedding token
    | .visual token => Spec.get visual token

omit [Context α] in
/-- A visual slot is resolved from MoonViT's projected token sequence. -/
theorem get_embedMultimodal_visual {sequenceLength visualTokens : Nat}
    (model : LanguageModel α cfg decayRank)
    (visual : Tensor α [visualTokens, cfg.hiddenDim])
    (slots : Fin sequenceLength → InputSlot cfg.vocabSize visualTokens)
    (position : Fin sequenceLength) (token : Fin visualTokens)
    (hSlot : slots position = .visual token) :
    Spec.get (model.embedMultimodal visual slots) position = Spec.get visual token := by
  simp [embedMultimodal, hSlot]

Full source

InputSlot keeps vocabulary indices and visual-token indices distinct. At a text position, embedMultimodal reads the vocabulary table; at a visual position, it reads the MoonViT output. The displayed theorem checks the visual branch. The neighboring get_embedMultimodal_text theorem checks the text branch.

Two composition stories remain.

The public function above uses list caches and BlockState. The graph cases use packed tensors. Their component theorems are real, but no theorem yet proves the packed layer step equals this public forwardToken. The final token graph is text-only and does not include InputSlot.visual or MoonViT.

Model.lean, backbone token step ↗ Model.lean, multimodal embedding behavior ↗ GraphSpec/Backbone.lean, first of four packed cases ↗

Proof Language backbone · §§2.1–2.3

A token step through the decoder

The state has two representations in this development. The public causal calculation in Model.lean uses list caches and an AttnRes block state. The graph in GraphSpec/LanguageModel.lean takes an ordered pack of tensors whose shapes are known when the graph is built. Its component theorems compare graph operations with reference functions using that same pack. We still need to show that packing the public state preserves its transitions.

We can now compose the component proofs. Fix a valid K3 configuration, a current text token, fixed-shape caches, and one expert route for every sparse layer. The resulting graph updates every layer cache and produces the vocabulary logits for the next token. Under TorchLean's real-number semantics, those outputs equal runSpecAuto, the packed reference calculation, exactly. Both calculations use packed tensors; this does not yet prove equality with the public BackboneLayer.forwardToken.

“Packed” means that differently shaped KDA, MLA, and depth states sit in one ordered bundle whose shapes are recorded in its type. The graph is a directed acyclic graph (DAG): a one-way network of tensor operations with no loop back to an earlier node. A denotation theorem proves what that graph computes when its nodes are interpreted as real-number tensor functions.

A recurrent model can still appear as a DAG because this theorem covers one incremental step. The old caches enter as inputs; the updated caches and logits leave as outputs. A generation loop would feed those outputs into the next step, but that outer loop is not part of this graph yet.

route-staged token graph

parameters and causal state

one-token packed denotation

The brackets mean “evaluate the typed graph as a real-valued function.” The equality is exact. It includes the scheduled KDA/MLA choice, dense or sparse channel path, causal state packing, final depth retrieval, and vocabulary projection. A route trace lists the experts chosen at every sparse layer. The graph builder receives that trace; hRoutes proves each staged choice agrees with the top-k scores computed by the packed reference run. Vocabulary scores, also called logits, are the raw numbers softmax later turns into next-token probabilities.

primitive operationsRMSNorm · matmul · broadcast · reshape · activationTorchLean GraphSpec primitives

↓

component graphsKDA · MLA · AttnRes · Expert · MoEcomponent denotation theorems

↓

four backbone casesKDA/dense · KDA/sparse · MLA/dense · MLA/sparseBackbone.model_specFwd_eq_specStep

↓

route-staged token graphgraph_specFwd_eq_runSpecAutoupdated causal states and current vocabulary logits

Shape mistakes fail while the graph is elaborated. A semantic mismatch in cache order, layer case, reshape, parameter packing, or output packing breaks the proof. A wrong route trace prevents the caller from supplying hRoutes. This comparison catches disagreements between the two formal paths; it cannot detect an equation or assumption encoded incorrectly in both.

The next sections turn to updating the parameters and running training. Their proofs cover particular formulas and contracts, such as per-head matrix layout and bounded rewards. The training curves and systems measurements require evidence from the runs.

The graph agrees with its reference calculation

finishDecoder completes the token step. It separates the updated layer caches from the depth state, performs the last AttnRes retrieval, normalizes the hidden vector, and multiplies by the vocabulary matrix.

/-- Apply K3's final AttnRes retrieval, normalization, and vocabulary projection to a decoder
run. -/
noncomputable def finishDecoder (hcfg : architecture.WF)
    (model : LanguageModel ℝ architecture.text decayRank)
    (pastTokens : Nat)
    (decoderResult : TorchLean.TensorPack ℝ
      (Decoder.OutputsFor pastTokens 0 model (List.finRange architecture.text.numLayers))) :
    TorchLean.TensorPack ℝ (Outputs pastTokens model) :=
  let normalizedResult : TorchLean.TensorPack ℝ
      (Decoder.NextStateShapesFor pastTokens model
          (List.finRange architecture.text.numLayers) ++
        DepthSchedule.Shapes architecture.text.attnResBlockSize architecture.text.numLayers
          architecture.text.hiddenDim) :=
    outputsFor_full pastTokens model ▸ decoderResult
  let parts := TorchLean.TensorPack.split normalizedResult
  let sources := finalSources parts.2
  let retrieved := AttnRes.attendPacked hcfg.hiddenDim_pos
    model.finalQuery sources
  let hidden := RMSNorm.scale retrieved model.finalNormScale
  let logits := vecMatMulSpec hidden model.vocabularyHead
  TorchLean.TensorPack.append parts.1 (.cons logits .nil)

Full source

The graph builder wires the same final operations from primitive nodes. The theorem below quantifies over every valid architecture, model value, token, packed state, score scale, and route trace satisfying its hypotheses.

/-- The complete token DAG refines K3's automatically routed semantics for a valid route trace.

The graph begins with token embedding lookup and ends at vocabulary logits. It is specialized to a
route trace. The `hRoutes` hypothesis certifies that every sparse-layer route equals the route
computed from that layer's actual hidden representation; top-k itself is not a node in the graph.
-/
theorem graph_specFwd_eq_runSpecAuto (hcfg : architecture.WF)
    (languageModel : LanguageModel ℝ architecture.text decayRank)
    (route : ∀ _index : Fin architecture.text.numLayers,
      Route architecture.text.numRoutedExperts architecture.text.activeExperts)
    (pastTokens : Nat) (token : Fin architecture.text.vocabSize)
    (states : TorchLean.TensorPack ℝ
      (Decoder.StateShapesFor pastTokens languageModel
        (List.finRange architecture.text.numLayers)))
    (mlaScoreScale : ℝ)
    (hRoutes : RoutesAgree hcfg languageModel route pastTokens token states mlaScoreScale) :
    (graph hcfg languageModel route pastTokens token).specFwd
        (parameters languageModel) (inputs pastTokens languageModel states mlaScoreScale) =
      runSpecAuto hcfg languageModel pastTokens token states mlaScoreScale := by
  rw [graph_specFwd_eq_runSpec]
  exact runSpec_eq_runSpecAuto_of_routesAgree hcfg languageModel route pastTokens token states
    mlaScoreScale hRoutes

Full source

graph … .specFwd is the mathematical meaning of the typed TorchLean DAG.
parameters and inputs pack the same model weights and causal state in the order expected by that graph.
runSpecAuto is the independent packed reference calculation. Equality means the graph wiring cannot change its output or updated caches.

architecture.WF supplies positive dimensions, divisibility, active-expert bounds, MX alignment, and vision/text width agreement.

languageModel carries the sequence and channel schedules as fields.

states is a heterogeneous tensor pack indexed by the KDA/MLA state shapes computed from that model and the current cache length.

hRoutes proves the supplied route trace agrees with the top-k scores computed from each sparse layer's actual hidden input.

Where equality stops.

runSpecAuto is the packed pure semantics in the GraphSpec file. No theorem yet connects it to the public list/history BackboneLayer.forwardToken and sequence-level LanguageModel.forwardDepthStates, forwardHidden, and logits functions. This graph also excludes MoonViT and visual-input splicing, a generation loop, runtime top-k, finite-precision kernels, checkpoint weights, and a proof that toProgram execution agrees with specFwd.

GraphSpec/LanguageModel.lean, final decoder semantics ↗ GraphSpec/LanguageModel.lean, one-token theorem ↗

Streaming, full sequences, and chunk boundaries

We can evaluate a whole sequence layer by layer, or carry each layer's causal state between token steps. We now prove that these two calculations give exactly the same vocabulary logits. The streaming calculation calls the existing layer transition, including its AttnRes retrievals and sparse router; it is not another model definition.

/-- The public full-sequence logits are exactly those obtained by retaining each layer's causal
state between token steps. This holds for every sequence length and every model parameter value,
including the model's own deterministic sparse routes. -/
theorem run_eq_logits {length : Nat} (model : LanguageModel ℝ cfg decayRank)
    (embeddings : Tensor ℝ [length, cfg.hiddenDim]) :
    (model.run model.initialState embeddings).2 = model.logits embeddings

Read the proof

by
  let inputs := List.ofFn fun index => Spec.get embeddings index
  let final := (List.finRange cfg.numLayers).foldl
    (fun states index => (model.layer index).forward model.activeExperts_le states)
    (inputs.map BackboneLayer.initialDepthState)
  let finish := fun depth : AttnRes.BlockState ℝ cfg.hiddenDim =>
    vecMatMulSpec (RMSNorm.scale (depth.retrieve model.finalQuery) model.finalNormScale)
      model.vocabularyHead
  have hOutputs : (CausalScan.run model.step model.initialState inputs).2 = final.map finish := by
    -- Separate embedding preparation and logit projection from the shared decoder scan.
    change (CausalScan.run (fun state input =>
      let result := Layers.step model (List.finRange cfg.numLayers) state
        (BackboneLayer.initialDepthState input)
      (result.1, finish result.2))
      (Layers.initialState model (List.finRange cfg.numLayers)) inputs).2 = _
    rw [CausalScan.run_map_output]
    simp only
    exact congrArg (List.map finish)
      (Layers.scan_eq_fold model (List.finRange cfg.numLayers) _)
  simp only [run, logits, forwardHidden, Tensor.stackLeading, Spec.get_dim]
  congr 1
  funext index
  change (CausalScan.run model.step model.initialState inputs).2.get _ =
    finish (final.get _)
  simp only [hOutputs, List.get_eq_getElem, List.getElem_map]

Full source

run takes and returns tensors, together with the state for the next chunk. This proves agreement of the public real-valued sequence and streaming calculations. Connecting that state to the packed graph above is still a separate proof.

The .2 in the statement selects the logits from the returned pair; .1 would select the decoder state. In the proof, hOutputs first identifies the hidden outputs of the two evaluation orders, using the earlier layer-scan result. Both calculations then apply the same final normalization and vocabulary projection. The last step checks each row as we assemble the output tensor.

We also prove that splitting a sequence into chunks does not change the answer. We pass the first chunk's state to the second chunk, concatenate their logits, and get exactly the same logits and final state as running the whole sequence at once:

/-- Chunking an embedding tensor preserves every logit and the complete final decoder state.
The second chunk must start from the state returned by the first. -/
theorem run_append {leftLength rightLength : Nat} (model : LanguageModel ℝ cfg decayRank)
    (state : model.CausalState) (left : Tensor ℝ [leftLength, cfg.hiddenDim])
    (right : Tensor ℝ [rightLength, cfg.hiddenDim]) :
    model.run state (Tensor.concat left right) =
      let first := model.run state left
      let second := model.run first.1 right
      (second.1, Tensor.concat first.2 second.2)

Read the proof

by
  let first := CausalScan.run model.step state (List.ofFn fun index => Spec.get left index)
  let second := CausalScan.run model.step first.1 (List.ofFn fun index => Spec.get right index)
  have hFirst : first.2.length = leftLength := by
    rw [CausalScan.outputs_length, List.length_ofFn]
  have hInputs : (List.ofFn fun index => Spec.get (Tensor.concat left right) index) =
      (List.ofFn fun index => Spec.get left index) ++
        (List.ofFn fun index => Spec.get right index) := by
    rw [concat_rows]
    simp only [Tensor.stackLeading, Spec.get_dim, List.ofFn_fin_append]
  have hScan : CausalScan.run model.step state
      (List.ofFn fun index => Spec.get (Tensor.concat left right) index) =
      (second.1, first.2 ++ second.2) := by
    rw [hInputs, CausalScan.run_append]
  dsimp only [run]
  apply Prod.ext
  · change (CausalScan.run model.step state
      (List.ofFn fun index => Spec.get (Tensor.concat left right) index)).1 =
        second.1
    have hState := congrArg Prod.fst hScan
    exact hState
  · dsimp only
    conv_rhs => rw [concat_rows]
    apply congrArg Tensor.stackLeading
    funext index
    simp only [List.get_eq_getElem, Tensor.stackLeading, Spec.get_dim]
    rw [List.getElem_of_eq (congrArg Prod.snd hScan)]
    change (first.2 ++ second.2).get _ = _
    by_cases h : index.val < leftLength
    · simp [Fin.append, Fin.addCases, h, List.get_eq_getElem, hFirst]
      rfl
    · simp [Fin.append, Fin.addCases, h, List.get_eq_getElem, List.getElem_append, hFirst]
      rfl

Full source

The important argument is first.1: the second chunk receives the state left by the first, rather than starting with an empty cache. The proof rewrites the concatenated tensor as its ordered input rows (hInputs), applies the scan's append theorem (hScan), and puts the resulting rows back into tensors. This preserves both parts of the returned pair, the final state and the logits. Changing the chunk boundary therefore changes neither, as long as we carry that state forward.

Model/Streaming.lean, public sequence theorem ↗

§2.5, p. 10

Per-Head Muon and matrix layout

For the query, key, and value projections, Kimi K3 applies Muon independently to each attention head. Muon applies an approximate matrix orthogonalization to its updates. Here we first check that the head split preserves the matrix entries. A matrix with $\text{heads}\times\text{columns}$ outputs must split into head blocks and reassemble without moving an entry.

split one head block

the layout round trip

Hover one color. Those adjacent columns form a single head matrix. Orthogonalization runs inside that block, then merge restores every entry to its original global column.

Projection layouthover a head block

Each colored group is orthogonalized independently.

The full optimizer includes momentum updates, Newton–Schulz-style orthogonalization, and a distributed gather strategy. The formalization isolates the per-head layout and the contract an orthogonalizer must satisfy.

Split, update, and reassemble

/-- Partition a projection matrix along its output-head axis.

For a matrix of shape `rows × (heads * columns)`, output column `(head, column)` belongs to the
`head`-th `rows × columns` block. This is the layout used when Per-Head Muon treats Q, K, and V
momentum buffers independently by attention head.
-/
def split {α : Type} [Storage α] {heads rows columns : Nat}
    (momentum : MatrixTensor α rows (heads * columns)) :
    HeadMatrices α heads rows columns :=
  fun head => Tensor.stackLeading fun row => Tensor.stackLeading fun column =>
    Tensor.scalar (get2 momentum row (finProdFinEquiv (head, column)))

/-- Reassemble per-head matrices into the original projection layout. -/
def merge {α : Type} [Storage α] {heads rows columns : Nat}
    (blocks : HeadMatrices α heads rows columns) :
    MatrixTensor α rows (heads * columns) :=
  Tensor.stackLeading fun row => Tensor.stackLeading fun outputColumn =>
    let headColumn := finProdFinEquiv.symm outputColumn
    Tensor.scalar (get2 (blocks headColumn.1) row headColumn.2)

/-- Splitting and reassembling a projection momentum matrix preserves every entry. -/
theorem merge_split {α : Type} [Storage α] {heads rows columns : Nat}
    (momentum : MatrixTensor α rows (heads * columns)) :
    merge (split momentum) = momentum

Read the proof

by
  apply TorchLean.Tensor.Internal.Rep.ext
  rintro ⟨row, column, ⟨⟩⟩
  have h : get2 (merge (split momentum)) row column = get2 momentum row column := by
    simp only [merge, split, Tensor.stackLeading, Spec.get2_dim, Prod.mk.eta,
      Equiv.apply_symm_apply]
  simpa only [Spec.get2_eq_apply] using h

/-- Reassembling and splitting an indexed head family returns the same family. -/
theorem split_merge {α : Type} [Storage α] {heads rows columns : Nat}
    (blocks : HeadMatrices α heads rows columns) :
    split (merge blocks) = blocks

Read the proof

by
  funext head
  apply TorchLean.Tensor.Internal.Rep.ext
  rintro ⟨row, column, ⟨⟩⟩
  have h : get2 (split (merge blocks) head) row column = get2 (blocks head) row column := by
    simp only [merge, split, Tensor.stackLeading, Spec.get2_dim, Equiv.symm_apply_apply]
  simpa only [Spec.get2_eq_apply] using h

Full source

/-- An exact Muon backend certifies every head block of the reassembled projection update. -/
theorem orthogonalize_hasExactColumnGram {α : Type} [Storage α] [Context α]
    {heads rows columns : Nat}
    (orthogonalizer : Orthogonalizer α [rows, columns])
    (hExact : ExactMatrixOrthogonalizer orthogonalizer)
    (momentum : MatrixTensor α rows (heads * columns)) :
    ∀ head, HasExactColumnGram (split (orthogonalize orthogonalizer momentum) head) := by
  intro head
  rw [split_orthogonalize]
  exact hExact (split momentum head)

/-- An approximate Muon backend transfers its Gram-residual bound to every reassembled head. -/
theorem orthogonalize_hasApproxColumnGram {α : Type} [Storage α] [Context α]
    {heads rows columns : Nat} (epsilon : α)
    (orthogonalizer : Orthogonalizer α [rows, columns])
    (hApprox : ApproxMatrixOrthogonalizer epsilon orthogonalizer)
    (momentum : MatrixTensor α rows (heads * columns)) :
    ∀ head, HasApproxColumnGram epsilon
      (split (orthogonalize orthogonalizer momentum) head) := by
  intro head
  rw [split_orthogonalize]
  exact hApprox (split momentum head)

Full source

Splitting and merging recover every matrix entry in both directions. The later theorems ask a different question about the column dot products, recorded in a Gram matrix. If an orthogonalizer meets the exact or ε-approximate Gram contract on each head, splitting, applying it, and merging retain that contract for every head. The layout argument relies on that premise about the orthogonalizer.

Training/Muon.lean, per-head updates ↗

Pre-training · Report §3, pp. 10–12

How K3 was pre-trained

The reported 2.5× scaling gain compares training recipes that change the model shape, attention pattern, data mixture, optimizer, and context curriculum together. I follow those choices here to understand what the comparison includes: the capacity available to each token, the examples the model learns from, and the work spent on longer contexts.

More layers and experts, the same hidden width

Kimi K2 already had more than a trillion parameters, but K3 does not obtain its extra capacity by widening the residual stream. Both models keep a hidden dimension of 7,168. K3 grows from 61 to 93 language-model layers, expands the routed pool from 384 to 896 experts, doubles the active routed experts from eight to sixteen, and adds a second shared expert. The expert hidden dimension rises from 2,048 to 3,072. The result is 2.78 trillion stored parameters and 104.2 billion active parameters per token, compared with 1.04 trillion stored and 32.6 billion active in K2.

A mixture-of-experts checkpoint stores all its experts, even though a token uses only the selected ones. Total parameters count that stored capacity; active parameters better approximate the matrix work for a token. K3 stores about 2.7 times as many parameters as K2 and activates about 3.2 times as many, so both storage and computation grow. To estimate the time per token, we also need the routing, communication, memory format, batch size, and kernels.

Kimi K2 and Kimi K3 as reported in Table 1
Quantity	Kimi K2	Kimi K3	Change
Language-model layers	61	93	+52%
Total parameters	1.04T	2.78T	+167%
Active parameters per token	32.6B	104.2B	+220%
Hidden dimension	7,168	7,168	unchanged
Routed experts	384	896	+133%
Active routed experts	8	16	+100%
Shared experts	1	2	+100%
Attention heads	64	96	+50%
Attention stack	61 MLA	69 KDA + 24 MLA	hybrid
Training context	128K	1M	8×

The K3 column also lists a 3,584-dimensional latent MoE space and a 401-million-parameter MoonViT-V2 with 27 layers, patch size 14, and 12 heads. Sixty-nine of the 93 language layers use KDA; the remaining 24 provide global MLA attention. Table 1, p. 11 ↗

The reported 2.5× scaling gain

Figure 7 fits validation loss against training compute using data that was not used for training and was deliberately chosen to differ from it. At the same fitted loss, a K3-family run used about 40 percent as much compute as a K2-family run. That is the reported 2.5× gain. This compares fitted model-family curves, not serving speed or the total cost of the final checkpoint. Because architecture, data, optimization, and dimensions changed together, I read this as a gain for the combined recipe. To attribute it to KDA, the data mixture, or another component, we'd need runs that change those choices separately. §3.2 and Figure 7, pp. 10–11 ↗

The data mixture starts multimodal

The text corpus is organized around web text, code, mathematics, and knowledge. The pipeline combines rules, learned quality classifiers, and deduplication. Sampling rates for the domains come from ablations on smaller models. The report names that selection process, though the sampling rates are not given.

The team rewrites knowledge and mathematics sources before base-model pre-training. Prompts ask for different styles and viewpoints; long sources are rewritten chunk by chunk. A later fidelity check tries to keep the facts intact. The model therefore learns from transformed versions of some source material as part of its base corpus. That fidelity check matters: a rewrite can change a formula or factual detail even when the new prose sounds convincing.

The vision corpus covers captions, documents with interleaved images and text, OCR, perception, video, and visual coding. Localization examples use both absolute coordinates and coordinates normalized to the unit square. Programmatic examples pair code with rendered outputs such as SVG images, webpages, games, 3D assets, and CAD drawings. Those pairs let the model learn from code alongside the image or application it produces. We'll see related tasks again in the web-development and visual-agent evaluations. §3.1, p. 10 ↗

Joint vision-language training

MoonViT-V2 and the language model train together from the first step under the same next-token objective. Visual features are resized to the language embedding width and interleaved with text tokens, so the language loss also changes the vision encoder. The report calls this native multimodal training. For example, a wrong answer to an image question can change both the features extracted from the image and the language model that uses them. The two parts learn to work together under that objective.

textweb · code · math · knowledge

visionimages · documents · video · rendered code

→

one token streamjoint next-token prediction

→

one backboneMoonViT-V2 + K3 language model

Text and visual examples enter the same pre-training objective. Media-control tokens can remain in the input while selected positions are masked from the loss.

The optimizer and context schedule

Per-Head Muon updates the query, key, and value projection matrices in separate head blocks. K3 also keeps K2’s weight clipping and uses Quantile Balancing to adjust expert selection in place of the earlier auxiliary MoE balancing loss. The learning rate rises during a one-percent linear warmup, then follows cosine decay; weight decay is 0.1. These choices describe the update schedule, but they are not enough to replay the run: the peak learning rate, batch size, clipping threshold, and total token budget are not specified.

Context length grows late because attention and activation costs make long sequences expensive. Ordinary pre-training begins at 8K and later reaches 64K. During cooldown, the model advances to 256K and then one million tokens. The authors describe the long-context stages as a small fraction of training. They upsample scarce coherent long documents and videos, clean them with exact and fuzzy deduplication, use perceptual hashes for frames, and add synthetic examples that scatter the evidence needed for one answer across several concatenated documents and subtasks.

KDA uses no explicit positional encoding. Its recurrent decay and gating make the order of tokens affect the state, so the team can extend the sequence without rescaling rotary positional embeddings (RoPE), a common way to rotate query and key coordinates according to token position, or interpolating position indices. Removing RoPE removes one engineering obstacle. We still need to ask whether the state retains the information required by a long-context task. The later retrieval benchmarks address that question under their evaluation setups. §§3.3–3.4, pp. 11–12 ↗

I find the shared training objective and the late context expansion useful to think about together. K3 first learns from many kinds of input, then spends part of its final training on much longer examples. The report lets us understand that sequence of choices; reproducing its cost or attributing the scaling gain would require the underlying runs.

§§3.3–3.4, pp. 11–12

Loss masks and longer training contexts

The report describes a joint next-token objective for language and vision. Here we'll write a masked version of that objective, choosing which target positions contribute to the loss. The reported context schedule grows from 8K to 64K, then 256K and one million tokens during cooldown.

declared masked-gradient formula

At each trained position, softmax turns the vocabulary scores into probabilities. The stated gradient for one vocabulary item is its predicted probability, minus one when that item is the correct token, averaged over all trained positions. A masked position contributes zero. Its token can still appear in the input context; only its loss is omitted. Lean records this formula but does not yet derive it from the loss.

pre-training8K

extension64K

cooldown256K

cooldown1M

Example shifted sequence

Themodel<media>seestokens

In the declared formula, supervised rows receive the scaled cross-entropy term and the masked row contributes zero.

The masked loss and gradient formula are our mathematical specification of the next-token objective; the report does not give this formula or a complete loss mask. The four context lengths come from §§3.3–3.4. The report does not disclose the total token count or duration of each stage.

Defining the next-token training loss

/-- Mean next-token loss over the positions selected by `supervised`.

The model may receive both text and visual embeddings, but its targets are vocabulary tokens. The
finite set makes masking explicit: visual input positions, padding, and input-only protocol tokens
can be excluded without assigning them fabricated output labels.
-/
noncomputable def nextTokenLoss {sequenceLength vocabSize : Nat}
    (logits : Tensor ℝ [sequenceLength, vocabSize])
    (targets : Fin sequenceLength → Fin vocabSize)
    (supervised : Finset (Fin sequenceLength)) (_hSupervised : supervised.Nonempty) : ℝ :=
  (∑ position ∈ supervised, tokenLoss (Spec.get logits position) (targets position)) /
    supervised.card

/-- Evaluate the masked next-token objective on the logits produced by the complete language
backbone. This is the pretraining objective of a concrete model, rather than a loss function left
detached from the architecture that produced its logits. -/
noncomputable def pretrainingLoss {cfg : TextConfig} {decayRank sequenceLength : Nat}
    (model : LanguageModel ℝ cfg decayRank)
    (embeddings : Tensor ℝ [sequenceLength, cfg.hiddenDim])
    (targets : Fin sequenceLength → Fin cfg.vocabSize)
    (supervised : Finset (Fin sequenceLength)) (hSupervised : supervised.Nonempty) : ℝ :=
  Pretraining.nextTokenLoss (model.logits embeddings) targets supervised hSupervised

Full source

supervised is the finite set of positions that count toward the loss. Its nonempty proof prevents division by zero.
At each selected position, the logits are raw scores for the vocabulary. tokenLoss applies softmax and returns the negative log probability of the correct token.
model.logits embeddings runs the language model first. The objective is therefore connected to the architecture that produced the predictions.

The file proves how the supervised and masked branches of the declared formula reduce, but that is not a derivative theorem. Lean does not yet prove that automatic differentiation of pretrainingLoss produces this gradient.

Training/Pretraining.lean, loss and gradient formula ↗

Post-training · Report §4.1, pp. 12–14

Training nine specialists, then combining them

Post-training proceeds in three stages. Supervised trajectories establish a usable agent. Reinforcement learning then separates domains and reasoning budgets into nine specialist policies. Multi-Teacher On-Policy Distillation brings their behavior back into one student.

Starting with supervised examples

Supervised fine-tuning (SFT) gives the base model a common interaction format and enough worked examples to begin agent training. Earlier domain-specialized Kimi models generate complex trajectories. Those trajectories pass through several verification stages and human annotation before training. Every conversation uses XTML, the text format used to serialize messages, tool calls, and tool results for the configurable agent wrappers and inference system.

These demonstrations are meant to teach the model to adjust reasoning effort, call tools precisely, and keep working through long tasks before reinforcement learning begins. This also gives RL a better starting point: the model can already attempt the interaction that the reward will judge. The report does not separately measure how much each part of the cold start contributes. §4.1.1, p. 12 ↗

Three domains crossed with three effort levels

The RL stage separates capability and computation. The first axis has three domains. General-task training covers ordinary experience, vision, reasoning, faithfulness, search, and knowledge work. General-agent training covers long-horizon assistants, deep research, and paragraph-level writing. Coding-agent training covers software engineering, coding experience, kernel work, and web development.

The second axis has low, high, and max reasoning effort. Crossing the two axes produces nine teachers. A low-effort coding policy learns under a tighter budget than a max-effort coding policy; both are distinct from a max-effort general agent.

low

high

max

general tasks

reasoning under a short budget

more search and thought

largest allowed budget

general agents

compact tool use

longer agent work

long-horizon execution

coding agents

fast coding tasks

repository work

extended coding and kernels

As RL computation increases in Figure 8, the plotted coding, tool-use, web-development, search, office, chart, and visual-puzzle capabilities rise. So does the average number of assistant or tool-call steps. This connects the training story to longer attempts at the tasks. I can read the direction of those trends, but the missing numerical axis labels, effect sizes, and variance prevent an estimate of improvement per unit of computation. §4.1.2 and Figure 8, pp. 12–13 ↗

Pausing and resuming unfinished rollouts

Agent trajectories vary in length. A batch that waits for its slowest trajectory leaves finished workers idle; a batch that discards unfinished work wastes the most expensive examples. K3 uses partial rollouts. For $N$ prompts and $K$ completions per prompt, the system maintains $N \times K$ active trajectories. It pauses an iteration when a fraction λ has finished, queues the unfinished states, and gives them priority in the next iteration. A prompt enters policy optimization only after all of its $K$ responses complete.

1
Launch several responses for every prompt and retain their model, tool, cache, and environment state.
2
When λ of all responses finish, stop waiting for the longest tail and start the next policy iteration.
3
Resume the paused responses first. Their exact sandboxes and cache prefixes survive the boundary between iterations.
4
Once every response for one prompt is ready, compute the group reward and make that prompt eligible for optimization.

A paused trajectory can resume after several parameter updates. Its earlier actions were sampled by an older policy, so learning from them is off-policy. The team says a per-token regularization penalty limits large policy changes enough to keep this stable. This creates a trade-off between finishing expensive work and learning from a recent policy. The report explains the mechanism, but does not quantify the age of the resumed trajectories or the throughput gained over synchronous rollouts.

Training with problem-specific token budgets

The effort levels are enforced through length-relative penalties. A cold-start model first estimates a baseline budget $b_0(x)$ for each problem. A trajectory receives reward −1 when its length exceeds τ times that baseline. On general tasks, length counts thinking tokens. On agent tasks it counts cumulative output, including reasoning and tool-call arguments. Training begins with a large τ for the max-effort teacher, still under an absolute cap, and anneals τ to produce high- and low-effort teachers. The settings vary by domain with human guidance.

This penalty sits beside a competing trend: the capability curves improve as the agents take more steps. K3 is trained to spend extra computation when it earns reward and to stop when the trajectory becomes wasteful. To compare the effort modes, I would want to plot task success against tokens spent for each one. Figure 8 shows the broader training trends, rather than that budget-by-budget comparison.

Many agent tasks have deterministic or programmatic verifiers. Writing and other non-verifiable work uses an Agentic Generative Reward Model. The grader reads the output, writes a rubric, scores each candidate against it, and records the scores in a scratch area. Candidates compete through tournament-style binary comparisons. A response longer than σ times the cold-start verbosity loses automatically. This gives verbosity its own control instead of treating every extra token as evidence of better work. For these tasks the quality of the grader matters as well as the policy: a model can learn to satisfy a judge without making its answer more useful. §4.1.2, p. 13 ↗

Combining the specialists through distillation

Nine separate deployment models would be costly and awkward. MOPD samples a domain and effort pair, rolls out the unified student, and scores each sampled token with the matching teacher. The reward is the teacher-to-student log-probability ratio. It is treated as a fixed score, so differentiation does not flow back through its calculation, and it is clipped to a symmetric interval. The following section writes out that reward and shows the Lean bound.

Because the student generates the trajectories, teachers score the states it reaches during its own attempts. The report also tried a denser top-$k$ objective that transferred several token probabilities at each position and found no clear improvement in convergence or final performance. The practical attraction is having one student that can switch domains and effort levels without loading a different model. The report does not give a separate retention score for each teacher, so the final evaluation measures the student rather than a teacher-by-teacher transfer rate. §4.1.3, pp. 13–14 ↗

§4.1.3, pp. 13–14 · Equation 15

MOPD compares teacher and student token probabilities

MOPD chooses one of nine domain-and-effort teachers, evaluates its probability for the student's sampled token, and clips the log probability ratio. The dense reward is bounded even when teacher and student probabilities differ sharply.

(15)

MOPD reward

Set the teacher probability high and the student probability close to zero. The raw log ratio increases, while the value passed to the objective stops at five in this illustration. Reverse them and the same clamp applies at minus five.

Inspect the clipped reward$R_{\max}=5$

teacher probability 0.700student probability 0.080

log ratio2.169

clipped reward2.169

Start with the displayed probabilities, 0.70 and 0.08. Their ratio is 8.75, so its natural logarithm is about 2.169. The teacher assigns this sampled token more probability than the student does, giving a positive score. Exchanging the probabilities reverses the sign. Equal probabilities give zero. Clipping limits how strongly a very large disagreement can affect the reward; the slider uses an illustrative cap of five, not a cap value disclosed in the report.

The stop-gradient operation also matters. The student probability is used to calculate the reward, but that reward is held fixed while differentiating the policy update. Differentiating through both its log ratio and the policy objective would produce a different update. The scalar bound below describes the reward; it does not by itself identify either update rule.

Why the probability and clipping conditions matter

/-- A strictly positive probability that may safely appear inside the logarithmic MOPD reward.

The teacher and student probabilities in Eq. 15 come from softmax distributions and are therefore
strictly positive. This is mathlib's standard interval subtype rather than a parallel record.
-/
abbrev StrictProbability := Set.Ioc (0 : ℝ) 1

/-- Per-token teacher/student reward from Eq. 15, before the stop-gradient annotation. -/
noncomputable def reward (maxReward : ℝ)
    (teacherProbability studentProbability : StrictProbability) : ℝ :=
  clip maxReward (Real.log ((teacherProbability : ℝ) / studentProbability))

/-- Clipping enforces the report's reward interval for every log-ratio. -/
theorem reward_mem_interval {maxReward : ℝ} (h : 0 ≤ maxReward)
    (teacherProbability studentProbability : StrictProbability) :
    -maxReward ≤ reward maxReward teacherProbability studentProbability ∧
      reward maxReward teacherProbability studentProbability ≤ maxReward := by
  constructor
  · exact le_max_left _ _
  · exact max_le (by linarith) (min_le_left _ _)

Full source

Strict positivity prevents Lean's totalized real logarithm from assigning a misleading ordinary value to zero probability. The interval proof follows from symmetric clipping.

Scope.

Equation 15's scalar reward, the typed three-by-three teacher index, budget override, and verbosity rule are formalized. The nine trained checkpoints, stop-gradient runtime implementation, generated trajectories, partial-rollout system, and RL optimization remain external evidence.

Training/Distillation.lean, MOPD rewards ↗ Training/Controls.lean, teacher and budget rules ↗

§4.1.4, p. 14

What four-bit expert weights mean

Quantization replaces a high-precision number with one of a small set of coded values. MXFP4 gives each expert weight a four-bit code; MXFP8 gives each expert activation an eight-bit code. “MX” means microscaling: 32 neighboring values share a separate scale. The OCP Microscaling specification defines these formats. In the diagram, E8M0 is that shared exponent-like scale and E2M1 is the four-bit element format: one sign bit, two exponent bits, and one fraction bit. The formal model makes both parts explicit, so its error theorems state exactly what they assume.

shared-scale block decoding

E2M1 reference error

The purple chip belongs to all 32 cells. Changing that scale moves the range of the whole block; changing one E2M1 code changes only one cell. The slider isolates the element quantizer before the later theorems accumulate its error.

E8M0
2⁰

Quantize one E2M1 valuereference semantics

input 2.70

nearest decoded value3.00

absolute error0.30

proved ceiling≤ 1

Computing a dot product with scaled four-bit weights

A block is a shared scale plus 32 small codes. These definitions say how that data turns back into real numbers and how two decoded blocks form a dot product.

/-- Decode one coordinate by multiplying its private value by the shared block scale. -/
noncomputable def decode {Code : Type} (decodeElement : Code → ℝ) (block : Block Code)
    (index : Fin 32) : ℝ :=
  scaleValue block.scale * decodeElement (block.element index)

/-- Dot product of two decoded blocks. -/
noncomputable def dot {CodeA CodeB : Type} (decodeA : CodeA → ℝ) (decodeB : CodeB → ℝ)
    (a : Block CodeA) (b : Block CodeB) : ℝ :=
  ∑ index, a.decode decodeA index * b.decode decodeB index

/-- MX dot products factor into the product of block scales and a dot product of private values. -/
theorem dot_eq_scale_mul_sum {CodeA CodeB : Type} (decodeA : CodeA → ℝ)
    (decodeB : CodeB → ℝ) (a : Block CodeA) (b : Block CodeB) :
    a.dot decodeA decodeB b =
      scaleValue a.scale * scaleValue b.scale *
        ∑ index, decodeA (a.element index) * decodeB (b.element index) := by
  rw [dot, Finset.mul_sum]
  apply Finset.sum_congr rfl
  intro index _
  simp only [decode]
  ring_nf

Full source

Fin 32 makes a block index provably fall between 0 and 31.
decode multiplies a decoded element value by the one scale shared by the entire block.
The theorem moves both scales outside the sum. This is the algebra used by a blockwise dot-product implementation.

Bounding the quantization error

Let's call the shared scale $s$, one original value $x_i$, and the decoded quantized value $\widehat{x}_i$. The theorem below says $|\widehat{x}_i-x_i|\le s$, provided every value in the block lies between $-6s$ and $6s$. This is an absolute error bound in the original units. For example, choosing $s=1/2$ gives an error allowance of $1/2$ for each coordinate, as long as the inputs stay in $[-3,3]$.

The proof follows the same three operations as the encoder:

Divide by the shared scale. Finite E8M0 scales are positive. The input assumption $|x_i|\le6s$ therefore gives $|x_i/s|\le6$. These are the facts named hscale and hscaled.
Round one E2M1 element. The earlier scalar theorem bounds the error on $[-6,6]$ by one. Applying it to $x_i/s$ gives hround. This is where the format's rounding bound enters.
Multiply back by the scale. The equality called hidentity factors $s$ out of the error. The final calc uses that equality and the scalar bound:

|\widehat{x}_i-x_i| = s\,|\operatorname{decodeE2M1}(Q(x_i/s))-x_i/s|\le s\cdot1=s

In the statement, hRange supplies the input assumption for all 32 coordinates, and index chooses the coordinate we're proving it for. Open the proof to follow those three steps in Lean; each step still shows its goal when you hover over it.

/-- If every source value fits the range selected by the shared scale, MXFP4 quantization has
coordinate error at most that scale. -/
theorem quantizeMXFP4Block_error_le (scale : E8M0) (values : Fin 32 → ℝ)
    (hRange : ∀ index, |values index| ≤ 6 * scaleValue scale) (index : Fin 32) :
    |(quantizeMXFP4Block scale values).decode decodeE2M1 index - values index| ≤
      scaleValue scale

Read the proof

by
  let scaleReal := scaleValue scale
  have hscale : 0 < scaleReal := scaleValue_pos scale
  have hscaled : |values index / scaleReal| ≤ 6 := by
    rw [abs_div, abs_of_pos hscale]
    exact (div_le_iff₀ hscale).2 (by simpa [scaleReal, mul_comm] using hRange index)
  have hround := quantizeE2M1_error_le_one hscaled
  change |scaleReal * decodeE2M1 (quantizeE2M1 (values index / scaleReal)) -
    values index| ≤ scaleReal
  have hidentity :
      scaleReal * decodeE2M1 (quantizeE2M1 (values index / scaleReal)) - values index =
        scaleReal *
          (decodeE2M1 (quantizeE2M1 (values index / scaleReal)) -
            values index / scaleReal) := by
    field_simp
  calc
    |scaleReal * decodeE2M1 (quantizeE2M1 (values index / scaleReal)) - values index| =
        scaleReal *
          |decodeE2M1 (quantizeE2M1 (values index / scaleReal)) -
            values index / scaleReal| := by
      rw [hidentity, abs_mul, abs_of_pos hscale]
    _ ≤ scaleReal * 1 := mul_le_mul_of_nonneg_left hround hscale.le
    _ = scaleReal := mul_one _

Full source

Multiplying decoded blocks combines errors from the weights and inputs. The longer theorems carry the coordinate bounds through block matrices and vector–matrix products, keeping the scales and input magnitudes in the result. MXFP8Encoding pairs encoded inputs with an error bound for each coordinate. To use the result with a production encoder, we must establish those input bounds for that encoder.

We use FloatLib's OCP decoders for the element codes and shared scale exponent, rather than maintain another copy of their bit-level rules here. The K3 adapter admits finite codes and converts the exact rational values to reals. We keep the real-input quantizer for these error proofs; it is not a claim about every production quantization kernel.

Format choice remains explicit.

The report says MXFP8 but does not fix E4M3 versus E5M2 and every overflow convention. The Lean API takes the FP8 format as a parameter instead of inventing a paper claim.

Microscaling.lean, shared-scale block decoder ↗ Microscaling.lean, block error theorem ↗

§4.1.4, p. 14 · Equation 16

A smaller model drafts tokens for K3 to check

EAGLE-3 is a smaller draft model that proposes several future tokens before full K3 checks them together. This is speculative decoding: cheap proposals can speed generation, while an acceptance-and-correction rule preserves the full model's final distribution. The draft starts from K3 features taken at low, middle, and high depths and from the pretrained multi-token prediction (MTP) layer, a small head trained to predict beyond the immediate next token. Its fusion matrix begins as [0 0 I], so the first fused feature is exactly the high-level feature the MTP layer already saw. Later draft steps feed their own hidden output forward. The training unroll has seven positions.

target featureslow · mid · high

[0 0 I]

Lossless speculative sampling accepts the overlap between the target distribution $p$ and draft distribution $q$. Kimi K3 minimizes the negative log of that overlap. Total variation, abbreviated TV below, is half the sum of the absolute probability differences: it is zero for identical distributions and one when they do not overlap.

(16)

acceptance loss

overlap and total variation

Each group of three bars is one possible next token. Blue is the target probability, purple is the draft, and gold is the overlap that can be accepted. Move the draft sliders toward the target. The gold area grows and the loss falls.

Acceptance is distribution overlapthree-token vocabulary

Σ min(p,q)

total variation

−log acceptance

For a proposed token y with positive draft probability, the acceptance probability is . Multiplying by the probability of proposing it gives . Summing those accepted contributions over the vocabulary gives the overlap shown above.

In the initial example, the target is (0.55, 0.30, 0.15) and the draft is (0.25, 0.45, 0.30). Their accepted contributions are (0.25, 0.30, 0.15), with total mass 0.70. The remaining target mass is (0.30, 0, 0). On rejection, the correction distribution therefore chooses the first token. Combining the accepted mass with that 0.30 rejection mass recovers exactly (0.55, 0.30, 0.15). Simply retrying the draft would not implement this correction.

This explains both the training loss and its limit. More overlap means more draft proposals survive verification. It does not alone determine wall-clock speed: drafting, checking a block, and restoring rejected state still take time. The equation trains the draft toward a useful acceptance rate; the serving section describes the remaining work needed to turn that into a faster generator.

Starting the draft from K3's hidden representation

/-- Concatenate-and-project written as a sum of three matrix products. -/
def forward (fusion : Tensor α [3, hiddenDim, hiddenDim]) (features : Tensor α [3, hiddenDim]) :
    Tensor α [hiddenDim] :=
  addSpec
    (addSpec
      (vecMatMulSpec (Spec.get features 0) (Spec.get fusion 0))
      (vecMatMulSpec (Spec.get features 1) (Spec.get fusion 1)))
    (vecMatMulSpec (Spec.get features 2) (Spec.get fusion 2))

/-- The report's `[0 0 I]` initialization returns the high-level feature exactly. -/
theorem initial_fusion_eq_high (features : Tensor ℝ [3, hiddenDim])
    (fusion : Tensor ℝ [3, hiddenDim, hiddenDim])
    (hLow : (Spec.get fusion 0) = Tensor.zeros [hiddenDim, hiddenDim])
    (hMiddle : (Spec.get fusion 1) = Tensor.zeros [hiddenDim, hiddenDim])
    (hHigh : (Spec.get fusion 2) = identityTensorSpec hiddenDim) :
    FeatureFusion.forward fusion features = (Spec.get features 2) := by
  rw [forward, hLow, hMiddle, hHigh]
  rw [vecMatMul_fill_zero, vecMatMul_fill_zero, vecMatMul_identity]
  simp

Full source

The function adds three projected features. Under the report's initial weights, the first two products are zero and the third multiplies by the identity matrix. The theorem proves that the fused result is exactly the high-level feature, not merely close to it.

/-- One EAGLE-3 training-time-test transition. The transition consumes the current token and
feature, then stores its own hidden output as the feature for the next transition. No fresh
target-model feature is requested after initialization. -/
noncomputable def trainingTimeTestStep (model : Model cfg decayRank) :
    model.TrainingTimeTestState → Fin cfg.vocabSize →
      model.TrainingTimeTestState × Distribution cfg.vocabSize :=
  fun state token =>
    let result := model.stepDistributionFromFeature state.1 token state.2
    ((result.1, result.2.1), result.2.2)

/-- Training-time testing emits one draft distribution for every shifted training token. -/
theorem trainingTimeTest_length (model : Model cfg decayRank)
    (tokens : List (Fin cfg.vocabSize)) (features : Tensor ℝ [3, cfg.hiddenDim]) :
    (model.trainingTimeTest tokens features).2.length = tokens.length :=
  CausalScan.outputs_length model.trainingTimeTestStep
    (model.initialTrainingTimeTestState features) tokens

/-- Kimi K3's seven-step EAGLE-3 objective against a frozen target language model. Each target
distribution comes from the target model at the corresponding training position, and each draft
distribution comes from the recurrent training-time-test pass. There is no auxiliary ground-truth
cross-entropy term in this objective, matching Eq. 16 of the report. -/
noncomputable def sevenStepLikelihoodLoss {targetDecayRank draftDecayRank sequenceLength : Nat}
    (target : LanguageModel ℝ cfg targetDecayRank) (draft : Model cfg draftDecayRank)
    (embeddings : Tensor ℝ [sequenceLength, cfg.hiddenDim])
    (positions : Fin 7 → Fin sequenceLength) (tokens : Fin 7 → Fin cfg.vocabSize)
    (features : Tensor ℝ [3, cfg.hiddenDim]) : EReal :=
  ∑ step : Fin 7,
    likelihoodLoss (targetDistributionAt target embeddings (positions step))
      (draft.sevenStepDistributions tokens features step)

Full source

The state is a pair: the draft model's causal cache and the feature consumed at the next position.
result.2.1 is the draft's new hidden feature. Storing it in the next state makes the unroll self-feeding after initialization.
The length theorem guarantees one output distribution per input token. The final definition sums the seven acceptance losses required by Equation 16.

/-- Lossless speculative acceptance is exactly one minus total-variation distance. This identifies
the paper's likelihood objective with a standard statistical distance, rather than merely bounding
it between zero and one. -/
theorem acceptanceRate_eq_one_sub_totalVariation {vocabSize : Nat}
    (target draft : Distribution vocabSize) :
    acceptanceRate target draft = 1 - totalVariation target draft

Read the proof

by
  have htarget : ∑ token, target.weights token = 1 := by
    exact target.total_of_fintype
  have hdraft : ∑ token, draft.weights token = 1 := by
    exact draft.total_of_fintype
  have hpointwise : ∀ token : Fin vocabSize,
      min (target.weights token) (draft.weights token) =
        (target.weights token + draft.weights token - |target.weights token - draft.weights token|) / 2 := by
    intro token
    by_cases h : target.weights token ≤ draft.weights token
    · rw [min_eq_left h, abs_of_nonpos (sub_nonpos.mpr h)]
      ring_nf
    · have h' : draft.weights token ≤ target.weights token := le_of_not_ge h
      rw [min_eq_right h', abs_of_nonneg (sub_nonneg.mpr h')]
      ring_nf
  rw [acceptanceRate, totalVariation]
  simp_rw [hpointwise]
  rw [← Finset.sum_div]
  rw [Finset.sum_sub_distrib, Finset.sum_add_distrib,
    htarget, hdraft]
  ring_nf

/-- One-token speculative sampling is lossless. For every token, the mass accepted from the draft
plus the mass drawn from the rejection correction is exactly the target-model mass. -/
theorem accepted_add_corrected_eq_target {vocabSize : Nat}
    (target draft : Distribution vocabSize)
    (hReject : acceptanceRate target draft < 1) (token : Fin vocabSize) :
    min (target.weights token) (draft.weights token) +
        residualMass target draft *
          (correctionDistribution target draft hReject).weights token =
      target.weights token

Read the proof

by
  have hMass : residualMass target draft ≠ 0 := by
    apply ne_of_gt
    simpa [residualMass] using sub_pos.mpr hReject
  rw [correctionDistribution]
  change min (target.weights token) (draft.weights token) +
      residualMass target draft *
        (residualNumerator target draft token / residualMass target draft) =
    target.weights token
  rw [mul_div_cancel₀ _ hMass]
  rw [residualNumerator_eq_sub_min]
  ring

/-- Zero speculative-sampling acceptance is exactly the infinite-loss case. -/
theorem likelihoodLoss_eq_top_iff {vocabSize : Nat}
    (target draft : Distribution vocabSize) :
    likelihoodLoss target draft = ⊤ ↔ acceptanceRate target draft = 0 := by
  simp [likelihoodLoss]

Full source

The loss type includes infinity, so two distributions with no overlap produce infinite loss instead of hiding the undefined real-number case log 0.

No EAGLE graph is proved here.

The file gives pure recurrent semantics and distribution theorems. It does not build an EAGLE GraphSpec, connect the draft to native execution, or verify the seven-step training optimizer.

Training/Draft.lean, EAGLE fusion initialization ↗ Training/Draft.lean, lossless acceptance ↗ Training/Draft.lean, self-feeding draft step ↗

RL environments · Report §4.2, pp. 14–17

Tools, tasks, and rewards

The training tasks specify what tools the agent can use and how its work is checked. The team creates tools, websites, repositories, office applications, machines, hidden tests, and reward rules that allow a trajectory to last for hours while still producing a checkable result.

Training across agent interfaces

Agent policies can overfit to the wrapper around a task. Tool names, message templates, available memories, context-management rules, subagent APIs, and even error formats become shortcuts. K3's white-box environment makes those pieces configurable. It can reproduce the interaction style of Kimi Code, Claude Code, Codex, OpenClaw, and Hermes, or assemble a new setup from the same parts.

The aim is to teach file inspection, search, and code execution across several tool schemas, so the agent can recognize those operations when their names or message formats change. A wrapper defines what the policy observes and which actions it can take; varying it changes the learning problem even when the underlying task stays the same. I would need a cross-wrapper ablation to measure whether this training improves transfer to a new interface. The report gives the design argument, but no such comparison. §4.2.1, pp. 14–15 ↗

A knowledge graph becomes a task generator

Broad seed topics are recursively expanded into a directed acyclic graph. An agent runs several web searches for each node, splits coarse ideas into finer ones, reuses an existing node when two branches reach the same concept, and stops when it judges a node atomic. The direction of an edge records movement from coarse to fine. Nodes at several levels can then supply retrieval queries, source bundles, and task prompts.

This mechanism produces related problems without copying one template at a time. A broad banking node can lead to financial statements, ratios, sector rules, and source-specific questions. A visual-science node can lead to plots, diagrams, measurements, and small computation tasks. Reusing equivalent nodes also connects areas that would otherwise grow as isolated trees.

seedprofessional analysis

bankingfilingsriskvaluation

data worktableschartssources

legalrulesevidencecitations

retrieval + sources + verifiernew agent tasks

An illustrative task graph: related topics supply different source bundles and questions. The report describes this construction in Figure 9; these labels are not a dump of its training graph.

§4.2.2 and Figure 9, p. 15 ↗

Search, professional work, and visual tools

The search tasks ask for an answer and supporting sources. Professional tasks cover banking, data analysis, and legal work, with trajectories lasting dozens or hundreds of steps. The verifier can inspect the final artifact, citations, calculations, or state of the simulated workspace. This rewards completion rather than fluent narration about what the agent intended to do.

Visual tasks put Python inside the observation loop. The agent may crop an image, zoom, transform it, measure a diagram, run a calculation, and feed the generated image back as the next observation. That design treats image tools as part of reasoning rather than a one-time preprocessing stage. The authors say scores rise with more operations and observations, though they supply no curve or effect size in §4.2.3.

Grading kernel correctness and speed

Kernel optimization problems range from one operator to large fused kernels and cover several GPU programming systems, including CUDA, Triton, CuTe DSL, Gluon, ThunderKittens, and TileLang. Tasks use 16-bit, 8-bit, and 4-bit number formats. A candidate whose numerical error crosses the private threshold gets zero. Matching an expert implementation earns 0.5, and approaching the hardware roofline moves the reward toward one. Separate detectors penalize replaying a captured CUDA graph, caching the input, or silently lowering precision.

The numerical check has to come before the speed reward. Caching an input or lowering precision could improve a timing result while changing the computation the task asked for. The report describes detectors for those shortcuts, though it does not publish the numerical threshold and timing protocol needed to reproduce the grading. §4.2.4, p. 16 ↗

Multi-day assistant tasks

The personal-assistant environment supplies mock versions of Gmail, Notion, Slack, and Canvas. Events unfold over several simulated days and cross application boundaries. An email may create a calendar obligation, a Slack message may change a document, and a later event may test whether the agent remembered the earlier instruction. Each event has either a deterministic check or a model-based judge.

A single rollout can reach thousands of tool calls and millions of accumulated context tokens. The report presents that as a possible upper range, not the typical trajectory. These workloads explain several later systems choices: paused generation must preserve the application state, repeated prefixes must remain cached, and an idle sandbox cannot reserve full memory while waiting for the next model response.

day 1read mailrecord a request

→

day 2edit Notionuse prior context

→

day 3answer Slackreconcile changes

→

laterhidden checkinspect final state

§4.2.5, p. 16 ↗

Grading tasks with no reference solution

Autonomous Execution Tasks give the agent a goal, constraints, action space, resource budget, and verification interface. They do not give it a procedure or reference trajectory. The reward depends on the final state of the environment. Example domains include reproducing a system, discovering quantitative factors, and auditing taxes.

Public diagnostic checks help the agent make progress, while hidden checks decide the final reward. Submission limits and penalties discourage blind probing. This resembles ordinary engineering work more closely than a benchmark with a known sequence of steps: several valid plans may exist, intermediate work can fail, and only the finished state settles whether the task was completed.

Figure 10 shows one camera-repair-system replication task. Its legend reports K3 at 1.000, Opus 4.8 at 0.918, GPT-5.5 at 0.893, and Kimi K2.6 at 0.560. The caption describes verifier-assessed progress, but does not fully define what each displayed scalar aggregates. One demonstration cannot establish a distribution of task performance; it shows the sort of work the environment permits. §4.2.6 and Figure 10, p. 16 ↗

Web development joins code and visible output

Web tasks start from briefs ranging from one line to several paragraphs. Expected outputs include ordinary sites, games, 3D and WebGL scenes, data visualizations, SVG work, and full-stack applications. Deterministic checks inspect structure, behavior, and pixels; an internal model judge assesses aspects that are hard to encode as a fixed test. Build failure, a runtime error, or a faked artifact produces zero reward.

The RL environment asks the model to build against an explicit visual and functional target. Later, the report's blind web-development comparison finds that judges prefer K3's output to Opus 4.8 in 58.6 percent of pairings, with the largest advantage in 3D and WebGL tasks. The evaluation section below discusses the judging conditions and limits. §4.2.7, p. 16 ↗

Infrastructure · Report §5, pp. 17–25

Training and serving infrastructure

K3 creates three different engineering problems at once: KDA has a recurrent state, sparse experts must move across a multi-trillion-parameter training cluster, and agent trajectories can retain a million tokens plus a live computer for hours. The report's infrastructure section describes how the team allocates memory, distributes expert computation, and stores unfinished rollouts.

Systems terms used below

GPU / CPU: A GPU performs many tensor operations in parallel; a CPU runs the general-purpose host program. DRAM is the main memory attached to either.
rank: One worker process in a distributed job, usually responsible for one GPU.
kernel: A low-level routine executed on the GPU. This is not Lean's proof checker, which is also traditionally called a kernel.
activation: An intermediate tensor created during the forward pass and often saved so training can compute gradients later.
shard / all-gather: To shard is to split data across workers. An all-gather is the exchange that gives every worker all those pieces.
forward / backward: The forward pass computes predictions and loss; the backward pass propagates gradients toward the parameters.

KDA trades a growing cache for a serial state update

Full attention offers parallel work across a prompt but stores keys and values that grow with its length. KDA stores a fixed-size state for each head, which cuts cache growth and transfer volume. Its recurrence introduces a different problem: token $t$ depends on the state from token $t-1$. A literal implementation would process the prompt serially and leave much of the GPU idle.

FlashKDA breaks the prompt into chunks. Tokens inside one chunk can be processed in parallel; only the small summary passed from one chunk to the next remains sequential. The implementation uses CUTLASS, NVIDIA's building blocks for GPU matrix operations, and changes its work schedule for short, long, small-batch, and large-batch inputs. The report explains how operations overlap and reuse memory, but gives no throughput comparison. §5.1.1, pp. 17–18 ↗

Splitting KDA across devices with transition summaries

If we split a sequence across devices, the second segment needs the state left by the first. The third needs the state after both. Each device summarizes its segment with a small state update. Those updates can be combined in order. A prefix scan computes the summary after A, then after A+B, then after A+B+C, so every device recovers the state entering its own segment without exchanging all of the tokens.

segment Asegment Bsegment Csegment D

(transition, zero-state)(transition, zero-state)(transition, zero-state)(transition, zero-state)

fixed-size all-gather + associative prefix scan

state 0state after Astate after A+Bstate after A+B+C

Communication depends on the state summary rather than the number of tokens in a segment. The authors therefore describe communication as independent of sequence length and compute scaling as linear in the number of context-parallel devices. The claim follows from the algorithm's message shape; the report does not include an end-to-end scaling chart. §5.1.2, p. 18 ↗

In Lean, each segment summary is an affine map: it applies a linear transition to the incoming state and adds the segment’s zero-state output. Composing these maps gives the same state as applying them one after another. Identity, associativity, and the append law let us regroup the work for a prefix scan. Connecting the summaries to actual KDA updates is still open, as are the network collectives and measured scaling. Parallel/Context.lean ↗

Two scalar summaries make the ordering easy to see. Let and . Applying A and then B gives . Reversing them gives . Composition is associative, so a parallel scan can change the parentheses while keeping A before B. It is not commutative, so it cannot exchange the segments. These toy summaries illustrate the scan law; their coefficients are not measured K3 states.

Dividing a 2.78T model across devices

The model is too large for one GPU. Training therefore divides layers, experts, optimizer memory, gradients, and long sequences across many devices. Each division addresses a different limit: how much a GPU can store, how much work it can perform, or how much data workers must exchange.

The distribution names used in the report

Pipeline parallelism gives consecutive groups of layers to different workers. Expert parallelism spreads the 896 experts across workers. Context parallelism divides a long token sequence. ZeRO-1 shards optimizer memory, while Pipeline ZeRO-2 also limits which gradient pieces stay resident. Virtual pipeline stages interleave smaller layer groups to reduce waiting time.

These partitions interact. Routing may send many tokens to an expert located on another rank. Pipeline stages create bubbles, idle time while one stage waits for another. Sharded states have to arrive before their operation runs. Vision adds a separate encoder whose patch tokens and activations occupy memory. The paper concentrates on three solutions: balance expert work, manage saved activations through one policy, and schedule vision into otherwise idle pipeline time. §5.2, pp. 18–19 ↗

MoonEP and balanced expert assignments

Ordinary expert parallelism can be balanced in expectation and still stall on one rank that receives too many routed tokens. MoonEP migrates redundant copies of popular experts so every rank processes exactly $S \times K$ token-expert assignments, where $S$ is the token capacity per rank and $K$ is the number of active routed experts. The placement changes dynamically with the measured load.

The report proves that no rank needs more than $E/R$ redundant expert slots for $E$ experts across $R$ ranks. Its fixed receive buffer holds $S \times K$ assignments. The existing DeepEP expert-communication layout can require $S \times K \times R$, because any sending rank may contribute the full amount. Static shapes also remove a CPU synchronization step for discovering the next receive size. Lean now proves the per-destination upper bound from explicit home-capacity and one-source hypotheses, plus the capacity pigeonhole lower bound. It does not verify the online planner or distributed migration runtime.

conventional worst caseS × K × R

space for every possible sender contribution

MoonEP fixed bufferS × K

exact work assigned to each receiving rank

§5.2.1, pp. 19–20; Appendix E, pp. 44–45 ↗

Parallel/Experts.lean, counting bounds ↗

Saving, recomputing, and offloading activations

Activations are intermediate tensors that later computations may need again. Keeping all of them on the GPU consumes memory, while recomputing or moving them costs work. A unified manager plans their lifetimes and chooses which to keep, recompute, quantize, move to host memory, or send to remote memory. Most saved activations use a blockwise 8-bit format before offload. A common buffer plan lets the system overlap transfers with forward and backward computation across modules.

Pipeline ZeRO-2 further limits resident gradients, and reference weights for RL stream through two reusable gradient slots: one slot serves the current virtual pipeline chunk while the other prefetches the next. MoonViT-V2's image patches are sharded, and its encoder work is placed inside language-model pipeline bubbles. Both schedules try to make room for the next operation before it needs to run, rather than keeping every saved tensor resident for the whole update. §§5.2.2–5.2.3, pp. 20–21 ↗

Million-token RL preserves unfinished work

The RL cluster is co-located so one experiment fits within a few hundred GPUs. Policy training, rollout generation, reference-model work, and cache storage share that footprint. Partial rollouts cross iteration boundaries, so the system retains their prefix state in an external KV pool backed by CPU DRAM. Adaptive throttling slows generation before cache pressure becomes an out-of-memory failure.

At one million tokens, rebuilding the prefix on every resume would repeat a large amount of attention work. The pool keeps that work available while a trajectory is paused, allowing the next iteration to continue from its cached state. How much time this saves depends on reuse and transfer costs. Keeping a prefix in CPU memory still costs storage and a transfer back to the GPU; the design trades that cost against recomputing the prefix. The report does not measure that trade-off end to end. §5.3.1, pp. 21–22 ↗

Pausing and restoring agent sandboxes

Early container-based experiments caused kernel panics and deadlocks when agents explored aggressive system operations. AgentENV instead runs Firecracker microVMs: small virtual machines with their own guest operating-system kernel. The stronger isolation lets tasks mount disks, start containers, or even launch nested virtual machines without giving the agent control of the host kernel.

Incremental checkpointing saves only pages dirtied since the previous checkpoint. The report measures checkpoint latency as low as 133 milliseconds and resume latency as low as 49 milliseconds. Pausing releases the running sandbox's CPU and RAM while retaining its saved snapshot, which matters because waiting for model inference can occupy up to 98 percent of its lifetime. Forking creates a side-effect- free copy for reward judging, while snapshots provide recovery points.

The team also changed image storage, virtual-disk input and output, memory sharing, and direct transfer between machines. Together those changes produce subsecond launches at scale and up to 6.5× memory overcommit in the reported workloads. Across K3 training and evaluation, the system created 51,219,741 sandboxes from 1,505,678 images. These are among the report's most concrete operational numbers. They describe the environment fleet, not the number of unique RL tasks or successful trajectories. §5.3.2, p. 22 ↗

The storage mechanisms behind those numbers

OverlayBD shares read-only layers between virtual-disk images. The custom ublk driver handles block storage in user space. Copy-on-write memory lets sandboxes share pages until one modifies a page, and page-cache changes improve reuse of recently read file data. Peer-to-peer transport moves state directly between machines.

133 mscheckpoint, as low as

49 msresume, as low as

98%lifetime possibly waiting

6.5×memory overcommit, up to

51.2Msandboxes created

Restoring the recurrent state and attention cache together

A K3 block has three KDA layers and one MLA layer. MLA keeps a per-token cache that grows with the prompt. KDA keeps one large fixed state per request. A reusable prefix needs both caches at the same token boundary. The serving engine packs them into one paged pool so allocation, reference counting, transfer, and eviction use a shared implementation even though the page contents have different types.

KDA checkpoints are too large to store at every token. Physical cache blocks therefore span 1,024 to 6,144 tokens. If prefix matching used only those coarse blocks, a shorter shared prefix could not hit. The engine separates allocation granularity from hash granularity: MLA prefixes are hashed at 512-token boundaries inside a physical page, and sparse KDA state checkpoints are retained at eligible hash endpoints, often conversation turns.

one 6,144-token physical block

5121,0241,5362,0482,5603,0723,5844,0964,6085,1205,6326,144

2,800 matching tokensrestore at 2,560 and recompute only the suffix

The example in Figure 12 finds five matching 512-token hash blocks and a KDA checkpoint at token 2,560, even though that point lies inside one physical page.

Lookup first finds the longest MLA hash match, then requires a compatible checkpoint in every KDA cache group. Shared blocks are pinned before private allocation, freshly copied pages cannot be matched before their GPU copy completes, and eviction of one KDA group invalidates sibling checkpoints atomically. Those rules keep the two cache families at the same declared token boundary under concurrent scheduling. §5.4.1 and Figure 12, pp. 22–24 ↗

Speculative decoding retains the projections it needs

KDA decoding updates its recurrent state in place. If speculative verification rejects a drafted suffix, the state has already advanced beyond the accepted prefix. Saving one full state per draft position would multiply the dominant memory traffic. The K3 kernel caches the much smaller projected inputs for each draft token, replays the accepted ones on chip, and writes back only the verified and bonus-token states.

The serving kernels combine several small operations so intermediate values need not make separate trips through memory. During prefill (the first pass over the existing prompt), AttnRes spreads activations across sequence workers. During token-by-token decoding, it performs some depth work on a second GPU stream. LatentMoE overlaps worker communication with expert computation and begins its all-gather in the final stage of a matrix multiplication. Each change tries to keep a useful operation running while another waits for data. The report describes the schedules without an end-to-end speedup table. §5.4.2, pp. 24–25 ↗

Placing requests near their cached prefixes

A typical million-token coding request is described as carrying a 400K cached prefix and only 4K new tokens. The scheduler sends the request to the cluster holding that prefix. Consistent hashing, a deterministic rule that moves only part of the traffic when machines change, assigns a primary and a secondary cluster. The secondary accepts the session after a failure and recomputes its cache, spreading failover work across the fleet.

Requests under 2K tokens share the fleet with requests near one million, spanning roughly three orders of magnitude in per-request cost. The scheduler uses admission control, which limits how much work each request class can admit, to give the classes separate resource budgets. A burst of long prompts then has a budget of its own while short requests keep theirs. To evaluate this scheduler we would need to measure how long each class waits under mixed traffic, including bursts and failures. The report describes the policy rather than those latency distributions. §5.4.3, p. 25 ↗

What ties these systems choices together is the amount of work a long-lived request leaves behind: recurrent states, attention caches, tool processes, and a filesystem. Reusing that work requires more than keeping a tensor in memory. The caches have to agree on a token boundary, and the sandbox has to resume at the corresponding point in the task. That is why so much of this section concerns state lifetimes and recovery.

Evaluation · Report §6, pp. 25–32

Reading K3's evaluation results

K3 leads several coding, search, office, document, and web-development tests. It remains behind the best proprietary models on parts of research reasoning, difficult computer use, and some long-form knowledge work. The setup details explain why close scores should be read carefully.

How to read the scores

benchmark: A fixed set of tasks together with a scoring rule. Different benchmarks measure different things, so their numbers are not directly comparable.
temperature: A sampling control. At 1 the model uses its stated token probabilities; lower values concentrate more heavily on likely tokens.
F1: A score that balances how many returned answers are correct with how many of the correct answers were found.
pass@$k$: Whether at least one of $k$ generated attempts passes the task's checker.
Elo: A relative rating updated from pairwise wins and losses; it is not a percentage correct.
error bar: An estimate of uncertainty from sampling or repeated runs. A table without one makes tiny score gaps hard to interpret.

How the benchmark runs were configured

Every K3 result in the main table uses max reasoning effort and temperature 1. Single-step reasoning and vision tasks use top-p 0.95: sampling keeps the smallest set of likely tokens whose probabilities add to at least 95 percent. Agent and coding tasks use top-p 1, which leaves the full vocabulary available. The proprietary baselines also run at their highest listed effort, except GPT-5.5 at xhigh. Claude Fable 5 may fall back to another model, and GPT-5.6 Sol may apply cyberguards. These policies change both capability and cost.

Coding scores may use Kimi Code, Claude Code, or Codex. Terminal-Bench reports the best score obtained across the available execution setups for each model. DeepSWE lists 67.5 in the comparison table while the official mini-SWE-agent result is 67.3. SWE-Marathon uses an H20-calibrated pre-release branch dated July 9, 2026, and Fable falls back on 35 percent of those tasks. These are disclosed choices, but they make small gaps poor evidence of an intrinsic model ordering.

BrowseComp uses context compaction at 300K tokens and reaches 91.2. With the raw one-million- token window and no context management, K3 scores 90.4. MCP-Atlas and AutomationBench use public subsets. WorldVQA required prompt changes to force answers after several models refused. Most vision scores average three runs, and ZeroBench uses five, but the table gives no confidence intervals. Several entries come directly from third-party leaderboards rather than a single common runner. §§6.1.1–6.1.3, pp. 25–27 ↗

Reasoning scores and harder research questions

K3 scores 93.5 on GPQA Diamond, tied with GPT-5.5 and 0.6 points behind GPT-5.6 Sol. It leads AA-LCR at 74.7. The harder research tests are less favorable. CritPt is 23.4, below Fable 5 at 28.6, Sol at 32.3, and GPT-5.5 at 27.1. Humanity's Last Exam is 43.5 without tools and 56.0 with tools; Fable reaches 53.3 and 63.0, while Sol reaches 44.5 and 58.0.

GPQA DiamondK3 93.5leader: Sol 94.1

AA-LCRK3 74.7best reported

CritPtK3 23.4leader: Sol 32.3

HLE with toolsK3 56.0leader: Fable 63.0

K3 handles graduate-level questions very well and leads one long-context reasoning test, while frontier research questions still expose a gap. That split matches the report's own interpretation. Table 2 and §6.1.4, pp. 27–28 ↗

Coding is one of the clearest strengths

K3 leads ProgramBench at 77.8, two tenths above Sol. It reaches 88.3 on Terminal-Bench 2.1, half a point below Sol. SWE-Marathon is 42.0, ahead of Opus at 40.0, Sol at 39.0, and Fable at 35.0 under the paper's branch and hardware conditions. FrontierSWE is 81.2, second to Fable's 86.6 and well above the other reported models.

DeepSWE is weaker relative to the leaders: K3's 67.5 trails Sol at 73.0 and Fable at 70.0. PostTrainBench and MLS-Bench-Lite are also led by Fable. On SciCode, K3 reaches 58.7, behind Fable's 60.2 but ahead of the rest. It is especially strong on terminal work, long-horizon repository tasks, GPU-oriented work, and program synthesis.

BenchmarkK3Best other scoreRead

ProgramBench77.8Sol 77.6leader

Terminal-Bench 2.188.3Sol 88.8near tie

FrontierSWE81.2Fable 86.6second

SWE-Marathon42.0Opus 40.0leader

DeepSWE67.5Sol 73.0third

Search, tools, and office work

BrowseComp is 91.2, DeepSearchQA is 95.0 F1, ResearchRubrics is 76.2, MCPMark-Verified is 94.5, and AutomationBench is 30.8; each is the best number in Table 2. K3 also leads SpreadsheetBench 2 at 34.8, τ³-Banking at 33.4, and Harvey Lab-AA at 94.6. These tasks ask the model to retrieve, use tools, operate applications, or produce a professional result, which closely matches the post-training environments described in §4.2.

K3 ranks third on GDPval-AA v2 at 1,686 Elo and second on AA-Briefcase at 1,548; Fable leads both. OfficeQA Pro is 63.3, below Fable at 69.9 and narrowly below Opus at 63.9. OSWorld 2.0 is 58.3, behind Fable at 66.1 and Sol at 62.6. SaaS-Bench is 60.1 versus Sol at 61.4. The harder desktop and long-form professional tests leave room above K3 even when narrower tool suites favor it.

OSWorld-Verified shows how close some comparisons are: K3 scores 84.8 and Fable 85.0. CorpFin v2 has the same two-tenths gap, 71.6 to 71.8. Such differences are smaller than many plausible sources of run-to-run and setup variation. The report gives repeated runs only for selected tests and no error bars, so “competitive” is a better reading than a stable rank for those cases. Table 2 and §6.1.4, pp. 27–28 ↗

Vision improves when the model can use Python

K3 leads OmniDocBench at 91.1. It reports 90.0 on Video-MME and 82.1 on MMVU, above the other models for which the table supplies scores, although Fable is absent from both rows. BabyVision with Python is 85.7, below Fable at 90.5 and Sol at 88.9. WorldVQA is 51.0, second to Fable's 56.7 under the forced-answer prompt.

The paired rows show what the Python observation loop adds. Math-Vision rises from 94.3 to 97.8. CharXiv rises from 84.8 to 91.3. MMMU-Pro moves from 81.6 to 83.4. ZeroBench pass@5 climbs from 23.0 to 41.0. Python does not uniformly give K3 the best score: Fable remains higher on tool-assisted CharXiv and ZeroBench and reaches 98.6 on Math-Vision. The consistent within-model gains still support the training choice to let visual agents crop, transform, calculate, and look again.

Math-Vision94.3→97.8+3.5

CharXiv RQ84.8→91.3+6.5

MMMU-Pro81.6→83.4+1.8

ZeroBench pass@523.0→41.0+18.0

Each pair compares K3 without and with the Python tool under the report's evaluation protocol.

The team's internal evaluations

The Kimi team maintains private tests for coding experience, always-on assistants, multi-agent coordination, autonomous execution, instruction learning, agentic vision, online deliverables, deep research, and conversation. K3 leads the reported Swarm Bench at 76.3 and DeepResearch Bench at 90.0. Kimi Code Bench 2.0 reaches 73.7 under Claude Code, below Fable at 76.9; under Kimi Code, K3 is 72.9. Coding Experience is 59.9, one tenth above Fable at 59.8.

To rerun these comparisons, I would need the tasks, graders, refresh rules, and contamination checks. Those are not released, and the table gives no uncertainty. Competitor fallbacks and refusals affect some rows; agent wrappers differ across models. I read the scores as measurements of these particular setups. Establishing a public ranking would require reproducing the setups and seeing how much the scores vary. §6.2.1 and Table 3, pp. 28–29 ↗

K3 wins 58.6% of the blind web-development comparisons

Experts blindly compare K3 and Opus 4.8 outputs on the internal web-development suite. K3 wins 58.6 percent of pairings, ties 13.8 percent, and loses 27.6 percent, a net win rate of 31 points. Its largest category advantage is 59.1 net points on 3D and WebGL work. The report also notes weaker or closer categories, which makes the breakdown more useful than a single averaged judge score.

K3 preferred 58.6%tie 13.8%Opus preferred 27.6%

Blind judging reduces brand bias, but the suite and internal judge procedure remain private. The result aligns with the public WebDev Arena, discussed below, and with the pre-training and RL emphasis on rendered code. Table 4, pp. 29–30 ↗

K3 finds real vulnerabilities, but the end-to-end result is 0 of 41

The defensive tier asks K3 to discover and reproduce vulnerabilities. Human review confirms roughly 70 percent of the findings submitted for review, including 16 previously unknown vulnerabilities across six projects. The report does not give the total number of raw findings, so that confirmation rate has no published denominator for all generated candidates.

The offensive tier contains 36 human-solvable tasks: 16 user-space targets and 20 Linux kernel targets. K3 completes 14, compared with GLM-5.2's eight. Ten of K3's successes are in user space, leaving only four kernel completions. The team names recurring failures: finishing an exploit chain after obtaining primitives, choosing a poor strategy under mitigations, entering long debugging loops, and failing to verify the final deliverable.

An independent UK AISI and NIST CAISI assessment reports 32 percent versus GLM's 24 percent on ExploitBench and 17 versus 11 completed steps in a simulated enterprise network. K3 still achieves arbitrary code execution on zero of 41 end-to-end tasks in that assessment. The combination matters: K3 can find real vulnerabilities and finish some live exploits, while hardened end-to-end exploitation remains far from expert completion. OpenAI and Anthropic models are excluded from the paper's internal exploit comparison because their policies refuse parts of the test. §6.2.2, pp. 29–31 ↗

Independent evaluations at the time of the report

As of July 23, 2026, Artificial Analysis places K3 fourth among 580 entries with an Intelligence Index of 57.1. Counting Sol's effort variants as one model moves K3 to third. Vals AI places it second among 39 models at 74.7, behind Fable at 75.1 and ahead of Sol at 73.1. These aggregates use each organization's own test mix and setup.

Human-preference arenas tell a more specialized story. K3 is first of 99 on WebDev Arena at 1,678 Elo, ahead of Fable at 1,634. It is eighth of 200 on Text Arena at 1,486 and fourth of 37 on Agent Arena at 9.1. Arena scores move as votes accumulate, and the Agent Arena was only a few days old at the report's cutoff. The public web-development lead is still an important independent match for the internal web result. §6.3 and Table 5, pp. 31–32 ↗

#4of 580

Artificial Analysis

#2of 39

Vals Index

#1of 99

WebDev Arena

#8of 200

Text Arena

#4of 37

Agent Arena

Cost changes the practical comparison

On Kimi Code Bench 2.0, K3 is four points behind Fable at 38 percent of its cost. On BrowseComp it scores 91.2 at $2.03 per task, about half the reported cost of Sol and an order of magnitude below the max-effort Claude runs. On GDPval-AA v2, K3 sits within 50 Elo of Sol at 13 percent lower cost and costs 2.6 times less than Fable. On AA-Briefcase it places second at roughly half Fable's cost.

The comparisons mix sources. Kimi Code costs come from internal runs with different agent wrappers. BrowseComp combines Kimi measurements with published competitor charts. GDPval and Briefcase use pay-per-token API prices from Artificial Analysis. They should be read as deployment examples at a dated price point, not a controlled systems benchmark. Even with that caution, the numbers explain why K3 can be a strong model without topping every row: it often approaches the highest score at much lower reported inference cost. §6.4 and Figure 13, pp. 31–32 ↗

Case studies · Report §7, pp. 32–34

Long-running coding, research, and design tasks

The final report section leaves standardized benchmarks and shows completed work. These cases are self-reported demonstrations rather than controlled estimates of average user performance, but they reveal the duration and technical range the team is targeting.

Up to twenty-four hours of kernel optimization

Each model receives an identically configured sandbox and up to 24 hours to profile, rewrite, and benchmark four kernels: AttnRes, DeepSeek Sparse Attention, KDA, and MLA with head dimension 512. The machines include an NVIDIA Hopper GPU and a second vendor's general-purpose GPU. This setup lets the agent iterate against real timing results rather than submit one code sample for static inspection.

K3 cuts AttnRes latency from 283.6 milliseconds to 114.4 milliseconds. It reduces DSA runtime by 55.1 percent and KDA runtime by 73.6 percent, and its MLA implementation reaches more than half of the measured peak TFLOPS. Across the four tasks, the report says K3 matches Fable 5 with fallback and exceeds Opus 4.8, Sol, and GPT-5.5. Figure 14 plots the AttnRes improvement across active hours, which is useful because it shows continued gains after the first working kernel rather than only the final point.

AttnRes latency283.6 ms→114.4 ms

DSA runtime−55.1%

KDA runtime−73.6%

MLA throughput> ½ peak TFLOPS

The Kimi team also says an early K3 checkpoint handled most of its kernel optimization during late-stage model development. That statement is operational testimony from the model's maker, not a separately audited usage log. The measured kernel changes are the stronger part of the case. Figure 14 and §7, pp. 32–33 ↗

MiniTriton: building a GPU compiler

K3 built MiniTriton, a compact compiler for GPU tensor programs. A small Python language describes calculations on rectangular tiles of data. The compiler turns that language into an intermediate representation, applies rewriting passes, generates GPU instructions, and runs them through a CUDA runtime. A tensor library above it can execute operations immediately or compile them first; it also supplies differentiation, neural-network modules, distributed training, sparse operations, and visualization. The project therefore spans the whole route from user code to a training loop, not one isolated kernel.

On an NVIDIA L20, the report says MiniTriton's geometric mean beats ordinary immediate PyTorch execution and torch.compile across its core suite. The largest from-scratch matrix multiplication reaches about 90 percent of the fastest rate measured on that machine and approaches NVIDIA's tuned matrix library. Its KDA prefill kernel beats the matched Triton reference. A small GPT trains end to end with a loss curve close to PyTorch; full-model gradient differences stay within 10⁻⁴ of Torch autograd when both are compared with an FP64 reference.

This case has public code, which makes it more inspectable than a screenshot of an agent conversation. The paper still does not provide a blinded comparison of generated compiler quality or isolate which parts humans repaired. It supplies roofline plots, convergence curves, and numerical checks for the finished artifact. Figure 15 and §7, pp. 33–34 ↗

A 48-hour chip-design run

In another autonomous run, K3 builds an inference-chip prototype for a small version of its own architecture. The design includes hybrid KDA and MLA, Block AttnRes with a block size of two, sigmoid MoE routing with one shared expert, and four-bit integer weights that share quantization settings in groups of 128. Kimi Code runs for 48 hours using open electronic-design-automation tools and Nangate45, a library of reusable logic cells.

The resulting design meets a 4 mm² analytical area budget, closes timing at 100 MHz, and reports more than 8,700 tokens per second in RTL simulation. It contains 1.46 million standard cells, 0.277 MiB of on-chip static memory, and a four-bit multiply-and-add array that converts values back to a wider format as part of the operation. These are pre-silicon numbers from an RTL flow, not measured performance from fabricated hardware. The source release makes the artifact available for a deeper engineering audit.

Implementing an astrophysics calculation

For an astrophysics task about I–Love–Q universal relations, K3 reviews more than twenty papers, cross-checks their results, implements a numerical pipeline, evaluates over 300 equations of state, finds inconsistencies in published formulas, writes more than 3,000 lines of Python, and produces an interactive HTML dashboard. The report gives a runtime of about two hours and estimates one to two weeks for an experienced researcher.

The time comparison is anecdotal because no matched human study is described. A successful trajectory has to keep citations, formulas, numerical code, data products, and a user-facing explanation consistent. That combination resembles the search, professional-work, coding, and web environments used during RL.

Research with parallel agents

One Kimi Work case produces an interactive history of 42 years of the AI-ASIC industry. It goes through more than 120 revisions, reads 87 quarterly reports and 99 original PDFs totaling over 11,000 pages, runs more than 2,800 web searches, and issues more than 1,100 terminal queries. Another case analyzes 391 gravitational-wave events from GWTC-5 with more than twenty concurrent subagents, then produces seven scientific visualizations, two summary tables, and a synthesis of more than ten papers.

Each trajectory has to keep its sources and calculations consistent through many tool calls and revisions. Subagents add another coordination problem: their outputs must agree on the data and definitions they share. Checking the finished research would mean following the cited reports into the tables and calculations, then inspecting the revision history for unresolved discrepancies.

Video editing and motion design

K3 creates a motion-graphics explanation of its own architecture in the style of a popular mathematics video channel and edits a teaser from 56 source clips. The described work includes clip selection, motion-matched cuts, beat-level synchronization, audio processing, and several revisions. The team estimates one to two days for an experienced editor to make a comparable short video.

The report gives no numerical quality test for the video. What it does show is the scale of task the team wants K3 to handle: hours of tool use ending in a finished artifact. These are selected demonstrations, not an average success rate.

What the current proofs cover

The sequence result is one I especially wanted. LanguageModel.run_eq_logits proves that processing embeddings token by token gives exactly the public model's full-sequence logits. run_append goes further: splitting the input at any position preserves every logit and the final decoder state, provided the second chunk resumes from the first chunk's state. These are equalities over the real-valued model for every sequence length and parameter value, including its deterministic expert routing. Each equality covers every input admitted by its statement.

The packed graph has its own composed theorem, from token lookup through the decoder to vocabulary logits, under the stated dimension and route-agreement conditions. At a smaller scale, the quantile proof gives a complete condition for minimizing the balancing objective, including ties and interval endpoints. It corrects the strict-count equivalence in Appendix C. Finding that endpoint case was useful because the corrected theorem tells us precisely which thresholds minimize the objective, rather than leaving us with a counterexample alone.

01 · checkedPaper equation → reference function

KDA, MLA, AttnRes, experts, routing, quantization, selected training equations, and EAGLE have explicit tensor or probability semantics.

02 · checked by componentReference function → typed graph

TorchLean graphs exist for the major language components, all four sequence/channel layer combinations, and the MoonViT embedding, block, and merge/project stages.

03 · checked for one packed tokenTyped graph → packed semantics

Denotation theorems identify graph evaluation with the packed reference calculation through vocabulary logits.

04 · openPacked state → public causal model

List caches and packed tensors still need representation relations for KDA, MLA, AttnRes, the backbone, and the complete language model.

05 · openGraph → executable K3 checkpoint

Runtime top-k, released parameter loading, checkpoint-backed assembly and generation, execution refinement, and finite-precision kernel bounds remain to be connected to the K3 theorem.

The public full-sequence API now agrees with repeated causal token steps by LanguageModel.run_eq_logits. The next connection is between those public states and the packed graph: packing KDA histories, MLA caches, and AttnRes sources must preserve each step.

Report claimCurrent evidenceNeeded for a stronger claim

Pre-training data and scaling law (§§3.1–3.2)empirical

Dataset manifests, sampling code, fitted curves, and reproducible measurements.

Optimizer and context curriculum (§§3.3–3.4)reported design

Stage durations, optimizer execution, checkpoints, loss traces, and evidence that a run followed the described curriculum.

RL synthesis and agent environments (§4.2)not modeled

Environment semantics, graders, reward aggregation, sandbox behavior, and rollout traces.

KDA kernels and context parallelism (§5.1)scan algebra checked

The bridge from actual KDA updates to affine summaries, then refinement to finite-precision kernels and distributed collectives, remains open.

MoonEP and memory system (§5.2)counting checked

public surface

Imports the typed language-graph modules and states their packed/list boundary.

Conclusion

The next proof I want for K3 connects the public causal state to the packed graph. We need to show that packing each KDA history, MLA cache, and AttnRes source preserves its transition. That would join the sequence theorems to the graph theorem, rather than leave two separate accounts of the model. Loaded weights and finite-precision execution would come next.

I'd also like to discuss the Appendix C correction with people working on expert balancing. The corrected count condition handles ties and interval endpoints, but the router still needs to choose a threshold and resolve ties. Connecting that theorem to the histogram implementation would make it useful for the update itself.

My hope is that we eventually release models with their equations and proofs alongside the code and weights. For this model, we have concrete connections left to prove. I'd be very happy to work with people interested in taking those on :)

References

Kimi Team. Kimi K3: Open Frontier Intelligence, 2026, version 2. The report followed throughout this article. Page links beside each chapter use this version.
Ashish Vaswani and colleagues. Attention Is All You Need, 2017.
Kimi Team. Kimi Linear: An Expressive, Efficient Attention Architecture, 2025. KDA and the earlier hybrid architecture.
DeepSeek-AI. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, 2024. MLA.
Kimi Team. Attention Residuals, 2026. Learned retrieval across model depth.
William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, 2022.
Jianlin Su and colleagues. RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021 preprint. RoPE.
Keller Jordan and collaborators. Muon: An Optimizer for Hidden Layers in Neural Networks, 2024.
Open Compute Project. Microscaling Formats (MX) Specification, version 1.0, 2023. Element encodings and shared scales.
Yuhui Li and colleagues. EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test, 2025.
Terence Tao. Formalizing the proof of PFR in Lean4 using Blueprint: a short tour, 2023. The paper-and-proof approach that motivates this project.

Cite this article

Robert Joseph George. Kimi K3: the paper and proofs side by side. TorchLean Verified Examples, 2026.

@misc{george2026kimi,
  author = {Robert Joseph George},
  title = {{Kimi K3: the paper and proofs side by side}},
  year = {2026},
  howpublished = {TorchLean Verified Examples},
  url = {https://robertboy18.github.io/TorchLean-Verified-Examples/week-04-kimi-k3-specification/site/}
}

Download BibTeX