Conversation
mergify
Bot
added
frontend
cpu
njhill
force-pushed
the
extensible-kvcache
branch
from
September 11, 2026 20:04
6a0dfc5 to
be4ca65
Compare
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #88410 for commit |
njhill
force-pushed
the
extensible-kvcache
branch
from
September 11, 2026 21:29
be4ca65 to
038f6d5
Compare
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #88434 for commit |
njhill
force-pushed
the
extensible-kvcache
branch
from
September 12, 2026 03:40
038f6d5 to
1fcb229
Compare
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #88465 for commit |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
njhill
force-pushed
the
extensible-kvcache
branch
from
September 12, 2026 21:04
1fcb229 to
0f24e2d
Compare
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #88567 for commit |
njhill
force-pushed
the
extensible-kvcache
branch
from
September 13, 2026 16:25
0f24e2d to
01bd8d0
Compare
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #88693 for commit |
njhill
force-pushed
the
extensible-kvcache
branch
from
September 14, 2026 21:50
01bd8d0 to
e7a1c2d
Compare
Member
Author
|
/ci run |
|
❌ This PR is 4 commits behind upstream |
njhill
force-pushed
the
extensible-kvcache
branch
from
September 15, 2026 00:17
e7a1c2d to
7cf8ad6
Compare
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #88959 for commit |
njhill
force-pushed
the
extensible-kvcache
branch
from
September 15, 2026 03:12
7cf8ad6 to
6a4ef96
Compare
|
✅ Buildkite CI #89994 failed during CI setup, so its test failure set is incomplete. Use |
Member
Author
|
/ci run all |
|
❌ This PR is 1 commit behind upstream |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
njhill
force-pushed
the
extensible-kvcache
branch
from
September 25, 2026 16:17
0f3d14d to
1e5b151
Compare
njhill
force-pushed
the
extensible-kvcache
branch
from
September 25, 2026 22:45
1e5b151 to
cc130c6
Compare
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #91337 for commit |
Port the library and configuration half of the extensible KV cache demo (#47363). ExtensibleTensor reserves a device virtual address range with CUDA VMM and commits physical pages incrementally, so a buffer can grow without moving its base pointer. The reservation is divided into equal segments that grow in lockstep, which backs KV cache layouts whose block dimension is not outermost, and `segment_view` exposes each committed prefix as a tensor whose storage covers exactly the backed bytes. `--enable-extensible-kv-cache` is plumbed through CacheConfig, EngineArgs and LLM but is not yet consumed; the V2 model runner integration follows. Co-authored-by: Zhuohan Li <zhuohan123@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
Wire the extensible KV cache into the V2 model runner so the cache is sized once, from the memory actually free after CUDA graph capture and kernel warmup, instead of from a profiling estimate that has to guess the capture footprint with a throwaway capture pass. Flow with `--enable-extensible-kv-cache`: - The engine core generates the KV cache config from the profiled budget as before; its `num_blocks` is now the reserved capacity. - `allocate_kv_cache` takes an allocator hook. The V2 runner passes `ExtensibleKVCache.allocate`, which reserves the whole backing allocation and returns a view over it, so layer views and kernel launches are unchanged. The single allocation splits into `size // (num_blocks * block_stride)` equal segments, each committed as a block prefix: one per layer for layer-outermost layouts, one overall for block-outermost. - Only block 0 is committed before capture (dummy block tables are zero). Warmup commits the prefix it addresses via `ensure_kv_cache_blocks`, and `set_dummy_context` wraps block IDs at the committed count. - After warmup each worker reports how many blocks fit in the memory left within its utilization budget (`measure_kv_cache_blocks`), less a margin and the granule-rounding bound. The engine core takes the minimum across workers, shrinks the configs' `num_blocks` (tensors keep describing the capacity), re-checks that `max_model_len` fits, and issues one `extend_kv_cache` RPC that commits the rest. Captured graphs stay valid because no address moves. - `profile_cudagraph_memory` is skipped on this path. Unsupported combinations (V1 runner, non-CUDA, kv_cache_memory_bytes, KV connectors) are rejected at config time; connectors follow separately. Validated on an RTX A1000 with Qwen3-0.6B: single capture pass, measured sizing (986 blocks vs. 1032 from the estimate at a 150 MiB margin), greedy outputs identical to the baseline in eager mode and for three of four prompts with cudagraphs; the fourth prompt is a near-tie that also differs between multiprocess and in-process baselines. Co-authored-by: Zhuohan Li <zhuohan123@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
Move the ctypes bindings behind a `VmmDriver` interface with CUDA (`cuMem*`) and HIP (`hipMem*`) implementations; the call signatures and struct layouts are identical, so `ExtensibleTensor` is driver-agnostic and the DLPack device type follows the driver. The HIP driver refuses runtimes older than ROCm 7.12, which reject `hipMemSetAccess` once a reservation holds several allocations (rocm-systems#2516), and access is granted per contiguous mapped run since ROCm rejects a range starting inside one. Other pieces carried over from the earlier prototype: a 2 MiB floor on the commit granule (ROCm reports 4 KiB), one `empty_cache` retry when a physical allocation fails because torch's caching allocator is holding idle memory, and `release_physical` so a later sleep-mode integration can drop pages while keeping the reservation and every view over it. `vmm_unavailable_reason` probes the driver (library load, entry points, a reserve/free round trip, known-defective runtimes). The engine core asks every worker before memory profiling and, if any cannot back the cache, logs a warning and disables the feature on all ranks so sizing falls back to profiling estimates. WSL2 and other VMM-less drivers take this path. Co-authored-by: Zhuohan Li <zhuohan123@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
The extensible KV cache is mapped by the driver, not the CuMem allocator pool, so sleep has to release it explicitly. `Worker.sleep` drops its physical pages before suspending the allocator, keeping the reservation so layer views and captured graphs stay valid; `Worker.wake_up` recommits zeroed pages for the same block count when the `kv_cache` tag wakes. Validated on an RTX A1000 with Qwen3-0.6B: sleep(level=1) frees 3.01 GiB including the 1.74 GiB KV cache, and greedy outputs after wake are identical. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
KV connectors register the KV cache memory and derive its extent from `untyped_storage().nbytes()`, which under a reservation would span uncommitted capacity. Rather than teach every connector about committed spans, make the storages honest: - The connector is created in `extend_kv_cache`, after the final commit, so it only ever sees the final block count and backed memory. - `ExtensibleKVCache.committed_views` rebuilds each layer's view over its segment's committed prefix (`ExtensibleTensor.segment_view`): same addresses, so captured graphs stay valid, but storage size equals the committed bytes. After the extend the runner is indistinguishable from one that allocated the final size up front, including its config's `num_blocks`. This needs at most one segment per layer, i.e. a block-compact layout (L or B outermost); the engine core narrows the backends' supported layouts to those when a connector is configured, and every backend declares at least one. - `commit(defragment=True)` remaps each segment as a single driver allocation before registration: UCX under NIXL cannot transfer a range that spans several `cuMemCreate` handles. Warmup contents are discarded. Validated on an RTX A1000 with Qwen3-0.6B and the ExampleConnector (shared storage): external cache hits on the second pass, outputs identical to the non-extensible run. NIXL and Mooncake need multi-GPU validation. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
… memory Make the extensible KV cache the default wherever it is supported (CUDA and ROCm with the V2 model runner, no manual kv_cache_memory_bytes). Left unset, `enable_extensible_kv_cache` resolves to that; requested explicitly where unsupported, it is an error. The driver probe still falls back to estimate-based sizing at startup where VMM is unavailable. The unused fraction of the device that `gpu_memory_utilization=0.92` reserves exists to cover what the profiling estimate misses. With the cache sized from memory measured after warmup that error is gone, so when the utilization is left unset it now defaults to 1.0 alongside the extensible cache (0.92 otherwise, and always restored on fallback). The utilization remains a cap on the fraction of total memory for co-tenancy; at 1.0 it can exceed what is free at startup, which only over-reserves address space, so the startup check tolerates that on the extensible path. What the measurement cannot see is kept as an explicit margin that scales with the workload instead of the device: max(256 MiB, 2% of the KV budget), plus the allocation-granule rounding bound. Both constants are provisional pending a sweep on long-running workloads. On a 6 GiB device (Qwen3-0.6B) the default now yields 3.64 GiB of KV cache where 0.92 gave 3.58 GiB, the floor dominating; on an 80 GiB device the same rule returns roughly 5 GiB of the former headroom to the cache. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
…bytes and scale the margin with it The headroom kept free after measured sizing had two parts: the transient peak (the larger of the profiled activation peak and the warmup steps' allocated peak), kept free in full, and a margin of 2% of the device headroom for allocations warmup did not exercise. On large devices the margin dominated (3.6 GiB of a 186 GiB GB200) while being unrelated to anything the model does. Measure the warmup transient as the allocator's reserved peak less the live allocations, so it includes the fragmentation and rounding the caching allocator actually needed from the device, and size the margin as a share of that measured peak (25%, floor 256 MiB) rather than of the device. Everything kept free is now either measured or a fraction of a measurement. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
Prompt logprobs materialize full-vocab logits in chunks of 1024 rows at runtime. Nothing else in warmup allocates that shape, so with measured KV cache sizing the first prompt-logprobs request took it out of the memory kept free after sizing; on a MIG slice that surfaced as CUBLAS failures late in evaluations once the allocator had fragmented around it. Run one chunk through the logits head and the top-k scorer during warmup, gated on the prompt-logprobs worker existing (last PP rank, generative models), logprobs being enabled, and the model exposing a logits head. The chunk width moves to a shared constant so warmup and runtime cannot drift. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
…s reserved peak free Measured sizing kept free the larger of the profiled activation peak and the warmup steps' peak. Validation on GB200 showed that was not enough: Qwen3-8B ran out of memory on its very first prefill, inside the compiled MLP, with the margin already spent, and Qwen3-0.6B came within 20 MiB of doing so. Two gaps combined. The profiling run is the only startup step at the full token budget, but it skips attention and is measured on allocated bytes, so neither attention workspaces nor the caching allocator's rounding and fragmentation are in the peak. And warmup never ran a full-budget step: the mixed prefill+decode warmup was only invoked by the FlashInfer sparse-MLA path, so the largest prefill the scheduler can issue was never exercised. Under the extensible cache, run the mixed prefill+decode warmup at max_num_batched_tokens after the kernel warmup so that step's reserved peak is measured, and take the profiling peak as the allocator's reserved peak less live allocations as well. With both, the measured transient matched the first real prefill on every model probed. Real steps still combine shapes warmup exercises separately (a prompt-logprobs chunk alongside a live decode, the rejection sampler over a full spec-decode batch) and were measured to allocate up to a quarter more than the warmup peak, so the margin becomes half the measured transient (floor unchanged at 256 MiB). Across Qwen3-0.6B, Qwen3-8B at TP1 and TP2, and Qwen3-1.7B with an EAGLE3 drafter, the true allocated need of an adversarial workload (long random prompts, sampling with logprobs, prompt logprobs) was 40-60% of the headroom kept, and the KV cache still ends up 7-8% larger than with the previous 0.92 utilization default. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
- ReplaySSM autotune bounds its dummy slots by the committable block count and commits them first; it ran before warmup committed anything, so its dummy run and cleanup wrote unmapped addresses. - A commit that fails part-way through mapping several runs no longer leaves earlier runs mapped without access: access is granted per run as it is mapped. - DLPack views no longer depend on module globals at interpreter exit: the DLManagedTensor is malloc'd for the process lifetime with a NULL deleter, and `_VirtualBuffer.__del__` binds `suppress` as a default argument. - The worker settles the DCP interleave size before the runner reads it in `initialize_kv_cache`; under the extensible cache the connector-time adjustment is deferred to `extend_kv_cache`. - When fewer blocks fit than warmup committed, the final commit remaps the smaller prefix instead of silently keeping the larger one, so the runner, scheduler and connector agree on the size. - `num_gpu_blocks_override` is treated like `kv_cache_memory_bytes`: the feature is off by default when it is set and an explicit request errors. - The external launcher reduces the enable/disable decision across ranks so no rank waits in the block-count all-reduce for one that dropped out. - Rebinding the KV cache keeps the ReplaySSM ring trackers when they still cover it: the extensible cache rebinds after CUDA graph capture, and the graphs address the tracker tensors (found on GB200 with Nemotron-H + FlashInfer ReplaySSM: illegal memory access on the first replay). Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ranks Warmup sizes its batches and decides whether the mixed prefill+decode step runs from the blocks this rank can commit, read from this rank's free memory. Ranks with different free memory would warm up different shapes, and a rank skipping a step the others take part in (the sparse-MLA autotune broadcasts afterwards) would leave them waiting. Reduce the count to the minimum over the world group before any caller uses it. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ors and exact fit sizing A KV connector registers each committed layer segment for RDMA, which cannot span several physical chunks. The defragmenting commit remapped segments sequentially, so a granule straddling a segment boundary was mapped by the previous segment's chunk whenever the capacity in bytes per segment was not a granule multiple. With a connector configured, the engine now rounds the reserved capacity and the final measured block count down to a count at which every segment starts and ends on a granule boundary, using the commit granule reported by the workers. Sizing also stops subtracting a per-segment rounding bound (segments x 2 MiB, which reached GiB for head-major layouts): the worker now measures available bytes and asks the cache for the most blocks whose rounded-up per-segment footprint fits. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
- Engine-args tests that exercise the feature skip off CUDA, where the platform check trips first. - The VMM probe does one create/map/set-access/unmap/release round trip of a granule, not just reserve/free. - Post-sizing "not enough KV cache memory" errors stop suggesting a higher `gpu_memory_utilization` under the extensible cache, where it is already 1.0, and point at max_model_len / max_num_seqs instead. - The CUDA driver calls cuInit and prefers the primary context of the requested device, so the probe works before torch touches the device. - The fallback driver-library lookup matches the mapped file name, not a substring: libcudart no longer stands in for libcuda. - Comments: hipMemAllocationProp layout, elastic-EP branch in the final sizing, usage-stats utilization value. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…engine, not by all-reduce Replace the worker-side all-reduce of the committable block count with the engine taking the minimum over workers. Workers return the count from `initialize_from_config`, the RPC that creates the cache, and the engine passes the minimum back as an argument to `compile_or_warm_up_model`, which warmup then uses as its bound. No new RPC and no collective inside warmup code. The external launcher reduces the value across its ranks the way it already does for memory and the final block count. Fold the commit-granule query into the driver probe the same way: the probe now returns the unsupported reason and the granule together, so the separate `kv_cache_commit_granule` RPC goes away. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…k agreement - The worker computes the committable block count once when it creates the cache and stores it as a plain attribute that warmup reads; the engine's minimum over ranks overwrites it. This drops the cap branch layered on the dynamic formula, and callers that warm up without the engine's value (elastic EP) keep the rank's own count. - Undo folding the commit granule into the driver probe: the separate `kv_cache_commit_granule` RPC is back, and the probe result type and its cross-worker merging are gone. Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
njhill
force-pushed
the
extensible-kvcache
branch
from
September 26, 2026 02:32
cc130c6 to
6ad605d
Compare
Member
Author
|
/ci retry |
|
✅ Triggered Buildkite CI #91369 for commit |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Productization of the demonstration #47363 from @zhuohan123.
WIP, replaces #50779