[None][feat] Contiguous primary KV cache (P0): per-sequence VMM arenas for KVCacheManagerV2

[thorjohnsen]

Summary

Prototype implementation (P0) of the contiguous primary KV cache for KVCacheManagerV2: each sequence's active KV blocks are contiguous in virtual memory, backed by physical pages mapped on demand via CUDA VMM, while reusable (stale) state lives in the host tier.

Flag-gated end to end: KvCacheConfig(use_contiguous_kv_arena=True) (prototype status). Default behavior is untouched — with the flag off, no code path changes.

Design

• One large VA reservation ("arena") per (pool group, pool); each sequence gets a contiguous, page-aligned block-index range sized for its per-request maximum, so slot_id = base + ordinal and the sequence's kernel-visible page indices are consecutive. Arena base = pool base pointer → zero kernel changes, CUDA-graph safe.
• Physical 2 MiB pages (configurable super-pages) map on demand against a level-wide page budget; freed ranges are parked still-mapped and adopted whole by the next admitted sequence (default ON), so steady state performs zero cuMemMap/cuMemUnmap; unmaps happen only under real pressure (spill) or teardown, behind a device quiesce.
• Active→stale: committed blocks copy out to the host tier on free, or eagerly via an opt-in write-through-on-commit policy that makes the free path copy-free (adoption of the pre-copied host slot).
• Stale→active: a reuse hit copies the matched prefix into the new sequence's range as sequence-private pages; canonical radix entries are untouched. Suspend/resume evacuates and re-onboards through the same paths.
• Prefix aliasing (opt-in, prefix_aliasing / TRTLLM_KV_ARENA_PREFIX_ALIASING=1): a closing canonical owner's fully-committed prefix pages are pinned (refcounted) in a canonical-span registry, and later same-prefix admissions cuMemMap the SAME physical pages into their own ranges — zero-copy, zero-budget reuse, with prefix-affine adoption (parked ranges carry an (span, extent) signature and are only handed to admissions matching it, since adoption overwrites page content). This recovers paged's refcount-sharing economics for the arena layout.
• Capacity planning (clamp_max_seq_len_for_mem, scheduler backpressure) moves from slot counts to physical pages; the v2 scheduler needed no changes.

Full design doc, VMM microbenchmarks (H100 + GH200), and per-commit work log are tracked in the development directory (contiguous_primary_kvcache/DESIGN.md); can be attached or moved into docs/ on request.

P0 scope gates (fail loudly, lifted in later phases)

Sliding-window/VSWA layers, SSM, SWA scratch reuse, partial-block reuse matches, seq rebasing, beam search (unsupported in v2 anyway), FlashInfer-style dense pool views (get_buffers), disaggregated serving.

Correctness campaign (8B scale): three distinct bugs, all root-caused and fixed

Scaling validation from 1B to Llama-3.1-8B exposed an intermittent illegal-memory-access that turned out to be three independent bugs (two of which also affect or inform main):

#	Bug	Class	Fix
a	KV block-offset staging race: the offset-gather kernel could read pinned staging buffers already overwritten by the next iteration under the overlap scheduler → stale offsets → IMA (arena) / silent wrong reads (classic)	CPU-side buffer reuse (nvbug-6293536 class); affects main	4d49f1d66e here; standalone main fix in PR #16088
b	Scheduler capacity deadlock: arena frees are event- and fence-gated, and the fence was only assigned after scheduling — pages freed by the scheduler's own evictions were unreclaimable in-pass → suspension cascade → guard raise	Arena-specific capacity model	4fefc34db5 — blocking fenced drain inside the pass + allocation retry
c	cuMemUnmap racing in-flight kernels: unmapping a VA range that running kernels concurrently read intermittently produces a Warp MMU fault	Driver-level; being escalated to the driver team with a repro matrix	b271d635fd (freed-range adoption removes steady-state unmaps entirely); TRTLLM_KV_ARENA_SYNC_BEFORE_UNMAP (default ON) quiesces the rare pressure/teardown unmaps

Evidence for (c), with (a) and (b) fixed: on a deterministic pressure reproducer (free_gpu_memory_fraction=0.42 ≈ 2.4× KV oversubscription @ conc 256), no-quiesce crashes ~4/7 across mapping margins (growth maps exonerated), while sync-before-unmap runs 6/6 clean (p ≈ 0.6% against the ~57% base rate). Two further ordering fixes landed from the same campaign: arena write-outs (suspend/close) are now ordered after the in-flight forward step via a caller-supplied prior event (52a311ed99 — silent stale-tail restore, proven both directions), and on-demand mapping (margin=1) was restored as the default (53f7167405) after pre-mapping proved to be a ~3-point stopgap for what was mostly bug (a).

