From: opencoti Subject: F5 M2 — M2-aware state I/O (bug-222 fix) Bench: perf/llamafile/advanced-kv-stack.bench.ts (Gemma 4 A4B + 31B), bug-222 fix Milestone: F5 M2 hardening — see docs/features/advanced_kv.md, .wolf/buglog.json bug-222 Pre-existing M2 head split bug surfaced on Gemma 4 iswa: state_write_data / state_read_data read/write the K and V cache assuming a single tensor per layer with full n_embd_k_gqa / n_embd_v_gqa row size. With M2 head split active the layer's tensor (k_stream[s] / v_stream[s]) holds ONLY the GPU half (n_embd_k_gpu width), and the CPU half lives in a separate k_cpu_per_stream[s] / v_cpu_per_stream[s] tensor. Reading buf_size = range_size * row_full from a tensor that has range_size * row_gpu bytes crashes with "tensor read out of bounds" in ggml_backend_tensor_get (ggml-backend.cpp:314) as soon as the per-range write exceeds the GPU half's nbytes — reproducible on Gemma 4 A4B & 31B at frac=0.5 with any prompt that crosses a 512-token microbatch boundary and triggers a context checkpoint (the cache_prompt server feature snapshots ≥ 634 cells after the second prompt batch). Fix: detect M2 split per layer (layer.gpu_heads > 0 && layer.k_cpu_per_stream[s]), emit GPU half (row_gpu) then CPU half (row_cpu) per range. Total bytes per range stay at range_size * row_full (header advertises the full row size unchanged) so io stream accounting is preserved; on-disk byte order within a range differs from upstream (GPU-first|CPU-first vs interleaved) but writer/reader are symmetric so round-trips are consistent. M2 inactive (no head_cpu) path is taken byte-identical to upstream — preserves Qwen / unified-mode behavior. V branch only fires in !v_trans (M2 is gated off when v_trans). Validated: - Gemma 4 A4B-98e-v6-coder-it-Q4_K_M: M5 stack bench 8/9 PASS (C5 byte-equality A2/A3 PASS, C9 perf threshold known cost). - Gemma 4 31B-it-Q4_K_M: M5 stack bench 8/9 PASS, same shape. - Qwen 2.5 1.5B Instruct: regression shield — M5 8/9, M2 6/6, M3 4/4, M0 5/5 all PASS unchanged. Bug-222 reproducer: /tmp/gemma-crash-probe.py (cycles reps 10..50 of the bench prompt against an A2 server; pre-fix crashes at reps=30, post-fix clean through reps=35). diff --git a/llama.cpp/src/llama-kv-cache.cpp b/llama.cpp/src/llama-kv-cache.cpp --- a/llama.cpp/src/llama-kv-cache.cpp +++ b/llama.cpp/src/llama-kv-cache.cpp @@ -3190,11 +3190,43 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t const uint64_t k_size_row = ggml_row_size(k->type, n_embd_k_gqa); io.write(&k_size_row, sizeof(k_size_row)); - // Read each range of cells of k_size length and write out - for (const auto & range : cr.data) { - const size_t range_size = range.second - range.first; - const size_t buf_size = range_size * k_size_row; - io.write_tensor(k, range.first * k_size_row, buf_size); + // opencoti bug-222 fix — F5 M2 split-aware state write. + // When M2 head split is active, `k` (k_stream[s]) holds only the + // GPU half of the layer's heads; the CPU half lives in + // k_cpu_per_stream[s]. The pre-fix path read `range_size * k_size_row` + // bytes from a tensor that only has `range_size * row_gpu` bytes → + // out-of-bounds crash on any range_size that pushes buf_size past + // the per-stream GPU tensor's nbytes. Fix: detect M2 split, emit + // GPU half then CPU half per range. Total bytes per range equals + // `range_size * k_size_row` (the header advertises this), so on-disk + // accounting matches; the layout differs only in within-range order + // (GPU-first|CPU-first vs interleaved). Reader (state_read_data) is + // symmetric so a save/load round-trip is consistent. When M2 is OFF + // (no head_cpu) the original path is taken byte-identical. + const bool m2_split_active = layer.gpu_heads > 0 + && !layer.k_cpu_per_stream.empty() + && layer.k_cpu_per_stream[cr.strm]; + if (m2_split_active) { + auto * k_cpu = layer.k_cpu_per_stream[cr.strm]; + const uint32_t n_embd_head_k = hparams.n_embd_head_k(il); + const uint32_t n_head_kv = hparams.n_head_kv(il); + const uint32_t n_embd_k_gpu = layer.gpu_heads * n_embd_head_k; + const uint32_t n_embd_k_cpu = (n_head_kv - layer.gpu_heads) * n_embd_head_k; + const uint64_t k_row_gpu = ggml_row_size(k->type, n_embd_k_gpu); + const uint64_t k_row_cpu = ggml_row_size(k_cpu->type, n_embd_k_cpu); + GGML_ASSERT(k_row_gpu + k_row_cpu == k_size_row); + for (const auto & range : cr.data) { + const size_t range_size = range.second - range.first; + io.write_tensor(k, range.first * k_row_gpu, range_size * k_row_gpu); + io.write_tensor(k_cpu, range.first * k_row_cpu, range_size * k_row_cpu); + } + } else { + // Read each range of cells of k_size length and write out + for (const auto & range : cr.data) { + const size_t range_size = range.second - range.first; + const size_t buf_size = range_size * k_size_row; + io.write_tensor(k, range.first * k_size_row, buf_size); + } } } @@ -3217,11 +3249,33 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t const uint64_t v_size_row = ggml_row_size(v->type, n_embd_v_gqa); io.write(&v_size_row, sizeof(v_size_row)); - // Read each range of cells of v_size length and write out - for (const auto & range : cr.data) { - const size_t range_size = range.second - range.first; - const size_t buf_size = range_size * v_size_row; - io.write_tensor(v, range.first * v_size_row, buf_size); + // opencoti bug-222 fix — F5 M2 split-aware V state write + // (same pattern as K above). M2 is gated off when v_trans, so + // this branch is reached only with the non-transposed V layout. + const bool m2_split_active = layer.gpu_heads > 0 + && !layer.v_cpu_per_stream.empty() + && layer.v_cpu_per_stream[cr.strm]; + if (m2_split_active) { + auto * v_cpu = layer.v_cpu_per_stream[cr.strm]; + const uint32_t n_embd_head_v = hparams.n_embd_head_v(il); + const uint32_t n_head_kv = hparams.n_head_kv(il); + const uint32_t n_embd_v_gpu = layer.gpu_heads * n_embd_head_v; + const uint32_t n_embd_v_cpu = (n_head_kv - layer.gpu_heads) * n_embd_head_v; + const uint64_t v_row_gpu = ggml_row_size(v->type, n_embd_v_gpu); + const uint64_t v_row_cpu = ggml_row_size(v_cpu->type, n_embd_v_cpu); + GGML_ASSERT(v_row_gpu + v_row_cpu == v_size_row); + for (const auto & range : cr.data) { + const size_t range_size = range.second - range.first; + io.write_tensor(v, range.first * v_row_gpu, range_size * v_row_gpu); + io.write_tensor(v_cpu, range.first * v_row_cpu, range_size * v_row_cpu); + } + } else { + // Read each range of cells of v_size length and write out + for (const auto & range : cr.data) { + const size_t range_size = range.second - range.first; + const size_t buf_size = range_size * v_size_row; + io.write_tensor(v, range.first * v_size_row, buf_size); + } } } } else { @@ -3433,7 +3487,36 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 } if (cell_count) { - if (sinfo.is_contiguous()) { + // opencoti bug-222 fix — F5 M2 split-aware state read (K). + // Symmetric to state_write_data: when M2 head split is active, + // read the GPU half then the CPU half and dispatch to the + // respective per-stream tensors. Total bytes equal + // `cell_count * k_size_row` (the header advertises this) so the + // sequential io stream matches what state_write_data wrote. + const bool m2_split_active = layer.gpu_heads > 0 + && !layer.k_cpu_per_stream.empty() + && layer.k_cpu_per_stream[strm]; + if (m2_split_active) { + auto * k_cpu = layer.k_cpu_per_stream[strm]; + const uint32_t n_embd_head_k = hparams.n_embd_head_k(il); + const uint32_t n_head_kv = hparams.n_head_kv(il); + const uint32_t n_embd_k_gpu = layer.gpu_heads * n_embd_head_k; + const uint32_t n_embd_k_cpu = (n_head_kv - layer.gpu_heads) * n_embd_head_k; + const uint64_t k_row_gpu = ggml_row_size(k->type, n_embd_k_gpu); + const uint64_t k_row_cpu = ggml_row_size(k_cpu->type, n_embd_k_cpu); + GGML_ASSERT(k_row_gpu + k_row_cpu == k_size_row); + if (sinfo.is_contiguous()) { + io.read_tensor(k, sinfo.head() * k_row_gpu, cell_count * k_row_gpu); + io.read_tensor(k_cpu, sinfo.head() * k_row_cpu, cell_count * k_row_cpu); + } else { + for (uint32_t i = 0; i < cell_count; ++i) { + io.read_tensor(k, sinfo.idxs[0][i] * k_row_gpu, k_row_gpu); + } + for (uint32_t i = 0; i < cell_count; ++i) { + io.read_tensor(k_cpu, sinfo.idxs[0][i] * k_row_cpu, k_row_cpu); + } + } + } else if (sinfo.is_contiguous()) { // Fast path: contiguous cells, single memcpy io.read_tensor(k, sinfo.head() * k_size_row, cell_count * k_size_row); } else { @@ -3476,7 +3559,33 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 } if (cell_count) { - if (sinfo.is_contiguous()) { + // opencoti bug-222 fix — F5 M2 split-aware state read (V), + // same pattern as K. M2 is gated off when v_trans, so this + // branch is only reached with the non-transposed V layout. + const bool m2_split_active = layer.gpu_heads > 0 + && !layer.v_cpu_per_stream.empty() + && layer.v_cpu_per_stream[strm]; + if (m2_split_active) { + auto * v_cpu = layer.v_cpu_per_stream[strm]; + const uint32_t n_embd_head_v = hparams.n_embd_head_v(il); + const uint32_t n_head_kv = hparams.n_head_kv(il); + const uint32_t n_embd_v_gpu = layer.gpu_heads * n_embd_head_v; + const uint32_t n_embd_v_cpu = (n_head_kv - layer.gpu_heads) * n_embd_head_v; + const uint64_t v_row_gpu = ggml_row_size(v->type, n_embd_v_gpu); + const uint64_t v_row_cpu = ggml_row_size(v_cpu->type, n_embd_v_cpu); + GGML_ASSERT(v_row_gpu + v_row_cpu == v_size_row); + if (sinfo.is_contiguous()) { + io.read_tensor(v, sinfo.head() * v_row_gpu, cell_count * v_row_gpu); + io.read_tensor(v_cpu, sinfo.head() * v_row_cpu, cell_count * v_row_cpu); + } else { + for (uint32_t i = 0; i < cell_count; ++i) { + io.read_tensor(v, sinfo.idxs[0][i] * v_row_gpu, v_row_gpu); + } + for (uint32_t i = 0; i < cell_count; ++i) { + io.read_tensor(v_cpu, sinfo.idxs[0][i] * v_row_cpu, v_row_cpu); + } + } + } else if (sinfo.is_contiguous()) { // Fast path: contiguous cells, single memcpy io.read_tensor(v, sinfo.head() * v_size_row, cell_count * v_size_row); } else {