Elastic K/V (--kv-grow, opt-in): the K/V takes VRAM as the context grows, the expert cache holds the rest (+68% cache on a 16 GB card at 262K)

[sergqwer]

At --max-context 262144 the K/V of the whole context is allocated at start: 3.35 GiB at int8 with the drafter's. The expert cache is sized around it. On a 16 GB card that is a quarter of the VRAM, and most requests never reach those cells.

With --kv-grow (or STRATA_KV_GROW=1), the K/V holds physical memory only for the cells the requests reach. The expert cache holds the rest. It is off by default: more slots move experts between the GPU and the CPU, which changes tokens, so it does not pass a byte-identical gate against main by design.

How it works:

• Virtual memory (core/vmm.hpp). The K/V pools and the cache arena are CUDA VMM ranges: addresses reserved once, physical memory mapped in 2 MiB chunks. The driver calls go through cudaGetDriverEntryPointByVersion, so nothing links cuda.lib.
  • Nothing moves, so no graph is captured again.
  • --max-context cells always fit: the window is never smaller.
• Growing. The K/V maps the chunks for 16384 cells at start and grows in steps of 8192.
  • It takes chunks from the slots just below the prompt path's loan.
  • A hotter expert in such a slot first moves (one D2D copy) into the slot of the coldest resident in the loan region, so the cache gives up its coldest experts, as a smaller cache would.
  • The CPU computes the experts given up, like any miss.
• Giving back. A short request after a long one hands the chunks back once the outgoing session is parked. The slots refill from the profile.
• Where it grows:
  • serve: at a request (its prompt, and a parked conversation before its restore) and in the decode rounds when a window would pass the mapped cells;
  • generate: once, for the prompt plus --max-new.
• When it applies: one GPU, a profile, the whole K/V in VRAM (no --kv-resident), every expert in RAM (not the low-RAM tier). Otherwise the flag is ignored.
• Test switches: STRATA_KV_GROW_INIT / _STEP / _FLOOR. tests/core/vmm_test maps, moves a chunk between ranges keeping its bytes, and releases.

One pitfall worth knowing: a device-to-device cudaMemcpy does not wait for the copy. The moves sync before their source chunks are unmapped. Without that, the copy faulted, and the error surfaced at the next cudaMemset as "no VRAM".

Same output when nothing moves

Setup: IQ2_XS (ISTA), RTX 5090, a 32K prompt and 64 tokens, fixed --expert-cache 14900. First-token logits were dumped with #276's STRATA_DUMP_FIRST_LOGITS in a test build.

These are bit-identical to 0.1.31, logits and tokens:

• without the flag;
• --kv-grow mapped whole (STRATA_KV_GROW_INIT=262144);
• --kv-grow grown from 4096 cells into new VRAM (_FLOOR above the cache).

Grown by giving up 266 slots (171 experts moved): the same first token and the same 64 tokens here. In general a given-up expert is computed on the CPU, which rounds differently.

Speed

Setup: IQ2_XS (ISTA), RTX 5090, 9950X3D, 128 GB, Windows 11, --max-context 262144 --kv int8 --prefill auto --spec 4, cache auto, 256 tokens. The runs alternated with and without the flag.

	0.1.31	--kv-grow
expert cache	15,119-15,207 slots	17,300-17,503
short chat	15.06-15.33 ms a round, 161-163 tok/s	14.71-14.83 ms, 164-167 tok/s
an emulated 16 GB card (--vram-reserve-mib 17884): expert cache	3,396 slots	5,719
16 GB: short chat	24.1-25.5 ms a round, 100 tok/s	21.1 ms, 115 tok/s
16 GB: a 32K prompt	7.44 s	6.36 s
16 GB: decode after it	24.6 ms a round	22.2 ms

On our NVFP4 fork (2.76 MB experts, 5090) the cache grows from ~7,060 to ~8,290 slots and a short chat from 111-115 to 126-133 tok/s.

Serve

A served run with --conversation-cache-mib 8192 behaved as expected:

• a chat;
• a 17.8K prompt and 4,000 tokens: grown 16K -> 20K at the request and 24K mid-decode;
• a 31.7K document: grown to 32K;
• two short chats: trimmed to 4K, 159 slots refilled;
• the document's follow-up: grown before its parked conversation was restored (31,749 tokens reused).

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[sergqwer]

Pushed 1d17e32, fixes from an audit of this PR (each checked against the code):