A fourth bug was root-caused and fixed after the table above (51e092e3ed): host-tier exhaustion during suspension. The host tier legitimately fills with already-suspended sequences' HELD pages (which LRU eviction correctly refuses to drop at the last level), and suspend() mutated sequence state before its fallible host allocation — so exhaustion either escaped as a process-fatal OutOfPagesError or, raced by the in-flight overlap step against the half-suspended state, poisoned the device as an IMA. Fix: an atomic pre-flight host reservation in suspend() (failure leaves the sequence fully intact) plus graceful degradation to v1-style drop-and-recompute through the adapter, scheduler, and executor (surviving canonical prefixes still shorten the re-prefill). Discriminator (Llama-3.1-8B, 0.42 quota, 4 GiB host ≈ 10× under suspension demand, conc 256): unfixed 4/4 dead (2 IMA + 2 clean FATAL — two faces of the same trigger); fixed 8/8 clean at 300/300 with 40+ drop-and-recompute cycles per run.

Endurance under production defaults on the pressure reproducer: 300-req ×3 + 2500-req runs all clean — zero IMA, zero deadlock.

Testing

• tests/unittest/kv_cache_manager_v2_tests/test_contiguous_arena.py: 108 tests — allocator/budget/int31 pure logic, sparse VirtMem, arena storage seam, manager-level end-to-end (growth, reuse round-trip with data verification, suspend/resume byte-identical restore, page-exhaustion backoff, page-based capacity planning), lazy-retention matrix, batched-sweep charging, freed-range adoption semantics, blocking-drain deadlock terminal state, host-exhaustion suspension atomicity (HELD-pinned raise + DROPPABLE churn), prefix aliasing (refcount matrix, zero-charge accounting, registry roundtrip, and two manager-level end-to-ends where zero-copy is proven by corrupting every host copy — including the partial-span shape where the owner's chain outlives a shorter adopter's match, and the pressure window where the registry spills between an admission's lookup and its first resume — dead spans fall back to the host-copy path), and write-out ordering (with explicit copy-kernel warmups — the first batchedCopy call pays ~200 ms module load and masks races otherwise), run under both write-through policies via class inheritance. Newest additions: span-spill protection (floor survives full-pressure spill while parked ranges die, excess spills oldest-first, fraction=0 restores unprotected spilling), map-ahead margin degradation (frontier-only mapping at the budget edge, deterministic batched-sweep charge replay, delta-degrade excess release), and owner-live aliasing (owner-charged span lifecycle at the pool-group seam; a manager-level end-to-end where the consumer admits while the owner is mid-generation and physical sharing is proven by writing through the owner's VA and reading the change through the consumer's).
• tests/unittest/_torch/executor/test_kv_cache_v2_contiguous_arena.py: 12 adapter tests (config field/validator, env knobs, real adapter construction, physical capacity accounting, request lifecycle, terminal-state allocation retry, drop-and-recompute fallback).
• tests/unittest/_torch/executor/test_kv_block_offset_overlap_race_v2.py: regression test for the offset-staging race (a).
• v2 regression suite: 219 passed / 12 skipped (211 inherited + 8 added), identical to main before/after every commit.
• api_stability 61/61; existing executor v2 tests 201/201.

Throughput benchmarks — Llama-3.1-8B (H100, solo runs on an idle node)

trtllm-bench throughput, 2500 requests, ISL/OSL 1024/1024, concurrency 256, v2 manager + block reuse + 64 GiB host tier, prefix-sharing dataset (15 shared system prompts, avg 512 tokens). Arena = production defaults (adoption ON, on-demand mapping, on_free, 16 MiB pages, linear kernels off).

Memory-equalized runs (both backends at ~52.5–52.9 GiB KV quota — a total_quota accounting bug that silently granted the arena ~8 GiB more than paged was found and fixed, c645628fcc):

config	tok/s	vs paged
paged v2 (baseline)	8,563	—
arena (adoption)	7,746	−9.5%
arena + prefix aliasing	7,816 / 7,844	−8.6%
arena + aliasing + span-spill protection (+ margin degrade, owner-live, dedup remap)	8,125–8,278 over 7 runs	−4.5% ± 0.5 (paged 8,617–8,631, ±0.05%)

At this quota the shared-prefix workload sits exactly at the arena's capacity edge even WITH aliasing ((3,385 pages − 60 registry)/13 charged pages/seq = 256), so pressure spills thrashed the canonical-span registry: alias hit rate dropped to ~9% and all 15 spans got spilled — aliasing recovered only ~0.9 points. Span-spill protection (463f45cee9) fixes this: the registry is exempt from pressure spills up to protected_span_fraction of the page budget (default 5%; the excess stays spillable LRU-first), and the resume-utilization gate stops counting protected pages as reclaimable. Re-run at the same equalized quota: zero spans spilled, 2,278 alias hits with zero lookup misses, recovering ~4 points (−8.6% → −4.5%). An error-bar note on that row: an interleaved A/B bisect across the night's runs showed the ARENA config wanders ~±0.7% across time windows on an otherwise idle node while paged holds ±0.05%, so the honest pooled number for the full stack is −4.5% ± 0.5 — single run-pairs at this scale cannot resolve sub-point deltas. Within that resolution the three follow-on commits are individually throughput-neutral at this benchmark shape, each by an understood mechanism: best-effort map-ahead margin (812ef73a32) converts capacity-edge growth failures into frontier-only mappings (this workload rarely hits that path); owner-live aliasing (35260d6afc) registers spans at context-end commit rather than close — at saturated concurrency admissions happen exactly at closes, so close-time registration was already timely, and its win shape is staggered/bursty same-prompt arrivals; pressure-time dedup remap (530da2a633) frees duplicate private prefix copies (commit-conflict losers, pre-span copy-onboards) by remapping their VA onto the canonical physical pages, as a reclaim-ladder rung between the parked-range spill and the registry-excess spill — one device quiesce per batch, paid only when the alternative is a preemption; the benchmark's steady state never reaches that rung (every admission aliases, so nothing is duplicated).

Capacity sensitivity (the same runs before the quota fix, i.e. arena inadvertently at 60.8 GiB — retained as the headroom data point):

config	tok/s	vs paged (52.5 GiB)
arena (adoption)	8,051 / 8,058	−6.1%
arena + prefix aliasing	8,471 / 8,461	−1.3% (87% alias hit rate)
arena + aliasing + linear kernels	8,434 / 8,439	−1.6%

Read together: the aliasing mechanism works — with ~8 GiB of headroom the spans stay resident and it recovers the full duplication gap (−1.3%, beating the −1.6% reuse-off bound) — but the current spill policy forfeits it exactly when memory is tight: the registry is only ~1 GB, and spilling it costs ~5 points. Attribution (reuse-off discriminator, identical footprint both modes): the arena-vs-paged gap is dominated by prefix duplication (paged refcount-shares and fits the quota eviction-free; arena's private copies oversubscribe it), with −1.6% of sharing-independent overhead underneath. Follow-ups: span-aware spill protection, map-ahead-margin degradation, and owner-live aliasing have all since landed (463f45cee9, 812ef73a32, 35260d6afc — see the equalized table above). Owner-live registration pins the span at the owner's context-end commit (matching paged, which shares context blocks the moment they commit), so same-prefix admissions alias while the owner is still generating; the entry stays owner-charged (no registry charge, spill-exempt) until the owner's free_sequence flips it to the registry-charged state. Pressure-time dedup remap (530da2a633) has also landed: a stateless scan over live sequences finds committed prefixes held as private duplicates of a registered span and re-backs their VA with the canonical physical pages (same-VA swap — slot ids, offsets, and page holders are untouched, since page objects are slot-indexed), freeing the duplicates' budget charge from LIVE sequences with zero information loss before the reclaim ladder resorts to spilling the registry or preempting.

Aliasing status: prototype, default OFF (ContiguousArenaConfig.prefix_aliasing / TRTLLM_KV_ARENA_PREFIX_ALIASING=1). Correctness-validated at production-shaped load (multiple 2×2500 runs, zero errors) and across a pressure discriminator matrix (adoption-off, tight-quota-only, registry-spill fallback). One rare first-wave IMA on the extreme pressure reproducer (2.4× KV oversubscription + 10× undersized host tier) remains under investigation — a dedicated 84-round multi-GPU grind reproduced it on BOTH arms, with the no-alias arm crashing at 5× the aliasing arm's rate (8/36 vs 2/48; the no-alias crashers fire with an empty registry and zero drop/OutOfPages events), confirming it is not the aliasing machinery but a rare residual of the pre-hardening IMA class (the driver-interference family being escalated separately).

Gap trajectory as the campaign's fixes landed (same workload; node-shared runs up to adoption, solo after; memory-equalized from c645628fcc): pre-map default −12.8% → margin=18 −9.4% → margin=1 −7.4% → adoption −6.8% → equalized adoption −9.5% → + prefix aliasing −8.6% → + span-spill protection (and follow-ons) −4.5% ± 0.5 (−1.3% with 8 GiB of KV headroom).

Throughput benchmarks — Llama-3.2-1B (H100)

5000 requests, ISL/OSL 1024/1024, concurrency 256, v2 manager + block reuse + 8 GiB host tier in all runs. Two datasets: shared (every prompt prefixed by one of 15 shared system prompts, avg 512 tokens — maximizes prefix reuse) and unique (fully unique prompts — isolates intrinsic overhead). Note both are cache-oversubscribed: ~320k committed blocks vs a ~65k-block GPU quota, so classic paged also offloads D2H at steady state (it runs copy-free only for the first ~15% of requests); copy volume is comparable between modes.

dataset	paged v2	arena on_free / 16 MiB	arena on_commit / 16 MiB	arena on_free / 2 MiB
unique	30,908 tok/s	30,785 (−0.4%)	28,147 (−8.9%)	—
shared	33,885 tok/s	31,987 (−5.6%)	30,423 (−10.2%)	28,689 (−15.3%)

Findings:

• Arena reaches parity with paged (−0.4%, noise) on the unique-prompt workload with the default on_free policy and 16 MiB super-pages. Demand paging, VA management, deferred reclaim, and write-out volume are all in the noise at this scale.
• The apparent "intrinsic" −8.9% was entirely the on_commit write-through's per-block dispatch granularity (~320k small migrations vs classic's large batched evictions) — the rate-limiting concern the design flags as risk Update TensorRT-LLM #4. on_free (the default) batches per request. A batched/async write-through dispatcher is the P1 fix before on_commit is recommendable.
• All runs completed 5000/5000 requests — sustained stress of demand paging, deferred reclaim, and the page budget at high concurrency.
• With policy and page size held at their best settings: arena is at parity (−0.4%) without prefix sharing and −4.4% to −5.6% behind with heavy sharing (−4.4% with 256 system prompts and a 64 GiB host tier); 2 MiB paging adds ~10 points of churn at this request rate (16 MiB super-pages recommended); on_commit write-through as currently implemented adds ~5-8 points of per-block dispatch (on_free is the default).

Attributing the residual gap (P2 + ablations)

The branch also implements P2 lazy GPU retention (§4.4 phase 2, opt-in): freed ranges stay mapped on a retained LRU, reuse hits copy D2D from the resident bytes (private copies re-register as fresh sources so hot prefixes survive LRU spills), and the page budget reclaims retained ranges only under pressure — plus the §4.2 batched map sweep (opt-in): growth maps charge the budget at resize time but execute in one batched pass per iteration. (Freed-range adoption, above, grew out of this parking machinery and is now default-ON in all modes.)

Three controlled ablations on the prefix-heavy workload all came back null, which bounds the residual precisely:

ablation	mechanism tested	result
lazy retention	reuse copy transfer (99.7% D2D hit rate measured)	no change
batched sweep	driver-call placement	no change
map-ahead	driver-call count (1 cuMemSetAccess/seq)	no change

Together with py-spy profiles (arena-vs-paged on identical workloads), the residual ~5% is distributed host-side Python bookkeeping across the arena's per-request paths (onboarding, close write-out, reclaim scans) plus associated GPU-wait — no concentrated sink. Both benchmark modes ran the pure-Python v2 package; the arena executes more Python per request, so mypyc compilation is the largest remaining lever, followed by moving the (now single-point) map flush to a helper thread. The sweep defaults off until then. The P1 linear-addressing kernels remain the upside this storage layout exists to enable. Profiles and per-run reports are archived with the work log.

Super-page size curve and follow-up optimizations

Sweeping TRTLLM_KV_ARENA_PHYS_PAGE_SIZE_MB on the prefix-heavy (shared) workload shows page count — not map-call count — scales several per-request costs at once (teardown unmap chunks, budget arithmetic, mapped-range scans):

super-page size	vs paged	mapped/seq (68-block range)	tail waste @ 256 seqs
16 MiB	−5.6%	80 MiB	~3 GB
32 MiB	−3.5%	96 MiB	~7 GB
64 MiB	−3.0%	128 MiB	~15 GB

Recommended default for this workload shape: 32 MiB (16→32 gains 2.1 points; 32→64 only 0.5 more while doubling tail waste). Long-context workloads can go larger, memory-tight deployments smaller. (With adoption ON, steady-state map traffic is zero, so page size matters mainly for spill/teardown and tail waste.)

Follow-ups landed since:

• Per-range mapped high-water frontier (Track per-range mapped frontier in contiguous KV arenas): ensure_mapped previously rescanned each range from its start on every growth call (O(blocks × pages) per sequence of pure-Python is_mapped checks — the inefficiency the page-size curve exposed). Ranges now track a mapped high-water frontier, making growth planning, retained-page counts, and reclaim O(1)/single-unmap. Recovers ~0.8 points at 16 MiB (−5.6% → −4.8% against a same-node paged baseline) without 32 MiB's tail-waste cost, and removes the two scanning helpers outright (net −18 lines).
• mypyc compilation (the biggest lever per profiling): main no longer mypyc-compiles at all; PR [None][fix] Restore KVCacheManagerV2 mypyc compilability #15990 restores compilability (pre-existing issues, separate PR), and this branch carries the arena-side typing fixes. With the package hybrid-compiled (20/22 modules; two excluded for mypyc codegen bugs documented in [None][fix] Restore KVCacheManagerV2 mypyc compilability #15990), both modes speed up (paged +2.3%, arena +2.8%) and the gap narrows to −4.3%. That is a lower bound on full compilation's effect: the single hottest module (_core/_kv_cache.py, resize/commit per token) is one of the two still interpreted.
• P1 linear-addressing kernels (opt-in, TRTLLM_KV_ARENA_LINEAR_KERNELS): XQA reads KV pages by linear arithmetic instead of table loads. E2e: +1.8% at short sequences (ISL/OSL 256/256, the µbench sweet spot), null at 1024/1024 — including under prefix aliasing in the KV-headroom regime (8,436 vs 8,466, −0.35%), so deduplicated, L2-friendlier prefix reads do not change the verdict. Graduation case is short-context serving regimes; default off.

Benchmarking also surfaced and fixed an arena-mode crash (commit "Handle intra-batch commit conflicts"): two in-flight requests committing the same system-prompt block asserted, since seq rebasing is disabled in arena mode; the loser now keeps its own pages as sequence-private committed copies referencing the existing tree block, and commit()'s block loop tolerates mid-loop virtual stops (hardening the classic path as well).

Hard-reclaim redesign + mappability-weighted canonical promotion (P3 v3)

Two commits (c9ac413fae, 7fd533b628) rework pressure-time reclaim and fix a pathology in how the reuse tree tracks residency.

The canonical-residency trap. Once a canonical block is evicted to host and its span pin dies, every later same-prefix admission copy-onboards from host, holds a sequence-private duplicate, and the tree never re-learns that a GPU copy exists — the prompt class copies from host forever. Fix: mappability-weighted promotion at both duplicate-creating seams (commit-conflict and onboarding). If the incumbent canonical is unmappable (no live registered span identity-covers it) and droppable, the GPU-resident challenger becomes the canonical instead of a private duplicate, and it adopts the incumbent's host slot as its pre-valid host backing — a future eviction is copy-free (verified byte-level: a corruption pattern written to the host copy survives the promoted owner's close, proving no D2H copy ran). The first descendant pays the copy once; its context-end commit registers a live span; every later descendant aliases.

