From: opencoti Subject: [PATCH 0020] F5 M2 — HeadInfer: head-wise GPU/CPU KV residency split Keep a configurable fraction of KV heads GPU-resident; offload the rest to host memory; reassemble the full head set via ggml_concat on the head axis at each attention step. Default off (frac = 1.0). Validated 6/6 on solidPC RTX 3090 with Qwen2.5-Coder-0.5B-IQ4_XS at -ngl 99: GPU KV halves (48 → 24 MiB at frac=0.5), CPU half appears (24 MiB), greedy decode byte-identical between baseline and split. See docs/features/advanced_kv.md (F5 M2). The CUDA backend's supports_op for GGML_OP_CONCAT was over-broad (claimed any non-I32/I16 type) while the kernel ggml_cuda_op_concat only implements F32. Tightened to match the kernel — F16 concat now correctly routes to the CPU backend (which has a real F16 path at ggml-cpu/ops.cpp:1980). Required for the reassembly to run on F16 KV caches. diff --git a/llama.cpp/common/arg.cpp b/llama.cpp/common/arg.cpp --- a/llama.cpp/common/arg.cpp +++ b/llama.cpp/common/arg.cpp @@ -1421,6 +1421,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.rest_kv_layer = value; } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REST_KV_LAYER")); + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + add_opt(common_arg( + {"--headinfer-gpu-heads-frac"}, "F", + string_format("headinfer: fraction of KV heads kept GPU-resident; the rest are offloaded to host memory and streamed back per step (default: %.2f, 1.0 = off). Only effective for GPU-offloaded layers with flash-attention.", (double) params.headinfer_gpu_heads_frac), + [](common_params & params, const std::string & value) { + params.headinfer_gpu_heads_frac = std::stof(value); + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_HEADINFER_GPU_HEADS_FRAC")); add_opt(common_arg( {"--chunks"}, "N", string_format("max number of chunks to process (default: %d, -1 = all)", params.n_chunks), diff --git a/llama.cpp/common/common.cpp b/llama.cpp/common/common.cpp --- a/llama.cpp/common/common.cpp +++ b/llama.cpp/common/common.cpp @@ -1632,6 +1632,8 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.slot_initial_ctx = params.slot_initial_ctx < 0 ? 0u : (uint32_t) params.slot_initial_ctx; // opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md cparams.slot_shrink_idle_ms = params.slot_shrink_idle_ms < 0 ? 0u : (uint32_t) params.slot_shrink_idle_ms; + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + cparams.headinfer_gpu_heads_frac = params.headinfer_gpu_heads_frac; cparams.n_rs_seq = params.speculative.need_n_rs_seq(); cparams.n_batch = params.n_batch; cparams.n_ubatch = params.n_ubatch; diff --git a/llama.cpp/common/common.h b/llama.cpp/common/common.h --- a/llama.cpp/common/common.h +++ b/llama.cpp/common/common.h @@ -435,6 +435,11 @@ struct common_params { // slot_initial_ctx and the over-cap pages are decommitted via // posix_madvise. 0 disables shrink-on-idle (Phase 2 behavior). int32_t slot_shrink_idle_ms = 0; + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + // Fraction of KV heads kept GPU-resident; the rest are offloaded to host + // memory and streamed back per step. 1.0 = off (no split). Only active for + // GPU-offloaded layers with flash-attention (non-transposed V cache). + float headinfer_gpu_heads_frac = 1.0f; int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) int32_t n_ubatch = 512; // physical batch size for prompt processing (must be >=32 to use BLAS) int32_t n_keep = 0; // number of tokens to keep from initial prompt diff --git a/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu --- a/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5308,8 +5308,16 @@ static bool GGML_CALL ggml_backend_cuda_device_supports_op(ggml_backend_dev_t de return op->type == GGML_TYPE_F32 && (op->src[0]->ne[2]*op->src[0]->ne[3]) <= (1 << 15); case GGML_OP_CONCAT: { - ggml_type src0_type = op->src[0]->type; - return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + // Tighten to F32-only to match ggml_cuda_op_concat's three + // hard asserts at ggml-cuda/concat.cu:158-160. The over-broad + // upstream filter ("any type except I32/I16") trusts the + // scheduler into dispatching F16 concats to CUDA, where the + // kernel asserts; with this fix the scheduler correctly routes + // F16 concat to the CPU backend, which has a real F16 path + // (ggml-cpu/ops.cpp:1980 concat_f16). Required for M2 get_k / + // get_v reassembly on F16 KV caches. + return op->src[0]->type == GGML_TYPE_F32; } break; case GGML_OP_CONV_TRANSPOSE_1D: { diff --git a/llama.cpp/include/llama.h b/llama.cpp/include/llama.h --- a/llama.cpp/include/llama.h +++ b/llama.cpp/include/llama.h @@ -353,6 +353,11 @@ extern "C" { // (Phase 2 behavior — soft cap grows monotonically). Requires // slot_initial_ctx > 0 to have any visible effect. uint32_t slot_shrink_idle_ms; + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + // Fraction of KV heads kept GPU-resident (rest offloaded to host and + // streamed back per step). 1.0 = off. Only effective for offloaded + // layers with flash-attention (non-transposed V cache). + float headinfer_gpu_heads_frac; uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing diff --git a/llama.cpp/src/llama-context.cpp b/llama.cpp/src/llama-context.cpp --- a/llama.cpp/src/llama-context.cpp +++ b/llama.cpp/src/llama-context.cpp @@ -62,6 +62,8 @@ llama_context::llama_context( cparams.slot_initial_ctx = params.slot_initial_ctx; // opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md cparams.slot_shrink_idle_ms = params.slot_shrink_idle_ms; + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + cparams.headinfer_gpu_heads_frac = params.headinfer_gpu_heads_frac; cparams.n_threads = params.n_threads; cparams.n_threads_batch = params.n_threads_batch; @@ -3353,6 +3355,7 @@ llama_context_params llama_context_default_params() { /*.n_seq_max =*/ 1, /*.slot_initial_ctx =*/ 0, /*.slot_shrink_idle_ms =*/ 0, + /*.headinfer_gpu_heads_frac =*/ 1.0f, /*.n_rs_seq =*/ 0, /*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default /*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS, diff --git a/llama.cpp/src/llama-cparams.h b/llama.cpp/src/llama-cparams.h --- a/llama.cpp/src/llama-cparams.h +++ b/llama.cpp/src/llama-cparams.h @@ -21,6 +21,8 @@ struct llama_cparams { // milliseconds: when a grown KV-cache is idle for at least this long, the // soft cap is reset to slot_initial_ctx and the pages beyond are decommitted. uint32_t slot_shrink_idle_ms; + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + float headinfer_gpu_heads_frac; uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing diff --git a/llama.cpp/src/llama-kv-cache-iswa.cpp b/llama.cpp/src/llama-kv-cache-iswa.cpp --- a/llama.cpp/src/llama-kv-cache-iswa.cpp +++ b/llama.cpp/src/llama-kv-cache-iswa.cpp @@ -22,6 +22,7 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( uint32_t kv_size, uint32_t kv_size_initial, uint32_t slot_shrink_idle_ms, + float headinfer_gpu_heads_frac, uint32_t n_seq_max, uint32_t n_ubatch, uint32_t n_pad, @@ -64,6 +65,7 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( kv_base = std::make_unique( model, type_k, type_v, v_trans, offload, unified, size_base, kv_size_initial, slot_shrink_idle_ms, + headinfer_gpu_heads_frac, n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, filter_base, reuse); LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa); @@ -71,6 +73,7 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( kv_swa = std::make_unique( model, type_k, type_v, v_trans, offload, unified, size_swa, kv_size_initial, slot_shrink_idle_ms, + headinfer_gpu_heads_frac, n_seq_max, n_pad, hparams.n_swa, hparams.swa_type, filter_swa, reuse); } diff --git a/llama.cpp/src/llama-kv-cache-iswa.h b/llama.cpp/src/llama-kv-cache-iswa.h --- a/llama.cpp/src/llama-kv-cache-iswa.h +++ b/llama.cpp/src/llama-kv-cache-iswa.h @@ -28,6 +28,8 @@ public: // opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md // Forwarded to both kv_base and kv_swa. 0 = shrink disabled. uint32_t slot_shrink_idle_ms, + // opencoti F5 M2 headinfer — forwarded to both kv_base and kv_swa. + float headinfer_gpu_heads_frac, uint32_t n_seq_max, uint32_t n_ubatch, uint32_t n_pad, 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 @@ -110,6 +110,7 @@ llama_kv_cache::llama_kv_cache( uint32_t kv_size, uint32_t kv_size_initial, uint32_t slot_shrink_idle_ms, + float headinfer_gpu_heads_frac, uint32_t n_seq_max, uint32_t n_pad, uint32_t n_swa, @@ -262,23 +263,67 @@ llama_kv_cache::llama_kv_cache( const bool has_k = true; const bool has_v = !is_mla; - ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr; - ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr; + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + // 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). + 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) { + 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); + if (n_head_kv > 1) { + uint32_t g = (uint32_t) std::lround((double) headinfer_gpu_heads_frac * (double) n_head_kv); + g = std::max(1u, std::min(g, n_head_kv - 1)); + gpu_heads = g; + n_embd_k_gpu = g * n_embd_head_k; + n_embd_v_gpu = g * n_embd_head_v; + n_embd_k_cpu = (n_head_kv - g) * n_embd_head_k; + n_embd_v_cpu = (n_head_kv - g) * n_embd_head_v; + LLAMA_LOG_INFO("%s: layer %3d: headinfer split heads %u GPU / %u CPU\n", + __func__, il, gpu_heads, n_head_kv - gpu_heads); + } + } + + 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); + } + 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_gqa, kv_size, k->nb[1], s*k->nb[2]) : nullptr); - v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr); + 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); } map_layer_ids[il] = layers.size(); - layers.push_back({ il, k, v, k_stream, v_stream, }); + layers.push_back({ il, k, v, k_stream, v_stream, k_cpu, v_cpu, gpu_heads }); } if (reuse) { @@ -440,6 +485,20 @@ void llama_kv_cache::ensure_cleared(uint32_t up_to_cells) { 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); + } + 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); + } } } @@ -610,6 +669,11 @@ void llama_kv_cache::shrink_if_idle() { 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); } // After decommit, treat cells [from_cells, to_cells) as "uncleared" — @@ -1645,15 +1709,41 @@ uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { const int32_t ikv = map_layer_ids.at(il); - auto * k = layers[ikv].k; + 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 + // dimension. The backend scheduler streams k_cpu to the compute backend. + if (layer.gpu_heads > 0 && layer.k_cpu) { + 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 kv_size = get_size(); const uint64_t n_embd_k_gqa = k->ne[0]; assert(n_embd_k_gqa == hparams.n_embd_k_gqa(il)); - const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; - 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)), @@ -1665,16 +1755,42 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k 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); - auto * v = layers[ikv].v; + 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) { + 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 kv_size = get_size(); const uint64_t n_embd_v_gqa = v->ne[0]; // [TAG_V_CACHE_VARIABLE] assert(n_embd_v_gqa >= hparams.n_embd_v_gqa(il)); - const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; - if (!v_trans) { // note: v->nb[1] <= v->nb[2] return ggml_view_4d(ctx, v, @@ -1699,7 +1815,8 @@ ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggm const int32_t ikv = map_layer_ids.at(il); - ggml_tensor * k = layers[ikv].k; + 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]; @@ -1711,6 +1828,24 @@ ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggm // TODO: add ggml helper function for this? GGML_ASSERT(ggml_row_size(k_cur->type, n_embd_head) == k_cur->nb[1]); + // 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) { + 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)); + } + 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]; @@ -1734,7 +1869,8 @@ ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggm const int32_t ikv = map_layer_ids.at(il); - auto * v = layers[ikv].v; + 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]; @@ -1745,6 +1881,20 @@ ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggm // we can merge dims 0 and 1 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) { + 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 int64_t n_stream = v->ne[2]; // take this branch when FA is enabled (the V cache is not transposed) 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 @@ -111,6 +111,11 @@ public: // back to kv_size_initial and the over-cap pages are // decommitted. 0 disables shrink-on-idle. uint32_t slot_shrink_idle_ms, + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + // Fraction of KV heads kept GPU-resident (rest offloaded + // to host, streamed back per step). 1.0 = off. Effective + // only for offloaded layers with !v_trans (flash-attn). + float headinfer_gpu_heads_frac, uint32_t n_seq_max, uint32_t n_pad, uint32_t n_swa, @@ -269,6 +274,14 @@ private: 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; + uint32_t gpu_heads = 0; }; bool v_trans = true; // the value tensor is transposed diff --git a/llama.cpp/src/llama-memory-hybrid-iswa.cpp b/llama.cpp/src/llama-memory-hybrid-iswa.cpp --- a/llama.cpp/src/llama-memory-hybrid-iswa.cpp +++ b/llama.cpp/src/llama-memory-hybrid-iswa.cpp @@ -48,6 +48,8 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( // Hybrid-iswa does not yet thread slot_shrink_idle_ms; pass 0 for // Phase 2 backward-compat (no shrink-on-idle). 0, + // opencoti F5 M2 headinfer — hybrid-iswa does not thread headinfer; 1.0 = off. + 1.0f, n_seq_max, n_ubatch, n_pad, diff --git a/llama.cpp/src/llama-memory-hybrid.cpp b/llama.cpp/src/llama-memory-hybrid.cpp --- a/llama.cpp/src/llama-memory-hybrid.cpp +++ b/llama.cpp/src/llama-memory-hybrid.cpp @@ -47,6 +47,8 @@ llama_memory_hybrid::llama_memory_hybrid( // Hybrid memory does not yet thread slot_shrink_idle_ms; pass 0 for // Phase 2 backward-compat (no shrink-on-idle). 0, + // opencoti F5 M2 headinfer — hybrid memory does not thread headinfer; 1.0 = off. + 1.0f, n_seq_max, n_pad, n_swa, diff --git a/llama.cpp/src/llama-model.cpp b/llama.cpp/src/llama-model.cpp --- a/llama.cpp/src/llama-model.cpp +++ b/llama.cpp/src/llama-model.cpp @@ -2077,6 +2077,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.slot_initial_ctx, // opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md cparams.slot_shrink_idle_ms, + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + cparams.headinfer_gpu_heads_frac, cparams.n_seq_max, cparams.n_ubatch, 1, @@ -2097,6 +2099,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.slot_initial_ctx, // opencoti F4 M3 Phase 3 — see docs/decisions/0001-lazy-slot-context.md cparams.slot_shrink_idle_ms, + // opencoti F5 M2 headinfer — see docs/features/advanced_kv.md + cparams.headinfer_gpu_heads_frac, cparams.n_seq_max, 1, hparams.n_swa,