Tiered kv cache implementation along with vattention (mad-lab-cache)

[kmbandy]

AMD ROCm Tiered KV Cache + NVMe Demand Weight Paging

Overview

This PR implements and stabilizes two major infrastructure features for running large models on AMD ROCm hardware with limited VRAM:

1. 3-tier KV cache — hot (VRAM), warm (secondary GPU VRAM or RAM), cold (SSD NVMe)
2. NVMe demand weight paging — 8-slot VRAM pool with on-demand weight loading for models that exceed VRAM

Both features are WIP infrastructure. Output correctness for weight paging is a known follow-up item.

---

Tiered KV Cache Fixes

Warm GPU tier OOM fix
When --kv-warm-device N is set, the pool previously allocated both RAM staging buffers and GPU VRAM buffers for all warm slots. On a 262144-slot warm tier this caused ~8GB of wasted RAM and an OOM kill. RAM staging buffers are now skipped when a GPU warm device is configured.

k_bytes/v_bytes quantization fix
Per-token KV byte size was computed as ne[0] * ggml_element_size(tensor). For block-quantized types (turbo4, Q4_K, etc.) ggml_element_size returns the block size in bytes rather than bytes-per-element, producing a 128× overestimate. Switched to tensor->nb[1] (actual row stride), which is correct for all types.

Hybrid model KV wiring
init_slot cast llama_get_memory() to llama_kv_cache *. Hybrid-architecture models (Qwen3-27B and similar) return llama_memory_hybrid or llama_memory_hybrid_iswa, causing silent metadata-only fallback. Added casts for both hybrid types so attention KV layers wire correctly into the tiered cache.

Tier init logging
Added startup log lines for warm GPU (ROCm%d KV warm buffer size = X MiB) and cold SSD (cold tier: N slots reserved, path: ...) to match the existing hot tier output.

---

Weight Pager Fixes

• Pool allocation via ggml buffer: replaced raw hipMalloc with ggml_backend_buft_alloc_buffer; sets tensor->buffer = pool.ggml_buf on page-in so ggml_cuda_mul_mat doesn't crash on src->buffer->buft dereference
• Async prefetch disabled by default: io_uring page-index tracking in complete_prefetch is broken; disabled until fixed
• fd dup moved outside prefetch gate: page_in uses pread regardless of prefetch state; fds must always be populated
• O_DIRECT cleared on dup fds: GGUF tensor offsets are not sector-aligned; O_DIRECT on dup fds caused silent misreads
• Padding zeroed after hipMemcpy: zeroes bytes past page.size up to slot_size so quantized kernel reads past tensor data are safe
• hipGraph disabled: GGML_CUDA_DISABLE_GRAPHS=1 via setenv; removed static from disable_cuda_graphs_due_to_env in common.cuh

---

Known Issues / Follow-ups

• Weight pager output correctness: gallocr allocates paged tensors (buffer=nullptr) in scratchpad, sets tensor->extra; callback overrides data and buffer but not extra; GPU kernel reads stale extra → coherent but wrong tokens. Fix: pre-seat paged tensors into pool.ggml_buf before ggml_backend_sched_alloc_graph.
• Async prefetch: re-enable after fixing page-index tracking in complete_prefetch / submit_prefetch
• Tiered KV eviction under load: watch for hot→warm→cold pressure as context fills up under concurrent agent workloads

---

Hardware / Test Config

• AMD Radeon AI PRO R9700 — gfx1201, 32GB VRAM (primary inference + hot KV tier)
• AMD RX 6900XT — 16GB VRAM (warm KV tier via --kv-warm-device 1)
• NVMe SSD (cold KV tier, weight paging source)

Models Tested

Model	Feature
Qwen3.6-27B-Q6_K	Tiered KV cache (hot/warm/cold), turbo4 KV quant
Qwen3.6-35B-A3B	Tiered KV cache
Qwen3.5-REAP-97B	Tiered KV cache
MiniMax-M2.7-229B-IQ3_S	NVMe demand weight paging (8-slot VRAM pool, 0.03 t/s)

Note: As context fills up under concurrent agent workloads, keep an eye on hot→warm→cold eviction pressure. The tiered eviction pipeline is functional but has not been stress-tested at full 512K context utilization.

---

Co-authored with Claude Sonnet 4.6