Hard reclaim. hard_reclaim_gpu(min_pages) replaces the ladder's three lower rungs: a sync-free exit when nothing is reclaimable anywhere, otherwise ONE device quiesce under which everything harvestable is taken in destruction order — dedup remap of ALL flagged sequences first (duplicates have zero retention value), then parked terminated ranges, then registry excess above the protection floor; pausing stays the caller's fallback. A per-request dirty set (flagged at the duplicate-creating sites, cleared on scan attempt) gates the scan at O(1). The old ladder short-circuited at the first progressing rung, so the dedup rung never fired on any real workload, and the parked-range spill re-quiesced per range; both are gone (a unit test pins the discipline: zero syncs when idle, exactly one per harvest).

Validation (Llama-3.1-8B, H100, solo, interleaved):

shape	config	throughput	dedup remapped pages
capacity edge (free_gpu_memory_fraction=0.7)	pre-v3 arena+aliasing	7763.6 / 7785.9	0 / 0
capacity edge (0.7)	v3 arena+aliasing	7852.9 / 7839.7 / 7885.3 (+0.9..1.4%)	187 / 184 / 150
saturated honest memory (0.85)	paged	8593.1 / 8583.1	—
saturated honest memory (0.85)	v3 arena+aliasing	8210.0 / 8198.9 (−4.47%, = pre-v3 pooled −4.5%±0.5)	31 / 24

