ggml decode: the KV cache was sized by two literals the write site ignored #372

Merged

dcharlot merged 1 commit from fix/kv-cache-attn-layer-count into main 2026-10-21 05:38:26 -04:00

Conversation 0 Commits 1 Files changed 3 +151 -11

2026-10-20

Claude

ggml decode: the KV cache was sized by two literals the write site ignored ...

`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]>

950634f200

• View at this point in history

2026-10-20 16:32:06 -04:00