Repository navigation
[Engine] shared KV pool for all lanes (--kv-pool-tokens): 256k lane moves 10x faster, 3 GiB less pinned RAM - #1614
noon-at-cgn wants to merge 19 commits into
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. |
- fix two PRs wrongly listed as already in main (Niko1221#1548/Niko1221#2224, Niko1221#1614) - correct Niko1221#1358's merge state (dirty, not cleanly mergeable) - add an implementation plan with one run command per phase for reading prompt/generation tok/s
|
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. |
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.
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
Branch
pr/kv-pool-141, head81a5c6ae(19 commits on upstream maine88ee39d=v0.2.0-beta-12, engine 0.1.45.3). Reopens this PR, rebased from v0.1.44.1.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.--kv-pool-tokens N(opt-in) pins 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.Measured on 0.2.0-beta, 2 x RTX 3080 20 GB, two lanes of about 256k tokens, the same binary in every arm and only the flag differing (
--kv-pool-tokens 524288present or removed):moved)copied)MemAvailable(mean of 3 starts each)MemAvailable)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 tokens) also fits: 13.1-13.2 GiB
MemAvailableidle, every needle answer correct.Short-context speed is unchanged. One stream 84.5 tok/s with the pool against 81.9 without, two streams 93.7 against 89.8 (means of the medians of three starts each), prompt reads within 1-3%. The +3-4% is smaller than the round-to-round spread (58-103 tok/s), so I make no speed claim. This is a memory and long-context lane-handling feature, not a decode-speed one.
The no-pool 256k numbers are one full run and two partial ones, because the host guard stopped two of the three no-pool starts (see Validation, Failed). They were also taken while the host had 1.1-1.8 GiB free and no free swap, which may add to the copy times.
Related discussion
96549a3e). Ours exists only with the pool: the chain follows the chunks when a conversation changes lane, and nothing is parked. It will conflict textually with [Engine] Preserve decoded batch-slot conversations before slot reuse #2526. I will rebase on whichever lands first.conversation_snapshot.cpptoo.pick_slot): the slot-choice change from our branch is not in this PR; see Approach.Approach
A session keeps a chunk table: a device copy for the kernels and a host copy for DMA and copies. Chunks the conversation has not reached point at one spare "trash" chunk. Without the flag the table is the identity layout and nothing changes. 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.What the rebase onto 0.2.0-beta changed (18 commits, 7 conflicts, all resolved by keeping main's code and re-applying the pool on top):
conversation_kv_copy(the direct device-to-device slot copy) andconversation_kv_restore_prefix(used byborrow).conversation_snapshot_testnow runs save, restore, verify, source, the direct copy and the prefix restore through a shared pool. With the direct copy's chunk branch disabled the test fails, and withrestore_prefixbypassing the chunk table it fails; both changes were reverted afterwards.kv_host_alloc, so it gets the WSL2 fallback and the pin budget.&st.host, which carries the chunk fields; I read every host-addressing site (append kernels, resolve and stage copies, verify, mtp, prefill, state hash).96549a3e). Without it a client that changes its last question finds no checkpoint below the change.9a5dcb4b: an active slot is never the source of a checkpoint move while the pool moves chunks. It matters only with the opt-inSTRATA_SLOT_CKPT_MOVE=1plus the pool. The benchmark binary was built fromeeaf7b3c, before this commit, so the guard was not in it and not exercised.pick_slot, 64e7f13) is split off to its own branchpr/pick-slot(4fcad010) because it changes every--batchsetup, pool or not, and main now has its own rule for the same family (Concurrent requests break prefix cache #2518). It will be a separate PR. The two touch the same spot inserve/server.py, so whichever lands second needs a small rebase.docs/BATCHING.mdanddocs/DETAILS.md(pool section, the 0.2.0-beta table, older figures labelled with their engine version). The "changed last question" figures (24.5 s / 13.3 s, 118,784 and 225,792 tokens reused) come from engine 0.1.44.1 with thepick_slotrule and are labelled as such.Validation
Environment: codebox, 2 x RTX 3080 20 GB, CUDA 13.0, sm_86, gcc, Ninja. Benchmark: lm-server, 2 x RTX 3080 20 GB at 220 W, a 90 GiB container. All raw outputs are kept and can be attached.
Passed:
ab83596b(full tree), thenstrata,kv_chunks_test,conversation_cache_test,conversation_snapshot_test,conversation_validation_testatfb493d3d,24695028,3ef8b4f2,596077c0,bd4eb343,96549a3e:ninjarc=0, 0 errors. Tip9a5dcb4bfull build rc=0, with warnings only in code this branch does not touch. Configure:cmake -DCMAKE_BUILD_TYPE=Release -DSTRATA_ENABLE_CUDA=ON -DSTRATA_NATIVE_EXPERTS=ON -DSTRATA_MMQ_KQUANTS=ON -DSTRATA_BUILD_TESTS=ON -DSTRATA_BUILD_CONVERSATION_TESTS=ON -DCMAKE_CUDA_ARCHITECTURES=86.CUDA_VISIBLE_DEVICES=-1): kv_chunks_test; conversation_cache_test (4791 checks); conversation_split_failure_test; conversation_memory_test (23); conversation_file_test (201); conversation_validation_test (972, host only); conversation_transfer_test (1548, host only); session_save_reclaim_test (119); batch_groups_test; batch_rows_test; memory_guard_test (134); aux_cpus_test.ctest -j1over the named list ateeaf7b3c: 28 of 28 passed (qsa, qsa_slots, qsa_topk, rope, verify, spec_verify, verify_batch and batch_rows parity; kv_q8, kv_stream, kv_window, kv_hybrid, kv_q4 and kv_steps parity; conversation_snapshot_test with 5007 checks including the new pooled variant; the conversation validation, transfer, file, memory, split_failure and cache tests; session_save_reclaim; draft_kv_plan; batch_rows and batch_groups; memory_guard; aux_cpus; kv_chunks).python -m unittest discover -s serve -t .on the tip: 897 tests, 14 skipped, passed. Theserve/test_kv_pool.pyandserve/test_slot_alloc.pymodules (23 tests) also pass on their own.strata-w14, sha25641fd4f0300ebc5c93f6acf9e87d8b5523cb8b347f4c93fb59cc60ab0a5a1daf5, built from the code ofeeaf7b3c; serve layer without thepick_slotrule; 8 engine starts, each from a clean restart; no kernel OOM kill; every needle check of a completed request correct):STRATA_BATCH_MTP_SERIAL=1,parallel 2,--max-context 262144 --kv int8 --kv-resident 32768 --spec 4 --batch-mtp --conversation-cache-mib 4096 --conversation-cache-slots 3 --prefill auto:16384. Only the pool flag differs per arm. 0.2.0 defaults (STRATA_SM86_PERF, the Linux file tier, the cgroup RAM check that counts clean page cache, zero-copy page-cached experts) are on in all arms.MemAvailable14.90 / 14.72 / 13.99 GiB against 10.93 / 11.66 / 11.54 GiB. Run minimum 10.4 GiB against 7.6 GiB (means). Start log:KV pool: 524288 cells shared by the sessions and the stages of the layer split, 6.24 GiB of pinned RAM.ab.py solo 12,ab.py c2 10,read96.pyafter every start; numbers in Problem and result.lanes_n.py 2 245000): pool, 2 starts, all steps completed. Without the pool, start 1 completed, starts 2 and 3 are in Failed.parallel 3, about 104k per lane, 2 starts): all answered,MemAvailable8.3-8.6 GiB at the run minimum, slot moves 89-326 ms.Failed:
guard_host.sh, stops the engine when hostMemAvailablefalls below 1.5 GiB) stopped the engine twice during the 245k run: startA1nopool2at 22:45:17 (hostMemAvailable1103 MiB, during the appended turn) and startA1nopool3at 23:22:59 (1403 MiB, during the second short question). It was the guard, not the kernel OOM killer (oom_killstayed 0). The 133k runs of those two starts have no data ("connection refused"). Never fired in a pool arm. This is a failure of the no-pool configuration on this box, recorded as a failure.serve.test_kv_pooltest_a_prompt_longer_than_the_context_is_a_400_not_a_503failed: main'sfit_max_tokensdefault turned the 400 into an engine call. The test now sets the strict mode.conversation_snapshot_testfailed underctestwhile another process held about 19 GB on both GPUs; it passed (5007 checks) when run alone a minute later.Not run:
cudaHostRegisterpath (the allocation goes throughkv_host_alloc, not run).STRATA_SLOT_CKPT_MOVE=1, and so the guard in9a5dcb4b.STRATA_SLOT_COPY_DIRECTend to end (unit test only).borrowend to end (main restores it only without a layer split; not applicable to the 2-GPU config). Session-fileSAVE/RESTOREwith a pool end to end (unit test only).pick_slotlayer (pr/pick-slot). The benchmark turns only append to one lane's history, and I did not send a changed-last-question request. That the results do not depend on it is my inference.MemAvailable, and about 3.1 GiB more per slot would have reached the 4 GiB abort line). Pool full paths (nokv_pool_fullin the logs). 4, 6 or 8 lanes, a larger pool, a single 512K conversation.ctest(150 tests) and the Python tests undertools/. Without the flag the default path is the identity layout and the parity tests pass, but I did not run a byte-for-byte output comparison against a build without this branch.ab.pywere identical configs and measure noise only.Limits
--kv-pool-tokens; pass it by hand inargs.STRATA_SLOT_CKPT_MOVEwith the pool and the SYCL, HIP, Metal and Windows paths are unverified.