From: opencoti Subject: [PATCH 0031] F4 M3 Phase 4 — per-stream KV tensor split (task #109) Resolves the long-deferred Phase 4 (`docs/decisions/0001-lazy-slot-context.md`) of F4 M3 — every layer's K/V (and M2's CPU subset) become OWNING per-stream tensor vectors rather than one 3-D tensor with a stream dim, with each (stream, buffer-type) pair owning its own backend buffer. Unified mode (`n_stream == 1`) collapses to the pre-Phase-4 layout and stays byte-identical; non-unified mode (`n_stream > 1`) is functional for the first time, unlocking the multi-tenant codepaths (M2 head-split now lifts its `n_stream == 1` construction-time gate). Sub-milestone shape (each landed against the M5-shipped tree at HEAD = `2373e69` and verified by the M5 stack + M0 turn-2 + M1 eviction + M2 residency benches before composing): - Phase 4-A — per-stream tensor split. `kv_layer` replaces `k`/`v`/`k_cpu`/`v_cpu` direct pointers with `std::vector k_per_stream` / `v_per_stream` / `k_cpu_per_stream` / `v_cpu_per_stream`; each per-stream tensor is `ggml_new_tensor_3d(ctx_s, type, n_embd, kv_size, 1)`. The trailing `dim 3 == 1` keeps stride math identical to the pre-Phase-4 `[n_embd, kv_size, n_stream]` layout under unified mode. The `k_stream` / `v_stream` 2-D view vectors are preserved (offset 0 within each per-stream tensor) so state I/O and cross-stream copies keep operating on the existing view handles. - Phase 4-B — per-stream backend buffers. `ctx_map` becomes a nested `std::vector> ctx_map_per_stream` indexed by stream, with a `ctx_for_buft(buft, s)` lambda that returns the get-or-create `(stream, buft)` ctx (mem_size budgeted at `4u*n_layer_kv* ggml_tensor_overhead()`). Allocation loops nest stream-outer × buft-inner so each (stream, buft) pair allocates one backend buffer. Phase 1 `ensure_cleared`, Phase 3 `opencoti_decommit_layer_range`, and the M2 CPU subset construction follow the same per-stream iteration shape. - Phase 4-C — reworked write/read graph ops. `cpy_k`/`cpy_v` ship option C-2 (N separate `ggml_set_rows`, one per `sinfo.n_stream()` stream, tied via `ggml_add` on 1-D views — unified-mode batch_ns==1 collapses to one `set_rows`, byte-identical). C-1 (concat-then- scatter) was rejected at impl time: `ggml_concat` materializes a new buffer rather than viewing inputs, so writes routed through the concat result would not land in the per-stream tensors. `get_k`/ `get_v` build per-cache-stream 4-D views and `ggml_concat` them along dim 3 (stream axis) — unified mode yields a single view, no concat. `set_input_k_idxs`/`set_input_v_idxs` emit LOCAL per-stream indices unconditionally (Phase 4-F fix; the previous gating on `sinfo.n_stream() > 1` left global indices in --parallel N servers serving single-session batches and crashed CUDA decode against the per-stream destination). `build_graph_shift` loops streams inside the layer loop, slicing `inp->k_shift` per stream. - Phase 4-D — M2 head-split is stream-aware. `get_k_gpu`/`get_v_gpu`/ `get_k_cpu`/`get_v_cpu` (and the M2 split branch inside `get_k`/ `get_v`/`cpy_k`/`cpy_v`) build per-stream 4-D views and concat them on dim 3 (stream axis); the construction-time gate at `headinfer_gpu_heads_frac < 1.0 && offload && !v_trans && n_stream == 1` drops the `n_stream == 1` clause so per-stream CPU subsets actually allocate for non-unified mode. - Phase 4-E — cross-stream copies in `update()`, `state_write_data`, `state_read_data`, and `seq_key_scores` continue to operate on the existing `k_stream[s]` / `v_stream[s]` 2-D views (now offset-0 views of per-stream tensors). Mechanically distinct from pre-Phase-4 (one buffer per stream rather than one shared buffer), semantically the same per the buffer transfer mechanism. - Phase 4-F — verification gates. Cosine ≥ 0.999 unified-mode equivalence captured indirectly via the M5 stack bench (B/A0/A1/ A2/A3/S configs all produce byte-identical KV MiB and `t2_cached_n=1271` against the M5-shipped baseline) plus M2's explicit greedy-decode byte-equality check (`split text === baseline text`, 6/6 PASS). Multi-stream RSS exercise comes via the M0 turn-2 scenarioB (`--parallel 2` with two sessions sharing a large system prefix — 5/5 PASS, validating per-stream KV reuse). M3 neo-pipeline O off-identity config crashes on both pre-Phase-4 and post-Phase-4 binaries with `concat.cu:165 GGML_ASSERT(src0-> type == GGML_TYPE_F32)` — pre-existing M3 bug, DEFERRED as task #278 and recorded as bug-207. Hook count stable at 20. The HARD gates that survived from earlier sub-milestones (the `GGML_ASSERT(n_stream == 1 && "Phase 4-A: get_k temporarily gated..."; "Phase 4-C lifts")` shape on get_k/get_v/cpy_k/cpy_v/build_graph_shift, and the parallel `Phase 4-D lifts` shape on the M2 accessors) all reach ship as lifted. Five `GGML_ASSERT(n_stream == 1)` lines deleted; the remaining `n_stream == 1` checks are real preconditions (e.g., v_trans + multi-stream is asserted out in cpy_v as not-a-real-workload). Unblocks F5 M6 (PolyKV) which needs per-stream lifetimes cleanly expressed for the shared compressed KV pool. 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 @@ -160,14 +160,35 @@ llama_kv_cache::llama_kv_cache( return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0; } }; - std::map ctx_map; - // create a context for each buffer type - auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * { + // opencoti F4 M3 Phase 4-B — per-stream backend buffer ownership. + // Each stream gets its own map of (buft → ctx) so each (stream, buft) + // pair allocates its own backend buffer. Phase 3's + // opencoti_decommit_layer_range then collapses to "decommit this + // stream's buffer range" cleanly, and M6 PolyKV's per-stream + // lifetimes (enroll/disenroll a stream into the pool) become + // mechanically straightforward. + // + // Unified mode (n_stream == 1): one stream, one inner map per + // buft, byte-identical layout to pre-Phase-4 (the buffer-allocation + // loop below ends up with exactly the same `(ctx, buf)` pairs in + // ctxs_bufs as before, in the same order, with the same sizes). + std::vector> ctx_map_per_stream; + ctx_map_per_stream.resize(n_stream); + + // ctx_for_buft(buft, s): get-or-create the ggml_context for + // (stream s, buft). Each per-stream ctx holds at most 4 tensors + // per layer (K + V owning, K + V views; doubled when M2's CPU + // subset is engaged on a layer). Budget conservatively at + // 4*n_layer_kv slots per stream per buft — address space is free, + // RSS still grows only on first write. + auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft, uint32_t s) -> ggml_context * { + GGML_ASSERT(s < ctx_map_per_stream.size()); + auto & ctx_map = ctx_map_per_stream[s]; auto it = ctx_map.find(buft); if (it == ctx_map.end()) { ggml_init_params params = { - /*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer_kv*ggml_tensor_overhead()), + /*.mem_size =*/ size_t(4u*n_layer_kv*ggml_tensor_overhead()), /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; @@ -255,10 +276,10 @@ llama_kv_cache::llama_kv_cache( LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); - ggml_context * ctx = ctx_for_buft(buft); - if (!ctx) { - throw std::runtime_error("failed to create ggml context for kv cache"); - } + // opencoti F4 M3 Phase 4-B: per-stream ctx selection has moved + // into the per-stream loop below. The layer-level `buft` is the + // device buffer-type for this layer; each stream allocates its + // own ctx for (buft, stream) via ctx_for_buft(buft, s). const bool has_k = true; const bool has_v = !is_mla; @@ -267,14 +288,21 @@ llama_kv_cache::llama_kv_cache( // Head-residency split: keep the first gpu_heads KV heads on this // (offloaded/device) layer's buffer and put the remaining heads on a // host buffer. Splitting on KV-head boundaries is inherently GQA-safe. - // Gated to offloaded, non-transposed-V, unified (n_stream==1) layers; - // frac >= 1.0 (or <= 1 head) means no split (byte-identical upstream). + // Gated to offloaded, non-transposed-V layers; frac >= 1.0 (or <= 1 + // head) means no split (byte-identical upstream). + // + // opencoti F4 M3 Phase 4-D: lifted the historical n_stream == 1 + // gate. With per-stream owning tensors and per-stream backend + // buffers (Phase 4-A/B), each stream's GPU/CPU half lives in its + // own buffer, so the split is well-defined for n_stream > 1. + // Unified mode (n_stream == 1) collapses to the pre-Phase-4 + // single-stream layout, byte-identical. uint32_t gpu_heads = 0; uint32_t n_embd_k_gpu = n_embd_k_gqa; uint32_t n_embd_v_gpu = n_embd_v_gqa; uint32_t n_embd_k_cpu = 0; uint32_t n_embd_v_cpu = 0; - if (headinfer_gpu_heads_frac < 1.0f && offload && !v_trans && n_stream == 1) { + if (headinfer_gpu_heads_frac < 1.0f && offload && !v_trans) { const uint32_t n_head_kv = hparams.n_head_kv(il); const uint32_t n_embd_head_k = hparams.n_embd_head_k(il); const uint32_t n_embd_head_v = hparams.n_embd_head_v(il); @@ -291,39 +319,119 @@ llama_kv_cache::llama_kv_cache( } } - ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gpu, kv_size, n_stream) : nullptr; - ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gpu, kv_size, n_stream) : nullptr; - - has_k && ggml_format_name(k, "cache_k_l%d", il); - has_v && ggml_format_name(v, "cache_v_l%d", il); - - // opencoti F5 M2 headinfer — the host-resident head subset lives in the - // CPU buffer-type context; the buffer-allocation loop below picks it up - // automatically (ctx_for_buft keys contexts by buffer type). - ggml_tensor * k_cpu = nullptr; - ggml_tensor * v_cpu = nullptr; - if (gpu_heads > 0) { - ggml_context * ctx_cpu = ctx_for_buft(ggml_backend_cpu_buffer_type()); - if (!ctx_cpu) { - throw std::runtime_error("failed to create CPU ggml context for headinfer kv cache"); - } - k_cpu = has_k ? ggml_new_tensor_3d(ctx_cpu, type_k, n_embd_k_cpu, kv_size, n_stream) : nullptr; - v_cpu = has_v ? ggml_new_tensor_3d(ctx_cpu, type_v, n_embd_v_cpu, kv_size, n_stream) : nullptr; - has_k && k_cpu && ggml_format_name(k_cpu, "cache_k_cpu_l%d", il); - has_v && v_cpu && ggml_format_name(v_cpu, "cache_v_cpu_l%d", il); - } - + // opencoti F4 M3 Phase 4 — per-stream tensor split. + // Each stream gets its own OWNING 3-D tensor of shape + // [n_embd, kv_size, 1]. Dim 3 == 1 keeps stride math + // identical to the pre-Phase-4 single-tensor layout, so the + // unified-mode (n_stream == 1) buffer layout is byte-identical + // to pre-Phase-4. Phase 4-B: each stream also gets its own + // (buft → ctx) entry, so each (stream, buft) pair allocates + // its own backend buffer in the loop below. + std::vector k_per_stream; + std::vector v_per_stream; + std::vector k_cpu_per_stream; + std::vector v_cpu_per_stream; std::vector k_stream; std::vector v_stream; for (uint32_t s = 0; s < n_stream; ++s) { - k_stream.push_back(has_k ? ggml_view_2d(ctx, k, n_embd_k_gpu, kv_size, k->nb[1], s*k->nb[2]) : nullptr); - v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gpu, kv_size, v->nb[1], s*v->nb[2]) : nullptr); + ggml_context * ctx_s = ctx_for_buft(buft, s); + if (!ctx_s) { + throw std::runtime_error("failed to create ggml context for kv cache"); + } + + ggml_tensor * k_s = has_k + ? ggml_new_tensor_3d(ctx_s, type_k, n_embd_k_gpu, kv_size, 1) + : nullptr; + ggml_tensor * v_s = has_v + ? ggml_new_tensor_3d(ctx_s, type_v, n_embd_v_gpu, kv_size, 1) + : nullptr; + + // Tensor naming preserves the pre-Phase-4 name in unified + // mode (no `_s%u` suffix) so external state-IO callers that + // grep by name find an unchanged layout. + if (k_s) { + if (n_stream == 1) { + ggml_format_name(k_s, "cache_k_l%d", il); + } else { + ggml_format_name(k_s, "cache_k_l%d_s%u", il, s); + } + } + if (v_s) { + if (n_stream == 1) { + ggml_format_name(v_s, "cache_v_l%d", il); + } else { + ggml_format_name(v_s, "cache_v_l%d_s%u", il, s); + } + } + + k_per_stream.push_back(k_s); + v_per_stream.push_back(v_s); + + // opencoti F5 M2 headinfer — per-stream host-resident head + // subset. Per-stream ctx_cpu so each stream's CPU half lives + // in its own host buffer (Phase 4-B). Vectors stay empty + // when gpu_heads == 0 so headinfer_split_active() can test + // by vector size. + if (gpu_heads > 0) { + ggml_context * ctx_cpu_s = ctx_for_buft(ggml_backend_cpu_buffer_type(), s); + if (!ctx_cpu_s) { + throw std::runtime_error("failed to create CPU ggml context for headinfer kv cache"); + } + ggml_tensor * k_cpu_s = has_k + ? ggml_new_tensor_3d(ctx_cpu_s, type_k, n_embd_k_cpu, kv_size, 1) + : nullptr; + ggml_tensor * v_cpu_s = has_v + ? ggml_new_tensor_3d(ctx_cpu_s, type_v, n_embd_v_cpu, kv_size, 1) + : nullptr; + if (k_cpu_s) { + if (n_stream == 1) { + ggml_format_name(k_cpu_s, "cache_k_cpu_l%d", il); + } else { + ggml_format_name(k_cpu_s, "cache_k_cpu_l%d_s%u", il, s); + } + } + if (v_cpu_s) { + if (n_stream == 1) { + ggml_format_name(v_cpu_s, "cache_v_cpu_l%d", il); + } else { + ggml_format_name(v_cpu_s, "cache_v_cpu_l%d_s%u", il, s); + } + } + k_cpu_per_stream.push_back(k_cpu_s); + v_cpu_per_stream.push_back(v_cpu_s); + } + + // 2-D views into per-stream 3-D tensor. Offset is 0 within + // the owning tensor — each stream owns the whole 3-D extent. + // + // opencoti F4 M3 Phase 4-E: keeping k_stream/v_stream as the + // existing 2-D view vectors means cross-stream copies in + // update() (ggml_backend_tensor_copy between views), state + // I/O paths (ggml_backend_tensor_set on the per-stream view), + // and seq_key_scores all work without per-stream logic at + // the consumer site — the view routes through to the + // per-stream buffer transparently. Unified-mode operations + // are byte-identical to pre-Phase-4 because the view has + // identical ne/nb and (in unified mode) the same underlying + // buffer layout. + k_stream.push_back(k_s + ? ggml_view_2d(ctx_s, k_s, n_embd_k_gpu, kv_size, k_s->nb[1], 0) + : nullptr); + v_stream.push_back(v_s + ? ggml_view_2d(ctx_s, v_s, n_embd_v_gpu, kv_size, v_s->nb[1], 0) + : nullptr); } map_layer_ids[il] = layers.size(); - layers.push_back({ il, k, v, k_stream, v_stream, k_cpu, v_cpu, gpu_heads }); + layers.push_back({ + il, + k_per_stream, v_per_stream, + k_stream, v_stream, + k_cpu_per_stream, v_cpu_per_stream, + gpu_heads, + }); } if (reuse) { @@ -359,31 +467,46 @@ llama_kv_cache::llama_kv_cache( // scheduler-visible tensor->buffer wiring intact (sched_reserve runs // at construction time and asserts buffer_id >= 0) while moving the // RSS commit to the first ensure_cleared() call. - for (auto & [buft, ctx] : ctx_map) { - ggml_backend_buffer_t buf; - if (model.hparams.no_alloc) { - buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer - for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) { - t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it + // + // opencoti F4 M3 Phase 4-B: per-stream backend buffers. Each + // (stream, buft) pair allocates its own backend buffer here. + // Iteration order is stable: outer loop over streams (s=0,1,...), + // inner loop over bufts (sorted by name via the ctx_map comparator). + // In unified mode (n_stream == 1) this collapses to exactly one + // pass per buft — byte-identical to pre-Phase-4 ctxs_bufs layout. + for (uint32_t s = 0; s < n_stream; ++s) { + for (auto & [buft, ctx] : ctx_map_per_stream[s]) { + ggml_backend_buffer_t buf; + if (model.hparams.no_alloc) { + buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer + for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) { + t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it + } + } else { + buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer (malloc'd but not zeroed yet) + } + if (!buf) { + throw std::runtime_error("failed to allocate buffer for kv cache"); + } + + if (n_stream == 1) { + LLAMA_LOG_INFO("%s: %10s KV buffer size = %8.2f MiB\n", __func__, + ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0); + } else { + LLAMA_LOG_INFO("%s: %10s KV buffer (stream %u) size = %8.2f MiB\n", __func__, + ggml_backend_buffer_name(buf), s, ggml_backend_buffer_get_size(buf)/1024.0/1024.0); } - } else { - buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer (malloc'd but not zeroed yet) - } - if (!buf) { - throw std::runtime_error("failed to allocate buffer for kv cache"); - } - LLAMA_LOG_INFO("%s: %10s KV buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0); + if (model.hparams.no_alloc) { + // Dummy buffer is size 0; clearing is a no-op but keeps the + // upstream invariant that constructed buffers are zeroed. + ggml_backend_buffer_clear(buf, 0); + } + // else: leave the buffer un-cleared. ensure_cleared() will zero it + // on the first inference entry point that touches tensor data. - if (model.hparams.no_alloc) { - // Dummy buffer is size 0; clearing is a no-op but keeps the - // upstream invariant that constructed buffers are zeroed. - ggml_backend_buffer_clear(buf, 0); + ctxs_bufs.emplace_back(std::move(ctx), buf); } - // else: leave the buffer un-cleared. ensure_cleared() will zero it - // on the first inference entry point that touches tensor data. - - ctxs_bufs.emplace_back(std::move(ctx), buf); } if (model.hparams.no_alloc) { @@ -471,33 +594,43 @@ void llama_kv_cache::ensure_cleared(uint32_t up_to_cells) { // memset on the malloc'd buffer; for offloaded buffers it's a backend // op. Either way only bytes for cells [a, b) get touched, so on Linux // only those pages get committed to RSS. + // + // opencoti F4 M3 Phase 4: per-stream tensors. Each layer's + // k_per_stream[s] / v_per_stream[s] is now its own owning tensor + // (shape [n_embd, kv_size, 1]) so the per-stream stride math + // collapses to a plain cell-offset memset within each stream's + // tensor. M2's CPU heads follow the same pattern via + // k_cpu_per_stream / v_cpu_per_stream (empty when no split). for (const auto & layer : layers) { for (uint32_t st = 0; st < n_stream; ++st) { - if (layer.k) { - const size_t off = (size_t) st * (size_t) layer.k->nb[2] - + (size_t) a * (size_t) layer.k->nb[1]; - const size_t sz = (size_t) (b - a) * (size_t) layer.k->nb[1]; - ggml_backend_tensor_memset(layer.k, 0, off, sz); - } - if (layer.v) { - const size_t off = (size_t) st * (size_t) layer.v->nb[2] - + (size_t) a * (size_t) layer.v->nb[1]; - const size_t sz = (size_t) (b - a) * (size_t) layer.v->nb[1]; - ggml_backend_tensor_memset(layer.v, 0, off, sz); - } - // opencoti F5 M2 headinfer — clear the host-resident head subset - // too (its own strides; same [a, b) cell range). - if (layer.k_cpu) { - const size_t off = (size_t) st * (size_t) layer.k_cpu->nb[2] - + (size_t) a * (size_t) layer.k_cpu->nb[1]; - const size_t sz = (size_t) (b - a) * (size_t) layer.k_cpu->nb[1]; - ggml_backend_tensor_memset(layer.k_cpu, 0, off, sz); + auto * k = layer.k_per_stream[st]; + if (k) { + const size_t off = (size_t) a * (size_t) k->nb[1]; + const size_t sz = (size_t) (b - a) * (size_t) k->nb[1]; + ggml_backend_tensor_memset(k, 0, off, sz); + } + auto * v = layer.v_per_stream[st]; + if (v) { + const size_t off = (size_t) a * (size_t) v->nb[1]; + const size_t sz = (size_t) (b - a) * (size_t) v->nb[1]; + ggml_backend_tensor_memset(v, 0, off, sz); + } + // M2 headinfer host-resident head subset (same [a, b) cells). + if (st < layer.k_cpu_per_stream.size()) { + auto * k_cpu = layer.k_cpu_per_stream[st]; + if (k_cpu) { + const size_t off = (size_t) a * (size_t) k_cpu->nb[1]; + const size_t sz = (size_t) (b - a) * (size_t) k_cpu->nb[1]; + ggml_backend_tensor_memset(k_cpu, 0, off, sz); + } } - if (layer.v_cpu) { - const size_t off = (size_t) st * (size_t) layer.v_cpu->nb[2] - + (size_t) a * (size_t) layer.v_cpu->nb[1]; - const size_t sz = (size_t) (b - a) * (size_t) layer.v_cpu->nb[1]; - ggml_backend_tensor_memset(layer.v_cpu, 0, off, sz); + if (st < layer.v_cpu_per_stream.size()) { + auto * v_cpu = layer.v_cpu_per_stream[st]; + if (v_cpu) { + const size_t off = (size_t) a * (size_t) v_cpu->nb[1]; + const size_t sz = (size_t) (b - a) * (size_t) v_cpu->nb[1]; + ggml_backend_tensor_memset(v_cpu, 0, off, sz); + } } } } @@ -542,22 +675,25 @@ bool llama_kv_cache::ensure_capacity(uint32_t cells_needed) { } // opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md -// Decommit pages of every layer's K/V tensors in the byte-range that -// corresponds to cells [from_cells, to_cells). Returns those pages to the OS -// — RSS drops by the size of the range, and the next access to those pages -// will fault them back in as fresh zero-filled pages. +// Decommit pages of one layer's K/V tensor for ONE STREAM in the byte- +// range that corresponds to cells [from_cells, to_cells). Returns those +// pages to the OS — RSS drops by the size of the range, and the next +// access to those pages will fault them back in as fresh zero-filled +// pages. // // Only safe to call after kv_size_alloc has been reset to from_cells (or -// lower) and n_cells_cleared has been reset accordingly, so the bookkeeping -// stays in sync with what's actually backed by physical memory. +// lower) and n_cells_cleared has been reset accordingly, so the +// bookkeeping stays in sync with what's actually backed by physical +// memory. // -// The decommit is per-stream because the KV tensors are 3-D -// [n_embd, kv_size, n_stream]; cells [a, b) within stream s occupy a -// contiguous byte-range whose start address is (base + s*nb[2] + a*nb[1]) -// and length is (b-a)*nb[1]. posix_madvise is page-grained — bytes that -// don't cover a full page are silently kept resident, which is fine for -// the cells that straddle page boundaries. -static void opencoti_decommit_layer_range(ggml_tensor * t, uint32_t n_stream, +// opencoti F4 M3 Phase 4: per-stream tensors. `t` is now ONE stream's +// owning 3-D tensor (dim 3 == 1). Cells [a, b) occupy contiguous bytes +// at (base + a*nb[1]) for length (b-a)*nb[1]. The stream-stride loop +// of the pre-Phase-4 implementation collapses; callers iterate streams +// at the call site instead. posix_madvise is page-grained — bytes that +// don't cover a full page are silently kept resident, which is fine +// for the cells that straddle page boundaries. +static void opencoti_decommit_layer_range(ggml_tensor * t, uint32_t from_cells, uint32_t to_cells) { if (!t || to_cells <= from_cells) { return; @@ -579,37 +715,32 @@ static void opencoti_decommit_layer_range(ggml_tensor * t, uint32_t n_stream, if (!base) { return; } - const size_t stride_stream = (size_t) t->nb[2]; - const size_t stride_cell = (size_t) t->nb[1]; - const size_t off_in_buf = (size_t) ((char *) t->data - (char *) base); - for (uint32_t s = 0; s < n_stream; ++s) { - const size_t off = off_in_buf - + (size_t) s * stride_stream - + (size_t) from_cells * stride_cell; - const size_t sz = (size_t) (to_cells - from_cells) * stride_cell; - // posix_madvise wants the address — work in absolute terms. - void * addr = (char *) base + off; - // Round up to a page boundary on the low side, round down on the - // high side, so we only decommit pages we fully own. The cell-aligned - // address is unlikely to be page-aligned in general. - const size_t page = (size_t) sysconf(_SC_PAGESIZE); - if (page == 0 || page == (size_t) -1) { - return; - } - uintptr_t lo = (uintptr_t) addr; - uintptr_t hi = lo + sz; - uintptr_t lo_aligned = (lo + page - 1) & ~(uintptr_t)(page - 1); - uintptr_t hi_aligned = hi & ~(uintptr_t)(page - 1); - if (hi_aligned <= lo_aligned) { - continue; - } - // madvise(MADV_DONTNEED) returns 0 on success, -1 on failure (with - // errno set). We treat any failure as advisory — the RSS win may - // be skipped, but correctness is unaffected (the K/V buffer is - // still mapped, just possibly still resident). - (void) madvise((void *) lo_aligned, (size_t) (hi_aligned - lo_aligned), - MADV_DONTNEED); + const size_t stride_cell = (size_t) t->nb[1]; + const size_t off_in_buf = (size_t) ((char *) t->data - (char *) base); + const size_t off = off_in_buf + (size_t) from_cells * stride_cell; + const size_t sz = (size_t) (to_cells - from_cells) * stride_cell; + // posix_madvise wants the address — work in absolute terms. + void * addr = (char *) base + off; + // Round up to a page boundary on the low side, round down on the + // high side, so we only decommit pages we fully own. The cell-aligned + // address is unlikely to be page-aligned in general. + const size_t page = (size_t) sysconf(_SC_PAGESIZE); + if (page == 0 || page == (size_t) -1) { + return; } + uintptr_t lo = (uintptr_t) addr; + uintptr_t hi = lo + sz; + uintptr_t lo_aligned = (lo + page - 1) & ~(uintptr_t)(page - 1); + uintptr_t hi_aligned = hi & ~(uintptr_t)(page - 1); + if (hi_aligned <= lo_aligned) { + return; + } + // madvise(MADV_DONTNEED) returns 0 on success, -1 on failure (with + // errno set). We treat any failure as advisory — the RSS win may + // be skipped, but correctness is unaffected (the K/V buffer is + // still mapped, just possibly still resident). + (void) madvise((void *) lo_aligned, (size_t) (hi_aligned - lo_aligned), + MADV_DONTNEED); } void llama_kv_cache::shrink_if_idle() { @@ -667,13 +798,20 @@ void llama_kv_cache::shrink_if_idle() { (unsigned long long) (now - last_active_ms_), evicted); for (const auto & layer : layers) { - opencoti_decommit_layer_range(layer.k, n_stream, from_cells, to_cells); - opencoti_decommit_layer_range(layer.v, n_stream, from_cells, to_cells); - // opencoti F5 M2 headinfer — the host-resident head subset IS - // host-pageable, so this is where the GPU-mode shrink actually returns - // RSS (the device k/v above are skipped by the is_host() guard). - opencoti_decommit_layer_range(layer.k_cpu, n_stream, from_cells, to_cells); - opencoti_decommit_layer_range(layer.v_cpu, n_stream, from_cells, to_cells); + for (uint32_t s = 0; s < n_stream; ++s) { + opencoti_decommit_layer_range(layer.k_per_stream[s], from_cells, to_cells); + opencoti_decommit_layer_range(layer.v_per_stream[s], from_cells, to_cells); + // opencoti F5 M2 headinfer — the host-resident head subset + // IS host-pageable, so this is where the GPU-mode shrink + // actually returns RSS (the device k/v above are skipped + // by the is_host() guard in the function). + if (s < layer.k_cpu_per_stream.size()) { + opencoti_decommit_layer_range(layer.k_cpu_per_stream[s], from_cells, to_cells); + } + if (s < layer.v_cpu_per_stream.size()) { + opencoti_decommit_layer_range(layer.v_cpu_per_stream[s], from_cells, to_cells); + } + } } // After decommit, treat cells [from_cells, to_cells) as "uncleared" — @@ -1683,11 +1821,14 @@ bool llama_kv_cache::get_has_shift() const { } ggml_type llama_kv_cache::type_k() const { - return layers[0].k->type; + // opencoti F4 M3 Phase 4: all per-stream tensors of a layer carry the + // same type — read it from stream 0. Works in both unified and + // non-unified modes. + return layers[0].k_per_stream[0]->type; } ggml_type llama_kv_cache::type_v() const { - return layers[0].v->type; + return layers[0].v_per_stream[0]->type; } uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { @@ -1710,104 +1851,166 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k const int32_t ikv = map_layer_ids.at(il); const auto & layer = layers[ikv]; - auto * k = layer.k; const uint64_t kv_size = get_size(); - const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; // opencoti F5 M2 headinfer — reassemble the full head set from the - // GPU-resident (k) and host-resident (k_cpu) subsets by concat on the head + // GPU-resident and host-resident subsets by concat on the head // dimension. The backend scheduler streams k_cpu to the compute backend. - if (layer.gpu_heads > 0 && layer.k_cpu) { + // + // opencoti F4 M3 Phase 4-D: stream-aware. For each cache stream covered + // by sinfo, build a 4-D head-axis-concat view (GPU half + CPU half) + // with that stream's per-stream tensors. Then concat the per-stream + // results along the stream axis (dim 3). Unified mode (sinfo.s1 == + // sinfo.s0 == 0) collapses to one 4-D head-concat view of stream 0's + // tensors — byte-identical to pre-Phase-4. + if (layer.gpu_heads > 0 && !layer.k_cpu_per_stream.empty() && layer.k_cpu_per_stream[0]) { const uint32_t n_embd_head_k = hparams.n_embd_head_k(il); const uint32_t n_head_gpu = layer.gpu_heads; const uint32_t n_head_cpu = hparams.n_head_kv(il) - layer.gpu_heads; const uint64_t n_embd_gpu = (uint64_t) n_head_gpu * n_embd_head_k; const uint64_t n_embd_cpu = (uint64_t) n_head_cpu * n_embd_head_k; - auto * kc = layer.k_cpu; - ggml_tensor * kg_v = ggml_view_4d(ctx, k, - n_embd_head_k, n_head_gpu, n_kv, ns, - ggml_row_size(k->type, n_embd_head_k), - ggml_row_size(k->type, n_embd_gpu), - ggml_row_size(k->type, n_embd_gpu*kv_size), - ggml_row_size(k->type, n_embd_gpu*kv_size)*sinfo.s0); - ggml_tensor * kc_v = ggml_view_4d(ctx, kc, - n_embd_head_k, n_head_cpu, n_kv, ns, - ggml_row_size(kc->type, n_embd_head_k), - ggml_row_size(kc->type, n_embd_cpu), - ggml_row_size(kc->type, n_embd_cpu*kv_size), - ggml_row_size(kc->type, n_embd_cpu*kv_size)*sinfo.s0); - return ggml_concat(ctx, kg_v, kc_v, 1); - } - - const uint64_t n_embd_k_gqa = k->ne[0]; + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * k_s = layer.k_per_stream[s_cache]; + ggml_tensor * kc_s = layer.k_cpu_per_stream[s_cache]; + ggml_tensor * kg_v = ggml_view_4d(ctx, k_s, + n_embd_head_k, n_head_gpu, n_kv, 1, + ggml_row_size(k_s->type, n_embd_head_k), + ggml_row_size(k_s->type, n_embd_gpu), + ggml_row_size(k_s->type, n_embd_gpu*kv_size), + 0); + ggml_tensor * kc_v = ggml_view_4d(ctx, kc_s, + n_embd_head_k, n_head_cpu, n_kv, 1, + ggml_row_size(kc_s->type, n_embd_head_k), + ggml_row_size(kc_s->type, n_embd_cpu), + ggml_row_size(kc_s->type, n_embd_cpu*kv_size), + 0); + return ggml_concat(ctx, kg_v, kc_v, 1); + }; + + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); + } + return result; + } + + // opencoti F4 M3 Phase 4-C: per-stream view + concat along stream axis + // (dim 3). In unified mode (sinfo.s1 == sinfo.s0 == 0) the loop does + // not fire, yielding a single 4-D view of layer.k_per_stream[0] — + // structurally and numerically byte-identical to pre-Phase-4 (same + // ne, same nb, same offset 0). In non-unified mode the result is N + // contiguous concat ops along dim 3 (the stream axis), one per + // covered cache stream. Cost: (ns - 1) concat nodes per layer per turn. + auto * k0 = layer.k_per_stream[sinfo.s0]; + const uint64_t n_embd_k_gqa = k0->ne[0]; assert(n_embd_k_gqa == hparams.n_embd_k_gqa(il)); - return ggml_view_4d(ctx, k, - hparams.n_embd_head_k(il), hparams.n_head_kv(il), n_kv, ns, - ggml_row_size(k->type, hparams.n_embd_head_k(il)), - ggml_row_size(k->type, n_embd_k_gqa), - ggml_row_size(k->type, n_embd_k_gqa*kv_size), - ggml_row_size(k->type, n_embd_k_gqa*kv_size)*sinfo.s0); + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * k_s = layer.k_per_stream[s_cache]; + return ggml_view_4d(ctx, k_s, + hparams.n_embd_head_k(il), hparams.n_head_kv(il), n_kv, 1, + ggml_row_size(k_s->type, hparams.n_embd_head_k(il)), + ggml_row_size(k_s->type, n_embd_k_gqa), + ggml_row_size(k_s->type, n_embd_k_gqa*kv_size), + 0); + }; + + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); + } + return result; } ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { const int32_t ikv = map_layer_ids.at(il); const auto & layer = layers[ikv]; - auto * v = layer.v; const uint64_t kv_size = get_size(); - const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; - // opencoti F5 M2 headinfer — concat the GPU + host head subsets. The split - // is gated off when v_trans at construction, so this branch is always the - // non-transposed layout (heads in dim 1, same as get_k). - if (layer.gpu_heads > 0 && layer.v_cpu) { + // opencoti F5 M2 headinfer — concat the GPU + host head subsets. The + // split is gated off when v_trans at construction, so this branch is + // always the non-transposed layout (heads in dim 1, same as get_k). + // + // opencoti F4 M3 Phase 4-D: stream-aware per-stream concat — same + // pattern as get_k's M2 branch. + if (layer.gpu_heads > 0 && !layer.v_cpu_per_stream.empty() && layer.v_cpu_per_stream[0]) { const uint32_t n_embd_head_v = hparams.n_embd_head_v(il); const uint32_t n_head_gpu = layer.gpu_heads; const uint32_t n_head_cpu = hparams.n_head_kv(il) - layer.gpu_heads; const uint64_t n_embd_gpu = (uint64_t) n_head_gpu * n_embd_head_v; const uint64_t n_embd_cpu = (uint64_t) n_head_cpu * n_embd_head_v; - auto * vc = layer.v_cpu; - ggml_tensor * vg_v = ggml_view_4d(ctx, v, - n_embd_head_v, n_head_gpu, n_kv, ns, - ggml_row_size(v->type, n_embd_head_v), - ggml_row_size(v->type, n_embd_gpu), - ggml_row_size(v->type, n_embd_gpu*kv_size), - ggml_row_size(v->type, n_embd_gpu*kv_size)*sinfo.s0); - ggml_tensor * vc_v = ggml_view_4d(ctx, vc, - n_embd_head_v, n_head_cpu, n_kv, ns, - ggml_row_size(vc->type, n_embd_head_v), - ggml_row_size(vc->type, n_embd_cpu), - ggml_row_size(vc->type, n_embd_cpu*kv_size), - ggml_row_size(vc->type, n_embd_cpu*kv_size)*sinfo.s0); - return ggml_concat(ctx, vg_v, vc_v, 1); - } - - const uint64_t n_embd_v_gqa = v->ne[0]; + + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * v_s = layer.v_per_stream[s_cache]; + ggml_tensor * vc_s = layer.v_cpu_per_stream[s_cache]; + ggml_tensor * vg_v = ggml_view_4d(ctx, v_s, + n_embd_head_v, n_head_gpu, n_kv, 1, + ggml_row_size(v_s->type, n_embd_head_v), + ggml_row_size(v_s->type, n_embd_gpu), + ggml_row_size(v_s->type, n_embd_gpu*kv_size), + 0); + ggml_tensor * vc_v = ggml_view_4d(ctx, vc_s, + n_embd_head_v, n_head_cpu, n_kv, 1, + ggml_row_size(vc_s->type, n_embd_head_v), + ggml_row_size(vc_s->type, n_embd_cpu), + ggml_row_size(vc_s->type, n_embd_cpu*kv_size), + 0); + return ggml_concat(ctx, vg_v, vc_v, 1); + }; + + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); + } + return result; + } + + // opencoti F4 M3 Phase 4-C: per-stream view + concat along stream + // axis. Same shape as get_k, but with v_trans branching. Unified mode + // (sinfo.s1 == sinfo.s0 == 0) keeps a single 4-D view → byte-identical. + auto * v0 = layer.v_per_stream[sinfo.s0]; + const uint64_t n_embd_v_gqa = v0->ne[0]; // [TAG_V_CACHE_VARIABLE] assert(n_embd_v_gqa >= hparams.n_embd_v_gqa(il)); if (!v_trans) { // note: v->nb[1] <= v->nb[2] - return ggml_view_4d(ctx, v, - hparams.n_embd_head_v(il), hparams.n_head_kv(il), n_kv, ns, - ggml_row_size(v->type, hparams.n_embd_head_v(il)), // v->nb[1] - ggml_row_size(v->type, n_embd_v_gqa), // v->nb[2] - ggml_row_size(v->type, n_embd_v_gqa*kv_size), // v->nb[3] - ggml_row_size(v->type, n_embd_v_gqa*kv_size)*sinfo.s0); + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * v_s = layer.v_per_stream[s_cache]; + return ggml_view_4d(ctx, v_s, + hparams.n_embd_head_v(il), hparams.n_head_kv(il), n_kv, 1, + ggml_row_size(v_s->type, hparams.n_embd_head_v(il)), + ggml_row_size(v_s->type, n_embd_v_gqa), + ggml_row_size(v_s->type, n_embd_v_gqa*kv_size), + 0); + }; + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); + } + return result; + } + + // v_trans branch: note v->nb[1] > v->nb[2] + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * v_s = layer.v_per_stream[s_cache]; + return ggml_view_4d(ctx, v_s, + n_kv, hparams.n_head_kv(il), hparams.n_embd_head_v(il), 1, + ggml_row_size(v_s->type, kv_size*hparams.n_embd_head_v(il)), + ggml_row_size(v_s->type, kv_size), + ggml_row_size(v_s->type, kv_size*n_embd_v_gqa), + 0); + }; + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); } - - // note: v->nb[1] > v->nb[2] - return ggml_view_4d(ctx, v, - n_kv, hparams.n_head_kv(il), hparams.n_embd_head_v(il), ns, - ggml_row_size(v->type, kv_size*hparams.n_embd_head_v(il)), // v->nb[1] - ggml_row_size(v->type, kv_size), // v->nb[2] - ggml_row_size(v->type, kv_size*n_embd_v_gqa), // v->nb[3] - ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0); + return result; } // opencoti F5 M3 neo — head-split-aware accessors. These factor the per-half @@ -1824,7 +2027,13 @@ bool llama_kv_cache::headinfer_split_active(int32_t il) const { return false; } const auto & layer = layers[it->second]; - return layer.gpu_heads > 0 && layer.k_cpu != nullptr; + // opencoti F4 M3 Phase 4: per-stream k_cpu vector is empty when no split, + // populated when gpu_heads > 0. Stream 0 always exists in unified mode + // (the only mode where split engages today; the M2 construction gate + // requires n_stream == 1). + return layer.gpu_heads > 0 + && !layer.k_cpu_per_stream.empty() + && layer.k_cpu_per_stream[0] != nullptr; } uint32_t llama_kv_cache::headinfer_gpu_heads(int32_t il) const { @@ -1840,93 +2049,131 @@ ggml_tensor * llama_kv_cache::get_k_gpu(ggml_context * ctx, int32_t il, uint32_t const auto & layer = layers[ikv]; GGML_ASSERT(layer.gpu_heads > 0 && "get_k_gpu called without an active headinfer split"); - auto * k = layer.k; - const uint64_t kv_size = get_size(); - const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; + // opencoti F4 M3 Phase 4-D: per-stream M2 GPU subset. In unified + // mode (sinfo.s1 == sinfo.s0 == 0) the loop runs once and yields + // a single 4-D view of layer.k_per_stream[0]'s GPU half — + // byte-identical to pre-Phase-4 (same ne, same nb, offset 0). + const uint64_t kv_size = get_size(); const uint32_t n_embd_head_k = hparams.n_embd_head_k(il); const uint32_t n_head_gpu = layer.gpu_heads; const uint64_t n_embd_gpu = (uint64_t) n_head_gpu * n_embd_head_k; - return ggml_view_4d(ctx, k, - n_embd_head_k, n_head_gpu, n_kv, ns, - ggml_row_size(k->type, n_embd_head_k), - ggml_row_size(k->type, n_embd_gpu), - ggml_row_size(k->type, n_embd_gpu*kv_size), - ggml_row_size(k->type, n_embd_gpu*kv_size)*sinfo.s0); + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * k_s = layer.k_per_stream[s_cache]; + return ggml_view_4d(ctx, k_s, + n_embd_head_k, n_head_gpu, n_kv, 1, + ggml_row_size(k_s->type, n_embd_head_k), + ggml_row_size(k_s->type, n_embd_gpu), + ggml_row_size(k_s->type, n_embd_gpu*kv_size), + 0); + }; + + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); + } + return result; } ggml_tensor * llama_kv_cache::get_k_cpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { const int32_t ikv = map_layer_ids.at(il); const auto & layer = layers[ikv]; - GGML_ASSERT(layer.gpu_heads > 0 && layer.k_cpu && "get_k_cpu called without an active headinfer split"); - - auto * kc = layer.k_cpu; - const uint64_t kv_size = get_size(); - const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; + GGML_ASSERT(layer.gpu_heads > 0 + && !layer.k_cpu_per_stream.empty() + && layer.k_cpu_per_stream[0] + && "get_k_cpu called without an active headinfer split"); + + // opencoti F4 M3 Phase 4-D: per-stream M2 CPU subset. Same pattern + // as get_k_gpu; unified mode collapses to a single 4-D view. + const uint64_t kv_size = get_size(); const uint32_t n_embd_head_k = hparams.n_embd_head_k(il); const uint32_t n_head_cpu = hparams.n_head_kv(il) - layer.gpu_heads; const uint64_t n_embd_cpu = (uint64_t) n_head_cpu * n_embd_head_k; - return ggml_view_4d(ctx, kc, - n_embd_head_k, n_head_cpu, n_kv, ns, - ggml_row_size(kc->type, n_embd_head_k), - ggml_row_size(kc->type, n_embd_cpu), - ggml_row_size(kc->type, n_embd_cpu*kv_size), - ggml_row_size(kc->type, n_embd_cpu*kv_size)*sinfo.s0); + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * kc_s = layer.k_cpu_per_stream[s_cache]; + return ggml_view_4d(ctx, kc_s, + n_embd_head_k, n_head_cpu, n_kv, 1, + ggml_row_size(kc_s->type, n_embd_head_k), + ggml_row_size(kc_s->type, n_embd_cpu), + ggml_row_size(kc_s->type, n_embd_cpu*kv_size), + 0); + }; + + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); + } + return result; } ggml_tensor * llama_kv_cache::get_v_gpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { const int32_t ikv = map_layer_ids.at(il); const auto & layer = layers[ikv]; GGML_ASSERT(layer.gpu_heads > 0 && "get_v_gpu called without an active headinfer split"); - // M2's gate is `!v_trans && n_stream == 1` — when split is active the - // V layout is always non-transposed (matches get_v's non-split branch - // shape modulo head count). Same head-dim-1 stride pattern as K. + // M2 requires non-transposed V (FA on). The split-active branch in + // get_v/cpy_v also requires this. Per-stream view + concat on stream + // axis. Phase 4-D: lifted the n_stream == 1 gate. GGML_ASSERT(!v_trans && "headinfer split requires non-transposed V"); - auto * v = layer.v; - const uint64_t kv_size = get_size(); - const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; + const uint64_t kv_size = get_size(); const uint32_t n_embd_head_v = hparams.n_embd_head_v(il); const uint32_t n_head_gpu = layer.gpu_heads; const uint64_t n_embd_gpu = (uint64_t) n_head_gpu * n_embd_head_v; - return ggml_view_4d(ctx, v, - n_embd_head_v, n_head_gpu, n_kv, ns, - ggml_row_size(v->type, n_embd_head_v), - ggml_row_size(v->type, n_embd_gpu), - ggml_row_size(v->type, n_embd_gpu*kv_size), - ggml_row_size(v->type, n_embd_gpu*kv_size)*sinfo.s0); + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * v_s = layer.v_per_stream[s_cache]; + return ggml_view_4d(ctx, v_s, + n_embd_head_v, n_head_gpu, n_kv, 1, + ggml_row_size(v_s->type, n_embd_head_v), + ggml_row_size(v_s->type, n_embd_gpu), + ggml_row_size(v_s->type, n_embd_gpu*kv_size), + 0); + }; + + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); + } + return result; } ggml_tensor * llama_kv_cache::get_v_cpu(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { const int32_t ikv = map_layer_ids.at(il); const auto & layer = layers[ikv]; - GGML_ASSERT(layer.gpu_heads > 0 && layer.v_cpu && "get_v_cpu called without an active headinfer split"); + GGML_ASSERT(layer.gpu_heads > 0 + && !layer.v_cpu_per_stream.empty() + && layer.v_cpu_per_stream[0] + && "get_v_cpu called without an active headinfer split"); GGML_ASSERT(!v_trans && "headinfer split requires non-transposed V"); - auto * vc = layer.v_cpu; - const uint64_t kv_size = get_size(); - const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; + // opencoti F4 M3 Phase 4-D: per-stream M2 CPU V subset. + const uint64_t kv_size = get_size(); const uint32_t n_embd_head_v = hparams.n_embd_head_v(il); const uint32_t n_head_cpu = hparams.n_head_kv(il) - layer.gpu_heads; const uint64_t n_embd_cpu = (uint64_t) n_head_cpu * n_embd_head_v; - return ggml_view_4d(ctx, vc, - n_embd_head_v, n_head_cpu, n_kv, ns, - ggml_row_size(vc->type, n_embd_head_v), - ggml_row_size(vc->type, n_embd_cpu), - ggml_row_size(vc->type, n_embd_cpu*kv_size), - ggml_row_size(vc->type, n_embd_cpu*kv_size)*sinfo.s0); + auto build_stream_view = [&](uint32_t s_cache) -> ggml_tensor * { + ggml_tensor * vc_s = layer.v_cpu_per_stream[s_cache]; + return ggml_view_4d(ctx, vc_s, + n_embd_head_v, n_head_cpu, n_kv, 1, + ggml_row_size(vc_s->type, n_embd_head_v), + ggml_row_size(vc_s->type, n_embd_cpu), + ggml_row_size(vc_s->type, n_embd_cpu*kv_size), + 0); + }; + + ggml_tensor * result = build_stream_view(sinfo.s0); + for (uint32_t s = sinfo.s0 + 1; s <= sinfo.s1; ++s) { + result = ggml_concat(ctx, result, build_stream_view(s), 3); + } + return result; } ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const { - GGML_UNUSED(sinfo); - const int32_t ikv = map_layer_ids.at(il); const auto & layer = layers[ikv]; - ggml_tensor * k = layer.k; const int64_t n_embd_head = k_cur->ne[0]; const int64_t n_head = k_cur->ne[1]; @@ -1941,46 +2188,116 @@ ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggm // opencoti F5 M2 headinfer — scatter the incoming heads into the GPU and // host subsets. The first gpu_heads heads (contiguous in k_cur's dim 0..1) // store into k; the rest store into k_cpu, both at the same cell indices - // (k_idxs). The single returned dep node ties both set_rows into the - // caller's one ggml_build_forward_expand (its value is never read). Split - // only happens for n_stream==1 (gated at construction) — no per-stream - // reshape needed. - if (layer.gpu_heads > 0 && layer.k_cpu) { + // (k_idxs). The single returned dep node ties all set_rows into the + // caller's one ggml_build_forward_expand (its value is never read). + // + // opencoti F4 M3 Phase 4-D: per-stream M2 scatter — for each batch + // stream s, slice k_cur and k_idxs by stream, split the head-axis + // slice into GPU + CPU halves, and emit 2 set_rows per stream against + // the per-stream tensors. In unified mode (sinfo.n_stream() == 1) this + // collapses to exactly 2 set_rows on stream 0's tensors — byte-identical + // to pre-Phase-4. + if (layer.gpu_heads > 0 + && !layer.k_cpu_per_stream.empty() + && layer.k_cpu_per_stream[0]) { const int64_t g = layer.gpu_heads; const int64_t c = n_head - g; - ggml_tensor * kc = layer.k_cpu; - ggml_tensor * cur_g = ggml_view_2d(ctx, k_cur, n_embd_head*g, n_tokens, k_cur->nb[2], 0); - ggml_tensor * cur_c = ggml_view_2d(ctx, k_cur, n_embd_head*c, n_tokens, k_cur->nb[2], (size_t) (g*k_cur->nb[1])); - ggml_tensor * store_g = ggml_set_rows(ctx, k, cur_g, k_idxs); - ggml_tensor * store_c = ggml_set_rows(ctx, kc, cur_c, k_idxs); - return ggml_add(ctx, ggml_view_1d(ctx, store_g, 1, 0), ggml_view_1d(ctx, store_c, 1, 0)); + const uint32_t batch_ns = sinfo.n_stream(); + GGML_ASSERT(batch_ns >= 1); + + const int64_t tokens_per_stream = (batch_ns == 1) ? n_tokens : (n_tokens / (int64_t) batch_ns); + GGML_ASSERT((int64_t) batch_ns * tokens_per_stream == n_tokens); + + ggml_tensor * tie = nullptr; + for (uint32_t s = 0; s < batch_ns; ++s) { + ggml_tensor * k_s = layer.k_per_stream[sinfo.strm[s]]; + ggml_tensor * kc_s = layer.k_cpu_per_stream[sinfo.strm[s]]; + + // k_cur is [head_dim, n_head, n_tokens]; nb[2] is the stride + // between tokens. The per-stream slice is the contiguous block + // of tokens_per_stream tokens at offset s*tokens_per_stream. + const size_t per_stream_off = (size_t) s * tokens_per_stream * k_cur->nb[2]; + ggml_tensor * cur_g = ggml_view_2d(ctx, k_cur, + n_embd_head*g, tokens_per_stream, + k_cur->nb[2], per_stream_off); + // CPU half: remaining c heads, offset by g heads (g rows in the head-merged dim 0) + ggml_tensor * cur_c = ggml_view_2d(ctx, k_cur, + n_embd_head*c, tokens_per_stream, + k_cur->nb[2], per_stream_off + (size_t)(g*k_cur->nb[1])); + ggml_tensor * idxs_s = (batch_ns == 1) + ? k_idxs + : ggml_view_1d(ctx, k_idxs, tokens_per_stream, + (size_t) s * tokens_per_stream * k_idxs->nb[0]); + ggml_tensor * store_g = ggml_set_rows(ctx, k_s, cur_g, idxs_s); + ggml_tensor * store_c = ggml_set_rows(ctx, kc_s, cur_c, idxs_s); + ggml_tensor * pair = ggml_add(ctx, + ggml_view_1d(ctx, store_g, 1, 0), + ggml_view_1d(ctx, store_c, 1, 0)); + if (!tie) { + tie = pair; + } else { + tie = ggml_add(ctx, + ggml_view_1d(ctx, tie, 1, 0), + ggml_view_1d(ctx, pair, 1, 0)); + } + } + return tie; } k_cur = ggml_view_2d(ctx, k_cur, n_embd_gqa, n_tokens, k_cur->nb[2], 0); - const int64_t n_stream = k->ne[2]; - - if (n_stream > 1) { - const int64_t kv_size = get_size(); - - assert(n_embd_gqa == k->ne[0]); - assert(kv_size == k->ne[1]); - - // merge the buffer across all streams because the idxs are global - k = ggml_reshape_2d(ctx, k, n_embd_gqa, kv_size*n_stream); + // opencoti F4 M3 Phase 4-C: per-stream scatter via N separate + // ggml_set_rows ops (option C-2). In unified mode (sinfo.n_stream() + // == 1) the fast path collapses to one ggml_set_rows against + // layer.k_per_stream[sinfo.strm[0]] — byte-identical to pre-Phase-4. + // In non-unified mode (sinfo.n_stream() > 1), set_input_k_idxs has + // emitted LOCAL indices (no global offset), so each per-stream slice + // of k_idxs is valid against its own k_per_stream[cache_strm] tensor. + // The N stores are tied into a single dep node via ggml_add so the + // caller's one ggml_build_forward_expand picks them all up. + // + // C-1 (concat-then-scatter) is not expressible with per-stream owning + // tensors in ggml: ggml_concat materializes a new buffer rather than + // creating a view into its inputs, so writes routed through the concat + // result would not land in the per-stream tensors. C-2 is the natural + // shape that preserves per-stream buffer ownership. + const uint32_t batch_ns = sinfo.n_stream(); + GGML_ASSERT(batch_ns >= 1); + + if (batch_ns == 1) { + ggml_tensor * k_s = layer.k_per_stream[sinfo.strm[0]]; + return ggml_set_rows(ctx, k_s, k_cur, k_idxs); + } + + const int64_t tokens_per_stream = n_tokens / (int64_t) batch_ns; + GGML_ASSERT((int64_t) batch_ns * tokens_per_stream == n_tokens); + + ggml_tensor * tie = nullptr; + for (uint32_t s = 0; s < batch_ns; ++s) { + ggml_tensor * k_s = layer.k_per_stream[sinfo.strm[s]]; + ggml_tensor * cur_s = ggml_view_2d(ctx, k_cur, + n_embd_gqa, tokens_per_stream, + k_cur->nb[1], + (size_t) s * tokens_per_stream * k_cur->nb[1]); + ggml_tensor * idxs_s = ggml_view_1d(ctx, k_idxs, + tokens_per_stream, + (size_t) s * tokens_per_stream * k_idxs->nb[0]); + ggml_tensor * store_s = ggml_set_rows(ctx, k_s, cur_s, idxs_s); + if (!tie) { + tie = store_s; + } else { + tie = ggml_add(ctx, + ggml_view_1d(ctx, tie, 1, 0), + ggml_view_1d(ctx, store_s, 1, 0)); + } } - - // store the current K values into the cache - return ggml_set_rows(ctx, k, k_cur, k_idxs); + return tie; } ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const { - GGML_UNUSED(sinfo); - const int32_t ikv = map_layer_ids.at(il); const auto & layer = layers[ikv]; - auto * v = layer.v; const int64_t n_embd_head = v_cur->ne[0]; const int64_t n_head = v_cur->ne[1]; @@ -1992,38 +2309,109 @@ ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggm GGML_ASSERT(ggml_row_size(v_cur->type, n_embd_head) == v_cur->nb[1]); // opencoti F5 M2 headinfer — head-scatter into the GPU + host subsets. - // Reached only when !v_trans and n_stream==1 (the split is gated off - // otherwise at construction), so this mirrors the cpy_k store exactly. - if (layer.gpu_heads > 0 && layer.v_cpu) { + // Reached only when !v_trans (the split is gated off at construction + // otherwise). + // + // opencoti F4 M3 Phase 4-D: per-stream M2 V scatter — same shape as + // cpy_k's M2 branch. Unified mode collapses to 2 set_rows on stream 0, + // byte-identical to pre-Phase-4. + if (layer.gpu_heads > 0 + && !layer.v_cpu_per_stream.empty() + && layer.v_cpu_per_stream[0]) { const int64_t g = layer.gpu_heads; const int64_t c = n_head - g; - ggml_tensor * vc = layer.v_cpu; - ggml_tensor * cur_g = ggml_view_2d(ctx, v_cur, n_embd_head*g, n_tokens, v_cur->nb[2], 0); - ggml_tensor * cur_c = ggml_view_2d(ctx, v_cur, n_embd_head*c, n_tokens, v_cur->nb[2], (size_t) (g*v_cur->nb[1])); - ggml_tensor * store_g = ggml_set_rows(ctx, v, cur_g, v_idxs); - ggml_tensor * store_c = ggml_set_rows(ctx, vc, cur_c, v_idxs); - return ggml_add(ctx, ggml_view_1d(ctx, store_g, 1, 0), ggml_view_1d(ctx, store_c, 1, 0)); + const uint32_t batch_ns = sinfo.n_stream(); + GGML_ASSERT(batch_ns >= 1); + + const int64_t tokens_per_stream = (batch_ns == 1) ? n_tokens : (n_tokens / (int64_t) batch_ns); + GGML_ASSERT((int64_t) batch_ns * tokens_per_stream == n_tokens); + + ggml_tensor * tie = nullptr; + for (uint32_t s = 0; s < batch_ns; ++s) { + ggml_tensor * v_s = layer.v_per_stream[sinfo.strm[s]]; + ggml_tensor * vc_s = layer.v_cpu_per_stream[sinfo.strm[s]]; + + const size_t per_stream_off = (size_t) s * tokens_per_stream * v_cur->nb[2]; + ggml_tensor * cur_g = ggml_view_2d(ctx, v_cur, + n_embd_head*g, tokens_per_stream, + v_cur->nb[2], per_stream_off); + ggml_tensor * cur_c = ggml_view_2d(ctx, v_cur, + n_embd_head*c, tokens_per_stream, + v_cur->nb[2], per_stream_off + (size_t)(g*v_cur->nb[1])); + ggml_tensor * idxs_s = (batch_ns == 1) + ? v_idxs + : ggml_view_1d(ctx, v_idxs, tokens_per_stream, + (size_t) s * tokens_per_stream * v_idxs->nb[0]); + ggml_tensor * store_g = ggml_set_rows(ctx, v_s, cur_g, idxs_s); + ggml_tensor * store_c = ggml_set_rows(ctx, vc_s, cur_c, idxs_s); + ggml_tensor * pair = ggml_add(ctx, + ggml_view_1d(ctx, store_g, 1, 0), + ggml_view_1d(ctx, store_c, 1, 0)); + if (!tie) { + tie = pair; + } else { + tie = ggml_add(ctx, + ggml_view_1d(ctx, tie, 1, 0), + ggml_view_1d(ctx, pair, 1, 0)); + } + } + return tie; } - const int64_t n_stream = v->ne[2]; + // opencoti F4 M3 Phase 4-C: per-stream V scatter. Two sub-branches + // depending on v_trans (FA vs non-FA layout). + // + // !v_trans (FA on): same shape as cpy_k — N separate ggml_set_rows on + // the 2-D-merged v_cur, sliced per stream. set_input_v_idxs emits + // LOCAL per-stream indices in non-unified mode. Unified mode (batch_ns + // == 1) keeps the simple one-set_rows fast path, byte-identical. + // + // v_trans (FA off, V transposed): each row is a single head element. + // The unified-mode path is preserved; multi-stream + v_trans is not a + // current workload (the v_trans layout requires no FA, while parallel + // multi-stream batches are always paired with FA). Asserted out as a + // safety net; future Phase 4-D/E may lift if a workload appears. + const uint32_t batch_ns = sinfo.n_stream(); + GGML_ASSERT(batch_ns >= 1); - // take this branch when FA is enabled (the V cache is not transposed) if (!v_trans) { v_cur = ggml_view_2d(ctx, v_cur, n_embd_gqa, n_tokens, v_cur->nb[2], 0); - if (n_stream > 1) { - const int64_t kv_size = get_size(); - - assert(n_embd_gqa == v->ne[0]); - assert(kv_size == v->ne[1]); - - // merge the buffer across all streams because the idxs are global - v = ggml_reshape_2d(ctx, v, n_embd_gqa, kv_size*n_stream); + if (batch_ns == 1) { + ggml_tensor * v_s = layer.v_per_stream[sinfo.strm[0]]; + return ggml_set_rows(ctx, v_s, v_cur, v_idxs); + } + + const int64_t tokens_per_stream = n_tokens / (int64_t) batch_ns; + GGML_ASSERT((int64_t) batch_ns * tokens_per_stream == n_tokens); + + ggml_tensor * tie = nullptr; + for (uint32_t s = 0; s < batch_ns; ++s) { + ggml_tensor * v_s = layer.v_per_stream[sinfo.strm[s]]; + ggml_tensor * cur_s = ggml_view_2d(ctx, v_cur, + n_embd_gqa, tokens_per_stream, + v_cur->nb[1], + (size_t) s * tokens_per_stream * v_cur->nb[1]); + ggml_tensor * idxs_s = ggml_view_1d(ctx, v_idxs, + tokens_per_stream, + (size_t) s * tokens_per_stream * v_idxs->nb[0]); + ggml_tensor * store_s = ggml_set_rows(ctx, v_s, cur_s, idxs_s); + if (!tie) { + tie = store_s; + } else { + tie = ggml_add(ctx, + ggml_view_1d(ctx, tie, 1, 0), + ggml_view_1d(ctx, store_s, 1, 0)); + } } - - return ggml_set_rows(ctx, v, v_cur, v_idxs); + return tie; } + // v_trans path: not a multi-stream workload; preserve unified-mode + // behavior exactly and assert multi-stream out. + GGML_ASSERT(batch_ns == 1 && "v_trans + multi-stream is unsupported; v_trans implies !FA which is mutually exclusive with parallel batches"); + ggml_tensor * v = layer.v_per_stream[sinfo.strm[0]]; + if (ggml_row_size(v_cur->type, n_embd_gqa) == v_cur->nb[2]) { // we can merge dims 0, 1 and 2 v_cur = ggml_reshape_2d(ctx, v_cur, n_embd_gqa, n_tokens); @@ -2119,11 +2507,23 @@ void llama_kv_cache::set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ub GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); int64_t * data = (int64_t *) dst->data; - for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { - const int64_t offs = sinfo.strm[s]*get_size(); + // opencoti F4 M3 Phase 4-C/F: emit LOCAL per-stream indices unconditionally. + // Phase 4 ALWAYS scatters to per-stream tensors (cpy_k routes to + // k_per_stream[sinfo.strm[s]] which is sized kv_size, not kv_size*n_stream), + // so the index must be local to that stream's tensor. In unified mode + // (n_stream == 1, sinfo.strm[0] == 0) pre-Phase-4 offs was 0 already, + // so this is byte-identical. + // + // Phase 4-F bug fix: the previous version conditioned the offset + // on sinfo.n_stream() (per-batch streams), which left it global + // when a --parallel N server served a single-session batch + // (sinfo.n_stream() == 1 but sinfo.strm[0] != 0). The resulting + // global indices wrote out-of-bounds against the per-stream + // tensor and crashed CUDA decode. + for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { for (uint32_t i = 0; i < sinfo.size(); ++i) { - data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i]; + data[s*sinfo.size() + i] = sinfo.idxs[s][i]; } } } @@ -2135,12 +2535,15 @@ void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ub GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); int64_t * data = (int64_t *) dst->data; + // opencoti F4 M3 Phase 4-C/F: same LOCAL-only fix as set_input_k_idxs. + // The v_trans branch's offset is also dropped — v_trans + multi-stream + // is asserted out in cpy_v, but the unified-mode v_trans path here is + // byte-identical regardless (sinfo.strm[0] == 0 makes offs == 0). + if (!v_trans) { for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { - const int64_t offs = sinfo.strm[s]*get_size(); - for (uint32_t i = 0; i < sinfo.size(); ++i) { - data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i]; + data[s*sinfo.size() + i] = sinfo.idxs[s][i]; } } } else { @@ -2150,11 +2553,9 @@ void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ub const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa_max(); for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { - const int64_t offs = sinfo.strm[s]*kv_size*n_embd_v_gqa; - for (uint32_t i = 0; i < sinfo.size(); ++i) { for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { - data[s*sinfo.size()*n_embd_v_gqa + i*n_embd_v_gqa + j] = offs + j*kv_size + sinfo.idxs[s][i]; + data[s*sinfo.size()*n_embd_v_gqa + i*n_embd_v_gqa + j] = j*kv_size + sinfo.idxs[s][i]; } } } @@ -2459,10 +2860,16 @@ size_t llama_kv_cache::total_size() const { } size_t llama_kv_cache::size_k_bytes() const { + // opencoti F4 M3 Phase 4: sum bytes across all per-stream tensors of + // each layer. In unified mode (n_stream == 1) this is byte-identical + // to ggml_nbytes(layer.k) of pre-Phase-4. M2's host subset is + // accounted separately (it lives in a different buffer-type). size_t size_k_bytes = 0; for (const auto & layer : layers) { - size_k_bytes += ggml_nbytes(layer.k); + for (auto * k : layer.k_per_stream) { + if (k) size_k_bytes += ggml_nbytes(k); + } } return size_k_bytes; @@ -2472,7 +2879,9 @@ size_t llama_kv_cache::size_v_bytes() const { size_t size_v_bytes = 0; for (const auto & layer : layers) { - size_v_bytes += layer.v ? ggml_nbytes(layer.v) : 0; + for (auto * v : layer.v_per_stream) { + if (v) size_v_bytes += ggml_nbytes(v); + } } return size_v_bytes; @@ -2570,6 +2979,19 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co const auto & cparams = lctx->get_cparams(); + // opencoti F4 M3 Phase 4-C: k-shift loops over per-stream tensors. + // With per-stream owning tensors, each stream's K is a separate buffer, + // so we build N rope-shift ops per layer (one per stream). The k_shift + // input tensor has length get_size()*n_stream, with stream s occupying + // indices [s*get_size(), (s+1)*get_size()) — set_input_k_shift writes + // them in that order. Each per-stream rope-shift consumes its + // contiguous slice via ggml_view_1d on k_shift. + // + // In unified mode (n_stream == 1), the loop runs once with stream 0, + // and the k_shift view degenerates to the full tensor (offset 0, + // length get_size()). Byte-identical to pre-Phase-4 — same rope-shift + // op shape and inputs. + for (const auto & layer : layers) { const uint32_t il = layer.il; @@ -2585,16 +3007,24 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); - ggml_tensor * k = - ggml_view_3d(ctx, layer.k, - n_rot, n_head_kv, get_size()*n_stream, - ggml_row_size(layer.k->type, n_embd_head_k), - ggml_row_size(layer.k->type, n_embd_k_gqa), - ggml_row_size(layer.k->type, n_embd_nope)); + for (uint32_t s = 0; s < n_stream; ++s) { + ggml_tensor * layer_k = layer.k_per_stream[s]; + ggml_tensor * k = + ggml_view_3d(ctx, layer_k, + n_rot, n_head_kv, get_size(), + ggml_row_size(layer_k->type, n_embd_head_k), + ggml_row_size(layer_k->type, n_embd_k_gqa), + ggml_row_size(layer_k->type, n_embd_nope)); + + ggml_tensor * k_shift_s = (n_stream == 1) + ? inp->k_shift + : ggml_view_1d(ctx, inp->k_shift, get_size(), + (size_t) s * get_size() * inp->k_shift->nb[0]); - ggml_tensor * cur = build_rope_shift(cparams, ctx, k, inp->k_shift, inp->k_rot, rope_factors, freq_base_l, freq_scale_l, il); + ggml_tensor * cur = build_rope_shift(cparams, ctx, k, k_shift_s, inp->k_rot, rope_factors, freq_base_l, freq_scale_l, il); - ggml_build_forward_expand(gf, cur); + ggml_build_forward_expand(gf, cur); + } } res->add_input(std::move(inp)); diff --git a/llama.cpp/src/llama-kv-cache.h b/llama.cpp/src/llama-kv-cache.h --- a/llama.cpp/src/llama-kv-cache.h +++ b/llama.cpp/src/llama-kv-cache.h @@ -286,19 +286,57 @@ private: // note: can be different from the layer index in the KV cache uint32_t il; - ggml_tensor * k; - ggml_tensor * v; - + // opencoti F4 M3 Phase 4 — per-stream tensor split. + // Each element is an OWNING 3-D tensor of shape + // [n_embd_k_gqa, kv_size, 1] (dim 3 == 1 keeps stride math + // identical to the pre-Phase-4 single-tensor layout). Vector + // size always equals n_stream. + // + // Unified mode (n_stream == 1): k_per_stream[0] is byte- + // identical to the pre-Phase-4 `k` field (same buft, same + // ctx, same offsets). Non-unified mode (n_stream > 1): each + // stream's tensor lives in its own (possibly per-stream) + // backend buffer (Phase 4-B); cross-stream operations are + // buffer-to-buffer copies. + std::vector k_per_stream; + std::vector v_per_stream; + + // k_stream[s] / v_stream[s] are 2-D views into + // k_per_stream[s] / v_per_stream[s]. Kept for backward compat + // with consumers that index by stream; same shape and offset + // semantics as before. With per-stream tensors the view's + // offset within its owning tensor is always 0 (each stream + // owns the whole 3-D extent of its tensor with dim 3 == 1). std::vector k_stream; std::vector v_stream; - // opencoti F5 M2 headinfer — head-residency split. When gpu_heads > 0, - // `k`/`v` hold the first gpu_heads KV heads on the layer's (GPU) buffer - // and `k_cpu`/`v_cpu` hold the remaining (n_head_kv - gpu_heads) heads - // on a host buffer. gpu_heads == 0 means no split (k/v hold all heads). - ggml_tensor * k_cpu = nullptr; - ggml_tensor * v_cpu = nullptr; + // opencoti F5 M2 headinfer — per-stream now. When gpu_heads > 0, + // k_per_stream[s] / v_per_stream[s] hold the first gpu_heads + // KV heads on the layer's device buffer, and + // k_cpu_per_stream[s] / v_cpu_per_stream[s] hold the remaining + // (n_head_kv - gpu_heads) heads on the layer's host buffer. + // Vectors are empty when no split (gpu_heads == 0). + std::vector k_cpu_per_stream; + std::vector v_cpu_per_stream; uint32_t gpu_heads = 0; + + // Convenience accessors for unified-mode (n_stream == 1) call + // sites. Returns nullptr in non-unified mode — that's the + // signal that the caller must use per-stream access instead. + // Phase 4-A uses these at sites that haven't yet been migrated + // to per-stream awareness; Phase 4-C/D lift them. + inline ggml_tensor * k_unified() const { + return k_per_stream.size() == 1 ? k_per_stream[0] : nullptr; + } + inline ggml_tensor * v_unified() const { + return v_per_stream.size() == 1 ? v_per_stream[0] : nullptr; + } + inline ggml_tensor * k_cpu_unified() const { + return k_cpu_per_stream.size() == 1 ? k_cpu_per_stream[0] : nullptr; + } + inline ggml_tensor * v_cpu_unified() const { + return v_cpu_per_stream.size() == 1 ? v_cpu_per_stream[0] : nullptr; + } }; bool v_trans = true; // the value tensor is transposed