Repository navigation
llama/cuda: tiered KV cache (--kv-vram-cells): VRAM head, pinned-host tail in one VMM range - #319
professorpalmer wants to merge 2 commits into
Conversation
…ff; drop --spec-draft-n-max-tail Review of PrismML-Eng#320: - The window cutoff was the minimum position of the whole batch and was applied to every slot, so a deep slot could remove the rows of a shallow slot that was not in the batch, and a shallow slot in the batch left extra rows in a deep one. Now each sequence in the batch gets its own cutoff, and sequences that are not in the batch keep their rows. - --spec-draft-n-max-tail needed n_kv_vram_cells from the tiered-KV change (PrismML-Eng#319), so this branch did not compile on its own. The tail flag, its per-slot fields and the n_max widening at load are removed; the tail returns stacked on the tiered-KV change, with the requested draft limits kept apart from the allocation size (the widening was also recaptured as the normal limit after sleep/wake). - --spec-draft-window help text: one sentence, no forced line break.
|
Retested on Linux (H200, driver 570): with --kv-vram-cells set, it still fails at startup with cuMemCreate(host): invalid argument, and the context then fails to load. The Linux driver doesn't accept CU_MEM_LOCATION_TYPE_HOST for cuMemCreate, but CU_MEM_LOCATION_TYPE_HOST_NUMA with id set to the GPU's CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID does. I tried that one change locally and it works: output is token-identical to an all-VRAM cache at 512 cells (staging on and off) and at a 29k-token prompt with 8192 cells in VRAM. Could you (1) use HOST_NUMA on non-Windows builds, and (2) fall back to the plain device buffer with a warning when the tier allocation fails, instead of failing context creation? Smaller points: please append n_kv_vram_cells at the end of llama_context_params rather than in the middle (ABI), and mention the cost in the help text. On an H200 with 72% of the cache in host memory, decode at 29k drops from about 250 to 60 t/s, which is the right trade on a 12 GB card but not on one where the cache fits. Also, have you run the staging path with CUDA graphs enabled? |
… tail in one VMM range --kv-vram-cells N (llama_context_params.n_kv_vram_cells) keeps cells [0, N) of each layer's K/V in device memory and the rest in pinned host memory, mapped into the same device virtual range with CUDA VMM (cuMemCreate with CU_MEM_LOCATION_TYPE_HOST), so a context can be larger than VRAM. Attention kernels are unchanged: they only touch host pages once a sequence is that deep. Output is bit-identical to an all-device cache. When an attention op's K/V range reaches the host tail, the used host rows are first copied into a VRAM staging buffer with the copy engine (a second VMM alias maps the same VRAM head pages plus the staging buffer), and the op reads the all-VRAM alias. Prefill reads the tail once per op instead of once per query tile; decode moves it by DMA instead of SMs reading host memory from inside the kernel. GGML_CUDA_KV_TIER_STAGING=0 disables staging. Bonsai 2 27B on an RTX 4070 12 GB: q8_0 K/V at the full 262,144-token window (9.1 GB of KV), first ~94k positions in VRAM. Staging: prefill in the tail 146 -> 280 tok/s at 180k, decode +28% at 180k. Plumbed through llama_cparams to the standard KV cache and the hybrid (attention + recurrent) memory via a trailing defaulted constructor parameter; other memory types are untouched.
… Windows (the Linux driver rejects CU_MEM_LOCATION_TYPE_HOST for cuMemCreate); fall back to the plain device buffer with a warning when the tiered buffer cannot be allocated; n_kv_vram_cells at the end of llama_context_params (ABI); the cost in the help text
…ff; drop --spec-draft-n-max-tail Review of PrismML-Eng#320: - The window cutoff was the minimum position of the whole batch and was applied to every slot, so a deep slot could remove the rows of a shallow slot that was not in the batch, and a shallow slot in the batch left extra rows in a deep one. Now each sequence in the batch gets its own cutoff, and sequences that are not in the batch keep their rows. - --spec-draft-n-max-tail needed n_kv_vram_cells from the tiered-KV change (PrismML-Eng#319), so this branch did not compile on its own. The tail flag, its per-slot fields and the n_max widening at load are removed; the tail returns stacked on the tiered-KV change, with the requested draft limits kept apart from the allocation size (the widening was also recaptured as the normal limit after sleep/wake). - --spec-draft-window help text: one sentence, no forced line break.
6fa6f2b to
b56e067
Compare
|
Thank you for the Linux test and the fix. All points are in
CUDA graphs: yes. The staging copy is captured with the graph and replayed at each step, so it copies the current host rows. RTX 4070, Bonsai 2 27B PTQ1_0,
0 CUDA errors. With prism's head-72 padding path, the conflict in I could not test the HOST_NUMA path here (Windows only). Could you retest it on the H200 with this head? |
One of the five pieces of #285, split per bri-prism's request on #317 so each can land on its own. Single commit on current
prism(eaecb50c7), applies without conflicts.What it does
llama_context_params.n_kv_vram_cells/--kv-vram-cells Nkeeps cells[0, N)of each layer's K/V in device memory and the rest in pinned host memory mapped into the same device range with CUDA VMM (cuMemCreatewithCU_MEM_LOCATION_TYPE_HOST);ggml_backend_cuda_tier_buffer_typeis exported through the registry. Kernels are unchanged and output is bit-identical to an all-device cache. When an attention op reaches the host tail, the used rows are copied into a VRAM staging buffer by the copy engine (a second VMM alias maps the same VRAM head plus the staging buffer) and the op reads the alias: prefill reads the tail once per op instead of once per query tile (146 -> 280 tok/s at 180k on the 4070), decode moves it by DMA (+28% at 180k). Plumbed throughllama_cparamsinto the standard KV cache and the hybrid memory with a trailing defaulted constructor parameter; other memory types untouched. Works on Windows/WDDM.Why: q8_0 at 262k is 9.1 GB of K/V for Bonsai 2 27B and does not fit next to the weights on 12 GB; #221 reached 262k with q4_0 (mean KLD 0.00218 vs f16 K/V, 97.9% top-1 agreement) where q8_0 measures 0.00017 / 99.4%.
Receipts (RTX 4070 12 GB, sm_89, PCIe 4.0 x16, GDDR6X +1500 MHz as in #221;
llama-server, one slot, greedy, MTP head, genuine prefill)The decode numbers above 32k include the MTP drafting-window commit (separate PR) and the FA GQA-decode commit that is being folded into #307; the tiered cache itself changes no arithmetic.
--kv-vram-cells 2048vs all-VRAM, q8_0, greedy, at 4096 / 20480 / 40960 tokens: 6/6 identical.test-backend-ops -b CUDA0on the stacked head and again on the rebased head:FLASH_ATTN_EXT2994/2994,MUL_MAT1516/1516,GATED_DELTA_NET39/39,GET_ROWS119/119 (MUL_MAT_ID1034/1035: the intermittentm=70, n=1case already inprism, fixed in tests: MUL_MAT_ID cases for the one-column mat-vec store (small K, several expert rows per block) #295).Serving recipe and all receipts: https://github.com/professorpalmer/bonsai-ada-surgery/blob/main/docs/Q8_FULL_CONTEXT.md