• VmmRange: a run whose cuMemSetAccess failed stayed counted as mapped, so a later map_range skipped it and the chunks were never made accessible. It is unmapped and freed again now (commit_run).
• A failed growth or trim could leave host_res, d_res and the slot range half-changed while serve went on to the next request. Such a failure now stops the engine, as other CUDA failures do. The residency uploads are checked.
• The trim maps the chunks it takes back as one run (one cuMemSetAccess instead of one per 2 MiB chunk). It refills the slots with async copies and one sync instead of a blocking cudaMemcpy per expert.

IQ2_XS: the same served scenario (grow at a request and mid-decode, trim and refill, regrow, a parked conversation's restore) runs as before; vmm_test passes.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[sergqwer]

Rebased on 0.1.32 (c499bd1). One conflict, in generate.cpp, and both sides are kept:

• Session start. Prefill regression on sm_86 (2x RTX A5000) in 0.1.30: ~-35% on ~2K-token prompts, not recovered by STRATA_PREFILL_STREAM_MIN / GDN_PIPELINE=0 / SCORES_MULTI=0 #340's STRATA_SPLIT_SMALL_OWN block stays, then kvg_start. The split block runs on a layer split only, and the elastic K/V is single-GPU (!multi_gpu), so they never act together.
• A request's start. Verify: the commit graph no longer waits; it overlaps the MTP draft #284's wait_commit comes first, then the K/V growth. Its quiesce synchronizes the device anyway.

Checked on 0.1.32: IQ2_XS (ISTA), RTX 5090, a fixed --expert-cache, --vram-reserve-mib 1500. The first generated token's logits (STRATA_DUMP_FIRST_LOGITS, now in main) are byte-identical to main's after a 2K and a 32K prompt. Every target builds with STRATA_BUILD_TESTS.

With --kv-grow, a cache budget below the free VRAM is needed for a byte-for-byte comparison:

• Why. --expert-cache N is a byte budget that the free VRAM caps. With --kv-grow the start leaves 3.1 GiB more free, so at --expert-cache 14900 the cache gets 15,628 slots instead of main's 15,445. That is the point of the change, but then the first token is computed differently.
• Check at --expert-cache 13000 (13,643 slots either way):
  • the logits are identical to main's;
  • the 256 greedy tokens are identical to main's;
  • the K/V stayed at 16,384 cells, with no growth needed.
• At 32K the K/V grew to 32,768 cells, and the cache gave 151 slots for it.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[BlueKingMuch]

Tested on Windows 11 (WDDM): RTX 4080 SUPER 32 GB (sm_89), Ryzen 7 5800X3D, 64 GB, PCIe 3.0, CUDA 13.3; IQ3_S at --max-context 262144, int8 K/V, --expert-cache auto with --expert-profile, every expert in RAM (the arena is pinned in 47 slices, 45 of 46.8 GiB). vmm_test passes. Build: 0.1.32 (+#363, #413, #346) with this PR merged together with #385, #372, #374 and #407, no conflicts; STRATA_KV_GROW=1 against unset in the same build.

• The auto sizing sees 24.85 instead of 21.76 GiB free: 12693 instead of 11055 slots (the K/V starts with 16384 cells, 0.25 of 3.35 GiB).
• Decode, 256-token requests, 3 per cell, ms per round / tokens per second:

	greedy 8K	greedy 64K	sampled 8K	sampled 64K
unset	21.4 / 124.6	21.7 / 127.3	20.7 / 118.2	21.8 / 130.5
STRATA_KV_GROW=1	20.4 / 132.0	21.6 / 128.3	20.6 / 123.9	21.7 / 129.4

At 8K the CPU computes 0.52 instead of 0.92 experts per layer (VRAM hits 98.3 instead of 96.8 %). For the 64K requests the K/V grew once to 65536 cells (330 slots given up, 12363 still hold experts); there the VRAM hit rate was 98.3 % without it already, and nothing changes.

• Prefill: fresh 32K / 64K prompts +1.1 / +0.8 %. Short prompts gain more, as fewer experts have to be streamed: a 2K prompt takes 2.45 instead of 2.68 s, and the second, 923-token piece of the 64K prompt (read here as 63035 + 923 tokens) 1.76 instead of 1.92 s (about 800 instead of 985 ms waiting for copies).
• With the conversation cache (--conversation-cache-mib 2816 --conversation-cache-slots 2), a main agent (60K), subagents (15K) and a second main agent (100K) taking turns: the K/V grows before each restore (to 65536 cells for the 60K conversation, 106496 for the 100K one) and is trimmed when a 15K subagent takes over (330 / 601 slots back, refilled from the profile). Time to first token was never longer (back at the main agent 3.0 -> 2.8 s and 2.8 -> 2.6 s, the other steps equal or 0.1-0.2 s shorter); parked and restored sizes and RAM as without.
• As you wrote, the tokens differ: the deterministic requests diverge from token 30, since experts the CPU computed before now run on the GPU.
• No errors in 4 engine starts and 44 requests with it.

I run with it now: STRATA_KV_GROW=1 in the config's env, so builds without this PR still start with the same config. Thanks!

[sergqwer]

@BlueKingMuch thank you for the thorough test, the conversation-cache run included.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[sergqwer]

Rebased on 0.1.37, plus a fix from a code review (7e5b7e7).

The bug: with --kv-grow the K/V is made elastic at session init and maps its first 16,384 cells. Growing it needs more than init checks: a device residency table, cache slots and a loan above the floor.

• Without them, kvg_start left the growth off and nothing mapped more. That happens with --no-pool, --no-token-graph, the dump modes, or --expert-cache auto resolving to 0.
• A prompt past 16K then wrote K/V into unmapped memory.
• Reproduced on IQ2_XS with --kv-grow --no-token-graph and a 32K prompt: prefill copy_i32: an illegal memory access was encountered.

Now: such a run maps the whole window from new memory at start ("elastic K/V off for this run ... mapped up front"), or stops with a message if it does not fit. The same 32K prompt reads through: 32,036 tokens in 4 chunks.

The growth's and the trim's residency-table uploads also wait for their own copy. A pageable cudaMemcpy may return before its DMA lands, and the streams that read the table are non-blocking. The same pattern elsewhere is #550's subject.

Without --kv-grow, nothing changes.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[Niko1221]

We'd take --kv-grow as an opt-in. Please rebase it on 0.1.38. Without the flag, the output must stay byte-identical to 0.1.38 (fixed cache, 10/10).

[sergqwer]

I rebased this onto 0.1.38 (c053829). The only conflict was the source list in CMakeLists.txt (peer_experts.cpp beside vmm.cpp). --kv-grow stays opt-in, and nothing changes without it.

Without the flag, compared with v0.1.38 built from the tag with the same toolchain:

• Setup: RTX 5090, IQ2_XS, fixed --expert-cache 14900, --max-context 262144 --kv int8 --spec 4. Five prompts (two short chats of 24 and 44 tokens, then 2K, 8K and 16K documents), each run twice, interleaved with v0.1.38. Every pair compared byte for byte: the first-token logits (STRATA_DUMP_FIRST_LOGITS) and 64 greedy tokens.
• Result: 10/10 byte-identical, logits and tokens both.
• Decode in the same runs:

prompt	v0.1.38	this branch
chat, 24	178.5 / 172.8	168.8 / 172.6
chat, 44	149.7 / 149.5	147.0 / 149.4
2K	164.5 / 164.6	165.4 / 164.3
8K	174.1 / 173.7	173.9 / 174.1
16K	182.8 / 183.2	183.5 / 182.8

Decode is in tokens/s, rounds 1 / 2. The one low cell, 168.8 on the 24-token chat, is run-to-run spread: the second round gives 172.6 against 172.8.

With --kv-grow, on the same rebased build (cache auto, a 32K document, 64 tokens):

• The run starts with 16,384 of 262,144 cells in VRAM.
• The expert cache holds 17,626 slots. With a fixed-size K/V it gets 15,342 at the same reserve (1500 MiB).
• For the prompt, the K/V grew to 32,768 cells (0.46 GiB). The cache gave 145 slots for it.
• Prefill ran at 6,783 tokens/s, then 64 tokens decoded at 185 tokens/s, with no errors.

[sergqwer]

Rebased onto 0.1.39 (d5ee0cc). 0.1.39's segmented expert cache (#533, --vram-elastic) also reserves the cache's address range through VMM, so the two can't run together:

• ExpertCache::open / close take the segments, else this PR's VMM range, else one cudaMalloc.
• --kv-grow stays off beside --vram-elastic and beside the batch slots (Feature request: serve concurrent requests (cross-request batching) #465, parallel; each slot's K/V would need the same growth), and says so: strata generate: --kv-grow is off beside ....

Without the flag, against v0.1.39 (IQ2_XS, fixed --expert-cache 14900, 64 greedy tokens): first-token logits and tokens byte-identical on a short chat, 2K and 16K (3/3).

[sergqwer]

Reopened as #1040, the same branch and commit. GitHub closed this one when my fork was made private by mistake, which took it out of the fork network.

[Niko1221]

In 0.1.40: --kv-grow is in (opt-in, off by default). It was ported over --vram-elastic: the cache arena tries the segmented backing first, then the VMM range, then cudaMalloc. It switches itself off with a warning under --batch, --vram-elastic and --peer-device.