Notes from an NVFP4 port of Kimi K3.1 Pro
We integrated Kimi K3.1 Pro into our own C++/CUDA inference engine (no PyTorch
at runtime). It is the same 2.8-trillion-parameter, ~104-billion-active
KimiLinear tower as K3 — 93 layers at hidden 7168, 69 linear-attention (KDA)
layers, 24 gated NoPE multi-head-latent-attention layers, an attention-residual
depth mixer every 12 layers, and a 896-expert latent MoE that runs 16 experts
per token in a 3584-wide latent — with one addition: the 24 full-attention
layers now carry a Lightning Indexer that keeps the best 2048 keys per query.
Of the 497 340 tensors in the release, 497 220 have exactly the names K3 uses
and the 120 new ones are the indexer — five on each full-attention layer; the
values themselves are re-trained, not copied. The checkpoint is 1.5 TB because
the routed experts are already 4-bit (MXFP4, E2M1 with one E8M0 exponent per 32
columns) while everything else is BF16. MXFP4 and NVFP4 share the E2M1 grid, so
the experts repack losslessly — nibbles verbatim, each E8M0 byte duplicated
into the two 16-wide halves of its 32-block, the exponent bias folded into one
per-tensor power-of-two global scale: 82 432 expert files, repacked as pure I/O
inside a 17-minute conversion, with no GPU in the expert path.
The indexer is the entire delta, and the checkpoint does not explain it. It
is 32 heads of 128 dimensions scoring relu(q·k) against the whole causal
history and keeping the best 2048 positions exactly; the query comes off the
normalised q-LoRA projection, the key from the attention input through a
128-wide LayerNorm that carries both a weight and a bias, and a per-head weight
vector sums the heads before the select. The release ships the tensors and adsa_config but no modeling code at all — the repository carries the
configuration classes and the encoding helpers and nothing else — so the
selection arithmetic had to be reconstructed from the shapes and the published
method. We filed that as open debt rather than presenting it as a verified port,
and we never quantize the five indexer tensors: they are a routing signal, and
4-bit rounding in a routing score does not blur an answer, it swaps the key set.
Details that silently change the answer. (1) The variant marker lives in the
wrapper config (model_type: "kimi_k31"), while text_config.model_type iskimi_linear — identical to K3's — and dsa_config is nested insidetext_config. A loader that reads either field alone turns K3.1 Pro into K3: it
loses the indexer and still answers fluently. We gate on all three (wrapper
type, dsa_config, indexer tensors present) in both directions. (2)dsa_config.dsa_layers is 1-based [4, 8, …, 92, 93] over 93 layers; the
0-based set is exactly K3's full-attention table {3, 7, …, 91, 92}, so this
variant indexes every full-attention layer, not a subset. (3) index_rope_dim
is 0 — no rotary on the indexer's queries or keys. (4) Crossing the top-k
boundary is a tie, and the selection emits positions strictly above the
threshold in ascending order first, then the boundary group from the lowest
position up — not a value-sorted list. (5) The K3 traps still hold: its MLA
layers are NoPE (the 64 "rope" dims are plain key dims), KDA's decay is
channel-wise from a 128-entry A_log, depth is mixed by attention residuals
every 12 layers, the activation is SiTU-GLU, and the model answers inside athink channel of the XTML frame.
Measured. SelfCheck on the converted tree (framed control prompt) PASS ×3
cycles, bit-exact — logits_hash 98ec1bbb372c9181 and top-1 id 40 on every
cycle — and a native tool call round-trips, get_weather {"city":"Paris"},
through the checkpoint's own XTML frames. The tree is 1.5 TB: 23 GB of pinned
weights (attention, dense MLP, norms, router, and the 120 indexer tensors kept
in 16-bit) plus 8.9 GB of latent/shared plane streamed per layer, with the 1.4 TB
of routed experts behind a host-side store. That is what buys 2.8 T parameters
on one 32 GB card — and it is also what sets the speed, because a token routes
into 1472 distinct experts (92 layers × top-16, 17.72 MB each, 25.5 GiB of
reads per token).
The first version of this port decoded at 6.95 s/token, and that was our bug,
not the model's. The routed-expert window was 512 against a working set of
1472; the layers are visited once each per token in ascending order, so with a
window smaller than one token the access pattern is a cyclic scan whose hit rate
is zero by construction — every expert was re-opened, re-mapped, re-parsed and
evicted. A sweep over the window, same binary, same weights, bit-identical
logits every time: 512 → 6.95 s/token, 1472 → 6.05, 2048 → 5.57,
65 536 → 2.98 (K3; the Pro tree carries the indexer on top and lands at
3.23–3.35). Two details mattered as much as the number. A large window had
looked expensive because the store eagerly pre-parses max_open experts at
load; those experts are chosen in path order, not decode order, so the
pre-parse cost 156 s and bought nothing — switching it off gives the same
rate with a 5-second load. And the obvious fix, prefetching each layer's
top-k into the page cache in parallel (serial cold reads run at 1.8 GB/s where
32 threads reach 39.8 GB/s, so it looked like an easy 20×), measured worse:
reading the files into a discard buffer cost 7.16 s/token because that traffic
is pure waste once a page is cached, and fadvise(WILLNEED) was a wash. We
removed it rather than ship it. The mechanism that pays is retention, not
prefetch — never unmap, and a re-read becomes a cache hit.
Eviction was the other half of the same loop. Every evicted expert used to
force a full stream synchronize, so the host waited for unrelated queued work
each time an expert aged out — 20.3 ms per eviction, measured against 20.2 ms of
unrelated traffic. That is now a CUDA event recorded on the copying stream: the
mapping is retired, not unmapped, and released once the copies queued before it
have executed, which measured 0.07 ms instead of 20.3 ms. At the 512-slot
window (1472 evictions per token) that alone took the decode from 6.95 to
6.26 s/token; at the shipped 65 536-slot window evictions are rare by design, so
nothing moved. The first version of that fix regressed another model by 6.6 %,
because it queried every retired event on every eviction; the events complete in
FIFO order on one stream, so only the front needs checking — a reminder that the
correct mechanism can still be shipped with expensive bookkeeping.
What that leaves honest. 3.2–3.35 s/token is still 7.6–8 GB/s against the2 s/token`
~40 GB/s the array can do, so the remaining lever is real but structural:
overlap the next layer's expert staging with the current layer's GEMMs, which
means replacing the synchronous per-expert residency-and-copy loop — shared
code, not a Kimi-local edit. It is filed as debt with the measurement attached,
not as a closed question. And for the record, the "4× slower than K3" figure
that this variant was first reported against does not exist: K3's `
was prose written minutes before its own run log recorded 6.7 s/token, and at
identical binary, geometry and window the two models measure 6.95 vs 6.98 — a
ratio of 1.00. A prose number in a report becomes a benchmark nobody measured;
that is what this entry cost us.
Not measured (and why). The source cannot execute on a single 32 GB card (it
needs fla-core plus a multi-GPU MXFP4 runner), so the FSI/CODE/DET goldens are
template-derived — payload-identical to K3's, the prompt set's own Expected
Template — and at the shipped 3.2–3.35 s/token one 2048-token prompt is ~1.8–1.9 h
while the full parity-plus-throughput sweep is ~14.7–15.3 h of exclusive GPU.
Those rows are declared debt with that measurement attached rather than quietly
skipped; the kernels are pinned instead by host-side reference tests against
independently computed arithmetic. And even a green row would not verify the
indexer: at the 2048-token protocol every key fits inside the 2048-key budget,
so the selection is dormant by construction and the sparse path is bit-identical
to the dense K3 path. The measured numbers therefore describe the dense path,
which is the honest thing to publish.