ggml decode: the KV cache was sized by two literals the write site ignored #372
Loading…
Reference in a new issue
No description provided.
Delete branch "fix/kv-cache-attn-layer-count"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
DecodeBuffers::newbuilt the KV cache from two literals:while
encode_decode_stepindexes that cache with the model's real attention-layer ordinal and the real position:Neither mismatch reported anything.
What went wrong
A seventh attention layer lost its KV writes. With
attn_idx == 6, the guardif attn_idx < bufs.kv_k_cache.len()was false, so that layer's K and V were never written. Attention then read a zero-filled cache — wrong output, no error anywhere. The guard that looks defensive is what made it silent._positionis never bounded. Not anywhere inencode_decode_step. At or past 4096,set_bufferreceives an offset past the end of the buffer.The fix
Both lengths come from the model now.
initresolves the layer types and passes the attention-layer count;max_seqcomes fromconfig.max_seq_len, falling back only when it is 0 — which would otherwise allocate a zero-length cache that fails every write.The capacity is recorded on
DecodeBuffersso the decode path can check instead of assuming, andencode_decode_stepchecks both once, up front. That placement is deliberate: the per-siteif idx < lenguards sit inside nested closures wherereturn Errcannot propagate, which is why the original failure had nowhere to surface.Tests
tests/ggml_kv_cache.rspins the two lengths:max_seq_lenfor 512, 4096 and 8192max_seq_len == 0still producing a usable cacheVerified by corruption, not just a green run: restoring the two literals fails exactly the two tests that should fail, with
left: 4096 / right: 512andleft: 6 / right: 1. The third passes either way, because the old code ignored the config entirely — it guards the new fallback, not the old bug.Also: 0 metal compile errors, 26/26 Metal kernel tests still passing,
clippy -- -D warningsclean.Scope and what is still assumed
This path is behind the
GGML_GRAPHenvironment variable, so it is opt-in rather than the default decode path.Left alone deliberately, and worth separate attention:
GRAPHandLAYERS_CACHEare process-globalOnceLocks, so a second model loaded in the same process reuses the first one's buffers and layer pattern.🤖 Generated with Claude Code
`DecodeBuffers::new` built the KV cache with let max_seq = 4096usize; let num_attn_layers = 6; // Nemotron: 6 attention layers while `encode_decode_step` indexes it with the model's real attention-layer ordinal and the real position: let attn_idx = layers.iter().take(layer.idx + 1) .filter(|l| l.layer_type == LayerType::Attention).count() - 1; let dst_offset = _position as u64 * kv_row_f16; Two consequences, neither of which reported anything: - A model with a seventh attention layer got `attn_idx == 6`, the `if attn_idx < bufs.kv_k_cache.len()` guard was false, and that layer's K and V were simply never written. Attention then read a zero-filled cache: wrong output, no error. - `_position` is never bounded anywhere in the function, so a position at or past 4096 gave `set_buffer` an offset past the end of the buffer. Both lengths now come from the model. `init` resolves the layer types and passes the attention-layer count; `max_seq` comes from `config.max_seq_len`, falling back only when it is 0, which would otherwise allocate a zero-length cache that fails every write. The capacity is recorded on `DecodeBuffers` so the decode path can check rather than assume, and `encode_decode_step` now checks both once, up front, where an error can propagate -- the per-site `if idx < len` guards are deep inside nested closures where a `return Err` cannot. tests/ggml_kv_cache.rs pins the two lengths: one buffer per attention layer for counts of 1, 6, 7 and 11 (6 being the literal that hid this), buffer length following max_seq_len for 512, 4096 and 8192, and a zero-length config still producing a usable cache. Scope: this path is behind the GGML_GRAPH environment variable, so it is opt-in rather than the default decode path. Still assumed, and worth separate attention: `GRAPH` and `LAYERS_CACHE` are process-global `OnceLock`s, so a second model loaded in the same process reuses the first one's buffers and layer pattern. Co-Authored-By: Claude Opus 5 <[email protected]>