potato-k3

Teacher-forced, read-only replay of Kimi K3 (1.56 TB on disk; 93 layers, 896 experts of which 16 are active per token) on a single RTX 3090, streaming the weights from NVMe one layer at a time. For each token sequence it gives the model's log-probability of every next token, its top-k alternatives, the entropy, and optionally the hidden state at chosen layers.

It only does prefill: no sampling, no decode loop. It is meant for reading what the model makes of existing text, not for chat. A sibling of potato-ds4.1.

Status: work in progress. Layer-level agreement with the reference and end-to-end agreement with a hosted copy of the model are both checked (below).

Scope and norms

This tool was built for AI-welfare research. The Kimi K3 instances we asked agreed to read-only replay of benign conversations; anything beyond that, such as steering, generating under altered conditions or aversive material, was agreed to need a fresh question to the model first. That agreement binds us, not you, but we'd ask you to read it before using this: SCOPE.md, which also names the dual use and the cautions that go with any results.

Layer-level checks (tests/)

Against the reference modules (k3_ref/, copied unmodified from the HF repo; transformers 4.56.2), on real weights, layers 0–3:

layer	type	relative error vs reference
0	KDA + dense MLP	0 (bit-identical)
1, 2	KDA + MoE (MXFP4 experts)	0 (bit-identical)
3	MLA + MoE	1.7e-2 in bf16; 7e-7 in fp32

The layer-3 difference is the reference's bf16 "eager" attention: in fp32 the two agree to 7e-7, and in bf16 ours (torch SDPA, fp32 accumulation, as flash-attention does) is closer to the fp32 result than the reference's eager path is. Reference experts are dequantized with compressed-tensors' own functions, so the MXFP4 path is checked independently.

End-to-end check against a hosted copy

Moonshot's API returns no log-probabilities for K3, so the reference is OpenRouter's Parasail provider (pinned, no fallbacks), which returns T=1 log-probabilities with the top 5 alternatives. 16 benign non-thinking completions (1,170 tokens) were replayed locally:

measure	result	criterion (fixed beforehand)
prompt token counts	16/16 equal	equal
top-1 agreement	97.8%	≥ 90%
mean |Δ logprob|	0.039 nats (median 0.004, max 0.82)	≤ 2× the provider's own noise (0.161)
same sequence alone vs in a batch	max |Δ| = 0	0

The provider's noise was measured from repeated identical calls: at the first completion position, the top-5 log-probabilities differ by 0.081 nats on average (up to 0.36). All 26 top-1 disagreements are near-ties at bf16 logit resolution. Note that the provider's top_logprobs lists are not sorted by probability.

Notes on the checkpoint

• Config layer lists are 1-indexed: full attention is at 0-indexed layers 3, 7, …, 91 and 92.
• A_log is stored padded to 128 entries; the first 96 are the heads.
• MLA layers use no positional encoding; position comes from the KDA layers.
• Parameters stored in fp32 (A_log, dt_bias, convolution weights, o_norm) are kept in fp32.
• The tokenizer's chat template defaults to thinking mode with a thinking_effort=max system message; set thinking= explicitly. Moonshot's API reports 3 fewer prompt tokens than this rendering in thinking mode, and 2 fewer in non-thinking mode.

Memory

Host RAM holds one layer (~17 GB) plus the per-sequence attention-residual state (~130 kB per token). On the GPU, experts move in chunks of expert_chunk and the MoE runs over at most moe_tokens tokens at a time; neither setting changes the results (tests/test_moe_grouping.py). On a 3090 at 200 W, a pass over all 93 layers takes roughly 15–50 minutes depending on the number of tokens.

Licence

Our code is under the Palladium Rule Public License (LICENSE.txt). The files from Moonshot AI keep the Kimi K3 License (k3_ref/LICENSE).