From ba0d8be2a62fe05f0280d296b9cfbac430d36de9 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 2 Nov 2026 11:10:00 -0800 Subject: [PATCH] [Bugfix] Pad extensible KV cache segments instead of rounding the block count With the extensible KV cache, any KV connector made the engine round the reserved capacity, and the count measured after warmup, down to a multiple of granule / gcd(granule, block stride) blocks. Each segment then starts on a commit granule, so a defragmenting commit can back it with one physical chunk for RDMA registration. For awkward strides that multiple is huge (2048 blocks, ~11.6 GB, for a 5,657,600 B stride), so engines with --kv-offloading-size failed to start, and registering connectors lost up to one alignment unit of KV cache. Instead, workers whose connector registers the KV cache pad each layer's region of a layer-outermost layout to a granule boundary. That costs address space only, since commits already map whole granules per segment. A defragmenting commit of any block count then backs every segment with one chunk, so the engine no longer rounds the count, or asks workers for their granule. Block-outermost layouts have a single segment at offset 0 and need no padding. Padding and the defragmenting commit apply only when a connector registers the device KV cache (registers_kv_cache_memory(): NIXL on device, Mooncake, MoRIIO and now UMBP, which registers GPU KV storage through MORI). CPU offloading copies blocks and needs neither. Fixes #65204. Supersedes #65210. Co-authored-by: Vegetog <110553275+Vegetog@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Xgy3B2ybgXQggEmtMSLHgG Signed-off-by: Nick Hill --- tests/v1/core/test_kv_cache_planning.py | 45 ++--- tests/v1/engine/test_engine_args.py | 1 + tests/v1/worker/test_extensible_kv_cache.py | 207 +++++++++++--------- tests/v1/worker/test_gpu_worker.py | 16 +- vllm/config/kv_transfer.py | 9 +- vllm/config/vllm.py | 2 +- vllm/utils/extensible_tensor.py | 25 +-- vllm/v1/core/kv_cache_planning.py | 57 +----- vllm/v1/engine/core.py | 10 - vllm/v1/worker/extensible_kv_cache.py | 148 +++++++++----- vllm/v1/worker/gpu/model_runner.py | 6 +- vllm/v1/worker/gpu_worker.py | 12 +- 12 files changed, 258 insertions(+), 280 deletions(-) diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py index 43221779a506..1612aa0ed304 100644 --- a/tests/v1/core/test_kv_cache_planning.py +++ b/tests/v1/core/test_kv_cache_planning.py @@ -2214,53 +2214,32 @@ def test_shrink_kv_cache_configs_keeps_tensors_and_checks_fit(): assert config.num_blocks == 5 -def test_extensible_kv_cache_granule_alignment_under_connector(): - """With a KV connector, capacity and the final measured count are rounded - down so every layer segment spans whole commit granules (one RDMA - registration each); without one the measured count is kept as is.""" +def test_finalize_extensible_kv_cache_keeps_measured_count(): + """The smallest count measured across ranks becomes every config's + num_blocks, not rounded to granules.""" from vllm.v1.worker.worker_base import CompilationTimes - spec = new_kv_cache_spec() # 1 head x 128 x fp16 x 2 x block 16 = 8 KiB - granule = 2 << 20 - blocks_per_granule = granule // spec.page_size_bytes + spec = new_kv_cache_spec() vllm_config = VllmConfig( model_config=ModelConfig(max_model_len=spec.block_size * 4) ) vllm_config.cache_config.kv_cache_layout = "LBNHC" groups = [KVCacheGroupSpec(["layer.0", "layer.1"], spec)] - capacity = 3 * blocks_per_granule + 7 + capacity = 1000 config = kv_cache_planning.get_kv_cache_config_from_groups( vllm_config, groups, available_memory=2 * spec.page_size_bytes * capacity ) scheduler_config = generate_scheduler_kv_cache_config([config]) - assert config.num_blocks == capacity - - kv_cache_planning.align_extensible_kv_cache_capacity( - vllm_config, [config], scheduler_config, granule - ) - aligned = 3 * blocks_per_granule - assert config.num_blocks == aligned - assert scheduler_config.num_blocks == aligned - assert vllm_config.cache_config.num_gpu_blocks == aligned - # The placements describe the smaller reservation to be allocated: the - # one two-layer tensor places its layers `aligned` blocks apart. - (tensor,) = config.kv_cache_tensors - assert tensor.layer_stride == aligned * spec.page_size_bytes - assert tensor.size == 2 * aligned * spec.page_size_bytes - - measured = [CompilationTimes(0.0, 0.0, num_kv_blocks=2 * blocks_per_granule + 5)] - assert ( - kv_cache_planning.finalize_extensible_kv_cache( - vllm_config, [config], scheduler_config, measured, commit_granule=granule - ) - == 2 * blocks_per_granule - ) - assert config.num_blocks == 2 * blocks_per_granule - measured = [CompilationTimes(0.0, 0.0, num_kv_blocks=blocks_per_granule + 5)] + measured = [ + CompilationTimes(0.0, 0.0, num_kv_blocks=num_blocks) + for num_blocks in (999, 997) + ] assert ( kv_cache_planning.finalize_extensible_kv_cache( vllm_config, [config], scheduler_config, measured ) - == blocks_per_granule + 5 + == 997 ) + assert config.num_blocks == scheduler_config.num_blocks == 997 + assert vllm_config.cache_config.num_gpu_blocks == 997 diff --git a/tests/v1/engine/test_engine_args.py b/tests/v1/engine/test_engine_args.py index c4c2963412c4..b04b49016885 100644 --- a/tests/v1/engine/test_engine_args.py +++ b/tests/v1/engine/test_engine_args.py @@ -424,6 +424,7 @@ def test_extensible_kv_cache_off_with_elastic_ep(monkeypatch, requested): ("MooncakeConnector", "cuda", True), ("MooncakeStoreConnector", "cuda", True), ("MoRIIOConnector", "cuda", True), + ("UMBPStoreConnector", "cuda", True), ("OffloadingConnector", "cuda", False), ("LMCacheConnectorV1", "cuda", False), ], diff --git a/tests/v1/worker/test_extensible_kv_cache.py b/tests/v1/worker/test_extensible_kv_cache.py index e90a2da79680..35004b5bd4f8 100644 --- a/tests/v1/worker/test_extensible_kv_cache.py +++ b/tests/v1/worker/test_extensible_kv_cache.py @@ -8,10 +8,6 @@ import torch from vllm.config import ModelConfig, VllmConfig -from vllm.utils.extensible_tensor import ( - granule_aligned_blocks, - granule_block_alignment, -) from vllm.utils.vmm_driver import vmm_unavailable_reason from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import ( @@ -25,6 +21,7 @@ from vllm.v1.worker.extensible_kv_cache import ( ExtensibleKVCache, measure_kv_cache_bytes, + pad_kv_cache_segments, ) from vllm.v1.worker.utils import allocate_kv_cache @@ -69,7 +66,7 @@ def test_segments_follow_layout(vmm, layout, expected_segments): """Layer-outermost layouts commit one prefix per layer; block-outermost layouts commit a single prefix of the whole allocation.""" config, spec = _make_config(layout) - kv_cache = ExtensibleKVCache(config, torch.device("cuda")) + kv_cache = ExtensibleKVCache(config, torch.device("cuda"), layout) try: assert kv_cache.buffer.num_segments == expected_segments assert kv_cache.bytes_per_block == NUM_LAYERS * spec.page_size_bytes @@ -86,7 +83,7 @@ def test_views_stay_valid_across_commits(vmm, layout): committed blocks read as zero.""" config, spec = _make_config(layout) device = torch.device("cuda") - kv_cache = ExtensibleKVCache(config, device) + kv_cache = ExtensibleKVCache(config, device, layout) try: kv_caches = allocate_kv_cache( config, device, layout, allocate=kv_cache.allocate @@ -191,7 +188,7 @@ def test_measured_sizing_caps_torch_memory(monkeypatch): reserved_headroom_bytes=gib, physical_bytes=gib, bytes_per_block=gib, - blocks_within=lambda num_bytes, aligned: num_bytes // gib, + blocks_within=lambda num_bytes: num_bytes // gib, physical_bytes_for=lambda num_blocks: num_blocks * gib, set_torch_memory_limit=limits.append, ) @@ -249,22 +246,6 @@ def test_freeing_the_cache_lifts_the_torch_memory_cap(monkeypatch): assert kv_cache.torch_memory_limit is None -def test_granule_aligned_blocks(): - """Blocks per stride must span whole granules; with several strides the - count is a multiple of every stride's requirement.""" - granule = 2 << 20 - # 64 KiB blocks: 32 per granule. - assert granule_block_alignment([64 << 10], granule) == 32 - assert granule_aligned_blocks(1000, [64 << 10], granule) == 992 - # 96 KiB blocks: 3 blocks cover 288 KiB; the granule is reached at 64. - assert granule_block_alignment([96 << 10], granule) == 64 - # Mixed strides: lcm(32, 64). - assert granule_aligned_blocks(1000, [64 << 10, 96 << 10], granule) == 960 - # A stride that is a granule multiple needs no alignment. - assert granule_aligned_blocks(1000, [4 << 20], granule) == 1000 - assert granule_aligned_blocks(31, [64 << 10], granule) == 0 - - def _make_mixed_config(num_blocks: int = NUM_BLOCKS): """One group whose layers differ in head size, so their block strides differ.""" specs = { @@ -290,13 +271,46 @@ def _make_mixed_config(num_blocks: int = NUM_BLOCKS): return config, specs +def test_pad_kv_cache_segments(): + """Layer-outermost regions are placed on granules: offsets, layer strides + and the size round up while block count and strides stay. Block-outermost + layouts have one segment at offset 0 and are left as they are.""" + granule = 2 << 20 + page = _make_config(KVCacheLayout.LBNHC)[1].page_size_bytes + # One block past a granule per layer. + num_blocks = granule // page + 1 + config, _ = _make_config(KVCacheLayout.LBNHC, num_blocks=num_blocks) + padded = pad_kv_cache_segments(config, KVCacheLayout.LBNHC, granule) + assert padded.num_blocks == num_blocks + (tensor,) = padded.kv_cache_tensors + assert tensor.block_stride == page + assert tensor.layer_stride == 2 * granule + assert tensor.offset == 0 + assert tensor.size == NUM_LAYERS * 2 * granule + assert config.kv_cache_tensors[0].layer_stride == num_blocks * page + + # Layers 0 and 2 share a stride, so one tensor places both ahead of layer 1. + config, _ = _make_mixed_config() + padded = pad_kv_cache_segments(config, KVCacheLayout.LBNHC, granule) + assert [t.layers for t in padded.kv_cache_tensors] == [ + ["layer.0", "layer.2"], + ["layer.1"], + ] + assert [t.offset for t in padded.kv_cache_tensors] == [0, 2 * granule] + assert all(t.layer_stride == granule for t in padded.kv_cache_tensors) + assert {t.size for t in padded.kv_cache_tensors} == {NUM_LAYERS * granule} + + config, _ = _make_config(KVCacheLayout.BLNHC) + assert pad_kv_cache_segments(config, KVCacheLayout.BLNHC, granule) is config + + @requires_cuda def test_mixed_block_strides_commit_per_layer(vmm): """Layers with different page sizes get segments of their own stride; a commit backs the same block count in each and the views stay valid.""" config, specs = _make_mixed_config() device = torch.device("cuda") - kv_cache = ExtensibleKVCache(config, device) + kv_cache = ExtensibleKVCache(config, device, KVCacheLayout.LBNHC) try: assert len(kv_cache.segment_strides) == NUM_LAYERS assert sorted(kv_cache.segment_strides) == sorted( @@ -331,7 +345,7 @@ def test_release_and_recommit_for_sleep(vmm): count under the same addresses, so existing views keep working.""" config, _ = _make_config(KVCacheLayout.LBNHC) device = torch.device("cuda") - kv_cache = ExtensibleKVCache(config, device) + kv_cache = ExtensibleKVCache(config, device, KVCacheLayout.LBNHC) try: views = list( allocate_kv_cache( @@ -362,7 +376,7 @@ def test_shrinking_commit_remaps_to_the_smaller_prefix(vmm, layout): smaller cache.""" config, _ = _make_config(layout) device = torch.device("cuda") - kv_cache = ExtensibleKVCache(config, device) + kv_cache = ExtensibleKVCache(config, device, layout) try: capacity_views = allocate_kv_cache( config, device, layout, allocate=kv_cache.allocate @@ -404,7 +418,7 @@ def test_committed_views_cover_exactly_the_committed_bytes(vmm, layout): views, so connectors deriving extents from storage size see backed memory.""" config, spec = _make_config(layout) device = torch.device("cuda") - kv_cache = ExtensibleKVCache(config, device) + kv_cache = ExtensibleKVCache(config, device, layout) try: capacity_views = allocate_kv_cache( config, device, layout, allocate=kv_cache.allocate @@ -442,7 +456,7 @@ def _segments_backed_by_one_chunk(kv_cache: ExtensibleKVCache) -> list[bool]: def _granule(layout: KVCacheLayout) -> int: - probe = ExtensibleKVCache(_make_config(layout)[0], torch.device("cuda")) + probe = ExtensibleKVCache(_make_config(layout)[0], torch.device("cuda"), layout) try: return probe.buffer.granularity finally: @@ -450,74 +464,80 @@ def _granule(layout: KVCacheLayout) -> int: @requires_cuda -def test_defragmenting_commit_backs_each_segment_with_one_allocation(vmm): - """Incremental commits map one driver allocation each; a defragmenting - commit of a granule-aligned count remaps every segment as a single - allocation (RDMA cannot span several) and still ends up zero-filled.""" - granule = _granule(KVCacheLayout.LBNHC) - stride = _make_config(KVCacheLayout.LBNHC)[1].page_size_bytes - alignment = granule_block_alignment([stride], granule) - # An unaligned capacity puts segment offsets inside granules. - capacity = 2 * alignment + alignment // 2 - config, _ = _make_config(KVCacheLayout.LBNHC, num_blocks=capacity) - kv_cache = ExtensibleKVCache(config, torch.device("cuda")) +@pytest.mark.parametrize("padded", [False, True]) +def test_defragmenting_commit_backs_each_segment_with_one_allocation(vmm, padded): + """Incremental commits map one driver allocation each. Unpadded, segments + after the first start mid-granule and share allocations even after a + defragmenting commit; padded, a defragmenting commit of any block count + remaps each segment as one allocation (RDMA cannot span several).""" + layout = KVCacheLayout.LBNHC + granule = _granule(layout) + page = _make_config(layout)[1].page_size_bytes + blocks_per_granule = granule // page + capacity = 2 * blocks_per_granule + blocks_per_granule // 2 + config, _ = _make_config(layout, num_blocks=capacity) + if padded: + config = pad_kv_cache_segments(config, layout, granule) + kv_cache = ExtensibleKVCache(config, torch.device("cuda"), layout) try: - assert kv_cache.block_alignment == alignment kv_cache.commit(1) - kv_cache.commit(alignment + 1) - num_segments = kv_cache.buffer.num_segments - assert kv_cache.buffer.num_physical_chunks == 2 * num_segments - - kv_cache.commit(capacity, defragment=True) - assert kv_cache.num_committed_blocks == capacity - assert torch.count_nonzero(kv_cache.buffer.full_view()) == 0 - # Segments after the first start mid-granule and share chunks. - assert not all(_segments_backed_by_one_chunk(kv_cache)) - finally: - kv_cache.free() - - config, _ = _make_config(KVCacheLayout.LBNHC, num_blocks=2 * alignment) - kv_cache = ExtensibleKVCache(config, torch.device("cuda")) - try: - kv_cache.commit(alignment + 1) - kv_cache.commit( - kv_cache.blocks_within(kv_cache.size, aligned=True), defragment=True - ) - assert kv_cache.num_committed_blocks == 2 * alignment + kv_cache.commit(blocks_per_granule + 1) buffer = kv_cache.buffer - assert buffer.num_physical_chunks == buffer.num_segments - assert all(_segments_backed_by_one_chunk(kv_cache)) + assert buffer.num_segments == NUM_LAYERS + assert buffer.num_physical_chunks == 2 * NUM_LAYERS + + kv_cache.commit(capacity - 3, defragment=True) + assert kv_cache.num_committed_blocks == capacity - 3 + for segment in range(buffer.num_segments): + assert torch.count_nonzero(buffer.segment_view(segment)) == 0 + assert all(_segments_backed_by_one_chunk(kv_cache)) == padded + if padded: + assert buffer.num_physical_chunks == NUM_LAYERS finally: kv_cache.free() @requires_cuda -@pytest.mark.parametrize("layout", [KVCacheLayout.LBNHC, KVCacheLayout.BLNHC]) -def test_aligned_capacity_config_allocates_and_defragments(vmm, layout): - """The engine-side alignment yields a config the cache can be built from, - with views over it, and whose full defragmenting commit backs each segment - with one allocation.""" - from vllm.v1.core.kv_cache_planning import align_extensible_kv_cache_capacity - from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config - +@pytest.mark.parametrize("case", ["LBNHC", "BLNHC", "mixed"]) +def test_padded_config_allocates_and_defragments(vmm, case): + """A padded config yields a cache with views over its reservation; a + defragmenting commit of an arbitrary count backs each segment with one + allocation, and the committed views and connector placements address the + same bytes as the capacity views.""" + layout = KVCacheLayout.LBNHC if case == "mixed" else KVCacheLayout[case] granule = _granule(layout) - config, spec = _make_config(layout, num_blocks=1000) - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.kv_cache_layout = layout.name - scheduler_config = generate_scheduler_kv_cache_config([config]) - align_extensible_kv_cache_capacity(vllm_config, [config], scheduler_config, granule) - assert config.num_blocks < 1000 + capacity = 1000 + if case == "mixed": + config, _ = _make_mixed_config(num_blocks=capacity) + else: + config, _ = _make_config(layout, num_blocks=capacity) + config = pad_kv_cache_segments(config, layout, granule) device = torch.device("cuda") - kv_cache = ExtensibleKVCache(config, device) + kv_cache = ExtensibleKVCache(config, device, layout) try: + assert all(offset % granule == 0 for offset in kv_cache.segment_offsets) views = allocate_kv_cache(config, device, layout, allocate=kv_cache.allocate) - assert all(v.shape[0] == config.num_blocks for v in views.values()) + assert all(v.shape[0] == capacity for v in views.values()) kv_cache.commit(3) - kv_cache.commit(config.num_blocks, defragment=True) + num_blocks = capacity - 3 + kv_cache.commit(num_blocks, defragment=True) assert all(_segments_backed_by_one_chunk(kv_cache)) - for view in views.values(): - view.fill_(1.0) + for i, view in enumerate(views.values()): + view[:num_blocks].fill_(float(i + 1)) torch.accelerator.synchronize() + + final = copy.copy(config) + final.num_blocks = num_blocks + committed = kv_cache.committed_views(final, layout) + assert committed is not None + for i, (name, view) in enumerate(views.items()): + assert committed[name].data_ptr() == view.data_ptr() + assert torch.all(committed[name] == float(i + 1)) + for tensor in kv_cache.committed_kv_cache_tensors(final, layout, num_blocks): + for layer_idx, name in enumerate(tensor.layers): + storage = committed[name].untyped_storage() + layer_start = tensor.offset + layer_idx * tensor.layer_stride + assert storage.data_ptr() + layer_start == committed[name].data_ptr() finally: kv_cache.free() @@ -525,24 +545,21 @@ def test_aligned_capacity_config_allocates_and_defragments(vmm, layout): @requires_cuda def test_blocks_within_accounts_for_granule_rounding(vmm): """The count that fits a byte budget is the most whose rounded-up - per-segment footprint fits, not the budget over the nominal block size; - aligned, it is the largest aligned count under the nominal footprint.""" + per-segment footprint fits, not the budget over the nominal block size.""" config, spec = _make_config(KVCacheLayout.LBNHC, num_blocks=4096) - kv_cache = ExtensibleKVCache(config, torch.device("cuda")) + kv_cache = ExtensibleKVCache(config, torch.device("cuda"), KVCacheLayout.LBNHC) try: granule = kv_cache.buffer.granularity - alignment = kv_cache.block_alignment + blocks_per_granule = granule // spec.page_size_bytes num_segments = kv_cache.buffer.num_segments - # One granule per segment holds `alignment` blocks exactly; one more - # block per segment needs a whole extra granule each. + # One granule per segment holds `blocks_per_granule` blocks exactly; + # one more block per segment needs a whole extra granule each. budget = num_segments * granule - assert kv_cache.blocks_within(budget) == alignment - assert kv_cache.blocks_within(budget + 1) == alignment - assert kv_cache.physical_bytes_for(alignment + 1) == 2 * budget - assert kv_cache.blocks_within(2 * budget - 1) == alignment - assert kv_cache.blocks_within(2 * budget) == 2 * alignment - nominal = (2 * alignment + 5) * kv_cache.bytes_per_block - assert kv_cache.blocks_within(nominal, aligned=True) == 2 * alignment + assert kv_cache.blocks_within(budget) == blocks_per_granule + assert kv_cache.blocks_within(budget + 1) == blocks_per_granule + assert kv_cache.physical_bytes_for(blocks_per_granule + 1) == 2 * budget + assert kv_cache.blocks_within(2 * budget - 1) == blocks_per_granule + assert kv_cache.blocks_within(2 * budget) == 2 * blocks_per_granule assert kv_cache.blocks_within(kv_cache.size * 2) == config.num_blocks finally: kv_cache.free() @@ -555,7 +572,7 @@ def test_committed_kv_cache_tensors_match_views(vmm, layout): the same bytes as the committed views, within storages of their size.""" config, _ = _make_config(layout) device = torch.device("cuda") - kv_cache = ExtensibleKVCache(config, device) + kv_cache = ExtensibleKVCache(config, device, layout) try: allocate_kv_cache(config, device, layout, allocate=kv_cache.allocate) committed = 5 diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 687b57a11fe6..1c90edda7f0c 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -526,6 +526,7 @@ def fail(*args, **kwargs): worker = SimpleNamespace( cache_config=SimpleNamespace(enable_extensible_kv_cache=True), model_config=SimpleNamespace(enable_return_routed_experts=False), + vllm_config=SimpleNamespace(kv_transfer_config=None), model_runner=runner, _maybe_get_memory_pool_context=lambda tag: nullcontext(), _v2_model_runner=lambda: runner, @@ -590,8 +591,8 @@ def get_mem_pool_context(): log.append("get_mem_pool_context") return connector_pool - def commit(num_blocks, **kwargs): - log.append(f"commit:{num_blocks}") + def commit(num_blocks, defragment, **kwargs): + log.append(f"commit:{num_blocks}:defragment={defragment}") raise Committed monkeypatch.setattr( @@ -614,14 +615,21 @@ def commit(num_blocks, **kwargs): commit=commit, ), cache_config=SimpleNamespace(get_resolved_kv_cache_layout=lambda: None), - vllm_config=SimpleNamespace(kv_transfer_config=object()), + # A connector that registers the KV cache needs one chunk per segment. + vllm_config=SimpleNamespace( + kv_transfer_config=SimpleNamespace(registers_kv_cache_memory=lambda: True) + ), ) kv_cache_config = SimpleNamespace(num_blocks=16, kv_cache_tensors=[]) if connector_pool is None: with pytest.raises(Committed): extensible_kv_cache.extend_kv_cache(runner, kv_cache_config, 8) - assert log == ["ensure_kv_transfer:8", "get_mem_pool_context", "commit:8"] + assert log == [ + "ensure_kv_transfer:8", + "get_mem_pool_context", + "commit:8:defragment=True", + ] else: with pytest.raises(ValueError, match="own memory pool"): extensible_kv_cache.extend_kv_cache(runner, kv_cache_config, 8) diff --git a/vllm/config/kv_transfer.py b/vllm/config/kv_transfer.py index 30482d3a22fe..68114357fb41 100644 --- a/vllm/config/kv_transfer.py +++ b/vllm/config/kv_transfer.py @@ -50,7 +50,12 @@ def hisparse_host_pool_gib( # Connectors that register the device KV cache with an RDMA engine whatever # their buffer device; NIXL is checked separately as it can stage through host. _KV_CACHE_REGISTERING_CONNECTORS = frozenset( - {"MooncakeConnector", "MooncakeStoreConnector", "MoRIIOConnector"} + { + "MooncakeConnector", + "MooncakeStoreConnector", + "MoRIIOConnector", + "UMBPStoreConnector", + } ) @@ -158,7 +163,7 @@ def uses_custom_mem_pool(self) -> bool: def registers_kv_cache_memory(self) -> bool: """Whether a connector registers the device KV cache itself with UCX or - RDMA (NIXL, Mooncake, MoRIIO), directly or in MultiConnector. NIXL with + RDMA (NIXL, Mooncake, MoRIIO, UMBP), directly or in MultiConnector. NIXL with a CPU kv_buffer_device registers only its host staging buffers.""" children = [{"kv_connector": self.kv_connector}] if self.kv_connector == "MultiConnector": diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7f381384dddc..c641ca022576 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3234,7 +3234,7 @@ def _resolve_extensible_kv_cache(self) -> None: # hsa_amd_portable_export_dmabuf accepts VMM ranges. unsupported = ( "on ROCm, UCX and RDMA cannot yet register driver-mapped memory, " - "which the NIXL, Mooncake and MoRIIO connectors require" + "which the NIXL, Mooncake, MoRIIO and UMBP connectors require" ) if requested is None: cache_config.enable_extensible_kv_cache = unsupported is None diff --git a/vllm/utils/extensible_tensor.py b/vllm/utils/extensible_tensor.py index 2a11b1032437..b816865644b8 100644 --- a/vllm/utils/extensible_tensor.py +++ b/vllm/utils/extensible_tensor.py @@ -6,9 +6,8 @@ import ctypes import itertools -import math import sys -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from contextlib import suppress from typing import Any @@ -21,28 +20,6 @@ def _round_up(value: int, multiple: int) -> int: return ((value + multiple - 1) // multiple) * multiple -def granule_block_alignment(block_strides: Iterable[int], granule: int) -> int: - """Block count multiple at which every stride's prefix is a granule multiple.""" - alignment = 1 - for stride in set(block_strides): - alignment = math.lcm(alignment, granule // math.gcd(granule, stride)) - return alignment - - -def granule_aligned_blocks( - num_blocks: int, block_strides: Iterable[int], granule: int -) -> int: - """Largest count up to num_blocks at which, for every block stride, that - many blocks span a whole number of granules. - - Segments hold num_blocks blocks of one stride each, so at such a count - every segment starts and ends on a granule boundary and one physical chunk - can back it exactly, as RDMA registration of a segment requires. - """ - alignment = granule_block_alignment(block_strides, granule) - return num_blocks // alignment * alignment - - # ROCm reports a 4 KiB granularity (CUDA 2 MiB); commit in larger units to # keep the granule bookkeeping small. _MIN_GRANULARITY = 2 << 20 diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py index ecb1cd44dd14..48545d5bc84a 100644 --- a/vllm/v1/core/kv_cache_planning.py +++ b/vllm/v1/core/kv_cache_planning.py @@ -11,7 +11,6 @@ from vllm.config import SpeculativeConfig, VllmConfig from vllm.logger import init_logger -from vllm.utils.extensible_tensor import granule_aligned_blocks from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import get_dtype_size @@ -1770,66 +1769,16 @@ def shrink_kv_cache_configs( kv_cache_config.num_blocks = num_blocks -def granule_aligned_kv_cache_blocks( - num_blocks: int, kv_cache_configs: list[KVCacheConfig], commit_granule: int -) -> int: - """Round num_blocks down so every worker's KV cache segments (one per - layer or one overall, num_blocks blocks of one tensor's block stride) - span whole commit granules. A KV connector registers each committed - segment for RDMA, which cannot span several physical chunks.""" - strides = { - tensor.block_stride - for kv_cache_config in kv_cache_configs - for tensor in kv_cache_config.kv_cache_tensors - } - return granule_aligned_blocks(num_blocks, strides, commit_granule) - - -def align_extensible_kv_cache_capacity( - vllm_config: VllmConfig, - kv_cache_configs: list[KVCacheConfig], - scheduler_kv_cache_config: KVCacheConfig, - commit_granule: int, -) -> None: - """Lower the reserved capacity to a granule-aligned block count, so that - segment offsets (multiples of the capacity) sit on granule boundaries.""" - num_blocks = granule_aligned_kv_cache_blocks( - scheduler_kv_cache_config.num_blocks, kv_cache_configs, commit_granule - ) - if num_blocks >= scheduler_kv_cache_config.num_blocks: - return - for kv_cache_config in kv_cache_configs + [scheduler_kv_cache_config]: - old_num_blocks = kv_cache_config.num_blocks - shrink_kv_cache_configs(vllm_config, [kv_cache_config], num_blocks) - # Unlike the post-warmup shrink, this happens before allocation: the - # placements must describe the smaller reservation. Layer-outermost - # tensors place layers `num_blocks` blocks apart. - kv_cache_config.kv_cache_tensors = [ - replace( - tensor, - layer_stride=num_blocks * tensor.block_stride, - offset=tensor.offset // old_num_blocks * num_blocks, - ) - if tensor.layer_stride == old_num_blocks * tensor.block_stride - else tensor - for tensor in kv_cache_config.kv_cache_tensors - ] - vllm_config.cache_config.num_gpu_blocks = num_blocks - - def finalize_extensible_kv_cache( vllm_config: VllmConfig, kv_cache_configs: list[KVCacheConfig], scheduler_kv_cache_config: KVCacheConfig, compilation_times: list["CompilationTimes"], - commit_granule: int | None = None, ) -> int: """Settle the KV cache size from the memory workers measured after warmup. Every rank must agree on one block count, so the minimum becomes the - scheduler's num_blocks; the caller commits it on the workers. With - commit_granule (set when a KV connector is configured) the count is - also granule-aligned, see `granule_aligned_kv_cache_blocks`. + scheduler's num_blocks; the caller commits it on the workers. """ # No measurements: the elastic-EP scale-up path skips warmup. Elastic EP # leaves the extensible KV cache off, so this branch is not reached; it @@ -1841,10 +1790,6 @@ def finalize_extensible_kv_cache( ) num_blocks = min(block_counts, default=scheduler_kv_cache_config.num_blocks) num_blocks = min(num_blocks, scheduler_kv_cache_config.num_blocks) - if commit_granule is not None: - num_blocks = granule_aligned_kv_cache_blocks( - num_blocks, kv_cache_configs, commit_granule - ) shrink_kv_cache_configs( vllm_config, kv_cache_configs + [scheduler_kv_cache_config], num_blocks ) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 05f9fda5c5f6..ae10644962c1 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -56,7 +56,6 @@ from vllm.v1.attention.backends.utils import resolve_kv_cache_layout from vllm.v1.core import kv_cache_utils from vllm.v1.core.kv_cache_planning import ( - align_extensible_kv_cache_capacity, finalize_extensible_kv_cache, get_global_kv_cache_groups, get_kv_cache_configs, @@ -379,14 +378,6 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: # With an extensible KV cache, `num_blocks` is the reserved capacity # until the post-warmup measurement below. extensible = vllm_config.cache_config.enable_extensible_kv_cache - # A KV connector may register each committed segment for RDMA, which - # cannot span physical chunks: keep segment bounds granule-aligned. - commit_granule: int | None = None - if extensible and kv_cache_groups and vllm_config.kv_transfer_config: - commit_granule = max(self.collective_rpc("kv_cache_commit_granule")) - align_extensible_kv_cache_capacity( - vllm_config, kv_cache_configs, scheduler_kv_cache_config, commit_granule - ) if kv_cache_groups and not extensible: update_kv_cache_capacity(vllm_config, scheduler_kv_cache_config) @@ -418,7 +409,6 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: kv_cache_configs, scheduler_kv_cache_config, compilation_times, - commit_granule, ) self.model_executor.extend_kv_cache(num_blocks) # Under the extensible KV cache, workers create the connector in diff --git a/vllm/v1/worker/extensible_kv_cache.py b/vllm/v1/worker/extensible_kv_cache.py index 05a0ba746073..6b2cd75bc7de 100644 --- a/vllm/v1/worker/extensible_kv_cache.py +++ b/vllm/v1/worker/extensible_kv_cache.py @@ -8,6 +8,7 @@ from collections.abc import Iterator from contextlib import contextmanager from dataclasses import replace +from itertools import pairwise from typing import TYPE_CHECKING import torch @@ -21,11 +22,8 @@ from vllm.model_executor.model_loader.reload.layerwise import max_layer_reload_bytes from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory from vllm.platforms import current_platform -from vllm.utils.extensible_tensor import ( - ExtensibleTensor, - granule_aligned_blocks, - granule_block_alignment, -) +from vllm.utils.extensible_tensor import ExtensibleTensor +from vllm.utils.math_utils import round_up from vllm.utils.mem_utils import format_gib from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -76,24 +74,27 @@ class ExtensibleKVCache: prefix of blocks on demand. Blocks keep fixed offsets, so views over the reservation and captured CUDA graphs stay valid as more are committed. The allocation splits into segments holding num_blocks blocks each: one per - layer for layer-outermost layouts (sized by that layer's block stride), one + layer for block-compact layer-outermost layouts (sized by that layer's + block stride, possibly padded apart by `pad_kv_cache_segments`), one overall for block-outermost. """ - def __init__(self, kv_cache_config: KVCacheConfig, device: torch.device): + def __init__( + self, + kv_cache_config: KVCacheConfig, + device: torch.device, + layout: KVCacheLayout, + ): tensors = kv_cache_config.kv_cache_tensors assert tensors, "Extensible KV cache requires at least one KV cache tensor." sizes = {tensor.size for tensor in tensors} assert len(sizes) == 1, "KV cache tensors must share one backing allocation." self.size = sizes.pop() self.capacity_blocks = kv_cache_config.num_blocks - if self.capacity_blocks <= 0 or self.size % self.capacity_blocks != 0: - raise ValueError( - f"KV cache allocation of {self.size} bytes is not a whole number " - f"of {self.capacity_blocks} blocks." - ) - self.bytes_per_block = self.size // self.capacity_blocks - self.segment_offsets, self.segment_strides = self._segments(tensors) + if self.capacity_blocks <= 0: + raise ValueError(f"Invalid KV cache capacity {self.capacity_blocks}.") + self.segment_offsets, self.segment_strides = self._segments(tensors, layout) + self.bytes_per_block = sum(self.segment_strides) self.num_committed_blocks = 0 # Memory to keep free while committing before the final sizing, e.g. # the profiled activation peak the warmup steps are about to hit. @@ -108,31 +109,41 @@ def __init__(self, kv_cache_config: KVCacheConfig, device: torch.device): self.size, device=device, segment_capacities=[ - self.capacity_blocks * stride for stride in self.segment_strides + end - start + for start, end in pairwise([*self.segment_offsets, self.size]) ], ) - def _segments(self, tensors: list[KVCacheTensor]) -> tuple[list[int], list[int]]: - """Segment offsets and block strides tiling the allocation. + def _segments( + self, tensors: list[KVCacheTensor], layout: KVCacheLayout + ) -> tuple[list[int], list[int]]: + """Segment offsets and block strides covering the allocation. - With one block stride the segments are implied by it; with several, each - tensor must give its layers contiguous num_blocks-block regions. + Block-compact layer-outermost layouts give each layer a num_blocks-block + region, which may be padded apart. Otherwise the segments are implied by + the one block stride. """ - block_strides = {tensor.block_stride for tensor in tensors} - if len(block_strides) == 1: + if not layout.is_block_compact or layout.is_block_outermost: + block_strides = {tensor.block_stride for tensor in tensors} + if len(block_strides) != 1: + raise ValueError( + "KV cache tensors with different block strides need a " + "block-compact layer-outermost layout for the extensible KV " + "cache." + ) stride = block_strides.pop() segment_bytes = self.capacity_blocks * stride + if self.size % segment_bytes != 0: + raise ValueError( + f"KV cache allocation of {self.size} bytes is not a whole " + f"number of {self.capacity_blocks}-block segments." + ) num_segments = self.size // segment_bytes return [i * segment_bytes for i in range(num_segments)], [ stride ] * num_segments segments: dict[int, int] = {} for tensor in tensors: - if tensor.layer_stride != self.capacity_blocks * tensor.block_stride: - raise ValueError( - "KV cache tensors with different block strides need a " - "layer-outermost layout for the extensible KV cache." - ) for layer_idx in range(len(tensor.layers)): offset = tensor.offset + layer_idx * tensor.layer_stride stride = segments.setdefault(offset, tensor.block_stride) @@ -143,12 +154,12 @@ def _segments(self, tensors: list[KVCacheTensor]) -> tuple[list[int], list[int]] ) offsets = sorted(segments) strides = [segments[offset] for offset in offsets] - end = 0 + end = offsets[0] for offset, stride in zip(offsets, strides): - if offset != end: - raise ValueError("KV cache tensors do not tile the allocation.") - end += self.capacity_blocks * stride - if end != self.size: + if offset < end: + raise ValueError("KV cache tensors overlap.") + end = offset + self.capacity_blocks * stride + if offsets[0] != 0 or end > self.size: raise ValueError("KV cache tensors do not tile the allocation.") return offsets, strides @@ -260,22 +271,9 @@ def physical_bytes_for(self, num_blocks: int) -> int: for stride in self.segment_strides ) - @property - def block_alignment(self) -> int: - """Block count multiple at which every segment ends on a granule.""" - return granule_block_alignment(self.segment_strides, self.buffer.granularity) - - def blocks_within(self, num_bytes: int, aligned: bool = False) -> int: - """Most blocks whose physical footprint fits in num_bytes. - - aligned restricts the count to multiples of block_alignment, so - that a defragmenting commit backs each segment with one chunk exactly. - """ + def blocks_within(self, num_bytes: int) -> int: + """Most blocks whose physical footprint fits in num_bytes.""" num_blocks = min(num_bytes // self.bytes_per_block, self.capacity_blocks) - if aligned: - return granule_aligned_blocks( - num_blocks, self.segment_strides, self.buffer.granularity - ) while num_blocks > 0 and self.physical_bytes_for(num_blocks) > num_bytes: num_blocks -= 1 return num_blocks @@ -350,6 +348,58 @@ def free(self) -> None: self.num_committed_blocks = 0 +def registers_kv_cache_memory(vllm_config: "VllmConfig") -> bool: + """Whether a KV connector registers the committed KV cache for RDMA, which + needs each segment backed by one physical chunk.""" + kv_transfer_config = vllm_config.kv_transfer_config + return ( + kv_transfer_config is not None + and kv_transfer_config.registers_kv_cache_memory() + ) + + +def pad_kv_cache_segments( + kv_cache_config: KVCacheConfig, layout: KVCacheLayout, granule: int +) -> KVCacheConfig: + """Start every layer's region on a granule, so that a defragmenting commit + of any block count can back each segment with one physical chunk. + + Only address space grows: commits already map whole granules per segment. + Block-outermost layouts have a single segment at offset 0 and are returned + as is. + """ + if layout.is_block_outermost: + return kv_cache_config + assert layout.is_block_compact, f"{layout.name} has no per-layer segments." + num_blocks = kv_cache_config.num_blocks + tensors = kv_cache_config.kv_cache_tensors + assert all(t.layer_stride == num_blocks * t.block_stride for t in tensors) + # Cache groups alias each other's layer regions. + regions = { + tensor.offset + layer_idx * tensor.layer_stride: tensor.layer_stride + for tensor in tensors + for layer_idx in range(len(tensor.layers)) + } + padded_offsets: dict[int, int] = {} + end = padded_end = 0 + for offset in sorted(regions): + assert offset == end, "KV cache tensors do not tile the allocation." + padded_offsets[offset] = padded_end + end += regions[offset] + padded_end += round_up(regions[offset], granule) + padded = copy.copy(kv_cache_config) + padded.kv_cache_tensors = [ + replace( + tensor, + size=padded_end, + offset=padded_offsets[tensor.offset], + layer_stride=round_up(tensor.layer_stride, granule), + ) + for tensor in tensors + ] + return padded + + def num_committable_kv_blocks(runner: "GPUModelRunner") -> int: """Blocks warmup may address: all of them, or, for an extensible cache, the count agreed across ranks so that every rank warms up the same shapes.""" @@ -409,7 +459,7 @@ def extend_kv_cache( ) kv_cache.commit( num_blocks, - defragment=runner.vllm_config.kv_transfer_config is not None, + defragment=registers_kv_cache_memory(runner.vllm_config), shrink=True, ) runner.kv_cache_config.num_blocks = kv_cache.num_committed_blocks @@ -512,9 +562,7 @@ def measure_kv_cache_blocks( vllm_config.model_config.multimodal_config, getattr(vllm_config.parallel_config, "_api_process_count", 1), ) - num_blocks = kv_cache.blocks_within( - kv_cache_bytes, aligned=vllm_config.kv_transfer_config is not None - ) + num_blocks = kv_cache.blocks_within(kv_cache_bytes) # Torch may grow into what stays free once the cache is extended, less the # memory kept for multimodal IPC and for allocations outside torch. left_free = ( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 665e8347329a..8d459b4a0f50 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -655,7 +655,11 @@ def initialize_kv_cache( self.kv_cache_config = kv_cache_config self.extensible_kv_cache: ExtensibleKVCache | None = None if extensible: - self.extensible_kv_cache = ExtensibleKVCache(kv_cache_config, self.device) + self.extensible_kv_cache = ExtensibleKVCache( + kv_cache_config, + self.device, + self.cache_config.get_resolved_kv_cache_layout(), + ) block_table_max_model_len = self.max_model_len if self.is_encoder_decoder: diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 6ae2160c701d..30d09e5621b9 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -113,6 +113,8 @@ extend_kv_cache, kv_cache_lent_for_weight_update, measure_kv_cache_blocks, + pad_kv_cache_segments, + registers_kv_cache_memory, ) from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel from vllm.v1.worker.startup_plan import ( @@ -955,6 +957,12 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: # related to kv cache connector (e.g. kv cache sharing layers). allocation_context = self._maybe_get_memory_pool_context(tag="kv_cache") if self.cache_config.enable_extensible_kv_cache: + if registers_kv_cache_memory(self.vllm_config): + kv_cache_config = pad_kv_cache_segments( + kv_cache_config, + self.cache_config.get_resolved_kv_cache_layout(), + granule_size(torch.accelerator.current_device_index()), + ) # The connector is created in `extend_kv_cache`, once the # memory it registers is final. self._kv_cache_config = kv_cache_config @@ -1027,10 +1035,6 @@ def extensible_kv_cache_unsupported_reason(self) -> str | None: return reason or "another rank cannot use it" return None - def kv_cache_commit_granule(self) -> int: - """Bytes one physical KV cache commit is rounded to on this device.""" - return granule_size(torch.accelerator.current_device_index()) - def disable_extensible_kv_cache(self) -> None: self.cache_config.enable_extensible_kv_cache = False self.requested_memory = request_memory(