Issues / #3351
#3351 [Engine] one shared KV pool for all lanes (--kv-pool-tokens): is it wanted, and in what shape? (scope agreement before I reopen #1614)
open · @noon-at-cgn · 0 comments · View on GitHub
BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsSecurityWindows
Description
### Area [Engine] ### What do you want to do? **User problem.** With KV streaming (`--kv-resident`) the main session and every `--batch` slot pin the K/V of the whole `--max-context` in RAM, however short the conversation: 3.09 GiB each at 262,144 cells with `--kv int8`. A conversation that moves between the main session and a slot is copied, and at 256k that takes 3.5-5.2 s. On a RAM-limited box with two long lanes the extra pinned RAM also brings the host close to its memory guard. Workaround today: a smaller `--max-context` or `--parallel`, or K/V in VRAM (fewer cached experts). **Proposed behaviour.** `--kv-pool-tokens N`, opt-in (absent means nothing changes). N cells per layer are pinned once, in 4,096-cell chunks, and a session holds the chunks its conversation has reached. A conversation moves between the main session and a slot by swapping chunk tables. When the pool is full an idle slot's cached conversation gives way, then a new request gets a retryable 503 (`kv_pool_full`), and a decoding lane that cannot grow ends with `length` and `"truncated": true`. It works on one GPU and across a layer split. The idea and the first four commits are xdbxdbx's #1011 (his authorship is kept); #1614 carried it on v0.1.44.1 and was closed by the history rewrite. It is now rebased on 0.2.0-beta (`e88ee39d`), branch `pr/kv-pool-144` on noon-at-cgn/Strata, 19 commits. **Tradeoffs.** - Memory and long-context lane handling only. I make no speed claim; short-context speed did not change. - N is chosen by the operator. Too small a pool means refused requests or shortened answers. - A new layer for K/V host addressing: every reader and writer of the host copy has to go through the chunk table. On this rebase two sites that were added to main after the old base (`conversation_kv_copy`, `conversation_kv_restore_prefix`) did not translate and are fixed, with a test that runs them through a shared pool. Later changes to those sites need the same care. - The MTP drafter ring is not pooled. CUDA only was built. **Measured on 0.2.0-beta** (engine 0.1.45.3, 2 x RTX 3080 20 GB, same binary with and without the flag, two lanes of about 256k tokens, appended turns): - Slot move 0.36-0.5 s with the pool against 3.5-5.2 s copied; first token of an appended short question 1.1-1.2 s against 2.3-6.7 s; both lanes decoding an appended turn 59-62 against 25-38 tok/s. - 3.2 GiB more idle `MemAvailable` (14.5 against 11.4 GiB, 3 starts each), 3.06 GiB less pinned shared memory. - No short-context speed effect (differences within the round-to-round spread). - Without the pool, the host memory guard stopped the engine twice in the 245k run (two of three no-pool starts), so the no-pool 256k timings are one full run and two partial ones. It never fired with the pool. - With the pool a third lane (three conversations of about 104k) fits at about 13 GiB `MemAvailable` idle. **Not measured:** three lanes without the pool; pool-full paths; more than 3 lanes; HIP, SYCL, Metal, Windows; `STRATA_SLOT_CKPT_MOVE=1` and `STRATA_SLOT_COPY_DIRECT` end to end; the slot-choice layer. One box only. The full list will be in the PR. **Overlap.** - #2526 (park an idle slot before reuse, a slot gives its checkpoints back with its conversation): the same idea and the same admission code as our checkpoint-chain commit. Ours exists only with the pool (the chain follows the chunks). I will rebase on whichever lands first. - #1962 (share identical KV prefix blocks): a different mechanism and complementary; the pool does not share prefixes. - #3251 (restore a session file without its K/V in RAM): needs the same chunk translation in `conversation_snapshot.cpp`. - #2518 (turn-start rule in `pick_slot`): our separate slot-choice change (`pr/pick-slot`) is not part of the pool PR; main's rule runs first there and ours is only a fallback. It would be its own PR, if you want it. ### What I am asking 1. Is an opt-in memory feature like this in scope? If not, it stays a branch. 2. If yes, in what shape: the whole branch as one PR (pool core plus the checkpoint-chain commit, which only exists with the pool), or only the pool core and the chain commit later (for example after #2526)? 3. Any evidence you want beyond the above. I am reopening #1614 as a draft with the rebased branch now, so the code and the evidence are in one place and nothing competes with it. If you want a different shape, I will change the PR rather than open another.
Related on this site
Quick links to related pages.