[Bugfix] Pad extensible KV cache segments instead of rounding the block count

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

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),
     ],

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)

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":

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

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

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
     )

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