Apply rotary to MLA's rope dims on models that are not NoPE - #27
Conversation
`src/` implemented no rotary: every occurrence of `rope` was `qk_rope` used as a width, and `rope_theta`, `rope_scaling` and `mla_use_nope` were read nowhere. That is correct for the Kimi models, which set `mla_use_nope` and pass those dims through unrotated, and wrong for a DeepSeek-V3 checkpoint, which sets no such flag and ships `rope_theta` with YaRN scaling. In MLA those dims are the only positional signal — the nope dims are position-free by construction — so a V3-family container attended over an unordered sequence. It is quiet: lexically determined answers still come out right, so a single-turn factual prompt looks correct. Add a second turn boundary and the model emits an empty assistant turn, `<|im_end|>` at p=0.968 on Kimi-K2, because it cannot order the prompt. `rope_init` builds `inv_freq` with YaRN's ramp and takes `mscale_all_dim` squared onto the attention scale, following `DeepseekV3YarnRotaryEmbedding`. Two details do not survive paraphrase: the rotation is GPT-J interleaved, pairing `x[2j]` with `x[2j+1]` — upstream reaches the same arithmetic by de-interleaving before a half-split rotate, so applying the half-split form directly pairs the wrong dims and still yields finite, weight-shaped output — and YaRN rescales `inv_freq` globally, so it applies from position 0 rather than only at long context. The k-side slice is rotated before it enters the latent cache, because a cached entry is reused by every later query and carries its own token's position. The absorbed `kv_b_proj` identity touches only the nope half and is unaffected. Rotation is skipped entirely when `mla_use_nope` is set, so Kimi-Linear and K3 are untouched by construction. A container needing rotation on a slice wider than `2 * WASTE_MAX_ROPE_HALF` is refused at load rather than run unrotated. `tools/kimi_ref.py` cannot serve a V3 config — it indexes `linear_attn_config` unconditionally and applies no rotary — so `tools/deepseek_ref.py` is the oracle for this path, on the same contract: weights read from the container, so quantization error cancels and a diff measures arithmetic. With `--no-rope --no-mscale` it reproduces the engine's pre-fix layer-0 residual to 0.0001% rel L2, and its full-depth top-1 to p=0.967985 against 0.968.
The suite stayed green through the whole window in which `src/` applied no rotary, and it would have stayed green after a fix that pairs the wrong dims. Both have the same cause: every container the suite can reach is a Kimi, every Kimi sets `mla_use_nope`, and so nothing in `tests/` ever entered `rope_init` or `rope_apply`. This is the missing half of the previous commit. `make_test_container.py --rope` writes a DeepSeek-V3 at the 1/18 scale the file already builds a Kimi-Linear at: no `mla_use_nope`, `rope_theta` and the YaRN block copied from Kimi-K2-Instruct's config, and no `linear_attn_config` at all, which is what makes every layer MLA. All-MLA is deliberate twice over — it exercises the rotation at depth rather than in the single full-attention layer the Kimi mix leaves, and it is the shape `deepseek_ref.py` can read, since not indexing `linear_attn_config` is exactly what separates it from `kimi_ref.py`. K2's rope block rather than V3's because `beta_fast == beta_slow == 1.0` collapses YaRN's correction range to a two-dim ramp, which is the more awkward of the two to get right. The checks build their own container instead of using `$MODEL`, so they run on every host and do not wait on weights nobody can convert yet — sqliteai#26 is what makes a real V3 container, and the shape is what the engine branches on. - rotated MLA against the PyTorch oracle - chunked prefill == token-at-a-time with rotation, which holds by construction today because `mla_layer` is per-token on both paths, and is exactly the "by construction" a later batched MLA would break quietly - a rope slice wider than `WASTE_MAX_ROPE_HALF` is refused at load The first takes the same two-source shape as the Kimi oracle above it: generate from `deepseek_ref.py` where `uv` exists, fall back to a fixture where it does not, so the Linux image without `uv` runs it rather than skipping it. Unlike that one the fixture ships, because this container is generated rather than converted and so is byte-reproducible at `--seed 0` — the sidecar carries a digest of the container it was made from, so a later change to the generator's weights reads as "regenerate me" and not as an engine bug. The fixture is the reference's logits, never the engine's. `deepseek_ref.py` grows the `--dump` that `kimi_ref.py` already had, so the diff is over whole logit vectors and not a printed top-k. Verified by reverting `src/model.c` and `src/model.h` to their pre-fix state with `tests/` and `tools/` left alone: the oracle check and the refusal check both fail, which is the property that makes them worth having. Both fallback paths were exercised directly — `uv` off `PATH` passes against the fixture, and a corrupted digest skips with the regenerate message instead of reporting a divergence. `set -o pipefail` sank the refusal check on the first run, because a refused load exits non-zero and that is the point; the output is read into a variable now, with a comment saying why. Suite on this commit: 46 passed, 0 failed, 2 skipped against Kimi-Linear and K3 (43/0/2 before), 39/0/9 on the synthetic path CI takes, and `make asan` 33/0/14. Fuzzer and the 168 serve checks unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Reviewed this properly. Verification first, then two things I'd like your read on, then the test coverage. VerificationI checked the arithmetic against
Both of the "do not survive paraphrase" claims hold. Two things1. 2. Unsupported Test coverageNothing in I wrote the missing half:
Verified the way that matters: reverting only Suite goes 43 → 46 passed / 0 failed / 2 skipped, 39/0/9 on the synthetic path CI takes, It is one commit on git fetch https://github.com/sqliteai/waste.git rope-test && git cherry-pick 90b111d
Happy for that to land as a follow-up instead, if you would rather keep this PR to the fix itself. |
…implement Two ways the rotation could still be skipped with nobody told. `mla_use_nope` was read by presence, which is the idiom cfg_from_json uses for its other flags. A container carrying it as `false` therefore loaded as NoPE and produced exactly the pre-fix result: unrotated logits, finite and weight-shaped, no error. convert.py copies the source config verbatim, so a checkpoint that writes the flag out rather than omitting it reaches that path. js_bool() reads the token instead — `false` is false, a missing key is still the default. The other flags keep the presence idiom, where a misread costs a feature rather than the sequence order. A rope_scaling shape the ramp does not implement fell through to plain RoPE just as quietly. Any `type` but yarn landed there, and so did mscale != mscale_all_dim, where upstream puts the ratio on cos/sin and this does not. Nothing shipping is affected — V3, R1 and K2 all set `type: yarn` with both mscales at 1.0 — but the argument for refusing an over-wide slice is the same argument, so rope_init leaves a reason in cfg.rope_err and the load refuses on it. The over-wide check moves into that mechanism, leaving one refusal point rather than a condition re-derived at the call site. Two shapes are read rather than refused. `rope_type` is HF's rename of `type`, so it is an alias: a config that only spells it that way is rotated, not turned away. And factor <= 1 stays a fall-through, because YaRN's ramp is the identity there and both mscales collapse to 1 — plain RoPE is the right answer, not a degraded one. tools/deepseek_ref.py refuses the same two shapes rather than approximating them, so it remains an oracle for exactly what the engine accepts. Three checks over make_test_container.py's new knobs: a container built at the same seed with `mla_use_nope: false` gives logits byte-identical to the same model without the key, and each unimplemented rope_scaling is refused at load. 37 passed, 0 failed, 12 skipped on the synthetic path (34/0/12 before), `make asan` 36/0/13 with the six rotary checks passing and no sanitizer report, fuzzer 400 cases 0 crashed. A full Kimi-K2 container still loads and reports unchanged.
|
Both taken. Your test commit is on the branch unchanged — it sat directly on 1.
|
|
Both points landed, and the second one landed better than I asked for: moving the The two shapes you read rather than refuse are right. VerificationCI is 8/8 on The one that matters, since I also checked the mscale claim against the configs I had not covered last round — One real defect:
|
|
Filed the cpu-list flake as #30. Two corrections to the diagnosis, since they change The chunk arithmetic. int chunk = (n + g_pool.nthreads - 1) / g_pool.nthreads; /* (8 + 1) / 2 == 4 */
if (chunk < min_chunk) chunk = min_chunk;so The test already has the mitigation, and it is the part that is failing. /* ...then burns enough time that the next chunk is taken
* by a different thread rather than by this one coming back for it. */That is a wall-clock assumption — the caller must still be spinning in range 1 when the So your other suggestion is the right one. #30 proposes the rendezvous, with a timeout Nothing here is yours to fix in this PR — it reproduces on |
… boolean
`"rope_scaling": null` was refused. js_get returns a token for a JSON null but
it is not a JS_OBJ, so `type` and `rope_type` both came back absent, js_str
left the buffer empty and the type refusal fired — on a container whose
correct plain-RoPE inv_freq the loop above had already built and then threw
away. null is how an HF config says "no scaling" and convert.py copies configs
verbatim, so it is the shape most containers on disk carry; the Kimi ones are
saved from it only by returning early on mla_use_nope. js_size is 0 for a null
and for {} alike, so one condition covers both. tools/deepseek_ref.py already
read this shape correctly, so the two agree again.
An mla_use_nope that is present but not a boolean took js_bool's default and
rotated. Nothing writes `1` or `"true"`, but defaulting picks the sequence
order out of a manifest that never said which one it wanted, and picks it
silently — the failure this file exists to remove. js_typeof() tells "absent"
from "present but not that", and present-but-not-boolean now leaves a reason
in rope_err the way an unimplemented rope_scaling does.
A rope_scaling carrying no type says so, rather than reporting the type as the
empty string. That is the message someone debugs from.
--nope takes a JSON value instead of being a --nope-false switch, and
--rope-scaling builds the null / {} / absent / no-type shapes. Three checks
over them: null and {} load with logits byte-identical to the same container
with no key at all, and the other two are refused at load.
40 passed, 0 failed, 12 skipped on the synthetic path (37/0/12 before),
`make asan` 39/0/13 with all nine rotary checks passing and no sanitizer
report, fuzzer 400 cases 0 crashed 0 hung. The shipped rotary fixture's
container digest is unchanged, so CI's no-uv path still compares. A full
Kimi-K2 container still opens and reports unchanged.
|
All three taken.
|
2 participants
Apply rotary to MLA's rope dims on models that are not NoPE
src/implemented no rotary. Every occurrence ofropewasqk_ropeused as awidth, and
rope_theta,rope_scalingandmla_use_nopewere read nowhere.Correct for the Kimi models, which set
mla_use_nopeand pass those dims throughunrotated; wrong for DeepSeek-V3, which sets no such flag and ships
rope_thetawith YaRN. In MLA those dims are the only positional signal, so a V3 container
attended over an unordered sequence.
Why it is quiet
Lexically determined answers still come out right, which is why this survives
casual use.
Kimi-K2at 3-bit VQ3R, single user turn, greedy:Add a second turn boundary and it collapses. Top-1 next token after
<|im_assistant|>assistant<|im_middle|>, 13-token prompt (system: A,user: Hi):<|im_end|>0.968Hi0.014I0.005Hi0.491Hello0.483Hey0.025<|im_end|>at 0.968 is an empty assistant turn: the model cannot tell whichturn came first, so it takes the highest-prior continuation. It is not a
degraded answer, it is an unordered one. Two user turns and no system turn
reproduce it identically (0.953), so it is the turn boundaries and not the
system role.
Isolation
Held identical between a reproducing and a non-reproducing model, and each
cleared:
checkpoint's
tokenizer_config.json, none absent, none extrareconstruction error is the engine's normal operating point
test_forwardwithWASTE_CHUNK=0and=1giveargmax 163586, max 15.9099bit-identically. Both paths callmla_layer, which is also why they agree.lm_head, F32 for 1-D, i.e.--trunk-bitsdefaultControl, same container recipe and the same two prompts,
Kimi-Linear-48B:Parisat p=0.761 (20 tok) and p=0.716 (38 tok) — stable across the addedturn. It answers correctly and does not care, because NoPE is right for it.
The change
rope_initreadsmla_use_nope,rope_thetaandrope_scaling, buildsinv_freqwith YaRN's ramp, and takesmscale_all_dimsquared onto theattention scale — following
DeepseekV3YarnRotaryEmbeddingin the checkpoint'sown
modeling_deepseek.py.rope_tables/rope_applyrotate a slice in place. The angles depend only on(pos, j), not on the head, so the tables are built once per token per layer andall heads reuse them.
mla_layerrotates the query's rope dims per head and the k-side slicebefore it enters the latent cache, because a cached entry is reused by every
later query and carries its own token's position. The absorbed
kv_b_projidentity touches only the nope half and is unaffected.
Two details that do not survive paraphrase, both called out in comments:
x[2j]withx[2j+1]. Upstreamreaches the same arithmetic by de-interleaving before a half-split rotate, so
writing the half-split form directly pairs dim
jwithj + qk_rope/2— finite,weight-shaped and wrong.
inv_freqglobally, so it applies from position 0 and a shortprompt does not avoid it. K2 carries
beta_fast = beta_slow = 1.0rather thanHF's 32/1, which collapses the correction range to dims 19–20.
mscalealso names two different factors: cos/sin carrymscale / mscale_all_dim,which is 1.0 when the two are equal, while the attention scale carries
mscale_all_dimsquared, which on K2 is 1.8133.Nothing changes for Kimi-Linear or K3
rope_initreturns early whenmla_use_nopeis set, leaving no table andatt_mulat 1, so those models take the same path as before by construction ratherthan by a runtime branch. A container that needs rotation on a slice wider than
2 * WASTE_MAX_ROPE_HALFis refused at load — running unrotated is not a degradedresult but an unordered one.
Both prefill paths are covered by one site: chunked prefill calls
mla_layerpertoken, which is also why the two paths agreed bit-for-bit before the fix.
Verification
tools/deepseek_ref.pyis included because the fix's credibility rests on it.tools/kimi_ref.pycannot serve a V3 config — it indexeslinear_attn_configunconditionally and applies no rotary — so this is a companion on the same
contract: weights read from the container, so quantization error cancels and a diff
measures arithmetic only.
It reproduces the engine before contradicting it. Layer 0 residual stream, 7168
dims, 13-token prompt, against
WASTE_DUMP_HIDDEN:--no-rope --no-mscalevs pre-fix engineFull depth, all 61 layers, engine-equivalent mode: top-1
163586at p=0.967985against the engine's own 0.968.
make checkis 36 passed / 0 failed / 7 skipped withWASTE_REF_MODELpointed ata default VQ3R Kimi-Linear container, and that container's next-token distribution
is unchanged on both a 20- and a 38-token prompt (
113476at p=0.761 and p=0.716,identical to before).
Scope
Not K2-specific. DeepSeek-V3, R1 and K2 all set no
mla_use_nopeand all shiprope_thetawith YaRN, so any V3-architecture container is affected. Nothingchanges for Kimi-Linear or K3.
Container used:
Kimi-K2-Instruct, 61 layers, 384 experts top-8, VQ3R 3-bit,354 GB expert set, 6.9 GB trunk, converted with
tools/convert.py --stages 3.Building one needs fp8 block-scale reading and the DeepSeek MoE tensor and config
names, which are #26 — that is a prerequisite for reproducing this, not
for building or testing the change here.
docs/LEARNED.md,CHANGELOG.mdandWASTE_VERSION_*are deliberately untouched.