diff --git a/tests/utils_/test_extensible_tensor.py b/tests/utils_/test_extensible_tensor.py new file mode 100644 index 000000000000..3bce9fc219e3 --- /dev/null +++ b/tests/utils_/test_extensible_tensor.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.utils.extensible_tensor import ExtensibleTensor + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def test_extensible_tensor_grows_without_moving() -> None: + buffer = ExtensibleTensor(4096, device="cuda") + try: + base_ptr = buffer.base_ptr + first_view = buffer.resize_(1024) + assert first_view.data_ptr() == base_ptr + first_view.fill_(7) + + second_view = buffer.resize_(2048) + assert second_view.data_ptr() == base_ptr + assert torch.equal(second_view[:1024], torch.full_like(second_view[:1024], 7)) + + second_view[1024:].fill_(3) + assert torch.equal(buffer.tensor, second_view) + + full_view = buffer.full_view() + assert full_view.data_ptr() == base_ptr + assert full_view.numel() == 4096 + finally: + buffer.free() + + +def test_extensible_tensor_rejects_shrink_and_overflow() -> None: + buffer = ExtensibleTensor(1024, device="cuda") + try: + buffer.resize_(512) + with pytest.raises(ValueError, match="grow-only"): + buffer.resize_(256) + with pytest.raises(ValueError, match="exceeds the segment capacity"): + buffer.resize_(1025) + finally: + buffer.free() + + +def test_segments_grow_in_lockstep_and_zero_new() -> None: + """Each segment's committed prefix grows in lockstep. + + Data written to a segment's committed prefix survives a grow; the newly + committed range of each segment is zeroed with `zero_new=True` while old + bytes are preserved. + """ + et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2) + try: + assert et.num_segments == 2 + assert et.segment_capacity_bytes == 4096 + + et.resize_per_segment_(256, zero_new=True) + assert et.bytes_per_segment == 256 + assert et.num_bytes == 512 + fv = et.full_view() + assert fv.shape == (8192,) + # Committed prefixes start zeroed. + assert torch.count_nonzero(fv[:256]) == 0 + assert torch.count_nonzero(fv[4096 : 4096 + 256]) == 0 + + pattern_a = torch.arange(256, device="cuda", dtype=torch.uint8) + pattern_b = 255 - pattern_a + fv[:256].copy_(pattern_a) + fv[4096 : 4096 + 256].copy_(pattern_b) + + et.resize_per_segment_(1024, zero_new=True) + fv2 = et.full_view() + assert fv2.data_ptr() == fv.data_ptr() + # Old bytes of both segments preserved; freshly committed ranges zeroed. + assert torch.equal(fv2[:256], pattern_a) + assert torch.equal(fv2[4096 : 4096 + 256], pattern_b) + assert torch.count_nonzero(fv2[256:1024]) == 0 + assert torch.count_nonzero(fv2[4096 + 256 : 4096 + 1024]) == 0 + finally: + et.free() + + +def test_segments_at_granularity_scale() -> None: + """Segments spanning multiple mapping granules commit correctly. + + Uses a segment capacity that is not a multiple of the allocation + granularity, so a granule straddles the segment boundary and is shared by + the first commit of one segment and a later commit of the other -- it must + be mapped exactly once. + """ + probe = ExtensibleTensor(max_num_bytes=1, device="cuda") + granularity = probe.capacity_bytes + probe.free() + # Two segments of 1.5 granules each; the middle granule straddles the + # boundary. + max_num_bytes = 3 * granularity + et = ExtensibleTensor(max_num_bytes=max_num_bytes, device="cuda", num_segments=2) + try: + seg = et.segment_capacity_bytes + assert seg == max_num_bytes // 2 + + step = granularity // 2 + et.resize_per_segment_(step, zero_new=True) + fv = et.full_view() + fv[:step].fill_(1) + fv[seg : seg + step].fill_(2) + + # Grow to the full segment capacity: previously mapped granules + # (including the boundary-straddling one) are reused, new ones are + # committed and zeroed. + et.resize_per_segment_(seg, zero_new=True) + fv2 = et.full_view() + assert torch.all(fv2[:step] == 1) + assert torch.all(fv2[seg : seg + step] == 2) + assert torch.count_nonzero(fv2[step:seg]) == 0 + assert torch.count_nonzero(fv2[seg + step :]) == 0 + finally: + et.free() + + +def test_multi_segment_invalid_usage_raises() -> None: + """Prefix-view APIs and invalid segment configs raise for multi-segment + buffers.""" + with pytest.raises(ValueError): + ExtensibleTensor(max_num_bytes=100, device="cuda", num_segments=3) + + et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2) + try: + with pytest.raises(ValueError): + _ = et.tensor + with pytest.raises(ValueError): + et.resize_(256) + + et.resize_per_segment_(256) + with pytest.raises(ValueError): + et.resize_per_segment_(128) # shrink + with pytest.raises(ValueError): + et.resize_per_segment_(et.segment_capacity_bytes + 1) # over capacity + finally: + et.free() diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py index e225666f8443..ac8d20b82f39 100644 --- a/tests/v1/core/test_encoder_cache_manager.py +++ b/tests/v1/core/test_encoder_cache_manager.py @@ -149,6 +149,30 @@ def test_has_cache_restores_from_freeable(): assert manager.num_freeable_slots == 6 +def test_make_profiling_reservation(): + assert ( + EncoderCacheManager.make_profiling_reservation( + cache_size=0, + embed_size=8, + dtype=torch.float16, + device="cpu", + ) + is None + ) + + reservation = EncoderCacheManager.make_profiling_reservation( + cache_size=7, + embed_size=8, + dtype=torch.float16, + device="cpu", + ) + + assert reservation is not None + assert reservation.shape == (7, 8) + assert reservation.dtype == torch.float16 + assert reservation.device.type == "cpu" + + def test_get_freed_mm_hashes_clears_freed_list(): manager = EncoderCacheManager(cache_size=10) req1 = MockRequest("reqA", ["a"], [5]) diff --git a/tests/v1/engine/test_engine_args.py b/tests/v1/engine/test_engine_args.py index c336f93c5fbb..3fb6a5d4d441 100644 --- a/tests/v1/engine/test_engine_args.py +++ b/tests/v1/engine/test_engine_args.py @@ -49,6 +49,18 @@ def test_prefix_caching_from_cli(): args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"]) +def test_extensible_kv_cache_from_cli(): + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + + args = parser.parse_args([]) + engine_args = EngineArgs.from_cli_args(args=args) + assert not engine_args.enable_extensible_kv_cache + + args = parser.parse_args(["--enable-extensible-kv-cache"]) + engine_args = EngineArgs.from_cli_args(args=args) + assert engine_args.enable_extensible_kv_cache + + @pytest.mark.skipif(_xxhash is None, reason="xxhash not installed") def test_prefix_caching_xxhash_from_cli(): parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) diff --git a/tests/v1/worker/test_extensible_kv_cache.py b/tests/v1/worker/test_extensible_kv_cache.py new file mode 100644 index 000000000000..d78ce79bd077 --- /dev/null +++ b/tests/v1/worker/test_extensible_kv_cache.py @@ -0,0 +1,369 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU integration tests for the extensible KV cache allocation paths. + +Drives `GPUModelRunner._allocate_kv_cache_tensors` / `_reshape_kv_cache_tensors` +/ `extend_kv_cache` directly with fake attention backends, covering the buffer +layouts the extensible flow supports: block-major (one committed prefix), +K/V-split (one prefix per half), Mamba (block-major per layer), and hybrid +attention + Mamba (attention re-strided to block-major). Buffer sizes exceed +the CUDA VMM allocation granularity so touching a block that the commit logic +missed would fault instead of silently passing. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, +) +from vllm.v1.worker.gpu_model_runner import GPUModelRunner +from vllm.v1.worker.utils import AttentionGroup + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + +BLOCK_SIZE = 16 +NUM_BLOCKS = 256 + + +class _SplitKVBackend(AttentionBackend): + """Fake backend with a K/V-split layout, like FlashAttention.""" + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + return (2, num_blocks, block_size, num_kv_heads, head_size) + + +class _BlockMajorBackend(AttentionBackend): + """Fake backend with a num-blocks-first layout, like FlashInfer.""" + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + return (num_blocks, 2, block_size, num_kv_heads, head_size) + + +class _StrideOrderBackend(AttentionBackend): + """Fake backend whose stride order makes a kv-first shape block-major.""" + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + return (2, num_blocks, block_size, num_kv_heads, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + assert not include_num_layers_dimension + return (1, 0, 2, 3, 4) + + +def _full_attention_spec() -> FullAttentionSpec: + # page_size_bytes = 2 (K+V) * 16 * 8 * 128 * 2 bytes = 64 KiB; 256 blocks + # = 16 MiB, several VMM granules per buffer. + return FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=8, + head_size=128, + dtype=torch.bfloat16, + ) + + +def _mamba_spec() -> MambaSpec: + # page_size_bytes = (8*128 + 16*64) * 4 bytes = 8 KiB per block per layer. + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=((8, 128), (16, 64)), + dtypes=(torch.float32, torch.float32), + ) + + +def _make_runner(kv_cache_config: KVCacheConfig, attn_groups) -> GPUModelRunner: + runner = object.__new__(GPUModelRunner) + runner.device = torch.device("cuda:0") + runner.kv_cache_config = kv_cache_config + runner.attn_groups = attn_groups + runner.runner_only_attn_layers = set() + runner.cache_config = SimpleNamespace(cache_dtype="auto") + return runner + + +def _attention_config(spec: FullAttentionSpec, backend) -> tuple[KVCacheConfig, list]: + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[ + KVCacheTensor(size=NUM_BLOCKS * spec.page_size_bytes, shared_by=["layer.0"]) + ], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)], + ) + attn_groups = [ + [ + AttentionGroup( + backend=backend, + layer_names=["layer.0"], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + ] + ] + return kv_cache_config, attn_groups + + +def _free_buffers(runner: GPUModelRunner) -> None: + for buffer, _ in getattr(runner, "_extensible_kv_cache_buffers", []): + buffer.free() + + +def test_kv_cache_num_segments_by_layer() -> None: + """Segment counts follow the physical layout of each layer's backend.""" + spec = _full_attention_spec() + for backend, expected in ( + (_SplitKVBackend, 2), + (_BlockMajorBackend, 1), + # kv-first logical shape but block-major physical order -> 1 segment. + (_StrideOrderBackend, 1), + ): + kv_cache_config, attn_groups = _attention_config(spec, backend) + runner = _make_runner(kv_cache_config, attn_groups) + assert runner._kv_cache_num_segments_by_layer() == {"layer.0": expected} + + +def test_extensible_split_layout_grows_both_halves() -> None: + """A K/V-split layer keeps its natural layout and both halves grow in + lockstep.""" + spec = _full_attention_spec() + kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend) + runner = _make_runner(kv_cache_config, attn_groups) + try: + raw_tensors = runner._allocate_kv_cache_tensors( + kv_cache_config, extensible=True + ) + kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE]) + kv_cache = kv_caches["layer.0"] + assert kv_cache.shape == (2, NUM_BLOCKS, BLOCK_SIZE, 8, 128) + [(buffer, bytes_per_block_per_segment)] = runner._extensible_kv_cache_buffers + assert buffer.num_segments == 2 + assert bytes_per_block_per_segment == spec.page_size_bytes // 2 + + # Only block 0 is committed -- in each half. + kv_cache[0, 0].fill_(1) # K, block 0 + kv_cache[1, 0].fill_(2) # V, block 0 + torch.cuda.synchronize() + + runner.extend_kv_cache(NUM_BLOCKS) + # Old data survives the grow; new blocks are usable in both halves and + # zeroed. + assert torch.all(kv_cache[0, 0] == 1) + assert torch.all(kv_cache[1, 0] == 2) + kv_cache[0, NUM_BLOCKS - 1].fill_(3) + kv_cache[1, NUM_BLOCKS - 1].fill_(4) + torch.cuda.synchronize() + assert torch.all(kv_cache[0, NUM_BLOCKS - 1] == 3) + assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 4) + assert torch.count_nonzero(kv_cache[:, 1 : NUM_BLOCKS - 1]) == 0 + finally: + _free_buffers(runner) + + +def test_extensible_block_major_layout() -> None: + """A layer whose physical layout is block-major uses a single segment.""" + spec = _full_attention_spec() + kv_cache_config, attn_groups = _attention_config(spec, _BlockMajorBackend) + runner = _make_runner(kv_cache_config, attn_groups) + try: + raw_tensors = runner._allocate_kv_cache_tensors( + kv_cache_config, extensible=True + ) + kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE]) + kv_cache = kv_caches["layer.0"] + assert kv_cache.shape == (NUM_BLOCKS, 2, BLOCK_SIZE, 8, 128) + [(buffer, bytes_per_block_per_segment)] = runner._extensible_kv_cache_buffers + assert buffer.num_segments == 1 + assert bytes_per_block_per_segment == spec.page_size_bytes + + kv_cache[0].fill_(1) + runner.extend_kv_cache(NUM_BLOCKS) + kv_cache[NUM_BLOCKS - 1].fill_(2) + torch.cuda.synchronize() + assert torch.all(kv_cache[0] == 1) + assert torch.all(kv_cache[NUM_BLOCKS - 1] == 2) + assert torch.count_nonzero(kv_cache[1 : NUM_BLOCKS - 1]) == 0 + finally: + _free_buffers(runner) + + +def test_legacy_split_layout_commits_everything() -> None: + """Without `extensible`, the full buffer is committed up front.""" + spec = _full_attention_spec() + kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend) + runner = _make_runner(kv_cache_config, attn_groups) + raw_tensors = runner._allocate_kv_cache_tensors(kv_cache_config, extensible=False) + kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE]) + kv_cache = kv_caches["layer.0"] + kv_cache[0, NUM_BLOCKS - 1].fill_(1) + kv_cache[1, NUM_BLOCKS - 1].fill_(2) + torch.cuda.synchronize() + assert torch.all(kv_cache[0, NUM_BLOCKS - 1] == 1) + assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 2) + with pytest.raises(RuntimeError, match="extensible"): + runner.extend_kv_cache(NUM_BLOCKS) + + +def test_extensible_mamba_grows_per_layer() -> None: + """Mamba per-layer buffers are block-major and grow with the KV cache.""" + spec = _mamba_spec() + num_blocks = 512 + layer_names = ["mamba.0", "mamba.1"] + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor(size=num_blocks * spec.page_size_bytes, shared_by=[name]) + for name in layer_names + ], + kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)], + ) + attn_groups = [ + [ + AttentionGroup( + backend=_BlockMajorBackend, + layer_names=layer_names, + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + ] + ] + runner = _make_runner(kv_cache_config, attn_groups) + try: + raw_tensors = runner._allocate_kv_cache_tensors( + kv_cache_config, extensible=True + ) + kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE]) + assert set(kv_caches) == set(layer_names) + assert len(runner._extensible_kv_cache_buffers) == len(layer_names) + for buffer, bytes_per_block_per_segment in runner._extensible_kv_cache_buffers: + assert buffer.num_segments == 1 + assert bytes_per_block_per_segment == spec.page_size_bytes + + # Write block 0 of every state of every layer (the committed + # prefixes), then grow. + for name in layer_names: + for state_tensor in kv_caches[name]: + state_tensor[0].fill_(1) + torch.cuda.synchronize() + runner.extend_kv_cache(num_blocks) + for name in layer_names: + for state_tensor in kv_caches[name]: + state_tensor[num_blocks - 1].fill_(2) + torch.cuda.synchronize() + for name in layer_names: + for state_tensor in kv_caches[name]: + assert torch.all(state_tensor[0] == 1) + assert torch.all(state_tensor[num_blocks - 1] == 2) + assert torch.count_nonzero(state_tensor[1 : num_blocks - 1]) == 0 + finally: + _free_buffers(runner) + + +def test_extensible_hybrid_attention_mamba() -> None: + """In hybrid models the attention cache is re-strided to block-major, so + its buffer must use a single segment.""" + attn_spec = _full_attention_spec() + mamba_spec = _mamba_spec() + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[ + KVCacheTensor( + size=NUM_BLOCKS * attn_spec.page_size_bytes, shared_by=["attn.0"] + ), + KVCacheTensor( + size=NUM_BLOCKS * mamba_spec.page_size_bytes, shared_by=["mamba.0"] + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["attn.0"], kv_cache_spec=attn_spec), + KVCacheGroupSpec(layer_names=["mamba.0"], kv_cache_spec=mamba_spec), + ], + ) + attn_groups = [ + [ + AttentionGroup( + backend=_SplitKVBackend, + layer_names=["attn.0"], + kv_cache_spec=attn_spec, + kv_cache_group_id=0, + ) + ], + [ + AttentionGroup( + backend=_BlockMajorBackend, + layer_names=["mamba.0"], + kv_cache_spec=mamba_spec, + kv_cache_group_id=1, + ) + ], + ] + runner = _make_runner(kv_cache_config, attn_groups) + try: + # The K/V-split attention layer is forced to one segment by the hybrid + # block-major re-stride. + assert runner._kv_cache_num_segments_by_layer() == {"attn.0": 1, "mamba.0": 1} + + raw_tensors = runner._allocate_kv_cache_tensors( + kv_cache_config, extensible=True + ) + kv_caches = runner._reshape_kv_cache_tensors( + raw_tensors, [BLOCK_SIZE, BLOCK_SIZE] + ) + attn_cache = kv_caches["attn.0"] + # `_update_hybrid_attention_mamba_layout` re-strides to interleave K/V + # per block: block b spans one contiguous page. + hidden_size = attn_cache.shape[2:].numel() + assert attn_cache.stride()[:2] == (hidden_size, 2 * hidden_size) + + attn_cache[0, 0].fill_(1) # K, block 0 + attn_cache[1, 0].fill_(2) # V, block 0 + for state_tensor in kv_caches["mamba.0"]: + state_tensor[0].fill_(3) + torch.cuda.synchronize() + + runner.extend_kv_cache(NUM_BLOCKS) + attn_cache[0, NUM_BLOCKS - 1].fill_(4) + attn_cache[1, NUM_BLOCKS - 1].fill_(5) + for state_tensor in kv_caches["mamba.0"]: + state_tensor[NUM_BLOCKS - 1].fill_(6) + torch.cuda.synchronize() + assert torch.all(attn_cache[0, 0] == 1) + assert torch.all(attn_cache[1, 0] == 2) + assert torch.all(attn_cache[0, NUM_BLOCKS - 1] == 4) + assert torch.all(attn_cache[1, NUM_BLOCKS - 1] == 5) + assert torch.count_nonzero(attn_cache[:, 1 : NUM_BLOCKS - 1]) == 0 + for state_tensor in kv_caches["mamba.0"]: + assert torch.all(state_tensor[0] == 3) + assert torch.all(state_tensor[NUM_BLOCKS - 1] == 6) + assert torch.count_nonzero(state_tensor[1 : NUM_BLOCKS - 1]) == 0 + finally: + _free_buffers(runner) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 2fb3358d55c3..3a8aa54f3542 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -174,6 +174,15 @@ class CacheConfig: gpu_memory_utilization. Note that kv_cache_memory_bytes (when not-None) ignores gpu_memory_utilization""" + enable_extensible_kv_cache: bool = False + """Use CUDA virtual memory to reserve the KV cache address range before + CUDA graph capture and commit the final size after capture. + + This makes automatic KV sizing account for the actual CUDA graph pool. + Supported for all V1 CUDA attention backends (block-major and K/V-split + KV cache layouts) and for Mamba / linear-attention models. + """ + kv_offloading_size: float | None = None """Size of the KV cache offloading buffer in GiB. When TP > 1, this is the total buffer size summed across all TP ranks. By default, this is set @@ -217,6 +226,8 @@ def compute_hash(self) -> str: "kv_cache_max_concurrency", # WIP feature toggle not impacting compiled graph shape "kv_sharing_fast_prefill", + # Runtime memory allocation strategy, not graph structure. + "enable_extensible_kv_cache", } from vllm.config.utils import get_hash_factors, hash_factors diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index eb13a379ff41..9c05befebb62 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2112,6 +2112,9 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: # Will be added by https://github.com/vllm-project/vllm/pull/35045 unsupported.append("KV sharing fast prefill") + if self.cache_config.enable_extensible_kv_cache: + unsupported.append("extensible KV cache") + if self.ec_transfer_config is not None: # Will be added by https://github.com/vllm-project/vllm/pull/38390 unsupported.append("EC transfer") diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index efdd7696fdce..72b14e7fa5dc 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -522,6 +522,7 @@ class EngineArgs: offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params") gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes + enable_extensible_kv_cache: bool = CacheConfig.enable_extensible_kv_cache max_num_batched_tokens: int | None = None max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills max_long_partial_prefills: int = SchedulerConfig.max_long_partial_prefills @@ -1152,6 +1153,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: cache_group.add_argument( "--kv-cache-memory-bytes", **cache_kwargs["kv_cache_memory_bytes"] ) + cache_group.add_argument( + "--enable-extensible-kv-cache", + **cache_kwargs["enable_extensible_kv_cache"], + ) cache_group.add_argument("--kv-cache-dtype", **cache_kwargs["cache_dtype"]) cache_group.add_argument( "--num-gpu-blocks-override", **cache_kwargs["num_gpu_blocks_override"] @@ -1869,6 +1874,7 @@ def create_engine_config( block_size=self.block_size, # type: ignore[arg-type] gpu_memory_utilization=self.gpu_memory_utilization, kv_cache_memory_bytes=self.kv_cache_memory_bytes, + enable_extensible_kv_cache=self.enable_extensible_kv_cache, cache_dtype=resolved_cache_dtype, # type: ignore[arg-type] is_attention_free=model_config.is_attention_free, num_gpu_blocks_override=self.num_gpu_blocks_override, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index a3ed94ee0aa9..bb8514c94332 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -119,6 +119,11 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): compared with using gpu_memory_utilization. Note that kv_cache_memory_bytes (when not-None) ignores gpu_memory_utilization + enable_extensible_kv_cache: Use CUDA virtual memory to reserve the KV + cache address range before CUDA graph capture and commit the final + cache size after capture. Supported by V1 CUDA workers for all + attention backends (block-major and K/V-split KV cache layouts) + and for Mamba / linear-attention models. cpu_offload_gb: The size (GiB) of CPU memory to use for offloading the model weights. This virtually increases the GPU memory space you can use to hold the model weights, at the cost of CPU-GPU data @@ -211,6 +216,7 @@ def __init__( profiler_config: dict[str, Any] | ProfilerConfig | None = None, attention_config: dict[str, Any] | AttentionConfig | None = None, kv_cache_memory_bytes: int | None = None, + enable_extensible_kv_cache: bool = False, compilation_config: int | dict[str, Any] | CompilationConfig | None = None, quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None, logits_processors: list[str | type[LogitsProcessor]] | None = None, @@ -320,6 +326,7 @@ def _make_config(value: Any, cls: type[_R]) -> _R: seed=seed, gpu_memory_utilization=gpu_memory_utilization, kv_cache_memory_bytes=kv_cache_memory_bytes, + enable_extensible_kv_cache=enable_extensible_kv_cache, cpu_offload_gb=cpu_offload_gb, offload_group_size=offload_group_size, offload_num_in_group=offload_num_in_group, diff --git a/vllm/utils/extensible_tensor.py b/vllm/utils/extensible_tensor.py new file mode 100644 index 000000000000..445fe0a14e5c --- /dev/null +++ b/vllm/utils/extensible_tensor.py @@ -0,0 +1,520 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Growable CUDA byte buffers backed by CUDA virtual memory management.""" + +from __future__ import annotations + +import ctypes +from contextlib import suppress + +import torch + +_CUDA_SUCCESS = 0 +_CU_MEM_ALLOCATION_TYPE_PINNED = 1 +_CU_MEM_LOCATION_TYPE_DEVICE = 1 +_CU_MEM_ALLOC_GRANULARITY_MINIMUM = 0 +_CU_MEM_ACCESS_FLAGS_PROT_READWRITE = 3 +_CU_MEM_ALLOCATION_COMP_NONE = 0 + + +class _CUmemLocation(ctypes.Structure): + _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)] + + +class _CUmemAllocFlags(ctypes.Structure): + _fields_ = [ + ("compressionType", ctypes.c_ubyte), + ("gpuDirectRDMACapable", ctypes.c_ubyte), + ("usage", ctypes.c_ushort), + ("reserved", ctypes.c_ubyte * 4), + ] + + +class _CUmemAllocationProp(ctypes.Structure): + _fields_ = [ + ("type", ctypes.c_int), + ("requestedHandleTypes", ctypes.c_int), + ("location", _CUmemLocation), + ("win32HandleMetaData", ctypes.c_void_p), + ("allocFlags", _CUmemAllocFlags), + ] + + +class _CUmemAccessDesc(ctypes.Structure): + _fields_ = [("location", _CUmemLocation), ("flags", ctypes.c_int)] + + +_CUdeviceptr = ctypes.c_ulonglong +_CUmemHandle = ctypes.c_ulonglong +_CUcontext = ctypes.c_void_p + +_libcuda: ctypes.CDLL | None = None + + +def _find_loaded_library(lib_name: str) -> str | None: + try: + with open("/proc/self/maps") as f: + for line in f: + if lib_name not in line: + continue + start = line.index("/") + return line[start:].strip() + except (OSError, ValueError): + return None + return None + + +def _load_libcuda() -> ctypes.CDLL: + for name in ("libcuda.so.1", "libcuda.so"): + try: + return ctypes.CDLL(name) + except OSError: + continue + if path := _find_loaded_library("libcuda"): + return ctypes.CDLL(path) + raise RuntimeError( + "Could not load libcuda. The CUDA driver library is required for " + "ExtensibleTensor." + ) + + +def _configure_signatures(lib: ctypes.CDLL) -> None: + pointer = ctypes.POINTER + lib.cuGetErrorString.argtypes = [ctypes.c_int, pointer(ctypes.c_char_p)] + lib.cuCtxGetCurrent.argtypes = [pointer(_CUcontext)] + lib.cuDevicePrimaryCtxRetain.argtypes = [pointer(_CUcontext), ctypes.c_int] + lib.cuCtxSetCurrent.argtypes = [_CUcontext] + lib.cuMemGetAllocationGranularity.argtypes = [ + pointer(ctypes.c_size_t), + pointer(_CUmemAllocationProp), + ctypes.c_int, + ] + lib.cuMemAddressReserve.argtypes = [ + pointer(_CUdeviceptr), + ctypes.c_size_t, + ctypes.c_size_t, + _CUdeviceptr, + ctypes.c_ulonglong, + ] + lib.cuMemCreate.argtypes = [ + pointer(_CUmemHandle), + ctypes.c_size_t, + pointer(_CUmemAllocationProp), + ctypes.c_ulonglong, + ] + lib.cuMemMap.argtypes = [ + _CUdeviceptr, + ctypes.c_size_t, + ctypes.c_size_t, + _CUmemHandle, + ctypes.c_ulonglong, + ] + lib.cuMemSetAccess.argtypes = [ + _CUdeviceptr, + ctypes.c_size_t, + pointer(_CUmemAccessDesc), + ctypes.c_size_t, + ] + lib.cuMemUnmap.argtypes = [_CUdeviceptr, ctypes.c_size_t] + lib.cuMemRelease.argtypes = [_CUmemHandle] + lib.cuMemAddressFree.argtypes = [_CUdeviceptr, ctypes.c_size_t] + + for fn in ( + lib.cuGetErrorString, + lib.cuCtxGetCurrent, + lib.cuDevicePrimaryCtxRetain, + lib.cuCtxSetCurrent, + lib.cuMemGetAllocationGranularity, + lib.cuMemAddressReserve, + lib.cuMemCreate, + lib.cuMemMap, + lib.cuMemSetAccess, + lib.cuMemUnmap, + lib.cuMemRelease, + lib.cuMemAddressFree, + ): + fn.restype = ctypes.c_int + + +def _cuda() -> ctypes.CDLL: + global _libcuda + if _libcuda is None: + lib = _load_libcuda() + _configure_signatures(lib) + _libcuda = lib + return _libcuda + + +def _check(result: int) -> None: + if result == _CUDA_SUCCESS: + return + msg = ctypes.c_char_p() + _cuda().cuGetErrorString(result, ctypes.byref(msg)) + detail = msg.value.decode() if msg.value else "unknown error" + raise RuntimeError(f"CUDA driver error {result}: {detail}") + + +def _ensure_context(device_index: int) -> None: + pctx = _CUcontext() + _check(_cuda().cuCtxGetCurrent(ctypes.byref(pctx))) + if pctx.value: + return + _check(_cuda().cuDevicePrimaryCtxRetain(ctypes.byref(pctx), device_index)) + _check(_cuda().cuCtxSetCurrent(pctx)) + + +def _make_alloc_prop(device_index: int) -> _CUmemAllocationProp: + prop = _CUmemAllocationProp() + prop.type = _CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = _CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = device_index + prop.allocFlags.compressionType = _CU_MEM_ALLOCATION_COMP_NONE + return prop + + +def _round_up(value: int, multiple: int) -> int: + return ((value + multiple - 1) // multiple) * multiple + + +class _VirtualBuffer: + """Own one device VA reservation and the physical chunks mapped into it. + + Physical memory is committed incrementally, at granularity-sized granules, + via `ensure_committed_range`; granules already mapped by an earlier + (possibly overlapping) range are skipped, so ranges may abut or overlap + freely. + """ + + def __init__(self, max_bytes: int, device_index: int) -> None: + _ensure_context(device_index) + self.device_index = device_index + + prop = _make_alloc_prop(device_index) + granularity = ctypes.c_size_t() + _check( + _cuda().cuMemGetAllocationGranularity( + ctypes.byref(granularity), + ctypes.byref(prop), + _CU_MEM_ALLOC_GRANULARITY_MINIMUM, + ) + ) + self.granularity: int = granularity.value + self.reserved_size: int = _round_up(max(max_bytes, 1), self.granularity) + + dptr = _CUdeviceptr() + _check( + _cuda().cuMemAddressReserve(ctypes.byref(dptr), self.reserved_size, 0, 0, 0) + ) + self.base_ptr: int = dptr.value + + # Granule indices (VA offset // granularity) that have physical + # memory mapped. + self._mapped_granules: set[int] = set() + # Each entry is (handle, va_offset, size) for one mapped physical chunk. + self._handles: list[tuple[int, int, int]] = [] + self._freed: bool = False + + @property + def committed_bytes(self) -> int: + """Total physically mapped bytes (a multiple of the granularity).""" + return len(self._mapped_granules) * self.granularity + + def ensure_committed(self, nbytes: int) -> None: + """Map physical pages so that at least the first `nbytes` are backed.""" + self.ensure_committed_range(0, nbytes) + + def ensure_committed_range(self, start: int, end: int) -> None: + """Map physical pages so that the byte range `[start, end)` is backed. + + The range is widened outward to granule boundaries; granules mapped by + earlier calls are skipped, so a granule shared by two requested ranges + is mapped once. + """ + if not 0 <= start <= end: + raise ValueError(f"Invalid range [{start}, {end}).") + if end > self.reserved_size: + raise ValueError( + f"Requested range end {end} exceeds reserved capacity " + f"{self.reserved_size}." + ) + if start == end: + return + first = start // self.granularity + last = (end + self.granularity - 1) // self.granularity # exclusive + run_start: int | None = None + for g in range(first, last + 1): + unmapped = g < last and g not in self._mapped_granules + if unmapped and run_start is None: + run_start = g + elif not unmapped and run_start is not None: + self._map_chunk_at( + run_start * self.granularity, (g - run_start) * self.granularity + ) + self._mapped_granules.update(range(run_start, g)) + run_start = None + + def _map_chunk_at(self, offset: int, size: int) -> None: + """Create one physical chunk of `size` bytes and map it at `offset`.""" + _ensure_context(self.device_index) + prop = _make_alloc_prop(self.device_index) + + handle = _CUmemHandle() + _check(_cuda().cuMemCreate(ctypes.byref(handle), size, ctypes.byref(prop), 0)) + + addr = self.base_ptr + offset + try: + _check(_cuda().cuMemMap(addr, size, 0, handle, 0)) + except RuntimeError: + _cuda().cuMemRelease(handle) + raise + + desc = _CUmemAccessDesc() + desc.location.type = _CU_MEM_LOCATION_TYPE_DEVICE + desc.location.id = self.device_index + desc.flags = _CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _check(_cuda().cuMemSetAccess(addr, size, ctypes.byref(desc), 1)) + + self._handles.append((handle.value, offset, size)) + + def free(self) -> None: + if self._freed: + return + self._freed = True + _ensure_context(self.device_index) + if self._handles: + torch.cuda.synchronize(self.device_index) + for handle, offset, size in self._handles: + _check(_cuda().cuMemUnmap(self.base_ptr + offset, size)) + _check(_cuda().cuMemRelease(handle)) + if self.base_ptr: + _check(_cuda().cuMemAddressFree(self.base_ptr, self.reserved_size)) + self._handles = [] + self._mapped_granules = set() + self.base_ptr = 0 + + def __del__(self) -> None: + with suppress(Exception): + self.free() + + +_K_DL_CUDA = 2 +_K_DL_UINT = 1 +_UINT8_BITS = 8 + + +class _DLDevice(ctypes.Structure): + _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int)] + + +class _DLDataType(ctypes.Structure): + _fields_ = [ + ("code", ctypes.c_uint8), + ("bits", ctypes.c_uint8), + ("lanes", ctypes.c_uint16), + ] + + +class _DLTensor(ctypes.Structure): + _fields_ = [ + ("data", ctypes.c_void_p), + ("device", _DLDevice), + ("ndim", ctypes.c_int), + ("dtype", _DLDataType), + ("shape", ctypes.POINTER(ctypes.c_int64)), + ("strides", ctypes.POINTER(ctypes.c_int64)), + ("byte_offset", ctypes.c_uint64), + ] + + +class _DLManagedTensor(ctypes.Structure): + pass + + +_DLDeleter = ctypes.CFUNCTYPE(None, ctypes.POINTER(_DLManagedTensor)) +_DLManagedTensor._fields_ = [ + ("dl_tensor", _DLTensor), + ("manager_ctx", ctypes.c_void_p), + ("deleter", _DLDeleter), +] + +_KEEPALIVE: dict[int, tuple[object, object, object]] = {} +_PyCapsule_New = ctypes.pythonapi.PyCapsule_New +_PyCapsule_New.restype = ctypes.py_object +_PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + + +def _uint8_tensor_from_ptr(ptr: int, num_bytes: int, device_index: int) -> torch.Tensor: + shape_arr = (ctypes.c_int64 * 1)(num_bytes) + + managed = _DLManagedTensor() + managed.dl_tensor.data = ctypes.c_void_p(ptr) + managed.dl_tensor.device = _DLDevice(_K_DL_CUDA, device_index) + managed.dl_tensor.ndim = 1 + managed.dl_tensor.dtype = _DLDataType(_K_DL_UINT, _UINT8_BITS, 1) + managed.dl_tensor.shape = ctypes.cast(shape_arr, ctypes.POINTER(ctypes.c_int64)) + managed.dl_tensor.strides = None + managed.dl_tensor.byte_offset = 0 + managed.manager_ctx = None + + key = ctypes.addressof(managed) + + def _deleter(_managed_ptr: object) -> None: + _KEEPALIVE.pop(key, None) + + deleter = _DLDeleter(_deleter) + managed.deleter = deleter + _KEEPALIVE[key] = (managed, shape_arr, deleter) + + capsule = _PyCapsule_New(ctypes.addressof(managed), b"dltensor", None) + return torch.from_dlpack(capsule) + + +class ExtensibleTensor: + """A 1-D CUDA byte buffer that can grow without moving its base pointer. + + With `num_segments > 1` the reservation is divided into that many equal + segments that grow in lockstep via `resize_per_segment_`: the committed + bytes form a prefix of each segment (segment `i` spans + `[i * segment_capacity_bytes, (i + 1) * segment_capacity_bytes)` of + `full_view()`). This backs layouts whose block dimension is not outermost, + e.g. a K/V-split KV cache (`num_segments=2`). `resize_` / `tensor` / + `append` assume a single contiguous prefix and are only valid when + `num_segments == 1`. + """ + + def __init__( + self, + max_num_bytes: int, + device: torch.device | str | int | None = None, + num_segments: int = 1, + ) -> None: + if max_num_bytes < 0: + raise ValueError("max_num_bytes must be non-negative.") + if num_segments < 1: + raise ValueError(f"num_segments must be positive, got {num_segments}.") + if max_num_bytes % num_segments != 0: + raise ValueError( + f"max_num_bytes ({max_num_bytes}) must be divisible by " + f"num_segments ({num_segments})." + ) + + if device is None: + device = torch.cuda.current_device() + dev = device if isinstance(device, torch.device) else torch.device(device) + if dev.type != "cuda": + raise ValueError(f"ExtensibleTensor requires a cuda device, got {dev}.") + self._device_index: int = ( + dev.index if dev.index is not None else torch.cuda.current_device() + ) + + torch.cuda.init() + + self._max_num_bytes: int = max_num_bytes + self._num_segments: int = num_segments + self._segment_capacity_bytes: int = max_num_bytes // num_segments + self._buffer: _VirtualBuffer = _VirtualBuffer(max_num_bytes, self._device_index) + self._bytes_per_segment: int = 0 + + @property + def tensor(self) -> torch.Tensor: + """Return a uint8 tensor view of the currently committed prefix.""" + if self._num_segments != 1: + raise ValueError( + "tensor (a single committed prefix) is only valid for " + "num_segments=1; use full_view() and index segments explicitly." + ) + return _uint8_tensor_from_ptr( + self._buffer.base_ptr, self._bytes_per_segment, self._device_index + ) + + def full_view(self) -> torch.Tensor: + """Return a uint8 tensor view spanning the requested maximum size.""" + return _uint8_tensor_from_ptr( + self._buffer.base_ptr, self._max_num_bytes, self._device_index + ) + + def resize_(self, num_bytes: int) -> torch.Tensor: + """Grow the buffer to `num_bytes` and return the committed-prefix view.""" + if self._num_segments != 1: + raise ValueError( + "resize_ (a single committed prefix) is only valid for " + "num_segments=1; use resize_per_segment_." + ) + self.resize_per_segment_(num_bytes) + return self.tensor + + def resize_per_segment_( + self, bytes_per_segment: int, zero_new: bool = False + ) -> None: + """Grow every segment's committed prefix to `bytes_per_segment` bytes. + + Existing bytes are preserved and the base pointer is unchanged. With + `zero_new=True` the newly committed byte range of each segment is + zeroed (bytes committed earlier are left intact). Raises if + `bytes_per_segment` is smaller than the current per-segment size + (shrink is unsupported) or larger than `segment_capacity_bytes`. + """ + old = self._bytes_per_segment + if bytes_per_segment < old: + raise ValueError( + f"ExtensibleTensor is grow-only: cannot resize from {old} " + f"to {bytes_per_segment} bytes per segment." + ) + if bytes_per_segment > self._segment_capacity_bytes: + raise ValueError( + f"Requested {bytes_per_segment} bytes per segment exceeds the " + f"segment capacity {self._segment_capacity_bytes}." + ) + if bytes_per_segment == old: + return + for i in range(self._num_segments): + start = i * self._segment_capacity_bytes + self._buffer.ensure_committed_range(start + old, start + bytes_per_segment) + self._bytes_per_segment = bytes_per_segment + if zero_new: + full = self.full_view() + for i in range(self._num_segments): + start = i * self._segment_capacity_bytes + full[start + old : start + bytes_per_segment].zero_() + + def append(self, num_bytes: int) -> torch.Tensor: + """Grow by `num_bytes` additional bytes and return the new view.""" + if num_bytes < 0: + raise ValueError("num_bytes to append must be non-negative.") + return self.resize_(self._bytes_per_segment + num_bytes) + + @property + def num_bytes(self) -> int: + """Current committed size in bytes, summed over all segments.""" + return self._bytes_per_segment * self._num_segments + + @property + def bytes_per_segment(self) -> int: + """Current committed prefix size of each segment in bytes.""" + return self._bytes_per_segment + + @property + def num_segments(self) -> int: + """Number of equal segments the reservation is divided into.""" + return self._num_segments + + @property + def segment_capacity_bytes(self) -> int: + """Maximum size of each segment (`max_num_bytes / num_segments`).""" + return self._segment_capacity_bytes + + @property + def capacity_bytes(self) -> int: + return self._buffer.reserved_size + + @property + def base_ptr(self) -> int: + return self._buffer.base_ptr + + @property + def device(self) -> torch.device: + return torch.device("cuda", self._device_index) + + def free(self) -> None: + self._buffer.free() + self._bytes_per_segment = 0 diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py index d479239e3b16..cff5d9086ac0 100644 --- a/vllm/v1/core/encoder_cache_manager.py +++ b/vllm/v1/core/encoder_cache_manager.py @@ -5,6 +5,8 @@ from collections.abc import Mapping from typing import TYPE_CHECKING +import torch + from vllm.logger import init_logger from vllm.v1.request import Request @@ -78,6 +80,17 @@ def __init__(self, cache_size: int): self.freeable: OrderedDict[str, int] = OrderedDict() self.freed: list[str] = [] + @staticmethod + def make_profiling_reservation( + cache_size: int, + embed_size: int, + dtype: torch.dtype, + device: torch.device | str, + ) -> torch.Tensor | None: + if cache_size <= 0: + return None + return torch.empty((cache_size, embed_size), dtype=dtype, device=device) + def reset(self) -> None: """Reset the encoder cache to its initial state. diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index f97f697dedca..e4c06158c565 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -89,6 +89,7 @@ logger = init_logger(__name__) HANDSHAKE_TIMEOUT_MINS = 5 +_CUDAGRAPH_MEMORY_BUFFER_BYTES = 150 * (1 << 20) _R = TypeVar("_R") # Return type for collective_rpc @@ -288,37 +289,73 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: assert len(kv_cache_specs) == len(available_gpu_memory) + use_extensible_kv_cache = ( + has_kv_cache and vllm_config.cache_config.enable_extensible_kv_cache + ) + if use_extensible_kv_cache: + from vllm.platforms import current_platform + + if vllm_config.cache_config.kv_cache_memory_bytes is not None: + raise ValueError( + "enable_extensible_kv_cache=True is not supported with " + "kv_cache_memory_bytes. The extensible path requires " + "automatic KV cache sizing." + ) + if not current_platform.is_cuda(): + raise ValueError( + "enable_extensible_kv_cache=True is only supported on CUDA." + ) + # Track max_model_len before KV cache config to detect auto-fit changes + # made by get_kv_cache_configs(). max_model_len_before = vllm_config.model_config.max_model_len - kv_cache_configs = get_kv_cache_configs( vllm_config, kv_cache_specs, available_gpu_memory ) + scheduler_kv_cache_config = self._apply_kv_cache_config( + vllm_config, + kv_cache_configs, + max_model_len_before, + ) - # If auto-fit reduced max_model_len, sync the new value to workers. - # This is needed because workers were spawned before memory profiling - # and have the original (larger) max_model_len cached. - max_model_len_after = vllm_config.model_config.max_model_len - if max_model_len_after != max_model_len_before: - self.collective_rpc("update_max_model_len", args=(max_model_len_after,)) - - scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) - vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks - kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups - if kv_cache_groups: - vllm_config.cache_config.block_size = min( - g.kv_cache_spec.block_size for g in kv_cache_groups + # Initialize KV cache and warm up execution. With extensible KV cache, + # this reserves the upper-bound address range, commits one block, and + # captures CUDA graphs before committing the post-capture KV size. + compilation_times = self.model_executor.initialize_from_config( + kv_cache_configs, + extensible=use_extensible_kv_cache, + ) + if use_extensible_kv_cache: + if len(compilation_times) != len(available_gpu_memory): + raise RuntimeError( + "Expected one CompilationTimes result per worker when " + "initializing extensible KV cache, but got " + f"{len(compilation_times)} results for " + f"{len(available_gpu_memory)} workers." + ) + final_available_gpu_memory = [ + max( + available_memory + - times.cuda_graph + - _CUDAGRAPH_MEMORY_BUFFER_BYTES, + 0, + ) + for available_memory, times in zip( + available_gpu_memory, compilation_times, strict=True + ) + ] + max_model_len_before = vllm_config.model_config.max_model_len + kv_cache_configs = get_kv_cache_configs( + vllm_config, + kv_cache_specs, + final_available_gpu_memory, ) - num_tokens, max_concurrency = get_kv_cache_capacity( - vllm_config, scheduler_kv_cache_config + scheduler_kv_cache_config = self._apply_kv_cache_config( + vllm_config, + kv_cache_configs, + max_model_len_before, ) - vllm_config.cache_config.kv_cache_size_tokens = num_tokens - vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency - - vllm_config.validate_block_size() - - # Initialize kv cache and warmup the execution - self.model_executor.initialize_from_config(kv_cache_configs) + self.model_executor.extend_kv_cache(scheduler_kv_cache_config.num_blocks) elapsed = time.time() - start compile_time = vllm_config.compilation_config.compilation_time @@ -347,6 +384,35 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: ) return scheduler_kv_cache_config + def _apply_kv_cache_config( + self, + vllm_config: VllmConfig, + kv_cache_configs: list[KVCacheConfig], + max_model_len_before: int, + ) -> KVCacheConfig: + # If auto-fit reduced max_model_len, sync the new value to workers. + # This is needed because workers were spawned before memory profiling + # and have the original (larger) max_model_len cached. + max_model_len_after = vllm_config.model_config.max_model_len + if max_model_len_after != max_model_len_before: + self.collective_rpc("update_max_model_len", args=(max_model_len_after,)) + + scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) + vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks + kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups + if kv_cache_groups: + vllm_config.cache_config.block_size = min( + g.kv_cache_spec.block_size for g in kv_cache_groups + ) + num_tokens, max_concurrency = get_kv_cache_capacity( + vllm_config, scheduler_kv_cache_config + ) + vllm_config.cache_config.kv_cache_size_tokens = num_tokens + vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency + + vllm_config.validate_block_size() + return scheduler_kv_cache_config + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: return self.model_executor.supported_tasks diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 4063844d469c..3420d7427fb6 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -115,12 +115,20 @@ def __init__( def _init_executor(self) -> None: raise NotImplementedError - def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None: + def initialize_from_config( + self, + kv_cache_configs: list[KVCacheConfig], + extensible: bool = False, + ) -> list[CompilationTimes]: """ Initialize the KV caches and begin the model execution loop of the underlying workers. """ - self.collective_rpc("initialize_from_config", args=(kv_cache_configs,)) + self.collective_rpc( + "initialize_from_config", + args=(kv_cache_configs,), + kwargs={"extensible": extensible} if extensible else None, + ) compilation_times: list[CompilationTimes] = self.collective_rpc( "compile_or_warm_up_model" ) @@ -135,6 +143,7 @@ def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None: self.vllm_config.compilation_config.encoder_compilation_time = max( t.encoder for t in compilation_times ) + return compilation_times def register_failure_callback(self, callback: FailureCallback): # noqa: B027 """ @@ -149,6 +158,9 @@ def determine_available_memory(self) -> list[int]: # in bytes def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]: return self.collective_rpc("get_kv_cache_spec") + def extend_kv_cache(self, num_blocks: int) -> None: + self.collective_rpc("extend_kv_cache", args=(num_blocks,)) + @overload def collective_rpc( self, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e8d472e9eb0b..dd40455a72d2 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4,6 +4,7 @@ import functools import gc import itertools +import math import threading import time from collections import defaultdict @@ -112,6 +113,7 @@ from vllm.tasks import GenerationTask, PoolingTask, SupportedTask from vllm.tracing import instrument from vllm.utils import length_from_prompt_token_ids_or_embeds +from vllm.utils.extensible_tensor import ExtensibleTensor from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks @@ -139,6 +141,7 @@ get_dcp_local_seq_lens, reorder_batch_to_split_decodes_and_prefills, ) +from vllm.v1.core.encoder_cache_manager import EncoderCacheManager from vllm.v1.core.sched.output import NewRequestData from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher from vllm.v1.kv_cache_interface import ( @@ -6238,6 +6241,7 @@ def _dummy_pooler_run( return self._dummy_pooler_run_task(hidden_states, max_task) def profile_run(self) -> None: + dummy_encoder_cache: torch.Tensor | None = None # Profile with multimodal encoder & encoder cache. if self.supports_mm_inputs: mm_config = self.model_config.multimodal_config @@ -6249,6 +6253,12 @@ def profile_run(self) -> None: else: mm_budget = self.mm_budget assert mm_budget is not None + dummy_encoder_cache = EncoderCacheManager.make_profiling_reservation( + mm_budget.encoder_cache_size, + self.inputs_embeds_size, + self.model_config.dtype, + self.device, + ) if (encoder_budget := mm_budget.get_encoder_budget()) > 0: if not mm_budget.mm_max_toks_per_item: @@ -6308,7 +6318,7 @@ def profile_run(self) -> None: else: output = None self._sync_device() - del hidden_states, output + del hidden_states, output, dummy_encoder_cache self.encoder_cache.clear() gc.collect() @@ -6387,6 +6397,11 @@ def _cleanup_profiling_kv_cache(self) -> None: self.attn_groups.clear() if hasattr(self, "kv_cache_config"): delattr(self, "kv_cache_config") + if hasattr(self, "_extensible_kv_cache_buffers"): + for buffer, _ in self._extensible_kv_cache_buffers: + buffer.free() + self._extensible_kv_cache_buffers = [] + self._extensible_kv_cache_enabled = False self.cache_config.num_gpu_blocks = None for layer in self.compilation_config.static_forward_context.values(): @@ -7028,19 +7043,78 @@ def may_reinitialize_input_batch( ) def _allocate_kv_cache_tensors( - self, kv_cache_config: KVCacheConfig + self, kv_cache_config: KVCacheConfig, extensible: bool = False ) -> dict[str, torch.Tensor]: """ Initializes the KV cache buffer with the correct size. The buffer needs to be reshaped to the desired shape before being used by the models. Args: - kv_cache_config: The KV cache config + kv_cache_config: The KV cache config; its `num_blocks` is the + declared capacity. + extensible: When True, reserve virtual address space for + `num_blocks` but commit only one block (per layout segment) + for CUDA graph capture; `extend_kv_cache` commits the rest + afterwards. When False, commit the full size up front. Returns: dict[str, torch.Tensor]: A map between layer names to their corresponding memory buffer for KV cache. """ kv_cache_raw_tensors: dict[str, torch.Tensor] = {} + if extensible: + if any(t.block_stride > 0 for t in kv_cache_config.kv_cache_tensors): + raise ValueError( + "enable_extensible_kv_cache=True is not supported with " + "packed KV cache tensor layouts." + ) + if kv_cache_config.num_blocks <= 0: + raise ValueError( + "enable_extensible_kv_cache=True requires at least one KV block." + ) + + # One CUDA virtual-memory byte buffer per KV cache tensor. Each + # buffer keeps its layers' physical layout and is committed as one + # prefix per layout segment (e.g. the K and V halves of a + # K/V-split layout) -- see `ExtensibleTensor`. + num_segments_by_layer = self._kv_cache_num_segments_by_layer() + self._extensible_kv_cache_buffers: list[tuple[ExtensibleTensor, int]] = [] + self._extensible_kv_cache_committed_blocks = 1 + self._extensible_kv_cache_enabled = True + for kv_cache_tensor in kv_cache_config.kv_cache_tensors: + bytes_per_block = kv_cache_tensor.size // kv_cache_config.num_blocks + assert bytes_per_block * kv_cache_config.num_blocks == ( + kv_cache_tensor.size + ) + segment_counts = { + num_segments_by_layer[layer_name] + for layer_name in kv_cache_tensor.shared_by + if layer_name in num_segments_by_layer + } + assert len(segment_counts) == 1, ( + "Layers sharing one KV cache tensor disagree on the buffer " + f"segmentation ({segment_counts}): {kv_cache_tensor.shared_by}" + ) + num_segments = segment_counts.pop() + assert bytes_per_block % num_segments == 0 + bytes_per_block_per_segment = bytes_per_block // num_segments + buffer = ExtensibleTensor( + max_num_bytes=kv_cache_tensor.size, + device=self.device, + num_segments=num_segments, + ) + self._extensible_kv_cache_buffers.append( + (buffer, bytes_per_block_per_segment) + ) + buffer.resize_per_segment_(bytes_per_block_per_segment, zero_new=True) + tensor = buffer.full_view() + for layer_name in kv_cache_tensor.shared_by: + kv_cache_raw_tensors[layer_name] = tensor + + return self._check_kv_cache_raw_tensors( + kv_cache_config, kv_cache_raw_tensors + ) + + self._extensible_kv_cache_enabled = False packed_backing: torch.Tensor | None = None for kv_cache_tensor in kv_cache_config.kv_cache_tensors: if kv_cache_tensor.block_stride > 0: @@ -7059,6 +7133,13 @@ def _allocate_kv_cache_tensors( for layer_name in kv_cache_tensor.shared_by: kv_cache_raw_tensors[layer_name] = tensor + return self._check_kv_cache_raw_tensors(kv_cache_config, kv_cache_raw_tensors) + + def _check_kv_cache_raw_tensors( + self, + kv_cache_config: KVCacheConfig, + kv_cache_raw_tensors: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: layer_names = set() for group in kv_cache_config.kv_cache_groups: for layer_name in group.layer_names: @@ -7070,6 +7151,49 @@ def _allocate_kv_cache_tensors( ) return kv_cache_raw_tensors + def _kv_cache_num_segments_by_layer(self) -> dict[str, int]: + """Number of equal contiguous segments of each layer's KV cache buffer + under its physical layout -- i.e. the product of the physical dims + preceding the block dim. Within each segment, block `b` occupies bytes + `[b * S, (b + 1) * S)` where `S = bytes_per_block / num_segments`, so + the extensible KV cache can commit a per-segment prefix of blocks. + """ + has_mamba = self.kv_cache_config.has_mamba_layers + num_segments_by_layer: dict[str, int] = {} + for group in self._kv_cache_spec_attn_group_iterator(): + kv_cache_spec = group.kv_cache_spec + if isinstance(kv_cache_spec, AttentionSpec) and not has_mamba: + attn_backend = group.backend + block_dim = attn_backend.get_kv_cache_block_dim( + kv_cache_spec.block_size, + kv_cache_spec.num_kv_heads, + kv_cache_spec.head_size, + cache_dtype_str=self.cache_config.cache_dtype, + ) + kv_cache_shape = attn_backend.get_kv_cache_shape( + 1, + kv_cache_spec.block_size, + kv_cache_spec.num_kv_heads, + kv_cache_spec.head_size, + cache_dtype_str=self.cache_config.cache_dtype, + ) + try: + stride_order = attn_backend.get_kv_cache_stride_order() + except (AttributeError, NotImplementedError): + stride_order = tuple(range(len(kv_cache_shape))) + num_segments = math.prod( + kv_cache_shape[dim] + for dim in stride_order[: stride_order.index(block_dim)] + ) + else: + # Mamba states are packed per block (block-major), and + # `_update_hybrid_attention_mamba_layout` re-strides attention + # caches of hybrid models to a block-major interleaved layout. + num_segments = 1 + for layer_name in group.layer_names: + num_segments_by_layer[layer_name] = num_segments + return num_segments_by_layer + def _attn_group_iterator(self) -> Iterator[AttentionGroup]: return itertools.chain.from_iterable(self.attn_groups) @@ -7224,8 +7348,35 @@ def _update_hybrid_attention_mamba_layout( stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]), ) + def extend_kv_cache(self, num_blocks: int) -> None: + """Commit physical pages so the KV cache holds `num_blocks` blocks. + + Grows the KV cache after CUDA graph capture, once the available memory + is known. No re-view is needed: the layers already view the full + capacity and each block stays at a fixed offset within its layout + segment, so captured graphs stay valid as more pages are mapped under + the stable base pointer. Newly committed blocks are zeroed. + """ + if not getattr(self, "_extensible_kv_cache_enabled", False): + raise RuntimeError("extend_kv_cache requires an extensible KV cache.") + if num_blocks <= self._extensible_kv_cache_committed_blocks: + return + + for buffer, bytes_per_block_per_segment in self._extensible_kv_cache_buffers: + # Zero only the freshly committed blocks; existing ones are left + # intact. + buffer.resize_per_segment_( + num_blocks * bytes_per_block_per_segment, zero_new=True + ) + + self._extensible_kv_cache_committed_blocks = num_blocks + logger.info("Extended KV cache to %d blocks.", num_blocks) + def initialize_kv_cache_tensors( - self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] + self, + kv_cache_config: KVCacheConfig, + kernel_block_sizes: list[int], + extensible: bool = False, ) -> dict[str, torch.Tensor]: """ Initialize the memory buffer for KV cache. @@ -7241,7 +7392,13 @@ def initialize_kv_cache_tensors( # Try creating KV caches optimized for kv-connector transfers cache_dtype = self.cache_config.cache_dtype - if self.use_uniform_kv_cache(self.attn_groups): + if extensible and self.use_uniform_kv_cache(self.attn_groups): + raise ValueError( + "enable_extensible_kv_cache=True is not supported with " + "cross-layer uniform KV cache layouts." + ) + + if not extensible and self.use_uniform_kv_cache(self.attn_groups): kv_caches, cross_layers_kv_cache, attn_backend = ( self.allocate_uniform_kv_caches( kv_cache_config, @@ -7256,7 +7413,10 @@ def initialize_kv_cache_tensors( else: # Fallback to the general case # Initialize the memory buffer for KV cache - kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config) + kv_cache_raw_tensors = self._allocate_kv_cache_tensors( + kv_cache_config, + extensible=extensible, + ) # Change the memory buffer to the desired shape kv_caches = self._reshape_kv_cache_tensors( @@ -7311,6 +7471,7 @@ def initialize_kv_cache( self, kv_cache_config: KVCacheConfig, is_profiling: bool = False, + extensible: bool = False, ) -> None: """ Initialize KV cache based on `kv_cache_config`. @@ -7343,7 +7504,9 @@ def initialize_kv_cache( # Reinitialize need to after initialize_attn_backend self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes) kv_caches = self.initialize_kv_cache_tensors( - kv_cache_config, kernel_block_sizes + kv_cache_config, + kernel_block_sizes, + extensible=extensible, ) if ( diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 28efee3dee8e..ef16c9e85939 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -479,6 +479,7 @@ def determine_available_memory(self) -> int: current_platform.is_cuda() and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE + and not self.cache_config.enable_extensible_kv_cache ): cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory() @@ -692,7 +693,11 @@ def update_max_model_len(self, max_model_len: int) -> None: logger.debug("Updated max_model_len to %d", max_model_len) @instrument(span_name="Allocate KV cache") - def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: + def initialize_from_config( + self, + kv_cache_config: KVCacheConfig, + extensible: bool = False, + ) -> None: """Allocate GPU KV cache with the specified kv_cache_config.""" # Update local config with adjusted num blocks after profiling, @@ -707,7 +712,12 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config) with self._maybe_get_memory_pool_context(tag="kv_cache"): - self.model_runner.initialize_kv_cache(kv_cache_config) + if extensible: + # Only the V1 GPU model runner implements the extensible flow; + # keep the default call signature untouched otherwise. + self.model_runner.initialize_kv_cache(kv_cache_config, extensible=True) + else: + self.model_runner.initialize_kv_cache(kv_cache_config) if self.model_config.enable_return_routed_experts: self.model_runner.init_routed_experts_capturer() @@ -720,6 +730,10 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: ): self.model_runner._init_kv_zero_meta() + def extend_kv_cache(self, num_blocks: int) -> None: + self.cache_config.num_gpu_blocks = num_blocks + self.model_runner.extend_kv_cache(num_blocks) + @instrument(span_name="Warmup (GPU)") def compile_or_warm_up_model(self) -> CompilationTimes: warmup_sizes: list[int] = [] @@ -894,6 +908,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes: return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time, + cuda_graph=cuda_graph_memory_bytes, ) def reset_mm_cache(self) -> None: diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 9381d71913dc..2c18732f547e 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -34,6 +34,7 @@ class CompilationTimes(NamedTuple): language_model: float encoder: float + cuda_graph: int = 0 class WorkerBase: @@ -99,11 +100,16 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: """Get specifications for KV cache implementation.""" raise NotImplementedError + def extend_kv_cache(self, num_blocks: int) -> None: + raise RuntimeError( + f"{self.__class__.__name__} does not support extensible KV cache." + ) + def compile_or_warm_up_model(self) -> CompilationTimes: """Prepare model for execution through compilation/warmup. Returns: - Compilation times (language_model, encoder) in seconds. + Compilation times in seconds and CUDA graph memory in bytes. """ raise NotImplementedError @@ -318,11 +324,20 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: # To make vLLM config available during worker initialization self.worker = worker_class(**kwargs) - def initialize_from_config(self, kv_cache_configs: list[Any]) -> None: + def initialize_from_config( + self, + kv_cache_configs: list[Any], + extensible: bool = False, + ) -> None: kv_cache_config = kv_cache_configs[self.global_rank] assert self.vllm_config is not None with set_current_vllm_config(self.vllm_config): - self.worker.initialize_from_config(kv_cache_config) # type: ignore + if extensible: + self.worker.initialize_from_config( # type: ignore + kv_cache_config, extensible=True + ) + else: + self.worker.initialize_from_config(kv_cache_config) # type: ignore def init_device(self): assert self.vllm_config is not None