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