Net: wins at the capacity edge where reclaim pressure is real, neutral at saturation; zero IMA / drops / alias misses across all runs. The second commit fixes a probe bug found during validation: the promotion mappability check originally reused the admission-path span lookup, whose verify loop invalidates a span on identity mismatch — probe chains legitimately diverge from the span owner's path past the shared prefix, so live spans were sporadically forfeited (47/2279 admissions missed). The probe is now a non-destructive O(1) identity check.

Mean-batch gap at saturation: root cause and fix (map-ahead margin)

Per-iteration instrumentation of the saturated 0.85 shape (scheduler batch, page-budget composition, and admission-refusal sites, correlated per iteration) decomposed the arena's −4.4% against paged into an admission-capacity story:

• New context requests are admitted through resume(), whose utilization gate (max_util_for_resume, default 0.95) refuses while (used − reclaimable)/total > 0.95. On refusing iterations the gate-visible utilization sits exactly at 0.951 — the system is a regulator pinned at its setpoint, holding the arena at ~234 active sequences while paged (whose evictable blocks reclaim via synchronous free-list ops and whose gate essentially never fires — 566 checks failed in a whole run vs ~197k) runs client-bound at 256.
• Of the 212 MB charged per active sequence at the pin, 187 MB is true KV capacity; the rest is the demand-paging map-ahead margin plus super-page rounding. The margin is charged to the page budget per sequence, so at the utilization ceiling it converts one-for-one into fewer admitted sequences — while its latency-hiding value does not measure, because at super-page granularity a sequence crosses a page boundary only every page_size/bytes_per_token tokens (128 tokens at 16 MiB for an 8B model).
• Steady-phase means (excluding the identical ramp/drain phases, which had masked the effect) on Llama-3.1-8B, H100, 2500×(1024/1024)@conc 256, 0.85 fraction, 64 GiB host, prefix aliasing on:

