Port of internal D110967544. The extensible KV cache flow (reserve KV
virtual address space up front, capture CUDA graphs first, then size and
commit the KV cache from post-capture free memory) previously required a
block-major attention backend and rejected Mamba models. This enables it
for every backend layout and for Mamba / linear attention:
- ExtensibleTensor gains num_segments: the reservation is divided into
equal segments that grow in lockstep, with committed bytes forming a
prefix of each segment. Physical pages are mapped at
allocation-granularity granules and deduped across overlapping ranges,
so a granule straddling a segment boundary is mapped exactly once.
resize_per_segment_(bytes, zero_new=True) zeroes only the newly
committed logical range of each segment.
- Each KV cache buffer keeps its layers' physical layout and is committed
as one prefix per layout segment. The segment count is derived from the
backend's get_kv_cache_shape / get_kv_cache_block_dim / stride order:
K/V-split layouts (e.g. FlashAttention) get one prefix per half,
block-major layouts (e.g. FlashInfer, MLA) a single prefix. Mamba state
pages are block-major per layer, and hybrid-model attention caches are
re-strided to block-major, so both use a single segment.
- Removed the supports_extensible_kv_cache gate plumbing from EngineCore,
Executor, Worker, WorkerBase and GPUModelRunner; a CUDA platform check
remains in EngineCore.
- enable_extensible_kv_cache is reported as unsupported by the V2 model
runner so V2-default models fall back to the V1 runner (which implements
the flow); also fixed initialize_kv_cache being called with the
extensible kwarg on runners that do not accept it, which broke every
default V2-runner boot on this branch.
Tested on H100:
- tests/utils_/test_extensible_tensor.py (5 passed, incl. new segmented
lockstep-grow/zero, granule-dedup and invalid-usage tests)
- tests/v1/worker/test_extensible_kv_cache.py (new, 6 passed: segment
derivation, split grows both halves, block-major, legacy full commit,
Mamba per-layer growth, hybrid attention+Mamba)
- E2E Qwen3-0.6B greedy with VLLM_ATTENTION_BACKEND=FLASH_ATTN (a K/V-split
backend the old gate rejected): extensible generations byte-identical to
the legacy path; log shows reserve then "Extended KV cache to 34663
blocks". V2->V1 auto-fallback path verified as well.
Zhuohan Li
0116e1cedc