Pull requests / #2201

#2201 --vram-elastic with --kv-grow: VRAM takes its chunks where the K/V grows into

open · @i-Vladimir-i · 0 commentaires · Sur GitHub

BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows

Description

**`--vram-elastic` with `--kv-grow`: VRAM takes its chunks where the K/V grows into**

Today the two switches exclude each other: with `--vram-elastic` the engine says `--kv-grow is off` and puts the whole
K/V in VRAM, because #533's cache is cut into segments and `--kv-grow` needs one VMM range of chunks. On our R9700 that
cost 1,413 cache slots for a 262K context (9,405 instead of 10,818) - more than the VRAM a game or an image encoder would
borrow. This lets both run.

**What it changes** (`src/program/generate.cpp`; docs in `DETAILS.md` and `VRAM_ELASTIC.md`)

- With both switches the cache keeps `--kv-grow`'s arena (one range of 2 MiB chunks), and `VRAM` works on the region
  the K/V grows into: slots `[lo, top)`, just below the prompt path's loan.
- **Shrink:** the same eviction the K/V's growth uses (hotter experts there move into the coldest loan-region slots
  first), now `kvg_take`, split out of `kvg_ensure` unchanged; the chunks go to the driver instead of to the K/V.
- **Grow:** new chunks mapped there and those slots refilled with the profile's hottest experts the GPU does not hold,
  now `kvg_give_back`, split out of `kvg_trim` unchanged. So after a grow the cache holds what a K/V trim would leave,
  not necessarily the experts the shrink took (the segments put the same ones back); the docs say so.
- The prompt path's loan and chunk are not touched. The K/V keeps growing while the cache is shrunk, from the cache's
  slots, not from the reserve. A small counter keeps the chunks given to the driver; it is settled against the arena's
  empty places, because the K/V's first cells are new VRAM and a trim below them maps those chunks into the arena too
  (without that the count drifted by ~100 MiB and a grow reported "no VRAM" for chunks that had no place).

**What it leaves alone:** either switch alone. `kvg_ensure` and `kvg_trim` do the same steps through the two helpers;
without `--kv-grow`, `VRAM` keeps the segments. The default path does not reach any of it.

**Checked** (on top of #2189 for HIP; the change itself does not depend on it)

- HIP build (gfx1201, ROCm 7.14.1) of this plus #2189: ctest 108/109, the one failure `ple_parity` as before; CUDA:
  this branch alone builds and links with nvcc 13.4.92, sm_120, 0 errors. Not run on an NVIDIA GPU.
- Live, R9700 32 GB, Unsloth UD-IQ4_XS, every expert in RAM, `--kv-grow --vram-elastic`, 262K context, a needle at 3/5
  of each prompt and a 320-token answer:

| step | VRAM call | slots holding experts (of 10,865) | free | needle / decode |
|---|---|---|---|---|
| 78K prompt, whole cache | - | 10,492 (the K/V took 373) | - | ok / 96.6 tok/s |
| `reserve 4000` | 106 ms | 8,872 | 4,000 MiB | - |
| 78K prompt (the K/V grows again while shrunk) | - | - | - | ok / 92.0 |
| `reserve 8000` | 113 ms | 7,094 | 8,000 MiB | - |
| 156K prompt (the K/V to 2.13 GiB) | - | - | - | ok / 87.8 |
| `reserve 0` | 418 ms | 10,865 | 453 MiB | - |
| 78K prompt, then a short one (the K/V trims) | - | - | - | ok / 96.1 |
| `reserve 4000`, `reserve 0` | 101 ms, 199 ms | 9,281, 10,865 | 4,001, 453 MiB | - |

  The driver's own counter followed (31.5 -> 30.5 -> 28.0 GiB at 4000, 24.1 GiB at 8000 with the 156K K/V, back to
  31.4); the engine's "MiB with the driver" went back to 0 at each full grow.
- `--kv-grow` alone with this build: 78K, a short request (trim), 156K, a short one, 78K - every needle found, the K/V
  grew and trimmed as before.
- Together with the elastic image encoder (#2196) on the same card: the first picture made `VRAM` keep 2 GiB free
  (52 ms) and the encoder started in 3.1 s; a 156K prompt then grew the K/V to 163,840 cells while the encoder held its
  share (needle found); after 60 s without a picture the encoder stopped and the cache was whole again (10,888 of
  10,888 slots, 0 MiB with the driver); a second cycle the same.

**Not tested:** NVIDIA at run time, Windows, more than one shrink-grow cycle with a parked conversation held by the K/V.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Sur le site

Liens install, modèles, releases.