config	steady mean batch	at 256 cap	throughput
paged	232.9	77.8%	8615–8621
arena+aliasing, margin 1 (old default)	218.5–218.8	13%	8238–8251
arena+aliasing, margin 0	228.3	14%	8383 (+1.7%)
arena+aliasing, margin 0 + max_util_for_resume: 0.98	232.9	77.9%	8473 / 8485 (+2.8%)

This PR flips the map-ahead default to 0 (TRTLLM_KV_ARENA_MAP_AHEAD_PAGES, ContiguousArenaConfig.map_ahead_pages; the knob remains). Margin 0 is already routinely exercised — the mapping planner degrades the margin to 0 under pressure before failing a charge — and the flip also reduces pressure-machinery activity (gate refusals drop 4×, −5.4% iterations to completion, allocation failures unchanged, zero IMA across all runs).

With the margin honest, raising max_util_for_resume to 0.98 in the serving config makes the arena client-bound with the same steady batch and cap-residency as paged, leaving −1.6%, which profiling attributes to per-step host bookkeeping (the mypyc item above), with no capacity component. Note the ordering matters: 0.98 without the margin fix ran the pressure machinery hot (5× dedup harvests, ~1000 iterations with blocking mid-scheduling drains) and measured slower (−1.4%); the gate default is therefore left at 0.95 and 0.98 is a recommended config for arena workloads rather than a default change (the field is shared with classic v2).

Status

