kv-cache : resolve the partial KV residency set once for the model

[Piggidragon]

Overview

--kv-gpu-layers keeps part of a host-resident attention cache on the devices. It was resolved
inside each cache, which went wrong three ways:

• Every cache made of several sub-caches ignored it. Only llama_kv_cache and
  llama_memory_hybrid override get_supports_partial_kv(); iSWA, DSA, DSV4 and MSA fall through to
  the base false, and create_memory handed them a specialized_placement with
  gpu_resident_layers = 0. Each sub-cache would otherwise have taken the full budget, so the option
  was disabled for all of them rather than divided.
• The layers were taken in layer order, which fills the device owning the first layers and leaves
  the free memory of the others unused.

Resolve one set of layer indices for the whole model instead, in create_memory, from the layers the
requesting context actually owns, and take them one per owning device in turn. get_supports_partial_kv
and the specialized_placement copy both go away.

An MTP context owns the nextn layers, which sit above hparams.n_layer(); the selection uses
n_layer_all and the context type so those still resolve.

Two filters exist only because the budget is now resolved once for the whole model rather than per
cache, where each cache applied its own. Neither fixes a defect on llama/dev; both prevent one
here. An MTP context owns the nextn layers, which sit above hparams.n_layer(), so the selection
runs over n_layer_all and filters by context type. A recurrent layer keeps its state outside the
attention cache, so it is skipped rather than counted against the budget.

Split out of #57. Independent of the host-cache correctness fix (#66).

Testing

Two GPUs, RTX 4070 (gen4 x16) + RTX 3060 (gen3 x4), CUDA, NCCL, stock clocks. A 15216-token prompt
from this repository's docs, -c 20480 -n 64 -ngl 99 -nkvo --kv-cpu-pinned -sm layer.

An iSWA cache could not use the option at all. gemma-4-26B-A4B:

--kv-gpu-layers	before	after
0	7.83 t/s	7.84 t/s
8	7.83 t/s (ignoring kv_gpu_layers)	10.68 t/s

Before, both sub-caches stay entirely in host memory (CUDA_Host KV buffer size = 400 MiB and
300 MiB, no device KV) and the context logs
partial GPU KV residency is not supported for this memory layout. After, the budget is divided
across the sub-caches and the devices. +36% generation, +14% prefill (1513 -> 1723 t/s).

The layers all landed on one device. Qwen3.8-27B-UD-Q5_K_M, --kv-gpu-layers 4:

	before	after
placement	CUDA0 KV 320 MiB, no CUDA1 KV	CUDA0 KV 160 MiB, CUDA1 KV 160 MiB
tg	4.56 t/s	5.37 t/s

Both runs allocate identical model and compute buffers and produce the same graph split count, so the
difference is only which device gives up which layers. Against the same model with
--kv-gpu-layers 0 (4.36 t/s), taking the layers in order is worth +4.6% and spreading them +23%.

tests/test-llama-archs.cpp gains test_mtp_kv_residency, which builds an MTP context with
--kv-gpu-layers 0 and 1 and asserts the host-resident context memory actually shrinks. It fails
if the budget resolves against the main context's layer range, which is disjoint from the MTP cache's.

test-llama-archs -s 1 at 1, 2, 3 and 4 virtual CUDA devices: passes. Built with
-DLLAMA_FATAL_WARNINGS=ON.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES - implemented by an agent on my instruction, see the Assisted-by: commit trailer.

[GenerelSchwerz]

Automated preliminary review by Codex; the repository owner plans a separate manual review.

Verdict: FAIL.

Blocking

• src/llama-model.cpp:2250-2252 applies a model-global recurrent-layer exclusion instead of each cache's ownership filter. Falcon H1's attention cache accepts every layer (src/llama-model.cpp:2560-2562), but Falcon H1 marks every layer recurrent, so a nonzero request can produce no resident attention KV. Nemotron H's attention cache accepts only non-recurrent, zero-FFN layers (src/llama-model.cpp:2563-2569), but the picker can spend the budget on non-recurrent FFN-only layers that the cache later discards, with no refill.
• src/llama-model.cpp:2292-2302 reports the picked set as device-resident and mutates cparams.kv_gpu_layers before memory construction. Null-memory architectures (src/llama-model.cpp:2305-2324) and filtered layouts can therefore over-report realized residency; src/llama-context.cpp:438-440 can then enable attention-compute offload when no attention KV actually became device-resident.

Will slow review

• The test is effectively single-device and covers only Qwen3.5 MTP (tests/test-llama-archs.cpp:1047-1067); it skips the relevant path when CPU CI has no GPU (tests/test-llama-archs.cpp:1056-1059). It does not cover multi-device selection, Falcon H1, Nemotron H ownership filtering, or a no-memory layout.
• The model-global picker duplicates cache-specific ownership policy (src/llama-model.cpp:2221-2281 versus src/llama-model.cpp:2560-2569), and create_memory now mutates cparams (src/llama-model.h:753-754), diverging from pristine upstream's immutable memory-creation contract.

Developmental progress

• Precursor scope, MTP scanning, and warning handling improved: selection now scans n_layer_all with context ownership (src/llama-model.cpp:2245-2249) and emits explicit empty/partial messages (src/llama-model.cpp:2295-2299). Those fixes deserve credit, but actual cache ownership and realized-state defects remain.

Reviewed head: e08ca85

[GenerelSchwerz]

Update for reviewed head 3f02d74765179b6313f80d2c42d141694a3215f8. Review, fixes and validation by Codex at the repository owner's request.

The blocking findings from the earlier review of e08ca85 are resolved: selection uses the cache ownership filters, and the realized count is reconciled after successful memory construction. The latest update also fixes router/all-nextn ownership and scopes the tensor/iSWA guard to KV layers owned by the requesting context. Documentation and existing tests now cover the resulting placement contract. No remaining PR-specific code blocker was found in the configurations tested.

Validation on the published source:

• CUDA Release build with fatal warnings enabled and -j20 passed.
• Residency tests passed with 1-4 logical CUDA devices: 98 placement observations and 46 matched-logit comparisons, all recorded NMSE values zero. CPU fallback and the phase/live-context workspace suites also passed.
• Thirty specialized-cache runs passed across Gemma4, DeepSeek32, MiniMax-M3, DeepSeek4 and Qwen3.5, including DeepSeek4 beyond its 128-token compression boundary. Prefill and decode logits matched the host baseline.
• Compute Sanitizer reported zero memory errors. The original PR library fails the new all-nextn regression; the updated library passes.

Before merging, finish and assess the fresh Linux, Windows, ARM64, server and WebGPU CI runs. Self-hosted checks are still queued. The previous ARM64 failure was an uninitialized-array warning in test-backend-ops.cpp; the updated branch includes the existing fix f114f91f9, pending confirmation from the fresh ARM64 run.

The CANN workflow failure occurs before any jobs run. Its unchanged workflow has all jobs commented out. This is an existing workflow issue outside this PR; repair or explicitly accept that CI exception before declaring the checks clear.

Coverage limit: local validation used tiny deterministic fixtures on one physical RTX 5070 Ti. The additional CUDA devices were virtual, so physical multi-GPU behavior is still unverified. The known tensor/iSWA mixed-residency combination remains guarded.

PR #68's separate capacity-accounting blocker does not block this PR. #67 remains open and unmerged.