forked from Karylab-cklius/vllm
Zhuohan Li
0116e1cedc
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.