Ready for review. Since the last body update: P3 v3 hard-reclaim redesign + mappability-weighted canonical promotion (c9ac413fae + probe fix 7fd533b628), and the saturation mean-batch root cause + map-ahead default flip (a4fa141fe7, sections above; the flip closes the batch gap's dominant term and combined with max_util_for_resume: 0.98 in serving configs reaches batch parity with paged). Remaining follow-ups tracked for later phases: batched on_commit dispatcher (its write-through prior-event is on the manager stream — needs the same forward-stream treatment as 52a311ed99 before default use), helper-thread map flush, partial-block reuse, VSWA/SSM support, disaggregated serving (P2), driver-team escalation of the unmap-vs-in-flight-kernels fault (repro matrix ready).

Summary by CodeRabbit

• New Features

  • Added experimental contiguous KV arena mode for KV cache management.
  • Introduced new arena-aware capacity, mapping, reclaim, and write-through behaviors.
  • Added public configuration options and APIs for arena-backed cache handling.

• Bug Fixes

  • Improved handling of cache lifecycle operations in arena mode, including resume, suspend, close, and capacity growth.
  • Tightened validation to prevent unsupported configuration combinations when arena mode is enabled.

• Tests

  • Added extensive CUDA unit coverage for contiguous arena allocation, mapping, reclamation, and end-to-end cache workflows.

[thorjohnsen]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces an experimental "contiguous KV cache arena" mode that reserves per-sequence contiguous virtual-address block ranges and maps physical pages on demand via a shared page budget, replacing scattered GPU slot allocation. Changes span configuration, CUDA virtual memory, storage layers, StorageManager/KVCacheManager APIs, per-sequence lifecycle logic, pyexecutor wiring, and new tests.

Changes

Contiguous KV Arena Mode

Layer / File(s)	Summary
Configuration and public exports tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/runtime/kv_cache_manager_v2/_config.py, tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py	Adds use_contiguous_kv_arena (forces use_kv_cache_manager_v2), ContiguousArenaConfig/WriteThroughPolicy, KVCacheManagerConfig.contiguous_arena with post-init compatibility checks, and re-exports.
Sparse VirtMem mapping tensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.py	Adds _page_map-backed sparse chunk mapping (map_range/unmap_range/is_mapped/num_sparse_chunks) exclusive of tail-stack mode, with teardown cleanup.
SequenceArena allocator tensorrt_llm/runtime/kv_cache_manager_v2/_sequence_arena.py, tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py	New PageBudget, BlockRangeAllocator, and SequenceArena implementing per-pool-group reservation, demand mapping, and deferred/immediate reclaim; module added to mypyc compile list.
Private committed page model tensorrt_llm/runtime/kv_cache_manager_v2/_page.py	Adds convert_to_private_committed and PrivateCommittedPage for sequence-private reuse copies, and refines eviction-exclusion logic in _PageHolder.__del__.
Arena pool group and cache-level storage tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py	Adds ArenaSlotPool, SequenceRange, ArenaPoolGroup (with reserve/take/ensure_mapped/queue_mapping/drain_reclaim/ghost registry), and GpuArenaCacheLevelStorage.
StorageManager arena API tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py	Adds arena-mode public API (reserve/take/ensure/free/drain/spill/queue/flush, ghosts, onboarding, offload, write-through), guards against scattered GPU allocation/eviction/rebalancing, and updates utilization/eviction logic.
KVCacheManager wiring tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py	Forwards arena_config, validates/threads max_capacity, disables rebasing/partial-match in arena mode, adds drain_gpu_reclaim/flush_gpu_mappings, and updates clamp_max_seq_len_for_mem for page-budget accounting.
_KVCache arena lifecycle tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py	Reworks close/suspend/resume/resize/_commit_block for arena mode: reserved-range mapping, write-through, onboarding, and private-committed-copy conflict handling.
pyexecutor integration tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py	Adds _arena_enabled detection, _build_arena_config() from environment variables, and arena-mode overrides across buffer access, free-block accounting, resource preparation, dummy requests, invalid-value checks, and shutdown.
Tests tests/unittest/_torch/executor/test_kv_cache_v2_contiguous_arena.py, tests/unittest/kv_cache_manager_v2_tests/test_contiguous_arena.py	New unit and integration test modules covering config wiring, adapter lifecycle, arena primitives, storage layers, and end-to-end KV manager behavior under arena mode.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant KVCacheManager
  participant KVCache as _KVCache
  participant StorageManager
  participant ArenaPoolGroup
  participant SequenceArena
  KVCacheManager->>KVCache: create_kv_cache(max_capacity)
  KVCache->>StorageManager: reserve_gpu_sequence(pg_idx)
  StorageManager->>ArenaPoolGroup: reserve_sequence(range)
  ArenaPoolGroup->>SequenceArena: allocate(range)
  KVCache->>StorageManager: ensure_gpu_mapped(range)
  StorageManager->>ArenaPoolGroup: ensure_mapped(range)
  ArenaPoolGroup->>SequenceArena: map physical pages
  KVCache->>KVCache: close()/suspend()
  KVCache->>StorageManager: free_gpu_sequence(range, last_consumer)
  StorageManager->>ArenaPoolGroup: free_sequence(range)
  ArenaPoolGroup->>SequenceArena: reclaim(range)

Loading

Suggested reviewers: joyang-nv, bo-nv, xinhe-nv, JunyiXu-nv, pengbowang-nv, yechank-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the main feature and follows the repository's ticket/type pattern.
Description check	✅ Passed	The description covers the summary, design, and testing well, with only the PR Checklist section left incomplete.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

tests/unittest/kv_cache_manager_v2_tests/test_contiguous_arena.py (1)

968-974: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consolidate duplicated _block_floats/_host_floats helpers.

_block_floats is redefined verbatim in TestArenaKVCacheManagerEndToEnd and TestArenaLazyRetention, and _host_floats duplicates the ctypes-buffer pattern already used inline in TestStorageManagerArenaMode._slot_floats. Extracting these into a shared module-level helper (or a small mixin) avoids drift if the addressing logic ever changes.

♻️ Proposed consolidation sketch

+def _gpu_block_floats(storage, page_index: int, n_floats: int) -> "torch.Tensor":
+    addr = int(storage.slot_address(GPU_LEVEL, PoolGroupIndex(0), page_index, PoolIndex(0)))
+    return TestSparseVirtMem._tensor_at(addr, n_floats)
+
+
+def _host_slot_floats(storage, slot_id: int, n_floats: int) -> "torch.Tensor":
+    import ctypes
+
+    addr = int(storage.slot_address(CacheLevel(1), PoolGroupIndex(0), slot_id, PoolIndex(0)))
+    buf = (ctypes.c_float * n_floats).from_address(addr)
+    return torch.frombuffer(buf, dtype=torch.float32)

Then have TestArenaKVCacheManagerEndToEnd._block_floats and TestArenaLazyRetention._block_floats/_host_floats delegate to these.

Also applies to: 1339-1357

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/kv_cache_manager_v2_tests/test_contiguous_arena.py` around
lines 968 - 974, Consolidate the duplicated `_block_floats` and `_host_floats`
logic used in `TestArenaKVCacheManagerEndToEnd`, `TestArenaLazyRetention`, and
`TestStorageManagerArenaMode` by extracting shared module-level helper(s) or a
small mixin. Keep the addressing/buffer creation behavior in one place so
`_block_floats` delegates to the shared helper and `_host_floats` reuses the
same ctypes-buffer pattern instead of reimplementing it inline. Update the test
methods to call the shared helper(s) and remove the verbatim duplicate
implementations.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/executor/test_kv_cache_v2_contiguous_arena.py`:
- Around line 103-152: Coverage for the arena adapter path is incomplete: add
tests in TestArenaAdapter that directly exercise get_buffers() raising
NotImplementedError when _arena_enabled, and check_invalid_values_in_kv_cache()
returning False without scanning in arena mode. Also add a test around
prepare_resources(), add_dummy_requests(), or extend_capacity_for_tokens() that
spies on self.mgr.impl to verify flush_gpu_mappings() and drain_gpu_reclaim()
are invoked. Finally, replace the manual kv_cache_map/index_mapper cleanup in
test_request_lifecycle_through_adapter with the public free_resources() path (or
add a separate test) so the full teardown flow and close() behavior are
validated through the production route.
- Around line 96-100: The defaults test for KVCacheManagerV2._build_arena_config
is not actually isolating the TRTLLM_KV_ARENA_* environment, so it can inherit
ambient values and become flaky. Update test_env_defaults in
test_kv_cache_v2_contiguous_arena to explicitly clear or override the relevant
arena env vars before calling KVCacheManagerV2._build_arena_config, rather than
using a no-op os.environ patch. Keep the assertions on phys_page_size and
write_through as the check for true default behavior.

---

Nitpick comments:
In `@tests/unittest/kv_cache_manager_v2_tests/test_contiguous_arena.py`:
- Around line 968-974: Consolidate the duplicated `_block_floats` and
`_host_floats` logic used in `TestArenaKVCacheManagerEndToEnd`,
`TestArenaLazyRetention`, and `TestStorageManagerArenaMode` by extracting shared
module-level helper(s) or a small mixin. Keep the addressing/buffer creation
behavior in one place so `_block_floats` delegates to the shared helper and
`_host_floats` reuses the same ctypes-buffer pattern instead of reimplementing
it inline. Update the test methods to call the shared helper(s) and remove the
verbatim duplicate implementations.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 0d56f573-44ed-4c74-add5-a96010929eb4

📥 Commits

Reviewing files that changed from the base of the PR and between b1ea26c and aa6d2eb.

📒 Files selected for processing (14)

• tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_config.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_cuda_virt_mem.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_sequence_arena.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py
• tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py
• tests/unittest/_torch/executor/test_kv_cache_v2_contiguous_arena.py
• tests/unittest/kv_cache_manager_v2_tests/test_contiguous_arena.py

[coderabbitai]

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Isolate arena env vars explicitly instead of relying on ambient environment.

mock.patch.dict("os.environ", {}, clear=False) is a no-op for isolation — it neither clears nor guarantees absence of TRTLLM_KV_ARENA_* vars that might already be set in the runner's environment, so this "defaults" test can silently validate the wrong state or flake depending on the ambient shell/CI environment.

🧪 Proposed fix: explicitly clear the relevant keys

     def test_env_defaults(self) -> None:
-        with mock.patch.dict("os.environ", {}, clear=False):
+        keys = (
+            "TRTLLM_KV_ARENA_PHYS_PAGE_SIZE_MB",
+            "TRTLLM_KV_ARENA_MAP_AHEAD_PAGES",
+            "TRTLLM_KV_ARENA_WRITE_THROUGH",
+            "TRTLLM_KV_ARENA_LAZY_RETENTION",
+            "TRTLLM_KV_ARENA_BATCHED_MAP",
+        )
+        with mock.patch.dict("os.environ", {}, clear=False):
+            for k in keys:
+                os.environ.pop(k, None)
             arena_cfg = KVCacheManagerV2._build_arena_config()

As per path instructions for tests/**, coverage/robustness of environment-dependent tests should be verified explicitly rather than relying on ambient state.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	def test_env_defaults(self) -> None:
	with mock.patch.dict("os.environ", {}, clear=False):
	arena_cfg = KVCacheManagerV2._build_arena_config()
	self.assertEqual(arena_cfg.phys_page_size, 2 * MiB)
	self.assertEqual(arena_cfg.write_through, WriteThroughPolicy.ON_FREE)
	def test_env_defaults(self) -> None:
	keys = (
	"TRTLLM_KV_ARENA_PHYS_PAGE_SIZE_MB",
	"TRTLLM_KV_ARENA_MAP_AHEAD_PAGES",
	"TRTLLM_KV_ARENA_WRITE_THROUGH",
	"TRTLLM_KV_ARENA_LAZY_RETENTION",
	"TRTLLM_KV_ARENA_BATCHED_MAP",
	)
	with mock.patch.dict("os.environ", {}, clear=False):
	for k in keys:
	os.environ.pop(k, None)
	arena_cfg = KVCacheManagerV2._build_arena_config()
	self.assertEqual(arena_cfg.phys_page_size, 2 * MiB)
	self.assertEqual(arena_cfg.write_through, WriteThroughPolicy.ON_FREE)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/executor/test_kv_cache_v2_contiguous_arena.py` around
lines 96 - 100, The defaults test for KVCacheManagerV2._build_arena_config is
not actually isolating the TRTLLM_KV_ARENA_* environment, so it can inherit
ambient values and become flaky. Update test_env_defaults in
test_kv_cache_v2_contiguous_arena to explicitly clear or override the relevant
arena env vars before calling KVCacheManagerV2._build_arena_config, rather than
using a no-op os.environ patch. Keep the assertions on phys_page_size and
write_through as the check for true default behavior.

Source: Path instructions

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Coverage gap: adapter-level arena integration points untested here.

This file is described as covering "KVCacheManagerV2 adapter plumbing," but the new adapter-level arena behavior isn't directly exercised:

• get_buffers() raising NotImplementedError in arena mode.
• check_invalid_values_in_kv_cache() returning False and skipping the scan.
• The flush_gpu_mappings()/drain_gpu_reclaim() calls newly added to prepare_resources(), add_dummy_requests(), and extend_capacity_for_tokens().

Additionally, test_request_lifecycle_through_adapter tears down via manual kv_cache_map.pop/index_mapper.remove_sequence instead of the public free_resources() path, so the arena close() behavior isn't validated through the same route production code uses.

Suggest adding to test_kv_cache_v2_contiguous_arena.py:

• A test asserting get_buffers() raises NotImplementedError when _arena_enabled.
• A test asserting check_invalid_values_in_kv_cache() returns False without raising in arena mode.
• A test that drives a request through prepare_resources()/add_dummy_requests() (mirroring _KVCache growth) to confirm flush_gpu_mappings/drain_gpu_reclaim are actually invoked (e.g. via mock/spy on self.impl).
• Replacing the manual teardown in test_request_lifecycle_through_adapter with mgr.free_resources(req) (or adding a second test that does), to validate the full production teardown path in arena mode.

As per path instructions, "Act as a QA engineer reviewing test changes and coverage... suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR": coverage here is insufficient for the adapter-level integration points that are the stated purpose of this file.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/executor/test_kv_cache_v2_contiguous_arena.py` around
lines 103 - 152, Coverage for the arena adapter path is incomplete: add tests in
TestArenaAdapter that directly exercise get_buffers() raising
NotImplementedError when _arena_enabled, and check_invalid_values_in_kv_cache()
returning False without scanning in arena mode. Also add a test around
prepare_resources(), add_dummy_requests(), or extend_capacity_for_tokens() that
spies on self.mgr.impl to verify flush_gpu_mappings() and drain_gpu_reclaim()
are invoked. Finally, replace the manual kv_cache_map/index_mapper cleanup in
test_request_lifecycle_through_adapter with the public free_resources() path (or
add a separate test) so the full teardown flow and close() behavior are
validated through the production route.

Source: Path instructions

[nvpohanh]

[by Codex] @lowsfer Could you review this PR? Thanks!

[nvpohanh]

[by Codex] @lowsfer Friendly reminder to review this PR as the primary KV-cache-manager reviewer when you have a chance. Thanks!