Repository navigation
[Engine] shared KV pool for all lanes (--kv-pool-tokens): at 2 x 256k the first token comes 2-5x sooner and lane moves ~10x faster; off = identical to main - #1614
Conversation
|
@xdbxdbx this is your #1011 rebased onto 0.1.41. GitHub closed #1011 automatically when |
noon-at-cgn
changed the title
--kv-pool-tokens): two 256k-context lanes in one 512k-token pool on two 20 GB cards, KV memory follows the tokens in use--kv-pool-tokens): two 256k lanes in one 512k pool, 3.1 GiB less pinned RAM at 2 lanes, no speed change
Niko1221#1637 and the four older ones), 0.1.41 batch-MTP, adaptive-tier and pool measurements
noon-at-cgn
changed the title
--kv-pool-tokens): two 256k lanes in one 512k pool, 3.1 GiB less pinned RAM at 2 lanes, no speed change--kv-pool-tokens), 3.1 GiB less pinned RAM at 2 lanes
400fc50 to
b65c6a8
Compare
|
Rebased onto current main (v0.1.44.1) and resolved the conflicts with the ported memory guard, batch MTP on a split and the adaptive tier. It is the same 14 commits in the same order; xdbxdbx's four keep their authorship. One extra change inside the existing reuse-image commit: main's new |
|
Result of running the pool end to end on this branch, and a fix I added (three commits on top, head 6ea1d2a). End to end (2 x RTX 3080 20 GB, layer split 23, A problem I found and the fix. An appended turn (same history, the answer, a new message) reused the cache on both lanes, from 30k to 256k tokens. A prompt with the same prefix but an edited last message reused nothing with the pool on: admitting a conversation into a slot dropped main's periodic checkpoint chain and kept only the last-turn checkpoint, which lies after the edit. The fix has two parts, in separate commits:
Costs and limits: each cached slot keeps up to 6 checkpoints, about 0.7 GB of host RAM per slot. An edit that lies below the first periodic checkpoint reuses nothing (30k tokens: 0 on both lanes; 60k: 31,232). Not measured: speed with the fix against without it, three or more lanes, pool-off at 256k. The CPU, GPU parity and serve tests pass; the C++ change was tested on the engine only in the runs above. |
|
A memory feature with no speed claim (the author measured no change); handed to the cache review. |
|
Thanks for the pull request. We're resetting the tracker for 0.2.0-beta. There are more than 1,300 open issues and pull requests, too many for us to review properly, so we're closing everything that isn't being worked on right now. Closed here does not mean rejected. To reopen it takes about two minutes:
The rules are in CONTRIBUTING.md. |
noon-at-cgn
changed the title
--kv-pool-tokens), 3.1 GiB less pinned RAM at 2 lanes
|
/reopen |
|
The check passed, but GitHub would not reopen it (Validation Failed: ). If the branch was deleted, restore it or open a new pull request and link this one. |
|
/reopen |
|
Reopened, thank you. A maintainer will pick it up from here. If it conflicts with main, rebase it first; that speeds up review. |
6ea1d2a to
81a5c6a
Compare
…PoolFull 503) — integ batch 19 [23 hunks via claude-strata; integ's take_reuse_for/retained_n flow supersedes PR reuse site; copy_to_slot 4-arg; NUMA binding ported into qsa_host_alloc; compile-verify pending]
- conversation_state.cpp: drop the integ-era checkpoint_validate copy and the duplicated snapshot_sources wrapper left side-by-side by the Niko1221#3341/Niko1221#3541 merges; restore the checkpoint_payload_sizes helper (Metal-aware) that the dedup slice clipped - conversation_cache.hpp: re-apply Niko1221#3615's min-tokens parser + READ gate, clobbered by a later header-side merge - generate.cpp: single parked Match (the Niko1221#3584 best_with_partial + Niko1221#3615 read gate wins; slot_source=-1 moves after it), carry_chain declared for Niko1221#1614's chain-carry in the integ move_cks block, one splice-aware read_from, 'moved' hoisted to the trace's scope - CMakeLists: the WIN32 ws2_32 link for stage_link.cpp sat inside the HIP-only guard Verified: cmake --build build-cu130 clean, strata.exe links
With KV streaming (--kv-resident) every session pins the K/V of the
whole context for each QSA layer, so N batch slots plus the main session
pin (N + 1) x the context even when the conversations together hold
little of it. --kv-pool-tokens N pins N cells per QSA layer once; a
session holds only the chunks (4096 cells) its conversation has reached.
- A per-session chunk table (device copy for the kernels, host copy for
the DMA movers) maps a session's blocks into the pool; unreserved
entries point at a trash chunk. Without a pool the layout is the
identity, as before.
- Lanes grow as they write. When the pool is full an idle slot gives its
K/V back (one holding nothing first, then the cached conversation used
longest ago); then a new request is refused ("ERR KV pool full", which
the server answers 503 with Retry-After); a decoding lane that cannot
grow ends ("length", with "truncated": true in the response).
- After a slot takes a conversation over, the main session gives its
copy back.
- The engine reports the pool on a POOL line; /metrics shows it as
live.kv_pool.
- The batch slots borrow the main session's RoPE table (64 MiB each at
262K).
Single GPU only: refused with a layer split.
Measured, IQ2_XS, RTX 4090 + 64 GB RAM, --resident-experts, 4 slots x
262144 cells, --kv-pool-tokens 524288: the pool pins 6.6 GB (5 x 3.3 GB
without it); 4 lanes of 512 tokens 83 t/s total; two lanes at 134K and
143K tokens deep at once; available RAM never below 15.1 GB.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Ported from upstream PR Niko1221#1011 by xdbxdbx (commit 66082ec), re-applied onto the current main line.
…ing pool chunks Admission and a slot giving a conversation back exchange the two sessions' chunk tables instead of copying the streamed K/V, so the pool never holds a conversation twice (a 104K-token admission: 52 ms instead of ~1.1 s copied). The prompt read that gives way to a waiting request (BYIELD) still copies. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Ported from upstream PR Niko1221#1011 by xdbxdbx (commit 28ed67b), re-applied onto the current main line.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Ported from upstream PR Niko1221#1011 by xdbxdbx (commit 3c3e736), re-applied onto the current main line.
…3 paths - A drafter streamed whole (its window covers the context) pins its own host copy; without one the conversation cache refused every snapshot. - A chunk-table upload completes before the next launch reads it (the kernels run on non-blocking streams). - pool_reserve gives nothing back when all the idle slots together would not cover the request, and a prompt read that gave way (BYIELD) gives way last instead of never. - The pool's refusal is a 503 with Retry-After on non-streamed /v1/responses too; in a stream it carries code kv_pool_full (Anthropic: overloaded_error). - A solo request that ends with finish "pool" ends there; it is not resent to a batch slot. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Ported from upstream PR Niko1221#1011 by xdbxdbx (commit 2040789), re-applied onto the current main line.
…k table; a session RESTORE reserves its chunks Three host-addressing sites that did not exist (or did not stream) when Niko1221#1011 was written: - kv_hybrid_k_half/kv_hybrid_v_half build a fresh KvHostPools from the arrays only, which dropped the chunk table: a pooled --kv k8v4 layer wrote its cells at their session block numbers, into other sessions' chunks. - conversation_kv_source (the SAVE <path> session file) read the authoritative pools linearly; it walks the chunks now (each_piece, as conversation_kv_save/restore/verify do). - RESTORE <path> restored the file's K/V into the main session before any chunk of it was reserved (everything past the reservation went to the trash chunk): it reserves the file's cells first and refuses, retryable, when the lanes hold the pool.
configuration the pool is for here, so a lane (the main session or a batch slot) now spans the stages: - KvChunkMap (kv_chunks.hpp/.cpp): the chunk bookkeeping on its own (free list, per-lane tables, reserve/shrink/swap), no CUDA, so it has a CPU test (tests/core/kv_chunks_test.cpp). - KvPool: init takes every stage's reference session (each stage pins its own QSA layers' arrays, in its GPU's context); attach opens a lane with its stage-0 session, attach_stage adds the session of a later stage with a copy of the same table in that GPU's memory; detach closes a lane before its sessions go. Reserve/shrink/release/swap act on the lane, named by any of its sessions, and upload the table to every stage. - generate.cpp: the pool is created once every stage's session exists; the slot sessions join their lane per stage; pool_move swaps the lane's tables once and copies/resets the indexer and residency of every stage; the slot copies skip the streamed layers of every stage (the pool carried them), not only CUDA0's.
…unk tables - kv_pool_policy.hpp: kv_pool_eviction_plan (lowest rank first, the longest unused, the lower slot; nobody is asked when all that can give way together would not cover the request); pool_reserve in generate.cpp acts on its answer. - KvHostPools::for_each_piece: the walk that cuts a byte range where a chunk ends, moved out of conversation_snapshot.cpp so the test can check every byte against the block mapping. - tests/core/kv_chunks_test.cpp (target kv_chunks_test, no CUDA): KvChunkMap rules and a random walk of reserve/shrink/swap that checks every chunk is free or held exactly once, the host mapping, the eviction plan.
… the header before <cstddef>)
…/test_kv_pool.py)
…s borrow RoPE only from a stage that has QSA layers --pipeline-windows launches windows ahead of each other and never reserved chunks per window like the serial loop does (everything past the prompt's reservation would have gone to the trash chunk): it reserves max-context up front and decodes serially when the pool cannot back it (it runs without --batch, so no slot has to give way).
…ation it describes With the shared KV pool, an admission move clears main's live conversation but left the reuse image retained at a cache restore. limit_reuse() kept it for the next request that resumed from another slot, and the next park reused its bytes for cells of a different conversation (silent wrong K/V in a parked entry). The reuse image now carries the token ids and image keys of its conversation; take_reuse() hands it back only for the prefix both conversations share, and the paths that move a conversation out of main drop it. Test: conversation_cache_test.
… and four lanes, what was not measured)
… (RAM, speed, restarts; what was not measured)
With --kv-pool-tokens an admission moved the K/V to the slot but left it only the deepest checkpoint (the turn checkpoint) and emptied the main session's chain; a give-back to the solo path dropped even that one. A later turn that changed the last question (or re-rendered the reply) has no checkpoint below the change except the periodic ones, so both lanes of two long conversations read everything again. The chain now travels with the chunks both ways, and an uncached slot drops its chain. Measured on 2 lanes x 133k with a changed last question: both lanes reuse 118,784 tokens (time to first token 24.5 and 13.3 s) where the unfixed binary reused 0 (52 and 142 s).
The pool section said a changed last question resumes from a checkpoint but not where that stops. Measured with two lanes: an edit below the first periodic checkpoint reuses nothing (30k: 0, 60k: 31,232); 133k per lane works with the default 4096 MiB conversation cache, 256k per lane needs at least 4528 MiB (4608 tested); each cached slot now keeps up to 6 checkpoints, about 0.7 GB of host RAM per slot.
…the direct copy and the prefix restore through a shared pool's chunk tables A second streamed state (its chunks interleaved with the first's) shares the host arrays; every K/V format the streamed test covered runs once more with the host copy in chunks, and checks that a neighbour's writes stay out of a state's chunks.
…e while the pool moves chunks STRATA_SLOT_CKPT_MOVE (opt-in) lets a request mount a checkpoint from a slot that is still decoding. With --kv-pool-tokens the slot's chunk table would move to the main session and be released under the slot's feet, so the pool keeps the idle-slot rule.
…s at 104k; what was not run) Same binary with and without the flag on 2 x RTX 3080: a ~256k conversation moves in 0.36-0.5 s instead of 3.5-5.2 s copied, 3.2 GiB more idle MemAvailable, no short-context speed effect. Older 0.1.41 / 0.1.44.1 figures are labelled as older. States that the binary was built from eeaf7b3 and that the serve layer had no pick_slot rule.
…kens Without the flag each slot keeps its own table again, as on main. Borrowing it saved 0.07 GiB of VRAM per slot, which --expert-cache auto turned into about 43 more cached experts per card, and so into different greedy output than main with the pool off.
81a5c6a to
1e1cd76
Compare
noon-at-cgn
changed the title
|
Rebased on main
The two red CI jobs ( |
Branch
pr/kv-pool-141, head1e1cd762: 20 commits on upstream main41cf05a0(v0.2.5-6, engine 0.1.50). Opt-in: without--kv-pool-tokensthe output is byte-identical to main (identity gate below). Scope issue: #3351.Problem and result
With KV streaming (
--kv-resident), the main session and every--batchslot pin the K/V of the whole--max-contextin 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, which at 256k takes seconds, and two long lanes push a RAM-limited host to its memory guard.--kv-pool-tokens Npins N cells once. Sessions hold only the 4,096-cell chunks their conversation has reached, and a move between lanes swaps chunk tables instead of copying.Two lanes of about 256k tokens (2 x RTX 3080 20 GB, UD-Q4_K_XL, same binary in both arms, only the flag differs; measured on 0.2.0-beta, engine 0.1.45.3):
moved)copied)MemAvailable)MemAvailable(mean of 3 starts each)At about 133k per lane the move took 0.1-0.25 s with the pool and 0.9-2.0 s without. With the pool a third lane (three conversations of about 104k) fits, every needle answered.
Short context is unchanged: one stream 84.5 vs 81.9 tok/s, two streams 93.7 vs 89.8 (means of three starts each, within the round-to-round spread). The gains are at long context with more than one lane.
On this PR's head (0.2.5) the pool runs the same workload end to end (2 x 256k, numbers under Validation). I did not repeat the no-pool arm on 0.2.5, so the comparison above is the 0.2.0-beta one.
Related discussion
conversation_snapshot.cpp.pick_slot): our slot-choice change is not in this PR.Approach
A session keeps a chunk table, a device copy for the kernels and a host copy for DMA and copies. Unreached chunks point at one spare "trash" chunk. Without the flag the table is the identity layout. A lane reserves chunks as it writes. When the pool is full, an idle slot's cached conversation gives way first, then a new request gets a retryable 503 (
kv_pool_full), and a decoding lane that cannot grow ends withlengthand"truncated": true. It works across a layer split: each card pins its own layers' share and a conversation uses the same chunks on every card. The MTP drafter ring is not pooled.KvHostPools::block,block_host,for_each_piece): append kernels, resolve and stage copies, verify, MTP, prefill, state hash,conversation_kv_copy,conversation_kv_restore_prefix, session files, K8V4 halves. I checked the 967 upstream commits since the last base for new host K/V sites and found none.1e1cd762; see Failed).kv_host_alloc(WSL2 fallback, pin budget).docs/BATCHING.md,docs/DETAILS.md,docs/MULTI_GPU.md.Validation
Head
1e1cd762unless a line says otherwise.ebb61958is the same series before the last two changes (the RoPE gate, and(void) cudaFreeinkv_pool.cpp; that file compiles to identical machine code either way).Passed:
41cf05a0vs this head without--kv-pool-tokens, production config otherwise (2 x RTX 3080 20 GB, layer split 23,parallel 2,--spec 4 --batch-mtp,--max-context 262144 --kv int8 --kv-resident). Greedy (temperature 0, top_k 1, seed 1234), 256 tokens, prompts of 989, 31,280 and 128,209 tokens, same bytes for both. The stock config is not reproducible even main against main (adaptive experts), so both arms ran the deterministic settings fromdocs/DETAILS.md(--adapt-every 100000 --adapt-async 0 --pcie-frac 0) withSTRATA_STATE_HASH=1. Two starts of each: all three outputs byte-identical to main's, the STATE_HASH fields equal, the expert cache at 4893 / 4161 slots like main's. Atebb61958the same held with--spec 2.--kv-pool-tokens 524288, one start):KV pool: 524288 cells shared by the sessions and the stages of the layer split, 6.24 GiB of pinned RAM. Two concurrent 256k prompts, then appended turns: 4/4 needles correct, slot moves 444-842 ms, appended-turn decode 53.0 tok/s aggregate, single-lane questions 43.5 and 50.5 tok/s, nokv_pool_full, hostMemAvailablenever below 17.2 GiB.-DSTRATA_ENABLE_CUDA=ON -DSTRATA_NATIVE_EXPERTS=ON -DSTRATA_MMQ_KQUANTS=ON): rc 0, same warning count as main.CUDA_VISIBLE_DEVICES=-1): kv_chunks_test; conversation_cache_test (4791); conversation_memory_test (31); conversation_file_test (201); conversation_validation_test (972); conversation_transfer_test (1548 host-only); session_save_reclaim_test (119); memory_guard_test (134).ctest -j1atebb61958: 161 of 168 passed, 4 skipped (CPU without AVX-512, ggml without trellis types), 3 failed, all pre-existing (Failed). At1e1cd762the pool-related subset again: 17/17 (conversation_snapshot_test with the pooled variant, kv_chunks, the kv_* parity tests, batch_rows, verify_batch_parity).ebb61958:python -m unittest discover -s serve -t .: 1050 tests, 15 skipped, OK.ebb61958(rocm/dev-ubuntu-24.04:7.2.4, AMD clang 22 / HIP 7.2, flags ofpackaging/build_hip_linux.sh):stratafor gfx1100 and for gfx1151;gfx1201, and all targets with tests on, rc 0 on the branch and on main. The branch adds one test (kv_chunks_test). The only new warning,cudaFreeresult unused inkv_pool.cpp, is fixed on this head.kv_chunks_testandconversation_cache_test(4791) pass on the host.ebb61958(oneAPI DPC++ 2026.1.1, flags ofpackaging/build_sycl_linux.sh):strataand all targets, rc 0, warnings identical to main. The SYCL tree compiles its own migrated copies and does not build the pool: of the touched files onlylayer.hpp,kv_stream.hppandconversation_cache.hppare compiled there, and they build cleanly.Failed:
ebb61958, flag off was not identical to main: the batch slots borrowed the RoPE table without the pool too, so each slot was 0.07 GiB smaller,--expert-cache autocached 43 more experts per card, and greedy output differed. With a VRAM reserve that equalised the expert count the outputs matched, which located the cause.1e1cd762borrows only with the pool; the gate above is from that head.ple_parity(its Q2_0 GGUF is not on this host),ggml_tq_overlay(cmake script "Unknown arguments", same on main),isa_guard_check(iq512x_rowsviolation, same on main's library). This branch does not touchsrc/kernels/cpu,cmakeortests/cmake.metal-linkfails inmetal/src/mtl_rt.mm(MTLGPUFamilyMetal4,MTLCompileOptions.mathMode: macOS 26 API, in main sincead9eb890) on themacos-14runner's SDK. Other PRs on current main fail the same job; this branch does not touchmetal/.Linux / runtime requirements onlyat1e1cd762: 1051 serve tests, 1 failure,test_lifeline.test_the_portable_watcher_ends_the_child_when_the_server_is_killed("the child outlived a killed server (watcher)"). It depends on how the tests are started: on the lm-server box it fails on main41cf05a0too, 5 of 5 runs from an ssh shell (python -m unittest serve.test_lifeline), while the whole serve suite above passed when run from a systemd unit. This branch does not touch the lifeline code or its test.Not run:
1e1cd762(the change sinceebb61958is 3 lines of host C++ ingenerate.cpp, which SYCL does not compile, plus the(void)cast); Metal; Windows and the WSL2cudaHostRegisterpath.1e1cd762(subset and CPU tests only); the no-pool 256k arm on 0.2.5.STRATA_SLOT_CKPT_MOVE=1with the pool;STRATA_SLOT_COPY_DIRECTandborrowend to end; session-fileSAVE/RESTOREwith a pool end to end (unit tests only); pool-full paths under load; more than 3 lanes; a single 512k conversation.All raw outputs (prompts, per-request JSON, engine logs, binaries' sha256) are kept and can be attached.
Limits
--kv-pool-tokens; pass it inargs.