[ZenCPU] Minimal zen image build on x86 premerge queue

Add the zen CPU docker image build (docker/Dockerfile.zen +
image_build_zen_cpu.sh) and an image-build-zen-cpu step pinned to the
small_cpu_queue_premerge queue, with no source_file_dependencies so it
always runs, to confirm the zen image builds on the generic x86 build
queue. Also drop the stale vllm-openai-zen reference from
docker/Dockerfile.cpu.

Image-build only: the zen5 hardware test and its runner are intentionally
excluded from this confirmation PR.

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Chinmay Kulkarni <Chinmay.Kulkarni@amd.com>

.buildkite/image_build/image_build.yaml

@@ -81,6 +81,29 @@ steps:
         - exit_status: -10  # Agent was lost
           limit: 2
 
+  - label: ":docker: Build Zen CPU image"
+    key: image-build-zen-cpu
+    # Minimal confirmation build: run the zen image build on the generic x86
+    # premerge queue to verify it succeeds. Reuses the shared `-cpu` image as its
+    # base, so it must run after the CPU image build.
+    #
+    # No source_file_dependencies: this step always runs for this validation PR
+    # so the build is exercised regardless of which files changed.
+    depends_on:
+      - image-build-cpu
+    agents:
+      queue: small_cpu_queue_premerge
+    commands:
+    - .buildkite/image_build/image_build_zen_cpu.sh $REGISTRY $REPO $BUILDKITE_COMMIT
+    env:
+      DOCKER_BUILDKIT: "1"
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2
+
   - label: ":docker: Build HPU image"
     soft_fail: true
     depends_on: []

.buildkite/image_build/image_build_zen_cpu.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Build the AMD Zen CPU image (vLLM + zentorch) as a two-step layered build:
+#   1. a CPU base image (vLLM installed) built from docker/Dockerfile.cpu
+#   2. docker/Dockerfile.zen --target vllm-zen-test -> zen image on top
+#
+# The image is (re)built from source every time, mirroring
+# .buildkite/scripts/hardware_ci/run-cpu-test.sh. It is not pushed to a registry.
+#
+# See docker/Dockerfile.zen for the build workflow this mirrors.
+set -e
+
+if [[ $# -lt 3 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+
+# Local image tags (not pushed).
+BASE_IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu-base-for-zen"
+IMAGE="$REGISTRY/$REPO:$BUILDKITE_COMMIT-zen-cpu"
+
+# ZENTORCH_VERSION is optional; when unset the Dockerfile falls back to
+# installing zentorch via `vllm[zen]`.
+ZENTORCH_VERSION=${ZENTORCH_VERSION:-}
+
+# Step 1: build the CPU base image that Dockerfile.zen layers on.
+echo "--- :docker: Building CPU base image"
+docker build --file docker/Dockerfile.cpu \
+  --platform linux/amd64 \
+  --build-arg max_jobs=16 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --build-arg VLLM_CPU_X86=true \
+  --tag "$BASE_IMAGE" \
+  --target vllm-openai \
+  --progress plain .
+
+# Step 2: build the zen test image on top of the CPU base.
+echo "--- :docker: Building Zen test image"
+docker build --file docker/Dockerfile.zen \
+  --platform linux/amd64 \
+  --build-arg BASE_IMAGE="$BASE_IMAGE" \
+  ${ZENTORCH_VERSION:+--build-arg ZENTORCH_VERSION="$ZENTORCH_VERSION"} \
+  --tag "$IMAGE" \
+  --target vllm-zen-test \
+  --progress plain .

docker/Dockerfile.cpu

@@ -9,7 +9,6 @@
 #
 # Build targets:
 #   vllm-openai (default): used for serving deployment
-#   vllm-openai-zen: vLLM from source + zentorch from PyPI via vllm[zen]
 #   vllm-test: used for CI tests
 #   vllm-dev: used for development
 #
@@ -305,19 +304,3 @@ LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
 COPY examples examples
 
 ENTRYPOINT ["vllm", "serve"]
-
-
-######################### ZEN CPU PYPI IMAGE #########################
-FROM vllm-openai AS vllm-openai-zen
-
-ARG TARGETARCH
-
-RUN if [ "$TARGETARCH" != "amd64" ]; then \
-        echo "ERROR: vllm-openai-amd only supports --platform=linux/amd64"; \
-        exit 1; \
-    fi
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install "vllm[zen]"
-
-ENTRYPOINT ["vllm", "serve"]

docker/Dockerfile.zen

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# This Dockerfile builds AMD Zen CPU images for vLLM with zentorch.
+# It layers on top of the CPU image produced by docker/Dockerfile.cpu.
+#
+# Supported platforms:
+#   - linux/amd64 (Zen2 and newer; Zen5 is the canonical CI target)
+#
+# Build targets:
+#   vllm-openai-zen (default): vLLM + zentorch for serving
+#   vllm-zen-test            : vllm-openai-zen + test deps for CI
+#
+# Build workflow (see .buildkite/image_build/image_build_zen_cpu.sh):
+#   1. Obtain a CPU base image that already has vLLM installed. In CI this is the
+#      shared `<repo>:<commit>-cpu` image (docker/Dockerfile.cpu --target
+#      vllm-test); either the vllm-test or vllm-openai target works as the base
+#      since both install vLLM into the same venv this stage extends.
+#   2. docker build -f docker/Dockerfile.zen --build-arg BASE_IMAGE=<base>:<tag> \
+#                    --target vllm-zen-test -t <out>:<tag> .
+
+ARG BASE_IMAGE
+ARG PYTHON_VERSION=3.12
+ARG ZENTORCH_VERSION
+
+######################### ZEN OPENAI IMAGE #########################
+FROM ${BASE_IMAGE} AS vllm-openai-zen
+
+ARG TARGETARCH
+ARG ZENTORCH_VERSION
+
+RUN if [ "${TARGETARCH}" != "amd64" ]; then \
+        echo "ERROR: Dockerfile.zen requires --platform=linux/amd64"; \
+        exit 1; \
+    fi
+
+WORKDIR /vllm-workspace
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -n "${ZENTORCH_VERSION}" ]; then \
+        uv pip install "zentorch==${ZENTORCH_VERSION}"; \
+    else \
+        uv pip install "vllm[zen]"; \
+    fi
+
+LABEL org.opencontainers.image.title="vLLM Zen CPU"
+LABEL org.opencontainers.image.description="vLLM with zentorch for AMD Zen CPUs"
+LABEL org.opencontainers.image.vendor="vLLM Project"
+LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm"
+LABEL ai.vllm.build.target-arch="${TARGETARCH}"
+LABEL ai.vllm.build.python-version="${PYTHON_VERSION}"
+LABEL ai.vllm.build.zentorch-version="${ZENTORCH_VERSION:-unpinned}"
+
+ENTRYPOINT ["vllm", "serve"]
+
+######################### ZEN TEST IMAGE #########################
+FROM vllm-openai-zen AS vllm-zen-test
+
+WORKDIR /vllm-workspace
+
+COPY requirements/test/cuda.in requirements/test/zen.in
+
+RUN sed -i '/mamba_ssm/d'                  requirements/test/zen.in && \
+    sed -i 's/^torch==.*/torch==2.11.0/g'  requirements/test/zen.in && \
+    sed -i 's/torchaudio.*/torchaudio/g'   requirements/test/zen.in && \
+    sed -i 's/torchvision.*/torchvision/g' requirements/test/zen.in && \
+    # zentorch parity: keep sentence-transformers pinned where vllm-test-deps did.
+    sed -i 's/^sentence-transformers.*/sentence-transformers==5.3.0/g' requirements/test/zen.in && \
+    uv pip compile requirements/test/zen.in \
+       -o requirements/test/zen.txt \
+       --index-strategy unsafe-best-match --torch-backend cpu
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install -r requirements/test/zen.txt
+
+# Reassert zentorch's torch in case test-deps moved it.
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -n "${ZENTORCH_VERSION}" ]; then \
+        uv pip install --force-reinstall --no-deps "zentorch==${ZENTORCH_VERSION}"; \
+    else \
+        uv pip install --force-reinstall --no-deps "vllm[zen]"; \
+    fi
+
+ADD ./tests/ ./tests/
+ADD ./examples/ ./examples/
+ADD ./benchmarks/ ./benchmarks/
+ADD ./vllm/collect_env.py .
+ADD ./docker/ ./docker/
+ADD ./.buildkite/ ./.buildkite/
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install -e tests/vllm_test_utils
+
+ENV HF_XET_HIGH_PERFORMANCE=1
+ENV HF_HUB_DOWNLOAD_TIMEOUT=60
+
+ENTRYPOINT []

docs/features/kv_offloading_usage.md

@@ -75,7 +75,7 @@ vllm serve <model> \
 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
-| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
+| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
 
 ## Secondary Tiers

tests/multimodal/media/test_connector.py

@@ -6,6 +6,7 @@ import mimetypes
 import os
 import shutil
 import time
+from io import BytesIO
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
 import aiohttp
@@ -111,6 +112,34 @@ async def test_fetch_image_base64(
         assert _image_equals(data_image_sync, data_image_async)
 
 
+@pytest.mark.asyncio
+async def test_fetch_image_keep_original_mode():
+    """media_io_kwargs can disable the default RGB conversion."""
+    # RGBA image: opaque black pixel on a fully transparent background
+    rgba_image = Image.new("RGBA", (4, 4), (0, 0, 0, 0))
+    rgba_image.putpixel((2, 2), (0, 0, 0, 255))
+    buffer = BytesIO()
+    rgba_image.save(buffer, "PNG")
+    data_url = (
+        f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
+    )
+
+    # Default behavior: RGBA is composited onto a white background
+    default_image = MediaConnector().fetch_image(data_url)
+    assert default_image.mode == "RGB"
+    assert default_image.getpixel((0, 0)) == (255, 255, 255)
+    assert default_image.getpixel((2, 2)) == (0, 0, 0)
+
+    # image_mode=None via media_io_kwargs: original mode is preserved
+    connector = MediaConnector(media_io_kwargs={"image": {"image_mode": None}})
+    image_sync = connector.fetch_image(data_url)
+    image_async = await connector.fetch_image_async(data_url)
+    for image in (image_sync, image_async):
+        assert image.mode == "RGBA"
+        assert image.getpixel((0, 0)) == (0, 0, 0, 0)
+        assert image.getpixel((2, 2)) == (0, 0, 0, 255)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
 async def test_fetch_image_local_files(image_url: str):

tests/multimodal/media/test_image.py

@@ -80,6 +80,29 @@ def test_image_media_io_rgba_custom_background(tmp_path):
     assert green_numpy[0][0][2] == 0  # B
 
 
+def test_image_media_io_no_mode_conversion(tmp_path):
+    """image_mode=None skips conversion and preserves the original mode."""
+    # RGBA image: opaque black pixel on a fully transparent background
+    rgba_image = Image.new("RGBA", (10, 10), (0, 0, 0, 0))
+    rgba_image.putpixel((5, 5), (0, 0, 0, 255))
+    test_image_path = tmp_path / "test_rgba.png"
+    rgba_image.save(test_image_path)
+
+    # Default behavior: RGBA is composited onto a white background
+    image_io_default = ImageMediaIO()
+    converted_default = image_io_default.load_file(test_image_path)
+    assert converted_default.media.mode == "RGB"
+    assert converted_default.media.getpixel((0, 0)) == (255, 255, 255)
+    assert converted_default.media.getpixel((5, 5)) == (0, 0, 0)
+
+    # image_mode=None: original mode and alpha channel are preserved
+    image_io_keep = ImageMediaIO(image_mode=None)
+    converted_keep = image_io_keep.load_file(test_image_path)
+    assert converted_keep.media.mode == "RGBA"
+    assert converted_keep.media.getpixel((0, 0)) == (0, 0, 0, 0)
+    assert converted_keep.media.getpixel((5, 5)) == (0, 0, 0, 255)
+
+
 def test_image_media_io_rgba_background_color_validation():
     """Test that invalid rgba_background_color values are properly rejected."""
 

tests/quantization/test_modelopt.py

@@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
     assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
 
 
+def test_modelopt_mixed_precision_composes_gemma4_mappers():
+    from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+    from vllm.model_executor.models.gemma4_mm import (
+        Gemma4ForConditionalGeneration,
+    )
+
+    config = _mixed_precision_config(
+        {
+            "model.language_model.layers.0.experts": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+            "model.language_model.layers.1.moe.experts.gate_up_proj": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+        }
+    )
+
+    config.apply_vllm_mapper(
+        Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
+    )
+    config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
+
+    expected_prefix = "language_model.model.layers.0.moe.experts"
+    assert set(config.quantized_layers) == {
+        expected_prefix,
+        "language_model.model.layers.1.moe.gate_up_proj",
+    }
+    assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
+
+
 def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
     from vllm.model_executor.layers.linear import LinearBase
 

tests/utils_/test_extensible_tensor.py

@@ -1,141 +0,0 @@
-# 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()

tests/v1/core/test_encoder_cache_manager.py

@@ -149,30 +149,6 @@ 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])

tests/v1/engine/test_engine_args.py

@@ -49,18 +49,6 @@ 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())

tests/v1/kv_connector/unit/offloading_connector/test_events.py

@@ -7,7 +7,13 @@ import torch
 
 from tests.v1.kv_connector.unit.utils import create_vllm_config
 from vllm.config import KVEventsConfig, KVTransferConfig
-from vllm.distributed.kv_events import MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored
+from vllm.distributed.kv_events import (
+    MEDIUM_CPU,
+    MEDIUM_FS,
+    MEDIUM_OBJ,
+    BlockRemoved,
+    BlockStored,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
     build_offloading_config,
 )
@@ -62,8 +68,9 @@ def _wire_hash(block_hash: BlockHash):
     return maybe_convert_block_hash(block_hash)
 
 
-def _request(*, block_hashes: list[BlockHash], token_count: int):
+def _request(*, block_hashes: list[BlockHash], token_count: int, req_id: str = "req"):
     req = MagicMock()
+    req.request_id = req_id
     req.block_hashes = block_hashes
     req.all_token_ids = list(range(1, token_count + 1))
     req.lora_request = None
@@ -104,10 +111,32 @@ def _record_chunks(
     return keys
 
 
+def _record_lookup_chunks(
+    tracker: OffloadingEventsTracker,
+    req,
+    group_config: GroupOffloadConfig,
+    num_chunks: int,
+) -> list[OffloadKey]:
+    keys: list[OffloadKey] = []
+    hbf = group_config.hashes_per_chunk
+    for chunk_idx in range(num_chunks):
+        tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
+        assert tail_hash is not None
+        key = make_offload_key(tail_hash, group_config.group_idx)
+        tracker.record_lookup(
+            req,
+            group_config,
+            chunk_idx,
+            key,
+        )
+        keys.append(key)
+    return keys
+
+
 def _stored_event(
     keys: list[OffloadKey],
-    locality: Locality | None = None,
     medium: str = _CPU_MEDIUM,
+    locality: Locality | None = None,
 ) -> OffloadingEvent:
     return OffloadingEvent(
         keys=keys,
@@ -119,8 +148,8 @@ def _stored_event(
 
 def _removed_event(
     keys: list[OffloadKey],
-    locality: Locality | None = None,
     medium: str = _CPU_MEDIUM,
+    locality: Locality | None = None,
 ) -> OffloadingEvent:
     return OffloadingEvent(
         keys=keys,
@@ -130,6 +159,21 @@ def _removed_event(
     )
 
 
+def _lookup_chunk() -> tuple[
+    OffloadingEventsTracker, MagicMock, GroupOffloadConfig, OffloadKey
+]:
+    tracker = _tracker()
+    req = _request(block_hashes=[_hash(0)], token_count=4)
+    group_config = _group_config()
+    key = _record_lookup_chunks(
+        tracker,
+        req,
+        group_config,
+        num_chunks=1,
+    )[0]
+    return tracker, req, group_config, key
+
+
 def test_take_events_forwards_locality_to_rich_store():
     tracker = _tracker()
     req = _request(block_hashes=[_hash(0)], token_count=4)
@@ -220,18 +264,37 @@ def test_take_events_publishes_routable_block_stored():
     assert len(tracker._pending_event_metadata) == 6
 
 
-def test_take_events_factor_gt_1_chunk_store_and_remove():
+def test_promotion_emits_full_cpu_stored_event():
+    tracker, _, _, key = _lookup_chunk()
+
+    [event] = tracker.take_events([_stored_event([key])])
+
+    assert isinstance(event, BlockStored)
+    assert event.medium == MEDIUM_CPU
+    assert event.block_hashes == [_wire_hash(_hash(0))]
+    assert event.parent_block_hash is None
+    assert event.token_ids == [1, 2, 3, 4]
+    assert event.block_size == 4
+    assert event.lora_id is None
+    assert event.lora_name is None
+    assert event.extra_keys is None
+    assert event.group_idx == 0
+    assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
+    assert event.kv_cache_spec_sliding_window is None
+
+
+def test_lookup_promotion_factor_gt_1_store_and_remove():
     block_size = 4
-    blocks_per_chunk = 3
+    blocks_per_chunk = 2
     tracker = _tracker()
     group_config = _group_config(
         block_size=block_size, blocks_per_chunk=blocks_per_chunk
     )
     req = _request(
-        block_hashes=[_hash(i) for i in range(6)],
+        block_hashes=[_hash(i) for i in range(4)],
         token_count=block_size * blocks_per_chunk * 2,
     )
-    keys = _record_chunks(tracker, req, group_config, num_chunks=2)
+    keys = _record_lookup_chunks(tracker, req, group_config, num_chunks=2)
 
     stored = list(tracker.take_events([_stored_event(keys)]))
     assert len(stored) == 2
@@ -293,6 +356,7 @@ def test_take_events_opt_out_keeps_placeholders():
     group_config = _group_config()
     req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
     keys = _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
 
     assert not tracker.self_describing_enabled
     assert not tracker._pending_event_metadata
@@ -315,11 +379,21 @@ def test_take_events_opt_out_keeps_placeholders():
     assert len(events[3].block_hashes) == 3
 
 
-def test_record_store_skips_sliding_window_group():
+@pytest.mark.parametrize(
+    "sliding_window_size_in_chunks",
+    [1, 2],
+    ids=["ssm", "sliding-window"],
+)
+def test_event_metadata_skips_non_full_attention_group(
+    sliding_window_size_in_chunks: int,
+):
     tracker = _tracker()
-    group_config = _group_config(sliding_window_size_in_chunks=2)
+    group_config = _group_config(
+        sliding_window_size_in_chunks=sliding_window_size_in_chunks
+    )
     req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
     keys = _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
 
     assert not tracker._pending_event_metadata
 
@@ -329,6 +403,57 @@ def test_record_store_skips_sliding_window_group():
     assert events[0].block_size == 0
 
 
+def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
+    tracker = _tracker()
+    block_hashes = [_hash(0), _hash(1)]
+    req = _request(block_hashes=block_hashes, token_count=8)
+    group_config = _group_config(blocks_per_chunk=2)
+    key = _record_chunks(tracker, req, group_config, num_chunks=1)[0]
+    confirmed_meta = tracker._pending_event_metadata[key]
+    lookup_req = _request(
+        block_hashes=block_hashes,
+        token_count=8,
+        req_id="new-request",
+    )
+
+    tracker.record_lookup(
+        lookup_req,
+        group_config,
+        0,
+        key,
+    )
+    assert tracker._pending_event_metadata[key] is confirmed_meta
+
+    removed = list(tracker.take_events([_removed_event([key])]))
+    assert len(removed) == 1
+    assert removed[0].block_hashes == [
+        _wire_hash(_hash(0)),
+        _wire_hash(_hash(1)),
+    ]
+
+    stored = list(tracker.take_events([_stored_event([key])]))
+    assert len(stored) == 1
+    assert stored[0].block_size == 0
+    assert stored[0].token_ids == []
+
+    tracker.record_lookup(lookup_req, group_config, 0, key)
+    removed = list(tracker.take_events([_removed_event([key])]))
+    assert removed[0].block_hashes == [
+        _wire_hash(_hash(0)),
+        _wire_hash(_hash(1)),
+    ]
+
+
+@pytest.mark.parametrize("medium", [MEDIUM_FS, MEDIUM_OBJ])
+def test_secondary_stored_event_does_not_mutate_cpu_metadata(medium: str):
+    tracker, _, _, key = _lookup_chunk()
+    expected_metadata = dict(tracker._pending_event_metadata)
+
+    stored = list(tracker.take_events([_stored_event([key], medium)]))
+    assert stored[0].token_ids == [1, 2, 3, 4]
+    assert tracker._pending_event_metadata == expected_metadata
+
+
 def test_take_events_groups_removed_hashes_by_kv_group():
     tracker = _tracker()
     group0_config = _group_config(group_idx=0, blocks_per_chunk=2)
@@ -378,7 +503,7 @@ def test_reset_cache_clears_side_table():
     tracker = _tracker()
     group_config = _group_config()
     req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
-    _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)
 
     assert tracker._pending_event_metadata
 
@@ -387,7 +512,7 @@ def test_reset_cache_clears_side_table():
     assert not tracker._pending_event_metadata
 
 
-def test_tiering_rejects_self_describing_kv_events():
+def test_tiering_accepts_self_describing_kv_events():
     vllm_config = create_vllm_config(
         block_size=4,
         max_num_batched_tokens=16,
@@ -423,5 +548,9 @@ def test_tiering_rejects_self_describing_kv_events():
         ],
     )
 
-    with pytest.raises(ValueError, match="TieringOffloadingSpec"):
-        TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
+    spec = TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
+    tracker = OffloadingEventsTracker(spec.kv_events_config)
+
+    assert spec.kv_events_config.enable_kv_cache_events
+    assert spec.kv_events_config.self_describing_kv_events
+    assert tracker.self_describing_enabled

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, call
 
 import pytest
 import torch
@@ -11,6 +11,7 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
     to_keys,
 )
 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
+from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
     OffloadingConnectorStats,
     _ConnectorMetricName,
@@ -19,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
     OffloadingConnectorScheduler,
     RequestOffloadState,
 )
+from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheGroupSpec,
@@ -26,6 +28,7 @@ from vllm.v1.kv_cache_interface import (
 )
 from vllm.v1.kv_offload.base import (
     LookupResult,
+    OffloadingEvent,
     OffloadingManager,
     OffloadPolicy,
     ReqContext,
@@ -143,6 +146,159 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
     assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0
 
 
+def test_max_offload_tokens_zero_does_not_record_pending_lookups(request_runner):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+    )
+    runner.manager.lookup.return_value = LookupResult.RETRY
+    runner.manager.take_events.return_value = []
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    runner.new_request(
+        token_ids=[1] * 12,
+        kv_transfer_params={"max_offload_tokens": 0},
+    )
+    runner.run(decoded_tokens=[])
+
+    tracker = runner.connector_scheduler._events_tracker
+    assert runner.manager.lookup.call_count == 3
+    assert not tracker._pending_event_metadata
+    assert list(runner.connector_scheduler.take_events()) == []
+
+    runner.manager.lookup.return_value = LookupResult.MISS
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+
+    assert not tracker._pending_event_metadata
+    assert list(runner.connector_scheduler.take_events()) == []
+
+
+def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal(
+    request_runner,
+):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+        blocks_per_chunk=2,
+    )
+    raw_events: list[OffloadingEvent] = []
+
+    def take_raw_events():
+        yield from raw_events
+        raw_events.clear()
+
+    runner.manager.lookup.return_value = LookupResult.RETRY
+    runner.manager.take_events.side_effect = take_raw_events
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+
+    runner.new_request(token_ids=[1] * 8)
+    runner.run(decoded_tokens=[])
+
+    tracker = runner.connector_scheduler._events_tracker
+    assert not tracker._pending_event_metadata
+    key = runner.manager.lookup.call_args.args[0]
+    req_id = str(runner.req_id)
+    req_status = runner.connector_scheduler._req_status[req_id]
+
+    runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
+
+    assert not tracker._pending_event_metadata
+
+    raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))
+    events = list(runner.connector_scheduler.take_events())
+    assert len(events) == 1
+    assert isinstance(events[0], BlockStored)
+    assert events[0].block_size == 0
+    assert events[0].token_ids == []
+
+    runner.manager.lookup.return_value = LookupResult.HIT
+    group_config = runner.connector_scheduler.config.kv_group_configs[0]
+    assert (
+        runner.connector_scheduler._maximal_prefix_lookup(
+            [key],
+            req_status.req_context,
+            req_status.req,
+            group_config,
+            0,
+        )
+        == 1
+    )
+    assert key in tracker._pending_event_metadata
+
+    raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=True))
+    [event] = runner.connector_scheduler.take_events()
+    assert isinstance(event, BlockRemoved)
+    assert event.medium == MEDIUM_CPU
+    assert len(event.block_hashes) == 2
+    assert key not in tracker._pending_event_metadata
+
+
+@pytest.mark.parametrize("blocks_per_chunk", [1, 2])
+def test_promotion_hit_precedes_stored_event_translation(
+    request_runner,
+    blocks_per_chunk: int,
+):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+        blocks_per_chunk=blocks_per_chunk,
+    )
+    token_ids = [1] * 4 * blocks_per_chunk
+
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.new_request(token_ids=token_ids)
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=tuple(range(blocks_per_chunk)),
+    )
+    runner.scheduler.reset_prefix_cache()
+    runner.connector_scheduler._events_tracker.reset()
+
+    raw_events: list[OffloadingEvent] = []
+
+    def lookup(key, req_context):
+        raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))
+        return LookupResult.HIT
+
+    def take_raw_events():
+        yield from raw_events
+        raw_events.clear()
+
+    runner.manager.lookup.side_effect = lookup
+    runner.manager.take_events.side_effect = take_raw_events
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+    publisher = MagicMock()
+    runner.scheduler.kv_event_publisher = publisher
+
+    runner.new_request(token_ids=token_ids)
+    runner.run(
+        decoded_tokens=[],
+        expected_loaded=tuple(range(blocks_per_chunk)),
+    )
+
+    events = [
+        event
+        for publish_call in publisher.publish.call_args_list
+        for event in publish_call.args[0].events
+        if isinstance(event, BlockStored) and event.medium == MEDIUM_CPU
+    ]
+    assert len(events) == 1
+    assert len(events[0].block_hashes) == blocks_per_chunk
+    assert events[0].block_size == 4
+    assert events[0].token_ids == token_ids
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_offloading_connector(request_runner, async_scheduling: bool):
     block_size = 4
@@ -241,7 +397,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2))
 
     # single block lookup with a hit in a middle block
@@ -249,7 +405,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))
 
 
@@ -307,7 +463,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
 
     # request should now return from preemption
     # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 3
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
     )
@@ -427,7 +583,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
     # start a request to load the first block, but don't complete
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -439,7 +595,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
 
     # start a new request to load the same first block
     runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -491,7 +647,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
     # start a request to load the first block, but don't complete
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(
         decoded_tokens=[],
         complete_transfers=False,
@@ -793,73 +949,144 @@ def _make_scheduler_with_lookup(
 
     scheduler = object.__new__(OffloadingConnectorScheduler)
     scheduler.manager = manager
+    scheduler._events_tracker = MagicMock()
     return scheduler
 
 
 _EMPTY_REQ_CTX = ReqContext(req_id="")
+_LOOKUP_REQ = MagicMock()
+_LOOKUP_REQ.request_id = "req"
+_LOOKUP_GROUP_CONFIG = MagicMock()
+
+
+def _maximal_lookup(sched, keys, start_chunk_idx: int = 0):
+    return sched._maximal_prefix_lookup(
+        keys,
+        _EMPTY_REQ_CTX,
+        _LOOKUP_REQ,
+        _LOOKUP_GROUP_CONFIG,
+        start_chunk_idx,
+    )
 
 
 class TestMaximalPrefixLookup:
     def test_all_hit(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 2
+
+    def test_records_absolute_chunk_indices(self):
+        keys = to_keys([1, 2])
+        sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
+
+        assert _maximal_lookup(sched, keys, start_chunk_idx=3) == 2
+        assert sched._events_tracker.record_lookup.call_args_list == [
+            call(
+                _LOOKUP_REQ,
+                _LOOKUP_GROUP_CONFIG,
+                3,
+                keys[0],
+            ),
+            call(
+                _LOOKUP_REQ,
+                _LOOKUP_GROUP_CONFIG,
+                4,
+                keys[1],
+            ),
+        ]
 
     def test_all_miss(self):
         sched = _make_scheduler_with_lookup({})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 0
+        sched._events_tracker.record_lookup.assert_not_called()
 
     def test_partial_prefix(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) == 2
 
     def test_miss_then_hit(self):
         sched = _make_scheduler_with_lookup({2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 0
 
     def test_single_hit(self):
         sched = _make_scheduler_with_lookup({1: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
+        assert _maximal_lookup(sched, to_keys([1])) == 1
 
     def test_empty(self):
         sched = _make_scheduler_with_lookup({})
-        assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, []) == 0
+
+    @pytest.mark.parametrize(
+        "pending_result",
+        [LookupResult.RETRY, LookupResult.HIT_PENDING],
+    )
+    def test_pending_result_is_not_recorded(
+        self,
+        pending_result: LookupResult,
+    ):
+        sched = _make_scheduler_with_lookup({1: pending_result})
+
+        assert _maximal_lookup(sched, to_keys([1])) is None
+        sched._events_tracker.record_lookup.assert_not_called()
 
     def test_retry_defers(self):
+        keys = to_keys([1, 2])
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.RETRY, 2: LookupResult.HIT}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
         assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            1,
+            keys[1],
+        )
 
     def test_retry_after_hit_defers(self):
+        keys = to_keys([1, 2])
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.HIT, 2: LookupResult.RETRY}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            0,
+            keys[0],
+        )
 
     def test_hit_pending_defers(self):
+        keys = to_keys([1, 2])
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
         assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            1,
+            keys[1],
+        )
 
     def test_hit_pending_does_not_stop_scan(self):
         """HIT_PENDING defers but does not break — scan continues until miss."""
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.HIT_PENDING, 2: LookupResult.MISS, 3: LookupResult.HIT}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None
         assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_not_called()
 
     def test_retry_stops_at_miss(self):
         """RETRY is treated as hit for iteration, but miss stops the scan."""
         sched = _make_scheduler_with_lookup(
             {1: LookupResult.RETRY, 2: LookupResult.MISS, 3: LookupResult.HIT}
         )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None
         # lookup should have been called for blocks 1 and 2 (stops at miss)
         assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_not_called()
 
 
 class TestSlidingWindowLookup:
@@ -1011,7 +1238,7 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling
 
     # New request with 2 offloaded chunks; first matches what's in CPU.
     runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(keys)
     )
@@ -1042,7 +1269,7 @@ def test_loads_do_not_populate_fence_index(request_runner):
         async_scheduling=False,
     )
     runner.new_request(token_ids=[0] * 12)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.run(decoded_tokens=[], complete_transfers=False)
     assert runner.connector_scheduler._block_id_to_pending_jobs == {}
 
@@ -1088,7 +1315,7 @@ def test_fence_at_update_state_after_alloc(request_runner):
 
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * 4)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1139,7 +1366,7 @@ def test_fence_at_build_store_jobs(request_runner):
 
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[1] * 4)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 0
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1365,7 +1592,7 @@ def test_reset_cache(request_runner, async_scheduling: bool):
     # Leave the load in-flight so that reset_cache must flush it.
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output([])
     )
@@ -1553,9 +1780,7 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner
     # preemption batch's ModelRunnerOutput is consumed by update_from_output().
     free_block_queue.num_free_blocks = num_free_blocks_empty
     assert runner.scheduler.reset_prefix_cache()
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len(
-        key
-    )
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: len(keys)
 
     readmit_output = runner.scheduler.schedule()
 
@@ -1669,7 +1894,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * num_tokens + [1])
     runner.manager.lookup.return_value = LookupResult.HIT
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
         # Group 0: full prefix lookup hits 2 offloaded chunks
@@ -1840,6 +2065,13 @@ class TestEagle:
         req.request_id = "test-req"
         req.num_tokens = num_tokens
         req.kv_transfer_params = None
+        num_hash_blocks = max(
+            len(hashes) * scheduler.config.kv_group_configs[idx].hashes_per_chunk
+            for idx, hashes in enumerate(offload_keys_per_group)
+        )
+        req.block_hashes = [BlockHash(str(i).encode()) for i in range(num_hash_blocks)]
+        req.all_token_ids = list(range(num_tokens))
+        req.lora_request = None
 
         state = RequestOffloadState(
             config=scheduler.config,

tests/v1/worker/test_extensible_kv_cache.py

@@ -1,715 +0,0 @@
-# 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.attn_utils import (
-    _allocate_extensible_kv_cache,
-    _kv_cache_num_segments_by_layer,
-    _reshape_kv_cache,
-    narrow_kv_caches_to_num_blocks,
-)
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-from vllm.v1.worker.gpu_worker import Worker
-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:
-    buffers = getattr(runner, "extensible_kv_buffers", None)
-    if buffers is not None:
-        buffers.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_buffers.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.accelerator.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.accelerator.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_buffers.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.accelerator.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.accelerator.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_buffers.buffers) == len(layer_names)
-        for buffer, bytes_per_block_per_segment in runner.extensible_kv_buffers.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.accelerator.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.accelerator.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.accelerator.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.accelerator.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)
-
-
-# ---------------------------------------------------------------------------
-# V2 model runner (vllm.v1.worker.gpu) extensible allocation
-# ---------------------------------------------------------------------------
-
-
-def _v2_allocate(kv_cache_config, attn_groups, kernel_block_sizes):
-    flat_groups = [g for groups in attn_groups for g in groups]
-    raw_tensors, buffers = _allocate_extensible_kv_cache(
-        kv_cache_config,
-        {},
-        torch.device("cuda:0"),
-        flat_groups,
-        kernel_block_sizes,
-        "auto",
-    )
-    kv_caches = _reshape_kv_cache(
-        attn_groups=flat_groups,
-        kv_cache_raw_tensors=raw_tensors,
-        cache_dtype="auto",
-        kernel_block_sizes=kernel_block_sizes,
-        shared_kv_cache_layers={},
-        kv_cache_config=kv_cache_config,
-    )
-    return kv_caches, buffers
-
-
-def test_v2_num_segments_by_layer() -> None:
-    """V2 segment counts follow the layer's physical layout, and hybrid
-    models force block-major (one segment)."""
-    spec = _full_attention_spec()
-    for backend, expected in (
-        (_SplitKVBackend, 2),
-        (_BlockMajorBackend, 1),
-        (_StrideOrderBackend, 1),
-    ):
-        _, attn_groups = _attention_config(spec, backend)
-        flat_groups = [g for groups in attn_groups for g in groups]
-        assert _kv_cache_num_segments_by_layer(
-            flat_groups, [BLOCK_SIZE], "auto", has_mamba=False
-        ) == {"layer.0": expected}
-        assert _kv_cache_num_segments_by_layer(
-            flat_groups, [BLOCK_SIZE], "auto", has_mamba=True
-        ) == {"layer.0": 1}
-
-
-def test_v2_extensible_split_layout_grows_incrementally() -> None:
-    """A K/V-split layer grows both halves in lockstep through the staged
-    commits the V2 flow performs (init -> warmup prefix -> final size)."""
-    spec = _full_attention_spec()
-    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
-    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
-    try:
-        kv_cache = kv_caches["layer.0"]
-        assert kv_cache.shape == (2, NUM_BLOCKS, BLOCK_SIZE, 8, 128)
-        assert buffers.num_blocks_committed == 1
-
-        kv_cache[0, 0].fill_(1)  # K, block 0
-        kv_cache[1, 0].fill_(2)  # V, block 0
-        torch.accelerator.synchronize()
-
-        # Warmup-style prefix commit, then the final post-warmup commit.
-        buffers.commit(8)
-        kv_cache[0, 7].fill_(3)
-        torch.accelerator.synchronize()
-        buffers.commit(NUM_BLOCKS)
-        # Shrink requests are ignored.
-        buffers.commit(1)
-        assert buffers.num_blocks_committed == NUM_BLOCKS
-
-        kv_cache[1, NUM_BLOCKS - 1].fill_(4)
-        torch.accelerator.synchronize()
-        assert torch.all(kv_cache[0, 0] == 1)
-        assert torch.all(kv_cache[1, 0] == 2)
-        assert torch.all(kv_cache[0, 7] == 3)
-        assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 4)
-        assert torch.count_nonzero(kv_cache[:, 1:7]) == 0
-        assert torch.count_nonzero(kv_cache[:, 8 : NUM_BLOCKS - 1]) == 0
-        assert buffers.physical_bytes >= NUM_BLOCKS * spec.page_size_bytes
-    finally:
-        buffers.free()
-
-
-def test_v2_extensible_hybrid_attention_mamba() -> None:
-    """V2 hybrid models re-stride attention to block-major; both the
-    attention and Mamba buffers grow as single-segment prefixes."""
-    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,
-            )
-        ],
-    ]
-    kv_caches, buffers = _v2_allocate(
-        kv_cache_config, attn_groups, [BLOCK_SIZE, BLOCK_SIZE]
-    )
-    try:
-        attn_cache = kv_caches["attn.0"]
-        # Re-strided to interleave K/V per block: block b spans one page.
-        hidden_size = attn_cache.shape[2:].numel()
-        assert attn_cache.stride()[:2] == (hidden_size, 2 * hidden_size)
-
-        attn_cache[0, 0].fill_(1)
-        attn_cache[1, 0].fill_(2)
-        for state_tensor in kv_caches["mamba.0"]:
-            state_tensor[0].fill_(3)
-        torch.accelerator.synchronize()
-
-        buffers.commit(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.accelerator.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:
-        buffers.free()
-
-
-def test_v2_extensible_packed_layout() -> None:
-    """A packed (block_stride) layout uses one shared block-major buffer;
-    per-layer pages within a block stay isolated across commits."""
-    spec = _full_attention_spec()
-    page_bytes = spec.page_size_bytes
-    block_stride = 2 * page_bytes  # two layers packed per block
-    layer_names = ["packed.0", "packed.1"]
-    kv_cache_config = KVCacheConfig(
-        num_blocks=NUM_BLOCKS,
-        kv_cache_tensors=[
-            KVCacheTensor(
-                size=NUM_BLOCKS * block_stride,
-                shared_by=[name],
-                offset=i * page_bytes,
-                block_stride=block_stride,
-            )
-            for i, name in enumerate(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,
-            )
-        ]
-    ]
-    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
-    try:
-        assert len(buffers.buffers) == 1
-        [(buffer, bytes_per_block)] = buffers.buffers
-        assert buffer.num_segments == 1
-        assert bytes_per_block == block_stride
-
-        cache0, cache1 = kv_caches["packed.0"], kv_caches["packed.1"]
-        assert cache0.shape == (NUM_BLOCKS, 2, BLOCK_SIZE, 8, 128)
-
-        cache0[0].fill_(1)
-        cache1[0].fill_(2)
-        torch.accelerator.synchronize()
-        buffers.commit(NUM_BLOCKS)
-        cache0[NUM_BLOCKS - 1].fill_(3)
-        torch.accelerator.synchronize()
-        assert torch.all(cache0[0] == 1)
-        assert torch.all(cache1[0] == 2)
-        assert torch.all(cache0[NUM_BLOCKS - 1] == 3)
-        # The other layer's page of the same block is untouched, and all
-        # middle blocks were zeroed on commit.
-        assert torch.count_nonzero(cache1[1:]) == 0
-        assert torch.count_nonzero(cache0[1 : NUM_BLOCKS - 1]) == 0
-
-        committed = NUM_BLOCKS // 2
-        narrowed = narrow_kv_caches_to_num_blocks(
-            kv_caches,
-            [g for groups in attn_groups for g in groups],
-            [BLOCK_SIZE],
-            "auto",
-            committed,
-            kv_cache_config,
-        )
-        narrowed0 = narrowed["packed.0"]
-        narrowed1 = narrowed["packed.1"]
-        assert narrowed0.untyped_storage().data_ptr() == buffer.base_ptr
-        assert (
-            narrowed0.untyped_storage().data_ptr()
-            == narrowed1.untyped_storage().data_ptr()
-        )
-        assert narrowed0.untyped_storage().nbytes() == committed * block_stride
-        assert narrowed0.stride() == cache0.stride()
-        assert narrowed1.stride() == cache1.stride()
-    finally:
-        buffers.free()
-
-
-def test_v2_extensible_release_and_recommit() -> None:
-    """Sleep/wake cycle: release_physical discards data but keeps VA and
-    views valid; recommit restores the committed size with zeroed pages."""
-    spec = _full_attention_spec()
-    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
-    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
-    try:
-        kv_cache = kv_caches["layer.0"]
-        base_ptr = buffers.buffers[0][0].base_ptr
-        buffers.commit(NUM_BLOCKS)
-        kv_cache.fill_(7)
-        torch.accelerator.synchronize()
-        assert buffers.physical_bytes > 0
-
-        buffers.release_physical()
-        assert buffers.physical_bytes == 0
-        assert buffers.num_blocks_committed == 0
-
-        buffers.recommit()
-        assert buffers.num_blocks_committed == NUM_BLOCKS
-        assert buffers.buffers[0][0].base_ptr == base_ptr
-        torch.accelerator.synchronize()
-        # Data was discarded; fresh pages are zeroed and writable through
-        # the original views.
-        assert torch.count_nonzero(kv_cache) == 0
-        kv_cache[1, NUM_BLOCKS - 1].fill_(9)
-        torch.accelerator.synchronize()
-        assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 9)
-    finally:
-        buffers.free()
-
-
-def test_v2_extensible_connector_sleep_fails_before_remapping() -> None:
-    """Connector registrations must not survive physical-page replacement."""
-    worker = object.__new__(Worker)
-    worker.model_runner = SimpleNamespace(extensible_kv_buffers=object())
-    worker.vllm_config = SimpleNamespace(kv_transfer_config=object())
-
-    with pytest.raises(RuntimeError, match="invalidates.*memory registration"):
-        worker.sleep()
-
-
-def test_v2_narrow_kv_caches_to_num_blocks() -> None:
-    """Connector-registration views are trimmed to the committed block count
-    along each layout's block dim, keeping base pointers and strides (so the
-    K and V segment prefixes are addressed exactly)."""
-    spec = _full_attention_spec()
-    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
-    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
-    try:
-        committed = 16
-        buffers.commit(committed)
-        narrowed = narrow_kv_caches_to_num_blocks(
-            kv_caches,
-            [g for groups in attn_groups for g in groups],
-            [BLOCK_SIZE],
-            "auto",
-            committed,
-            kv_cache_config,
-        )
-        full = kv_caches["layer.0"]
-        trimmed = narrowed["layer.0"]
-        assert trimmed.shape == (2, committed, BLOCK_SIZE, 8, 128)
-        assert trimmed.stride() == full.stride()
-        # K prefix starts at the buffer base; V prefix at the segment offset.
-        assert trimmed[0].data_ptr() == full[0].data_ptr()
-        assert trimmed[1].data_ptr() == full[1].data_ptr()
-        # The narrowed views cover only committed memory.
-        trimmed[0, committed - 1].fill_(1)
-        trimmed[1, committed - 1].fill_(2)
-        torch.accelerator.synchronize()
-        assert torch.all(full[0, committed - 1] == 1)
-        assert torch.all(full[1, committed - 1] == 2)
-    finally:
-        buffers.free()
-
-
-def test_v2_extensible_defragment_on_commit() -> None:
-    """commit(defragment=True) re-maps each segment prefix as ONE physical
-    chunk (required for KV-transfer registration), discarding prior data."""
-    spec = _full_attention_spec()
-    kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
-    kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
-    try:
-        kv_cache = kv_caches["layer.0"]
-        # Staged commits spanning multiple VMM granules -> multiple physical
-        # chunks per segment.
-        buffers.commit(8)
-        buffers.commit(NUM_BLOCKS // 2)
-        kv_cache[0, 0].fill_(1)
-        torch.accelerator.synchronize()
-        [(buffer, _)] = buffers.buffers
-        assert len(buffer._buffer._handles) > 2
-
-        buffers.commit(NUM_BLOCKS, defragment=True)
-        # One chunk per segment; data discarded (zeroed); views still work.
-        assert len(buffer._buffer._handles) == 2
-        assert buffers.num_blocks_committed == NUM_BLOCKS
-        torch.accelerator.synchronize()
-        assert torch.count_nonzero(kv_cache) == 0
-        kv_cache[1, NUM_BLOCKS - 1].fill_(3)
-        torch.accelerator.synchronize()
-        assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 3)
-    finally:
-        buffers.free()

vllm/config/cache.py

@@ -177,18 +177,6 @@ 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 driver virtual memory to reserve the KV cache address range up
-    front, run warmup and CUDA graph capture with only a small block prefix
-    physically committed, and commit the final size afterwards.
-
-    This makes automatic KV sizing account for the memory that warmup and
-    CUDA graph capture actually consume (including worst-case activation
-    working sets, e.g. with speculative decoding), and avoids warmup-time
-    OOMs. Requires driver VMM support (CUDA or ROCm; falls back to standard
-    allocation with a warning where unavailable, e.g. WSL2).
-    """
-
     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
@@ -234,8 +222,6 @@ class CacheConfig:
             "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

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -255,14 +255,6 @@ class KVConnectorBase_V1(ABC):
 
         Args:
             kv_caches: dictionary of layer names, kv cache
-
-        Note:
-            The views' shapes/strides/numel are the authoritative source of
-            the KV cache geometry; do not derive block sizes or extents from
-            `untyped_storage().nbytes()`. With the extensible KV cache, the
-            underlying storage spans the reserved virtual-address capacity,
-            of which only each view's per-segment block prefix is physically
-            committed (and safe to access or register).
         """
         return
 

vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py

@@ -61,9 +61,9 @@ def get_offloading_event_group_spec(
 
 @dataclass(slots=True)
 class _OffloadEventMetadata:
-    """BlockStored payload snapshot for one OffloadKey, captured at store
-    time and kept until the matching eviction event. ``medium`` is forwarded
-    from the OffloadingEvent."""
+    """BlockStored payload snapshot for one OffloadKey, captured while the
+    Request is available and kept until the matching eviction event. ``medium``
+    is forwarded from the OffloadingEvent."""
 
     # The chunk's constituent block hashes; the last one is the OffloadKey.
     block_hashes: tuple[BlockHash, ...]
@@ -81,10 +81,11 @@ class _OffloadEventMetadata:
 class OffloadingEventsTracker:
     """Tracks offloaded chunks' KV event payloads from store to eviction.
 
-    The scheduler calls :meth:`record_store` from ``_build_store_jobs``
-    while the ``Request`` is available, and routes the manager's raw
-    :class:`OffloadingEvent` stream through :meth:`take_events`. All state
-    is bounded by the CPU pool capacity and cleared by :meth:`reset`.
+    The scheduler calls :meth:`record_store` from ``_build_store_jobs`` and
+    :meth:`record_lookup` for ready primary-tier hits while the ``Request`` is
+    available. Deferred and missing lookups add no state. Under the connector's
+    supported success-only transfer model, entries follow primary allocations
+    until CPU removal translation or :meth:`reset`.
     """
 
     def __init__(self, config: OffloadingKVEventsConfig):
@@ -93,8 +94,7 @@ class OffloadingEventsTracker:
             config.enable_kv_cache_events and config.self_describing_kv_events
         )
 
-        # OffloadKey -> payload snapshot, kept until the eviction event so
-        # BlockRemoved can fan out. Bounded: one entry per offloaded chunk.
+        # OffloadKey -> payload snapshot, kept until CPU removal or reset.
         self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}
 
     def record_store(
@@ -116,6 +116,23 @@ class OffloadingEventsTracker:
         meta = self._build_event_metadata(req, group_config, chunk_idx)
         self._pending_event_metadata[offload_key] = meta
 
+    def record_lookup(
+        self,
+        req: Request,
+        group_config: "GroupOffloadConfig",
+        chunk_idx: int,
+        offload_key: OffloadKey,
+    ) -> None:
+        """Snapshot metadata for a ready primary-tier lookup hit."""
+        if not self.self_describing_enabled:
+            return
+        if group_config.sliding_window_size_in_chunks is not None:
+            return
+        if offload_key not in self._pending_event_metadata:
+            self._pending_event_metadata[offload_key] = self._build_event_metadata(
+                req, group_config, chunk_idx
+            )
+
     def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
         """Translate raw OffloadingEvents into self-describing KV events.
 
@@ -165,7 +182,7 @@ class OffloadingEventsTracker:
         assert len(chunk_hashes) == hbf
 
         if group_config.sliding_window_size_in_chunks is not None:
-            # record_store filters these out before calling this helper.
+            # The recording methods filter these out before calling this helper.
             raise AssertionError("self-describing events only support full attention")
 
         parent_block_hash: BlockHash | None
@@ -232,7 +249,8 @@ class OffloadingEventsTracker:
                         "OffloadingEventsTracker: no event metadata for "
                         "offload key during BlockStored emission; emitting a "
                         "placeholder payload. Expected for non-full-attention "
-                        "groups; otherwise indicates a missing populate path."
+                        "groups and promotions not observed as a primary-tier "
+                        "hit before translation."
                     )
                 yield self._placeholder_stored(key, event.medium, locality)
                 continue

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -463,15 +463,27 @@ class OffloadingConnectorScheduler:
             del self._req_status[req_id]
 
     def _maximal_prefix_lookup(
-        self, keys: Iterable[OffloadKey], req_context: ReqContext
+        self,
+        keys: Iterable[OffloadKey],
+        req_context: ReqContext,
+        req: Request,
+        group_config: GroupOffloadConfig,
+        start_chunk_idx: int,
     ) -> int | None:
         """Return the number of consecutive offloaded chunks from the start,
         or None if the backend deferred a lookup."""
         hit_count = 0
         defer_lookup = False
-        for key in keys:
-            match self.manager.lookup(key, req_context):
+        for local_idx, key in enumerate(keys):
+            result = self.manager.lookup(key, req_context)
+            match result:
                 case LookupResult.HIT:
+                    self._events_tracker.record_lookup(
+                        req,
+                        group_config,
+                        start_chunk_idx + local_idx,
+                        key,
+                    )
                     hit_count += 1
                 case LookupResult.HIT_PENDING:
                     defer_lookup = True
@@ -616,7 +628,11 @@ class OffloadingConnectorScheduler:
                 num_hit_chunks: int | None
                 if sliding_window_size_in_chunks is None:
                     num_hit_chunks = self._maximal_prefix_lookup(
-                        offload_keys, req_status.req_context
+                        offload_keys,
+                        req_status.req_context,
+                        req_status.req,
+                        group_config,
+                        start_chunk_idx,
                     )
                 else:
                     required_window = sliding_window_size_in_chunks

vllm/engine/arg_utils.py

@@ -525,7 +525,6 @@ 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_scheduled_tokens: int | None = None
     max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills
@@ -1166,10 +1165,6 @@ class EngineArgs:
         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"]
@@ -1910,7 +1905,6 @@ class EngineArgs:
             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,

vllm/entrypoints/llm.py

@@ -119,11 +119,6 @@ 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
@@ -216,7 +211,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
         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,
@@ -315,7 +309,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
             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,

vllm/model_executor/models/gemma4.py

@@ -84,6 +84,12 @@ from .utils import (
 
 logger = init_logger(__name__)
 
+_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
+    orig_to_new_regex={
+        re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
+    }
+)
+
 
 def _remap_gemma4_expert_weight_name(name: str) -> str:
     return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
@@ -1508,7 +1514,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
 class Gemma4ForCausalLM(
     nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
 ):
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
         orig_to_new_prefix={
             # Gemma4ForConditionalGeneration already loads the text stack
             # from `model.language_model.*`. We reuse that same checkpoint

vllm/model_executor/models/gemma4_mm.py

@@ -40,7 +40,10 @@ from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+from vllm.model_executor.models.gemma4 import (
+    _GEMMA4_EXPERT_PARENT_MAPPER,
+    Gemma4ForCausalLM,
+)
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.transformers.utils import recursive_replace_linear
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -998,7 +1001,7 @@ class Gemma4ForConditionalGeneration(
     }
 
     # Maps checkpoint prefixes to vLLM module paths.
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
         orig_to_new_prefix={
             # vision tower
             "model.vision_tower": "vision_tower",
@@ -1010,7 +1013,7 @@ class Gemma4ForConditionalGeneration(
             "model.language_model.": "language_model.model.",
             "lm_head.": "language_model.lm_head.",
             "model": "language_model.model",
-        }
+        },
     )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

vllm/model_executor/warmup/qwen_triton_warmup.py

@@ -174,12 +174,6 @@ def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool:
     if not callable(zero_block_ids):
         return False
 
-    # With the extensible KV cache (V2), only a prefix of the blocks is
-    # physically committed; make sure the blocks zeroed below are backed.
-    ensure_kv_cache_blocks = getattr(runner, "ensure_kv_cache_blocks", None)
-    if callable(ensure_kv_cache_blocks):
-        ensure_kv_cache_blocks(max(_ZERO_KV_N_BLOCKS))
-
     for n_blocks in _ZERO_KV_N_BLOCKS:
         zero_block_ids(list(range(n_blocks)))
     return True

vllm/multimodal/media/connector.py

@@ -476,15 +476,17 @@ class MediaConnector:
         self,
         image_url: str,
         *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
     ) -> Image.Image:
         """
         Load a PIL image from an HTTP or base64 data URL.
 
-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
         """
         image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
 
         try:
@@ -501,15 +503,17 @@ class MediaConnector:
         self,
         image_url: str,
         *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
     ) -> Image.Image:
         """
         Asynchronously load a PIL image from an HTTP or base64 data URL.
 
-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
         """
         image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
 
         try:
@@ -526,14 +530,14 @@ class MediaConnector:
         self,
         video_url: str,
         *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
         video_processor: str | None = None,
     ) -> tuple[npt.NDArray, dict[str, Any]]:
         """
         Load video from an HTTP or base64 data URL.
         """
         image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
         video_io_kwargs = dict(self.media_io_kwargs.get("video", {}))
         if "video_backend" not in video_io_kwargs and (
@@ -552,16 +556,18 @@ class MediaConnector:
         self,
         video_url: str,
         *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
         video_processor: str | None = None,
     ) -> tuple[npt.NDArray, dict[str, Any]]:
         """
         Asynchronously load video from an HTTP or base64 data URL.
 
-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
         """
         image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
         )
         video_io_kwargs = dict(self.media_io_kwargs.get("video", {}))
         if "video_backend" not in video_io_kwargs and (

vllm/multimodal/media/image.py

@@ -25,9 +25,11 @@ class ImageMediaIO(MediaIO[Image.Image]):
     error handling.
     """
 
-    def __init__(self, image_mode: str = "RGB", **kwargs) -> None:
+    def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None:
         super().__init__()
 
+        # Target mode for loaded images; `None` keeps the original mode
+        # (i.e. no conversion, alpha channel is preserved as-is).
         self.image_mode = image_mode
         # `kwargs` contains custom arguments from
         # --media-io-kwargs for this modality, merged with
@@ -62,7 +64,7 @@ class ImageMediaIO(MediaIO[Image.Image]):
         """Convert image mode with custom background color."""
         if isinstance(image, MediaWithBytes):
             image = image.media
-        if image.mode == self.image_mode:
+        if self.image_mode is None or image.mode == self.image_mode:
             return image
         elif image.mode == "RGBA" and self.image_mode == "RGB":
             return rgba_to_rgb(image, self.rgba_background_color)

vllm/multimodal/utils.py

@@ -58,13 +58,14 @@ def encode_audio_url(
 def encode_image_base64(
     image: Image.Image,
     *,
-    image_mode: str = "RGB",
+    image_mode: str | None = "RGB",
     format: str = "PNG",
 ) -> str:
     """
     Encode a pillow image to base64 format.
 
     By default, the image is converted into RGB format before being encoded.
+    Pass `image_mode=None` to keep the original image mode.
     """
     image_io = ImageMediaIO(image_mode=image_mode)
     return image_io.encode_base64(image, image_format=format)
@@ -73,13 +74,14 @@ def encode_image_base64(
 def encode_image_url(
     image: Image.Image,
     *,
-    image_mode: str = "RGB",
+    image_mode: str | None = "RGB",
     format: str = "PNG",
 ) -> str:
     """
     Encode a pillow image as a data URL.
 
     By default, the image is converted into RGB format before being encoded.
+    Pass `image_mode=None` to keep the original image mode.
     """
     image_b64 = encode_image_base64(image, image_mode=image_mode, format=format)
     mimetype = mimetypes.types_map.get("." + format.lower(), "image")

vllm/utils/extensible_tensor.py

@@ -1,451 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Growable GPU byte buffers backed by driver virtual memory management."""
-
-from __future__ import annotations
-
-import ctypes
-from contextlib import suppress
-
-import torch
-
-from vllm.logger import init_logger
-from vllm.utils.vmm_driver import get_vmm_driver
-
-logger = init_logger(__name__)
-
-
-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, shareable: bool = False
-    ) -> None:
-        self._driver = get_vmm_driver()
-        self._driver.ensure_context(device_index)
-        self.device_index = device_index
-        self._shareable = shareable
-
-        self.granularity: int = self._driver.granularity(device_index)
-        self.reserved_size: int = _round_up(max(max_bytes, 1), self.granularity)
-        self.base_ptr: int = self._driver.reserve(self.reserved_size)
-
-        # 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`."""
-        driver = self._driver
-        driver.ensure_context(self.device_index)
-        if self._shareable:
-            try:
-                handle = driver.create(size, self.device_index, shareable=True)
-            except RuntimeError as e:
-                logger.warning_once(
-                    "Failed to allocate shareable (IPC/RDMA-capable) memory "
-                    "(%s); falling back to standard allocation. KV transfers "
-                    "from this memory may fail.",
-                    e,
-                )
-                self._shareable = False
-                handle = driver.create(size, self.device_index)
-        else:
-            handle = driver.create(size, self.device_index)
-
-        addr = self.base_ptr + offset
-        try:
-            driver.map(addr, size, handle)
-        except RuntimeError:
-            driver.release(handle)
-            raise
-        driver.set_access(addr, size, self.device_index)
-
-        self._handles.append((handle, offset, size))
-
-    def release_physical(self) -> None:
-        """Unmap and release all physical memory, keeping the VA reservation.
-
-        The base pointer (and any tensor views over it) stays valid but
-        unbacked; `ensure_committed_range` maps fresh physical pages again.
-        """
-        driver = self._driver
-        driver.ensure_context(self.device_index)
-        if self._handles:
-            torch.accelerator.synchronize(self.device_index)
-        for handle, offset, size in self._handles:
-            driver.unmap(self.base_ptr + offset, size)
-            driver.release(handle)
-        self._handles = []
-        self._mapped_granules = set()
-
-    def free(self) -> None:
-        if self._freed:
-            return
-        self._freed = True
-        self.release_physical()
-        if self.base_ptr:
-            self._driver.free_reserved(self.base_ptr, self.reserved_size)
-        self.base_ptr = 0
-
-    def __del__(self) -> None:
-        with suppress(Exception):
-            self.free()
-
-
-_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)
-    device_type = get_vmm_driver().dlpack_device_type
-    managed.dl_tensor.device = _DLDevice(device_type, 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,
-        shareable: bool = False,
-    ) -> 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.accelerator.current_device_index()
-        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.accelerator.current_device_index()
-        )
-
-        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, shareable=shareable
-        )
-        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 physical_bytes(self) -> int:
-        """Physically mapped bytes (committed size rounded up to granules)."""
-        return self._buffer.committed_bytes
-
-    def release_physical(self) -> None:
-        """Release all physical memory while keeping the VA reservation.
-
-        Existing tensor views stay pointer-valid but must not be accessed
-        until the buffer is committed again; the data is discarded.
-        """
-        self._buffer.release_physical()
-        self._bytes_per_segment = 0
-
-    @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
-
-
-class ExtensibleKVCacheBuffers:
-    """Grow-only physical backing for the KV cache: one CUDA virtual-memory
-    buffer per KV cache tensor, committed as a per-segment prefix of blocks.
-
-    `commit` maps (and zeroes) physical pages for additional blocks while
-    keeping every buffer's base pointer, existing data, and the logical views
-    built over the full reserved capacity stable.
-    """
-
-    def __init__(
-        self,
-        buffers: list[tuple[ExtensibleTensor, int]],
-        num_blocks_capacity: int,
-    ) -> None:
-        # Each entry is (buffer, bytes_per_block_per_segment).
-        self.buffers = buffers
-        self.num_blocks_capacity = num_blocks_capacity
-        self.num_blocks_committed = 0
-        self._num_blocks_to_recommit = 0
-
-    def commit(self, num_blocks: int, defragment: bool = False) -> None:
-        """Grow the committed prefix of every buffer to `num_blocks` blocks.
-
-        With `defragment=True`, all previously committed physical chunks are
-        released first so each segment's prefix is re-mapped as one physical
-        allocation. Existing contents are DISCARDED, so this is only valid
-        before real KV data exists (e.g. right after warmup). It is required
-        before KV-transfer registration: UCX cannot transfer memory regions
-        that span multiple VMM allocation handles.
-        """
-        if defragment and self.num_blocks_committed > 0:
-            self.release_physical()
-        if num_blocks <= self.num_blocks_committed:
-            return
-        for buffer, bytes_per_block_per_segment in self.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.num_blocks_committed = num_blocks
-
-    @property
-    def physical_bytes(self) -> int:
-        return sum(buffer.physical_bytes for buffer, _ in self.buffers)
-
-    def release_physical(self) -> None:
-        """Discard all physical memory (sleep), keeping VA and views valid."""
-        self._num_blocks_to_recommit = self.num_blocks_committed
-        for buffer, _ in self.buffers:
-            buffer.release_physical()
-        self.num_blocks_committed = 0
-
-    def recommit(self) -> None:
-        """Re-commit the pre-release block count with freshly zeroed pages."""
-        self.commit(self._num_blocks_to_recommit)
-
-    def free(self) -> None:
-        for buffer, _ in self.buffers:
-            buffer.free()
-        self.buffers = []
-        self.num_blocks_committed = 0

vllm/utils/vmm_driver.py

@@ -1,354 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""ctypes bindings for GPU virtual-memory-management (VMM) driver APIs.
-
-Exposes a uniform driver interface over the CUDA driver's ``cuMem*`` entry
-points and HIP's mirrored ``hipMem*`` entry points, used by
-:class:`vllm.utils.extensible_tensor.ExtensibleTensor`: reserve a virtual
-address range, create physical memory handles, map/unmap them into the
-reservation, and set access permissions.
-"""
-
-from __future__ import annotations
-
-import ctypes
-from functools import cache
-from typing import Any
-
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-_SUCCESS = 0
-_MEM_ALLOCATION_TYPE_PINNED = 1
-_MEM_LOCATION_TYPE_DEVICE = 1
-_MEM_ALLOC_GRANULARITY_MINIMUM = 0
-_MEM_ACCESS_FLAGS_PROT_READWRITE = 3
-_MEM_ALLOCATION_COMP_NONE = 0
-_MEM_HANDLE_TYPE_POSIX_FD = 1
-
-DevicePtr = ctypes.c_ulonglong
-MemHandle = ctypes.c_ulonglong
-_Context = ctypes.c_void_p
-
-
-class _MemLocation(ctypes.Structure):
-    _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)]
-
-
-class _MemAllocFlags(ctypes.Structure):
-    _fields_ = [
-        ("compressionType", ctypes.c_ubyte),
-        ("gpuDirectRDMACapable", ctypes.c_ubyte),
-        ("usage", ctypes.c_ushort),
-        ("reserved", ctypes.c_ubyte * 4),
-    ]
-
-
-class _MemAllocationProp(ctypes.Structure):
-    # Layout shared by CUmemAllocationProp and hipMemAllocationProp.
-    _fields_ = [
-        ("type", ctypes.c_int),
-        ("requestedHandleTypes", ctypes.c_int),
-        ("location", _MemLocation),
-        ("win32HandleMetaData", ctypes.c_void_p),
-        ("allocFlags", _MemAllocFlags),
-    ]
-
-
-class _MemAccessDesc(ctypes.Structure):
-    _fields_ = [("location", _MemLocation), ("flags", ctypes.c_int)]
-
-
-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
-
-
-class VmmDriver:
-    """Uniform interface over a GPU driver's virtual memory management API.
-
-    Subclasses supply the driver library candidates and symbol names; the
-    call signatures and struct layouts are shared between CUDA and HIP.
-    """
-
-    # DLPack device type for tensors viewing driver-mapped memory.
-    dlpack_device_type: int
-    _lib_candidates: tuple[str, ...]
-    _lib_search_name: str
-    # Logical name -> library symbol.
-    _symbols: dict[str, str]
-
-    def __init__(self) -> None:
-        self._lib = self._load_library()
-        self._fns: dict[str, Any] = {}
-        for logical, symbol in self._symbols.items():
-            self._fns[logical] = getattr(self._lib, symbol)
-        self._configure_signatures()
-
-    def _load_library(self) -> ctypes.CDLL:
-        for name in self._lib_candidates:
-            try:
-                return ctypes.CDLL(name)
-            except OSError:
-                continue
-        if path := _find_loaded_library(self._lib_search_name):
-            return ctypes.CDLL(path)
-        raise RuntimeError(
-            f"Could not load {self._lib_candidates[0]}. The GPU driver "
-            "library is required for VMM-backed tensors."
-        )
-
-    def _configure_signatures(self) -> None:
-        pointer = ctypes.POINTER
-        fns = self._fns
-        fns["get_granularity"].argtypes = [
-            pointer(ctypes.c_size_t),
-            pointer(_MemAllocationProp),
-            ctypes.c_int,
-        ]
-        fns["address_reserve"].argtypes = [
-            pointer(DevicePtr),
-            ctypes.c_size_t,
-            ctypes.c_size_t,
-            DevicePtr,
-            ctypes.c_ulonglong,
-        ]
-        fns["create"].argtypes = [
-            pointer(MemHandle),
-            ctypes.c_size_t,
-            pointer(_MemAllocationProp),
-            ctypes.c_ulonglong,
-        ]
-        fns["map"].argtypes = [
-            DevicePtr,
-            ctypes.c_size_t,
-            ctypes.c_size_t,
-            MemHandle,
-            ctypes.c_ulonglong,
-        ]
-        fns["set_access"].argtypes = [
-            DevicePtr,
-            ctypes.c_size_t,
-            pointer(_MemAccessDesc),
-            ctypes.c_size_t,
-        ]
-        fns["unmap"].argtypes = [DevicePtr, ctypes.c_size_t]
-        fns["release"].argtypes = [MemHandle]
-        fns["address_free"].argtypes = [DevicePtr, ctypes.c_size_t]
-        for fn in fns.values():
-            fn.restype = ctypes.c_int
-
-    def error_string(self, code: int) -> str:
-        raise NotImplementedError
-
-    def ensure_context(self, device_index: int) -> None:
-        """Make sure a driver context for `device_index` is current."""
-        raise NotImplementedError
-
-    def _check(self, result: int) -> None:
-        if result == _SUCCESS:
-            return
-        raise RuntimeError(f"GPU driver error {result}: {self.error_string(result)}")
-
-    def _make_alloc_prop(
-        self, device_index: int, shareable: bool = False
-    ) -> _MemAllocationProp:
-        prop = _MemAllocationProp()
-        prop.type = _MEM_ALLOCATION_TYPE_PINNED
-        prop.location.type = _MEM_LOCATION_TYPE_DEVICE
-        prop.location.id = device_index
-        prop.allocFlags.compressionType = _MEM_ALLOCATION_COMP_NONE
-        if shareable:
-            # KV transfer engines access this memory from other processes:
-            # intra-node CUDA IPC needs an exportable (POSIX FD) handle type,
-            # and NIC RDMA needs the GPU-direct-RDMA-capable flag.
-            prop.requestedHandleTypes = _MEM_HANDLE_TYPE_POSIX_FD
-            prop.allocFlags.gpuDirectRDMACapable = 1
-        return prop
-
-    def granularity(self, device_index: int) -> int:
-        prop = self._make_alloc_prop(device_index)
-        granularity = ctypes.c_size_t()
-        self._check(
-            self._fns["get_granularity"](
-                ctypes.byref(granularity),
-                ctypes.byref(prop),
-                _MEM_ALLOC_GRANULARITY_MINIMUM,
-            )
-        )
-        return granularity.value
-
-    def reserve(self, size: int) -> int:
-        """Reserve a virtual address range and return its base pointer."""
-        dptr = DevicePtr()
-        self._check(self._fns["address_reserve"](ctypes.byref(dptr), size, 0, 0, 0))
-        return dptr.value
-
-    def free_reserved(self, ptr: int, size: int) -> None:
-        self._check(self._fns["address_free"](ptr, size))
-
-    def create(self, size: int, device_index: int, shareable: bool = False) -> int:
-        """Create a physical memory handle of `size` bytes."""
-        prop = self._make_alloc_prop(device_index, shareable)
-        handle = MemHandle()
-        self._check(
-            self._fns["create"](ctypes.byref(handle), size, ctypes.byref(prop), 0)
-        )
-        return handle.value
-
-    def map(self, ptr: int, size: int, handle: int) -> None:
-        self._check(self._fns["map"](ptr, size, 0, handle, 0))
-
-    def set_access(self, ptr: int, size: int, device_index: int) -> None:
-        desc = _MemAccessDesc()
-        desc.location.type = _MEM_LOCATION_TYPE_DEVICE
-        desc.location.id = device_index
-        desc.flags = _MEM_ACCESS_FLAGS_PROT_READWRITE
-        self._check(self._fns["set_access"](ptr, size, ctypes.byref(desc), 1))
-
-    def unmap(self, ptr: int, size: int) -> None:
-        self._check(self._fns["unmap"](ptr, size))
-
-    def release(self, handle: int) -> None:
-        self._check(self._fns["release"](handle))
-
-
-class CudaVmmDriver(VmmDriver):
-    dlpack_device_type = 2  # kDLCUDA
-    _lib_candidates = ("libcuda.so.1", "libcuda.so")
-    _lib_search_name = "libcuda"
-    _symbols = {
-        "get_granularity": "cuMemGetAllocationGranularity",
-        "address_reserve": "cuMemAddressReserve",
-        "create": "cuMemCreate",
-        "map": "cuMemMap",
-        "set_access": "cuMemSetAccess",
-        "unmap": "cuMemUnmap",
-        "release": "cuMemRelease",
-        "address_free": "cuMemAddressFree",
-    }
-
-    def __init__(self) -> None:
-        super().__init__()
-        lib = self._lib
-        lib.cuGetErrorString.argtypes = [
-            ctypes.c_int,
-            ctypes.POINTER(ctypes.c_char_p),
-        ]
-        lib.cuGetErrorString.restype = ctypes.c_int
-        lib.cuCtxGetCurrent.argtypes = [ctypes.POINTER(_Context)]
-        lib.cuCtxGetCurrent.restype = ctypes.c_int
-        lib.cuDevicePrimaryCtxRetain.argtypes = [
-            ctypes.POINTER(_Context),
-            ctypes.c_int,
-        ]
-        lib.cuDevicePrimaryCtxRetain.restype = ctypes.c_int
-        lib.cuCtxSetCurrent.argtypes = [_Context]
-        lib.cuCtxSetCurrent.restype = ctypes.c_int
-
-    def error_string(self, code: int) -> str:
-        msg = ctypes.c_char_p()
-        self._lib.cuGetErrorString(code, ctypes.byref(msg))
-        return msg.value.decode() if msg.value else "unknown error"
-
-    def ensure_context(self, device_index: int) -> None:
-        pctx = _Context()
-        self._check(self._lib.cuCtxGetCurrent(ctypes.byref(pctx)))
-        if pctx.value:
-            return
-        self._check(
-            self._lib.cuDevicePrimaryCtxRetain(ctypes.byref(pctx), device_index)
-        )
-        self._check(self._lib.cuCtxSetCurrent(pctx))
-
-
-class HipVmmDriver(VmmDriver):
-    """HIP mirrors the CUDA driver's VMM API (``hipMem*``) with identical
-    call signatures, struct layouts, and constants; PyTorch's
-    expandable-segments allocator uses the same entry points on ROCm.
-    """
-
-    dlpack_device_type = 10  # kDLROCM
-    _lib_candidates = (
-        "libamdhip64.so",
-        "libamdhip64.so.7",
-        "libamdhip64.so.6",
-        "libamdhip64.so.5",
-    )
-    _lib_search_name = "libamdhip64"
-    _symbols = {
-        "get_granularity": "hipMemGetAllocationGranularity",
-        "address_reserve": "hipMemAddressReserve",
-        "create": "hipMemCreate",
-        "map": "hipMemMap",
-        "set_access": "hipMemSetAccess",
-        "unmap": "hipMemUnmap",
-        "release": "hipMemRelease",
-        "address_free": "hipMemAddressFree",
-    }
-
-    def __init__(self) -> None:
-        super().__init__()
-        lib = self._lib
-        lib.hipGetErrorString.argtypes = [ctypes.c_int]
-        lib.hipGetErrorString.restype = ctypes.c_char_p
-        lib.hipGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)]
-        lib.hipGetDevice.restype = ctypes.c_int
-        lib.hipSetDevice.argtypes = [ctypes.c_int]
-        lib.hipSetDevice.restype = ctypes.c_int
-
-    def error_string(self, code: int) -> str:
-        msg = self._lib.hipGetErrorString(code)
-        return msg.decode() if msg else "unknown error"
-
-    def ensure_context(self, device_index: int) -> None:
-        # The HIP runtime manages contexts implicitly; just make sure the
-        # buffer's device is current on this thread.
-        device = ctypes.c_int()
-        self._check(self._lib.hipGetDevice(ctypes.byref(device)))
-        if device.value != device_index:
-            self._check(self._lib.hipSetDevice(device_index))
-
-
-@cache
-def get_vmm_driver() -> VmmDriver:
-    import torch
-
-    if torch.version.hip is not None:
-        return HipVmmDriver()
-    return CudaVmmDriver()
-
-
-@cache
-def vmm_unavailable_reason() -> str | None:
-    """Probe VMM support; returns None if usable, else a reason string.
-
-    Checks that the driver library loads, exposes the VMM entry points, and
-    can actually reserve (and release) a virtual address range on the current
-    device. Notably returns a reason on platforms whose driver lacks VMM
-    support (e.g. WSL2) and on non-CUDA/ROCm builds.
-    """
-    try:
-        import torch
-
-        if not torch.accelerator.is_available():
-            return "no CUDA/ROCm device is available"
-        torch.cuda.init()
-        driver = get_vmm_driver()
-        device_index = torch.accelerator.current_device_index()
-        driver.ensure_context(device_index)
-        granularity = driver.granularity(device_index)
-        ptr = driver.reserve(granularity)
-        driver.free_reserved(ptr, granularity)
-    except Exception as e:
-        return str(e)
-    return None

vllm/v1/core/encoder_cache_manager.py

@@ -5,8 +5,6 @@ from collections import OrderedDict
 from collections.abc import Mapping
 from typing import TYPE_CHECKING
 
-import torch
-
 from vllm.logger import init_logger
 from vllm.v1.request import Request
 
@@ -80,17 +78,6 @@ class EncoderCacheManager:
         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.
 

vllm/v1/engine/core.py

@@ -91,7 +91,6 @@ logger = init_logger(__name__)
 
 
 HANDSHAKE_TIMEOUT_MINS = 5
-_WARMUP_MEMORY_BUFFER_BYTES = 150 * (1 << 20)
 
 _R = TypeVar("_R")  # Return type for collective_rpc
 
@@ -292,89 +291,37 @@ class EngineCore:
 
         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:
-            if (
-                vllm_config.kv_transfer_config is not None
-                and not vllm_config.use_v2_model_runner
-            ):
-                raise ValueError(
-                    "enable_extensible_kv_cache=True with KV connectors "
-                    "requires the V2 model runner (which defers connector "
-                    "registration until the final KV cache size is committed)."
-                )
-            # The workers' drivers must support virtual memory management
-            # (e.g. WSL2 and non-GPU platforms do not); fall back gracefully.
-            reasons: list[str | None] = self.collective_rpc(
-                "extensible_kv_cache_unsupported_reason"
-            )
-            if reason := next((r for r in reasons if r), None):
-                logger.warning(
-                    "Disabling extensible KV cache; falling back to standard "
-                    "KV cache allocation: %s",
-                    reason,
-                )
-                use_extensible_kv_cache = False
-
         # 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,
-        )
 
-        # Initialize KV cache and warm up execution. With extensible KV cache,
-        # this reserves the upper-bound address range, commits only the block
-        # prefix warmup needs, and runs warmup / CUDA graph capture before the
-        # post-warmup KV size is committed.
-        compilation_times = self.model_executor.initialize_from_config(
-            kv_cache_configs,
-            extensible=use_extensible_kv_cache,
-        )
-        if use_extensible_kv_cache:
-            if vllm_config.cache_config.kv_cache_memory_bytes is None:
-                # Automatic sizing: re-derive the KV cache size from the
-                # memory actually consumed by warmup and CUDA graph capture.
-                # With an explicit kv_cache_memory_bytes, the requested size
-                # is committed as-is (the extensible path still defers the
-                # commit until after warmup).
-                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.warmup_memory
-                        - _WARMUP_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,
-                )
-                scheduler_kv_cache_config = self._apply_kv_cache_config(
-                    vllm_config,
-                    kv_cache_configs,
-                    max_model_len_before,
-                )
-            self.model_executor.extend_kv_cache(kv_cache_configs)
+        # 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()
+
+        # Initialize kv cache and warmup the execution
+        self.model_executor.initialize_from_config(kv_cache_configs)
 
         elapsed = time.time() - start
         compile_time = vllm_config.compilation_config.compilation_time
@@ -403,35 +350,6 @@ class EngineCore:
             )
         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, ...]:
         supported_tasks = self.model_executor.supported_tasks
         self._log_pooler_config(supported_tasks)

vllm/v1/executor/abstract.py

@@ -115,20 +115,12 @@ class Executor(ABC):
     def _init_executor(self) -> None:
         raise NotImplementedError
 
-    def initialize_from_config(
-        self,
-        kv_cache_configs: list[KVCacheConfig],
-        extensible: bool = False,
-    ) -> list[CompilationTimes]:
+    def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None:
         """
         Initialize the KV caches and begin the model execution loop of the
         underlying workers.
         """
-        self.collective_rpc(
-            "initialize_from_config",
-            args=(kv_cache_configs,),
-            kwargs={"extensible": extensible} if extensible else None,
-        )
+        self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
         compilation_times: list[CompilationTimes] = self.collective_rpc(
             "compile_or_warm_up_model"
         )
@@ -143,7 +135,6 @@ class Executor(ABC):
             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
         """
@@ -158,10 +149,6 @@ class Executor(ABC):
     def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
         return self.collective_rpc("get_kv_cache_spec")
 
-    def extend_kv_cache(self, kv_cache_configs: list[KVCacheConfig]) -> None:
-        """Commit the final KV cache size on all workers (extensible flow)."""
-        self.collective_rpc("extend_kv_cache", args=(kv_cache_configs,))
-
     @overload
     def collective_rpc(
         self,

vllm/v1/kv_cache_interface.py

@@ -972,6 +972,23 @@ class KVCacheConfig:
     def has_mamba_layers(self) -> bool:
         return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)
 
+    @property
+    def has_mixed_precision_kv_cache(self) -> bool:
+        """Whether attention groups store their KV cache at more than one precision."""
+        kv_cache_precisions = {
+            (group.kv_cache_spec.dtype, group.kv_cache_spec.kv_quant_mode)
+            for group in self.kv_cache_groups
+            if isinstance(group.kv_cache_spec, AttentionSpec)
+        }
+        return len(kv_cache_precisions) > 1
+
     @property
     def needs_kv_cache_zeroing(self) -> bool:
-        return self.has_mamba_layers
+        """Whether newly allocated KV cache blocks must be zeroed before use.
+
+        Required for Mamba layers, whose state is read before it is fully written
+        (#35219), and for mixed-precision caches, where a block reused across
+        groups can be reinterpreted under a different precision and decode stale
+        bytes to NaN/Inf. Uniform-precision caches skip zeroing.
+        """
+        return self.has_mamba_layers or self.has_mixed_precision_kv_cache

vllm/v1/kv_offload/tiering/spec.py

@@ -136,14 +136,6 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
         super().__init__(config)
         # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it
         self._manager: OffloadingManager | None = None
-        if self.kv_events_config.self_describing_kv_events:
-            raise ValueError(
-                "self_describing_kv_events is not supported by "
-                "TieringOffloadingSpec. Tier promotions can emit primary-tier "
-                "store events that do not correspond to GPU store jobs, so the "
-                "current self-describing side table cannot describe them "
-                "correctly."
-            )
 
         # Parse secondary tier configurations
         self.secondary_tier_configs = self.extra_config.get("secondary_tiers", [])

vllm/v1/simple_kv_offload/worker.py

@@ -99,14 +99,12 @@ class SimpleCPUOffloadWorker:
         num_blocks = self.kv_cache_config.num_blocks
 
         # Deduplicate: multiple layers may share the same backing storage.
-        seen_ptrs: dict[
-            int, tuple[str, torch.Tensor, torch.Tensor | list[torch.Tensor]]
-        ] = {}
+        seen_ptrs: dict[int, tuple[str, torch.Tensor]] = {}
         for name, value in kv_caches.items():
             tensor = _repr_tensor(value)
             ptr = tensor.untyped_storage().data_ptr()
             if ptr not in seen_ptrs:
-                seen_ptrs[ptr] = (name, tensor, value)
+                seen_ptrs[ptr] = (name, tensor)
 
         # Build [num_blocks, block_bytes] int8 views from each unique
         # storage so that stride(0) gives block_bytes for the copy op.
@@ -114,51 +112,28 @@ class SimpleCPUOffloadWorker:
         # The physical layout varies across attention backends:
         #   FlashAttn/ROCm:  (2, num_blocks, ...) -> K/V outermost, 2 segments
         #   FlashInfer/MLA:  (num_blocks, ...)    -> blocks outermost, 1 segment
-        # We derive the per-block data size from the registration view rather
-        # than storage.nbytes(): with the extensible KV cache, the storage
-        # spans the reserved capacity while only the view's (per-segment
-        # prefix) extent is physically committed. Packed layouts keep the
-        # storage-based size (their bounded storage holds every layer's data
-        # per block, of which each layer's view only covers a slice). Dims
-        # whose byte-stride exceeds the per-block size are outer segment dims
-        # (e.g. the K/V dim of size 2); each segment's committed blocks form
-        # a prefix of that segment.
-        layer_is_packed: dict[str, bool] = {
-            ln: kv_tensor.block_stride > 0
-            for kv_tensor in self.kv_cache_config.kv_cache_tensors
-            for ln in kv_tensor.shared_by
-        }
+        # We derive page_size_bytes = storage.nbytes() // num_blocks, then
+        # classify dims: any dim whose byte-stride exceeds page_size_bytes
+        # must be an outer segment dim (e.g. the K/V dim of size 2). A less
+        # hacky way is to update the interface with the layout.
         unique_gpu_caches: dict[str, torch.Tensor] = {}
-        for name, tensor, value in seen_ptrs.values():
+        for name, tensor in seen_ptrs.values():
             storage = tensor.untyped_storage()
             raw = torch.empty(0, dtype=torch.int8, device=self.device).set_(
                 storage, 0, (storage.nbytes(),)
             )
             el = tensor.element_size()
-            if layer_is_packed.get(name, False):
-                # Bounded packed storage: every layer's data for all
-                # committed blocks.
-                page_size_bytes = storage.nbytes() // num_blocks
-            else:
-                # Sum over all state tensors of the layer (Mamba layers pack
-                # several per block); attention layers have a single tensor.
-                tensors = [value] if isinstance(value, torch.Tensor) else value
-                data_bytes = sum(t.numel() * t.element_size() for t in tensors)
-                page_size_bytes = data_bytes // num_blocks
+            page_size_bytes = storage.nbytes() // num_blocks
             outer_dims = [
                 d for d in range(tensor.ndim) if tensor.stride(d) * el > page_size_bytes
             ]
             if not outer_dims:
-                unique_gpu_caches[name] = raw[: num_blocks * page_size_bytes].view(
-                    num_blocks, -1
-                )
+                unique_gpu_caches[name] = raw.view(num_blocks, -1)
             else:
-                n_outer = tensor.shape[outer_dims[0]]
                 seg_stride = tensor.stride(outer_dims[0]) * el
-                seg_block_bytes = page_size_bytes // n_outer
-                for idx in range(n_outer):
+                for idx in range(tensor.shape[outer_dims[0]]):
                     offset = idx * seg_stride
-                    chunk = raw[offset : offset + num_blocks * seg_block_bytes]
+                    chunk = raw[offset : offset + seg_stride]
                     unique_gpu_caches[f"{name}.{idx}"] = chunk.view(num_blocks, -1)
 
         # Compute per-tensor bytes_per_block. Tensors may have different

vllm/v1/worker/gpu/attn_utils.py

@@ -16,11 +16,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.multimodal.inputs import MultiModalFeatureSpec
-from vllm.utils.extensible_tensor import (
-    ExtensibleKVCacheBuffers,
-    ExtensibleTensor,
-    uint8_tensor_from_ptr,
-)
 from vllm.utils.torch_utils import get_dtype_size
 from vllm.v1.attention.backend import (
     AttentionCGSupport,
@@ -192,15 +187,6 @@ def _allocate_kv_cache(
         for layer_name in kv_cache_tensor.shared_by:
             kv_cache_raw_tensors[layer_name] = tensor
 
-    _check_layer_coverage(kv_cache_config, kv_cache_raw_tensors, shared_layers)
-    return kv_cache_raw_tensors
-
-
-def _check_layer_coverage(
-    kv_cache_config: KVCacheConfig,
-    kv_cache_raw_tensors: dict[str, torch.Tensor],
-    shared_layers: dict[str, str],
-) -> None:
     layer_names = set()
     for group in kv_cache_config.kv_cache_groups:
         for layer_name in group.layer_names:
@@ -208,251 +194,7 @@ def _check_layer_coverage(
     assert layer_names == (kv_cache_raw_tensors.keys() | shared_layers.keys()), (
         "Some layers are not correctly initialized"
     )
-
-
-def _kv_cache_num_segments_by_layer(
-    attn_groups: Sequence[AttentionGroup],
-    kernel_block_sizes: list[int],
-    cache_dtype: str,
-    has_mamba: bool,
-) -> dict[str, int]:
-    """Number of equal contiguous segments of each layer's KV cache buffer
-    under its physical layout -- 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.
-    """
-    num_segments_by_layer: dict[str, int] = {}
-    for group in attn_groups:
-        if group.kv_cache_group_id >= len(kernel_block_sizes):
-            continue
-        kv_cache_spec = group.kv_cache_spec
-        if isinstance(kv_cache_spec, AttentionSpec) and not has_mamba:
-            if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
-                kernel_block_size = kv_cache_spec.storage_block_size
-            else:
-                kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
-            # Mirror the per-layer dtype selection of _reshape_kv_cache.
-            layer_cache_dtype = (
-                "auto"
-                if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
-                and not isinstance(kv_cache_spec, TQFullAttentionSpec)
-                else cache_dtype
-            )
-            block_dim = group.backend.get_kv_cache_block_dim(
-                kernel_block_size,
-                kv_cache_spec.num_kv_heads,
-                kv_cache_spec.head_size,
-                cache_dtype_str=layer_cache_dtype,
-            )
-            kv_cache_shape = group.backend.get_kv_cache_shape(
-                1,
-                kernel_block_size,
-                kv_cache_spec.num_kv_heads,
-                kv_cache_spec.head_size,
-                cache_dtype_str=layer_cache_dtype,
-            )
-            try:
-                stride_order = group.backend.get_kv_cache_stride_order()
-            except (AttributeError, NotImplementedError):
-                stride_order = tuple(range(len(kv_cache_shape)))
-            num_segments = 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_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 narrow_kv_caches_to_num_blocks(
-    kv_caches: dict[str, Any],
-    attn_groups: Sequence[AttentionGroup],
-    kernel_block_sizes: list[int],
-    cache_dtype: str,
-    num_blocks: int,
-    kv_cache_config: KVCacheConfig,
-) -> dict[str, Any]:
-    """Return views of the KV caches narrowed to the first `num_blocks` blocks.
-
-    With the extensible KV cache, the layer views span the full reserved
-    capacity while only a block prefix is physically committed. KV connectors
-    must only see (and register) backed memory, so hand them views whose block
-    dimension is trimmed to the committed count. Since committed blocks form a
-    prefix of each layout segment, a narrow along the block dim covers exactly
-    the committed bytes of every segment.
-    """
-    narrowed: dict[str, Any] = dict(kv_caches)
-    for group in attn_groups:
-        if group.kv_cache_group_id >= len(kernel_block_sizes):
-            continue
-        kv_cache_spec = group.kv_cache_spec
-        for layer_name in group.layer_names:
-            kv_cache = kv_caches.get(layer_name)
-            if kv_cache is None:
-                continue
-            if isinstance(kv_cache_spec, AttentionSpec):
-                if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
-                    kernel_block_size = kv_cache_spec.storage_block_size
-                else:
-                    kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
-                num_blocks_per_kv_block = (
-                    kv_cache_spec.storage_block_size // kernel_block_size
-                )
-                layer_cache_dtype = (
-                    "auto"
-                    if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
-                    and not isinstance(kv_cache_spec, TQFullAttentionSpec)
-                    else cache_dtype
-                )
-                block_dim = group.backend.get_kv_cache_block_dim(
-                    kernel_block_size,
-                    kv_cache_spec.num_kv_heads,
-                    kv_cache_spec.head_size,
-                    cache_dtype_str=layer_cache_dtype,
-                )
-                narrowed[layer_name] = kv_cache.narrow(
-                    block_dim, 0, num_blocks * num_blocks_per_kv_block
-                )
-            elif isinstance(kv_cache_spec, MambaSpec):
-                narrowed[layer_name] = [
-                    state.narrow(0, 0, num_blocks) for state in kv_cache
-                ]
-
-    _bound_packed_kv_cache_storages(narrowed, kv_cache_config, num_blocks)
-    return narrowed
-
-
-def _bound_packed_kv_cache_storages(
-    kv_caches: dict[str, Any],
-    kv_cache_config: KVCacheConfig,
-    num_blocks: int,
-) -> None:
-    """Rebase packed views onto storage ending at the committed block prefix."""
-    packed_by_storage: dict[int, tuple[int, list[str]]] = {}
-    for tensor_config in kv_cache_config.kv_cache_tensors:
-        if tensor_config.block_stride <= 0:
-            continue
-        committed_bytes = num_blocks * tensor_config.block_stride
-        for layer_name in tensor_config.shared_by:
-            cache = kv_caches.get(layer_name)
-            if not isinstance(cache, torch.Tensor):
-                continue
-            storage_ptr = cache.untyped_storage().data_ptr()
-            previous = packed_by_storage.get(storage_ptr)
-            if previous is None:
-                packed_by_storage[storage_ptr] = (committed_bytes, [layer_name])
-            else:
-                previous_bytes, layer_names = previous
-                if previous_bytes != committed_bytes:
-                    raise ValueError(
-                        "Packed KV cache views sharing storage disagree on the "
-                        f"committed size: {previous_bytes} != {committed_bytes}."
-                    )
-                layer_names.append(layer_name)
-
-    for storage_ptr, (committed_bytes, layer_names) in packed_by_storage.items():
-        first_cache = kv_caches[layer_names[0]]
-        assert isinstance(first_cache, torch.Tensor)
-        device_index = first_cache.device.index
-        assert device_index is not None
-        bounded_storage = uint8_tensor_from_ptr(
-            storage_ptr, committed_bytes, device_index
-        )
-        for layer_name in layer_names:
-            cache = kv_caches[layer_name]
-            assert isinstance(cache, torch.Tensor)
-            typed_storage = bounded_storage.view(cache.dtype)
-            kv_caches[layer_name] = torch.as_strided(
-                typed_storage,
-                size=cache.shape,
-                stride=cache.stride(),
-                storage_offset=cache.storage_offset(),
-            )
-
-
-def _allocate_extensible_kv_cache(
-    kv_cache_config: KVCacheConfig,
-    shared_layers: dict[str, str],
-    device: torch.device,
-    attn_groups: Sequence[AttentionGroup],
-    kernel_block_sizes: list[int],
-    cache_dtype: str,
-    shareable: bool = False,
-) -> tuple[dict[str, torch.Tensor], ExtensibleKVCacheBuffers]:
-    """Reserve virtual address space for the full KV cache capacity but commit
-    only one block per buffer. The returned raw tensors view the full capacity;
-    `ExtensibleKVCacheBuffers.commit` maps physical pages for more blocks.
-    """
-    num_blocks = kv_cache_config.num_blocks
-    if num_blocks <= 0:
-        raise ValueError(
-            "enable_extensible_kv_cache=True requires at least one KV block."
-        )
-
-    num_segments_by_layer = _kv_cache_num_segments_by_layer(
-        attn_groups,
-        kernel_block_sizes,
-        cache_dtype,
-        kv_cache_config.has_mamba_layers,
-    )
-    kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
-    buffers: list[tuple[ExtensibleTensor, int]] = []
-    packed_view: torch.Tensor | None = None
-    for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
-        bytes_per_block = kv_cache_tensor.size // num_blocks
-        assert bytes_per_block * num_blocks == kv_cache_tensor.size
-        if kv_cache_tensor.block_stride > 0:
-            # Packed layout: one backing shared by all layers, with block b
-            # occupying the b-th `block_stride`-byte row (holding every
-            # layer's page). The backing is block-major by construction, so
-            # one shared single-segment buffer commits a prefix of blocks.
-            # One packed row per logical block: _reshape_kv_cache, NIXL's
-            # packed registration, and _bound_packed_kv_cache_storages all
-            # rely on this equality.
-            assert bytes_per_block == kv_cache_tensor.block_stride
-            if packed_view is None:
-                packed_buffer = ExtensibleTensor(
-                    max_num_bytes=kv_cache_tensor.size,
-                    device=device,
-                    num_segments=1,
-                    shareable=shareable,
-                )
-                buffers.append((packed_buffer, bytes_per_block))
-                packed_view = packed_buffer.full_view()
-            tensor = packed_view
-        else:
-            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() if segment_counts else 1
-            assert bytes_per_block % num_segments == 0
-            buffer = ExtensibleTensor(
-                max_num_bytes=kv_cache_tensor.size,
-                device=device,
-                num_segments=num_segments,
-                shareable=shareable,
-            )
-            buffers.append((buffer, bytes_per_block // num_segments))
-            tensor = buffer.full_view()
-        for layer_name in kv_cache_tensor.shared_by:
-            kv_cache_raw_tensors[layer_name] = tensor
-
-    _check_layer_coverage(kv_cache_config, kv_cache_raw_tensors, shared_layers)
-    extensible_buffers = ExtensibleKVCacheBuffers(buffers, num_blocks)
-    extensible_buffers.commit(1)
-    return kv_cache_raw_tensors, extensible_buffers
+    return kv_cache_raw_tensors
 
 
 def _reshape_attention_kv_cache(
@@ -784,27 +526,12 @@ def init_kv_cache(
     cache_dtype: str,
     kernel_block_sizes: list[int],
     vllm_config: VllmConfig,
-    extensible: bool = False,
-) -> tuple[dict[str, Any], ExtensibleKVCacheBuffers | None]:
+) -> dict[str, Any]:
     shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
+    kv_cache_raw_tensors = _allocate_kv_cache(
+        kv_cache_config, shared_kv_cache_layers, device
+    )
     flattened_attn_groups = list(group for groups in attn_groups for group in groups)
-    extensible_buffers = None
-    if extensible:
-        kv_cache_raw_tensors, extensible_buffers = _allocate_extensible_kv_cache(
-            kv_cache_config,
-            shared_kv_cache_layers,
-            device,
-            flattened_attn_groups,
-            kernel_block_sizes,
-            cache_dtype,
-            # KV connectors export this memory for cross-process access
-            # (CUDA IPC intra-node, GPU-direct RDMA across nodes).
-            shareable=vllm_config.kv_transfer_config is not None,
-        )
-    else:
-        kv_cache_raw_tensors = _allocate_kv_cache(
-            kv_cache_config, shared_kv_cache_layers, device
-        )
     kv_caches = _reshape_kv_cache(
         attn_groups=flattened_attn_groups,
         kv_cache_raw_tensors=kv_cache_raw_tensors,
@@ -822,7 +549,7 @@ def init_kv_cache(
         else 1
     )
     bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module)
-    return kv_caches, extensible_buffers
+    return kv_caches
 
 
 def build_slot_mappings_by_layer(

vllm/v1/worker/gpu/model_runner.py

@@ -45,7 +45,6 @@ from vllm.model_executor.model_loader import get_model_loader
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
-from vllm.utils.extensible_tensor import ExtensibleKVCacheBuffers
 from vllm.utils.math_utils import cdiv
 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE
@@ -60,7 +59,6 @@ from vllm.v1.worker.gpu.attn_utils import (
     get_kv_cache_spec,
     init_attn_backend,
     init_kv_cache,
-    narrow_kv_caches_to_num_blocks,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.buffer_utils import (
@@ -253,7 +251,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         # KV Connector if configured.
         self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
-        self.extensible_kv_buffers: ExtensibleKVCacheBuffers | None = None
 
         # For transferring state from execute_model to subsequent sample_tokens call.
         self.execute_model_state: ExecuteModelState | None = None
@@ -411,12 +408,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
     def get_kv_cache_spec(self):
         return get_kv_cache_spec(self.vllm_config)
 
-    def initialize_kv_cache(
-        self, kv_cache_config: KVCacheConfig, extensible: bool = False
-    ) -> None:
-        if self.extensible_kv_buffers is not None:
-            self.extensible_kv_buffers.free()
-        self.extensible_kv_buffers = None
+    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         kv_cache_config = deepcopy(kv_cache_config)
         self.kv_cache_config = kv_cache_config
 
@@ -509,7 +501,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.speculator.init_cudagraph_manager(cudagraph_mode)
 
         self.kv_caches: list[torch.Tensor] = []
-        kv_caches_dict, self.extensible_kv_buffers = init_kv_cache(
+        kv_caches_dict = init_kv_cache(
             self.kv_caches,
             self.compilation_config.static_forward_context,
             self.kv_cache_config,
@@ -518,67 +510,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.cache_config.cache_dtype,
             self.kernel_block_sizes,
             self.vllm_config,
-            extensible=extensible,
         )
-        self._kv_caches_dict = kv_caches_dict
-        # With the extensible flow, KV transfer init is deferred until the
-        # final KV cache size is committed, so this yields a no-op connector
-        # that init_deferred_kv_connector later replaces.
         self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
 
-    def init_deferred_kv_connector(self) -> None:
-        """Create and register the KV connector after `extend_kv_cache`.
-
-        With the extensible KV cache, connectors must not register the cache
-        before its final size is physically committed. Registration views are
-        narrowed to the committed block count so connectors only see (and
-        register with e.g. RDMA) backed memory.
-        """
-        assert self.extensible_kv_buffers is not None
-        kv_caches = narrow_kv_caches_to_num_blocks(
-            self._kv_caches_dict,
-            [g for groups in self.attn_groups for g in groups],
-            self.kernel_block_sizes,
-            self.cache_config.cache_dtype,
-            self.kv_cache_config.num_blocks,
-            self.kv_cache_config,
-        )
-        self.kv_connector = get_kv_connector(self.vllm_config, kv_caches)
-
-    def ensure_kv_cache_blocks(self, num_blocks: int) -> None:
-        """Commit at least `num_blocks` KV blocks when the extensible KV cache
-        is in use (no-op otherwise). Warmup paths call this before executing
-        batches that write to a prefix of real block IDs.
-        """
-        if self.extensible_kv_buffers is not None:
-            self.extensible_kv_buffers.commit(
-                min(num_blocks, self.extensible_kv_buffers.num_blocks_capacity)
-            )
-
-    def extend_kv_cache(self, num_blocks: int, defragment: bool = False) -> None:
-        """Commit physical pages so the KV cache holds `num_blocks` blocks.
-
-        Grows the KV cache after warmup and CUDA graph capture, once the
-        actual available memory is known. No re-view is needed: the layers
-        already view the full reserved 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. `defragment` discards the warmup-time
-        commits so each segment is backed by a single physical allocation
-        (required before KV-transfer registration).
-        """
-        if self.extensible_kv_buffers is None:
-            raise RuntimeError("extend_kv_cache requires an extensible KV cache.")
-        self.extensible_kv_buffers.commit(num_blocks, defragment=defragment)
-        self.kv_cache_config.num_blocks = num_blocks
-        logger.info("Extended KV cache to %d blocks.", num_blocks)
-
-    @property
-    def kv_cache_committed_bytes(self) -> int:
-        """Physically committed KV cache bytes (0 without extensible KV)."""
-        buffers = getattr(self, "extensible_kv_buffers", None)
-        return buffers.physical_bytes if buffers is not None else 0
-
     def _init_kv_zero_meta(self) -> None:
         """Build KV-block zeroing metadata; invoked from gpu_worker."""
         self.kv_block_zeroer = KVBlockZeroer(
@@ -1662,9 +1596,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.attn_groups.clear()
         if hasattr(self, "kv_cache_config"):
             del self.kv_cache_config
-        if self.extensible_kv_buffers is not None:
-            self.extensible_kv_buffers.free()
-            self.extensible_kv_buffers = None
         free_before_shutdown(self.vllm_config)
         if hasattr(self, "model_state"):
             del self.model_state

vllm/v1/worker/gpu/warmup.py

@@ -73,10 +73,6 @@ def run_mixed_prefill_decode_warmup(
         )
         return False
 
-    # With the extensible KV cache, only a prefix of the blocks is physically
-    # committed so far; commit the prefix this warmup writes to.
-    model_runner.ensure_kv_cache_blocks(1 + required_blocks)
-
     next_block_id = 1
 
     def _alloc_blocks(num_blocks: int) -> list[int]:
@@ -225,10 +221,6 @@ def warmup_kernels(
         max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req),
     )
 
-    # With the extensible KV cache, only a prefix of the blocks is physically
-    # committed so far; commit the prefix this warmup writes to.
-    model_runner.ensure_kv_cache_blocks(1 + num_reqs * max_blocks_per_req)
-
     req_ids = [f"_warmup_{i}_" for i in range(num_reqs)]
 
     # SamplingParams exercising all sampling features.

vllm/v1/worker/gpu_model_runner.py

@@ -4,7 +4,6 @@
 import functools
 import gc
 import itertools
-import math
 import threading
 import time
 from collections import defaultdict
@@ -119,7 +118,6 @@ from vllm.sequence import IntermediateTensors
 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 ExtensibleKVCacheBuffers, 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
@@ -151,7 +149,6 @@ from vllm.v1.attention.backends.utils import (
     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 (
@@ -563,7 +560,6 @@ class GPUModelRunner(
         # self.model: nn.Module  # Set after load_model
         # Initialize in initialize_kv_cache
         self.kv_caches: list[torch.Tensor] = []
-        self.extensible_kv_buffers: ExtensibleKVCacheBuffers | None = None
         # Initialize in initialize_kv_cache_tensors
         self.cross_layers_kv_cache: torch.Tensor | None = None
         self.cross_layers_attn_backend: type[AttentionBackend] | None = None
@@ -6366,7 +6362,6 @@ class GPUModelRunner(
         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
@@ -6378,12 +6373,6 @@ class GPUModelRunner(
             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:
@@ -6443,7 +6432,7 @@ class GPUModelRunner(
         else:
             output = None
         self._sync_device()
-        del hidden_states, output, dummy_encoder_cache
+        del hidden_states, output
         self.encoder_cache.clear()
         gc.collect()
 
@@ -6526,9 +6515,6 @@ class GPUModelRunner(
             self.attn_groups.clear()
         if hasattr(self, "kv_cache_config"):
             delattr(self, "kv_cache_config")
-        if self.extensible_kv_buffers is not None:
-            self.extensible_kv_buffers.free()
-            self.extensible_kv_buffers = None
         self.cache_config.num_gpu_blocks = None
 
         for layer in self.compilation_config.static_forward_context.values():
@@ -7250,76 +7236,19 @@ class GPUModelRunner(
         )
 
     def _allocate_kv_cache_tensors(
-        self, kv_cache_config: KVCacheConfig, extensible: bool = False
+        self, kv_cache_config: KVCacheConfig
     ) -> 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; 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.
+            kv_cache_config: The KV cache config
         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()
-            buffers: list[tuple[ExtensibleTensor, int]] = []
-            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
-                buffer = ExtensibleTensor(
-                    max_num_bytes=kv_cache_tensor.size,
-                    device=self.device,
-                    num_segments=num_segments,
-                )
-                buffers.append((buffer, bytes_per_block // num_segments))
-                tensor = buffer.full_view()
-                for layer_name in kv_cache_tensor.shared_by:
-                    kv_cache_raw_tensors[layer_name] = tensor
-
-            self.extensible_kv_buffers = ExtensibleKVCacheBuffers(
-                buffers, kv_cache_config.num_blocks
-            )
-            self.extensible_kv_buffers.commit(1)
-            return self._check_kv_cache_raw_tensors(
-                kv_cache_config, kv_cache_raw_tensors
-            )
-
-        self.extensible_kv_buffers = None
         packed_backing: torch.Tensor | None = None
         for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
             if kv_cache_tensor.block_stride > 0:
@@ -7338,13 +7267,6 @@ class GPUModelRunner(
             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:
@@ -7356,49 +7278,6 @@ class GPUModelRunner(
         )
         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)
 
@@ -7597,31 +7476,8 @@ class GPUModelRunner(
                     stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]),
                 )
 
-    def extend_kv_cache(self, num_blocks: int, defragment: bool = False) -> 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 self.extensible_kv_buffers is None:
-            raise RuntimeError("extend_kv_cache requires an extensible KV cache.")
-        self.extensible_kv_buffers.commit(num_blocks, defragment=defragment)
-        logger.info("Extended KV cache to %d blocks.", num_blocks)
-
-    @property
-    def kv_cache_committed_bytes(self) -> int:
-        """Physically committed KV cache bytes (0 without extensible KV)."""
-        buffers = getattr(self, "extensible_kv_buffers", None)
-        return buffers.physical_bytes if buffers is not None else 0
-
     def initialize_kv_cache_tensors(
-        self,
-        kv_cache_config: KVCacheConfig,
-        kernel_block_sizes: list[int],
-        extensible: bool = False,
+        self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
     ) -> dict[str, torch.Tensor]:
         """
         Initialize the memory buffer for KV cache.
@@ -7637,13 +7493,7 @@ class GPUModelRunner(
 
         # Try creating KV caches optimized for kv-connector transfers
         cache_dtype = self.cache_config.cache_dtype
-        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):
+        if self.use_uniform_kv_cache(self.attn_groups):
             kv_caches, cross_layers_kv_cache, attn_backend = (
                 self.allocate_uniform_kv_caches(
                     kv_cache_config,
@@ -7658,10 +7508,7 @@ class GPUModelRunner(
         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,
-                extensible=extensible,
-            )
+            kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
 
             # Change the memory buffer to the desired shape
             kv_caches = self._reshape_kv_cache_tensors(
@@ -7716,7 +7563,6 @@ class GPUModelRunner(
         self,
         kv_cache_config: KVCacheConfig,
         is_profiling: bool = False,
-        extensible: bool = False,
     ) -> None:
         """
         Initialize KV cache based on `kv_cache_config`.
@@ -7749,9 +7595,7 @@ class GPUModelRunner(
         # 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,
-            extensible=extensible,
+            kv_cache_config, kernel_block_sizes
         )
 
         if (

vllm/v1/worker/gpu_worker.py

@@ -175,10 +175,6 @@ class Worker(WorkerBase):
         # pending non-blocking PP send work from the previous iteration
         self._pp_send_work: list[Handle] = []
 
-        # Set by initialize_from_config when the extensible KV cache defers
-        # KV transfer init until the final cache size is committed.
-        self._deferred_kv_transfer_init = False
-
         # Resolved lazily on first sleep/wake; persists worker-process state.
         self._sleep_mode_backend: SleepModeBackend | None = None
 
@@ -194,19 +190,6 @@ class Worker(WorkerBase):
         return self._sleep_mode_backend
 
     def sleep(self, level: int = 1) -> None:
-        extensible_kv_buffers = getattr(
-            self.model_runner, "extensible_kv_buffers", None
-        )
-        if (
-            extensible_kv_buffers is not None
-            and self.vllm_config.kv_transfer_config is not None
-        ):
-            raise RuntimeError(
-                "Sleep mode with an extensible KV cache and a KV connector is "
-                "not supported: waking remaps physical pages and invalidates "
-                "the connector's memory registration."
-            )
-
         torch.accelerator.synchronize()
         free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]
 
@@ -224,11 +207,6 @@ class Worker(WorkerBase):
 
         self._get_sleep_mode_backend().suspend(level)
 
-        # The extensible KV cache lives outside the torch/CuMem allocators;
-        # discard its physical memory directly (VA and views stay valid).
-        if extensible_kv_buffers is not None:
-            extensible_kv_buffers.release_physical()
-
         torch.accelerator.synchronize()
         deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
         while True:
@@ -266,11 +244,6 @@ class Worker(WorkerBase):
             self._sleep_rebuild_draft_metadata_buffers = False
 
         if tags is None or "kv_cache" in tags:
-            extensible_kv_buffers = getattr(
-                self.model_runner, "extensible_kv_buffers", None
-            )
-            if extensible_kv_buffers is not None:
-                extensible_kv_buffers.recommit()
             self.model_runner.post_kv_cache_wake_up()
 
     def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
@@ -530,7 +503,6 @@ class Worker(WorkerBase):
                 current_platform.is_cuda_alike()
                 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()
 
@@ -743,11 +715,7 @@ class Worker(WorkerBase):
         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,
-        extensible: bool = False,
-    ) -> None:
+    def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
         """Allocate GPU KV cache with the specified kv_cache_config."""
 
         # Update local config with adjusted num blocks after profiling,
@@ -759,20 +727,10 @@ class Worker(WorkerBase):
         # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
         # because `initialize_kv_cache` will inject kv cache groups not
         # related to kv cache connector (e.g. kv cache sharing layers).
-        # With the extensible KV cache, connectors must not register the KV
-        # cache memory before its final size is committed, so KV transfer
-        # init is deferred to `extend_kv_cache` (which receives the final,
-        # pristine kv_cache_config).
-        self._deferred_kv_transfer_init = (
-            extensible and self.vllm_config.kv_transfer_config is not None
-        )
-        if not self._deferred_kv_transfer_init:
-            ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
+        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, extensible=extensible
-            )
+            self.model_runner.initialize_kv_cache(kv_cache_config)
 
         if self.model_config.enable_return_routed_experts:
             self.model_runner.init_routed_experts_capturer()
@@ -785,27 +743,6 @@ class Worker(WorkerBase):
         ):
             self.model_runner._init_kv_zero_meta()
 
-    def extend_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
-        """Commit the final KV cache size after warmup (extensible flow)."""
-        num_blocks = kv_cache_config.num_blocks
-        self.cache_config.num_gpu_blocks = num_blocks
-        # Defragment when a connector will register the memory: UCX cannot
-        # transfer regions spanning multiple VMM allocation handles.
-        self.model_runner.extend_kv_cache(
-            num_blocks, defragment=self._deferred_kv_transfer_init
-        )
-        if self._deferred_kv_transfer_init:
-            # The final size is committed; now the connector may register the
-            # (physically backed) KV cache memory.
-            ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
-            assert hasattr(self.model_runner, "init_deferred_kv_connector")
-            self.model_runner.init_deferred_kv_connector()
-
-    def extensible_kv_cache_unsupported_reason(self) -> str | None:
-        from vllm.utils.vmm_driver import vmm_unavailable_reason
-
-        return vmm_unavailable_reason()
-
     @instrument(span_name="Warmup (GPU)")
     def compile_or_warm_up_model(self) -> CompilationTimes:
         warmup_sizes: list[int] = []
@@ -947,31 +884,6 @@ class Worker(WorkerBase):
             else:
                 self.model_runner._dummy_sampler_run(hidden_states=last_hidden_states)
 
-        warmup_memory_bytes = cuda_graph_memory_bytes
-        if self.cache_config.enable_extensible_kv_cache and hasattr(
-            self, "available_kv_cache_memory_bytes"
-        ):
-            # With the extensible KV cache, only a small prefix of the KV cache
-            # is committed so far, so the current memory usage reflects
-            # everything else at its post-warmup state: weights, CUDA graphs,
-            # NCCL buffers, and the allocator segments retained from the
-            # worst-case warmup batches (which can far exceed the profiled
-            # activation peak, e.g. with speculative decoding). Report the
-            # measured excess over the profiling estimate so the final KV cache
-            # size is computed from actual usage.
-            torch.accelerator.synchronize()
-            free_memory, _ = torch.accelerator.get_memory_info()
-            non_kv_used_memory = (
-                self.init_snapshot.free_memory
-                - free_memory
-                - self.model_runner.kv_cache_committed_bytes
-            )
-            post_warmup_available = int(self.requested_memory) - non_kv_used_memory
-            warmup_memory_bytes = max(
-                cuda_graph_memory_bytes,
-                int(self.available_kv_cache_memory_bytes) - post_warmup_available,
-            )
-
         # Reset the seed to ensure that the random state is not affected by
         # the model initialization and profiling.
         set_random_seed(self.model_config.seed)
@@ -1007,7 +919,6 @@ class Worker(WorkerBase):
         return CompilationTimes(
             language_model=self.compilation_config.compilation_time,
             encoder=self.compilation_config.encoder_compilation_time,
-            warmup_memory=warmup_memory_bytes,
         )
 
     def reset_mm_cache(self) -> None:

vllm/v1/worker/worker_base.py

@@ -34,10 +34,6 @@ _R = TypeVar("_R")
 class CompilationTimes(NamedTuple):
     language_model: float
     encoder: float
-    # GPU memory (bytes) consumed by warmup and CUDA graph capture beyond the
-    # profiled baseline; used by the extensible KV cache flow to compute the
-    # final KV cache size from actual usage.
-    warmup_memory: int = 0
 
 
 class WorkerBase:
@@ -103,20 +99,11 @@ class WorkerBase:
         """Get specifications for KV cache implementation."""
         raise NotImplementedError
 
-    def extend_kv_cache(self, kv_cache_config: Any) -> None:
-        raise RuntimeError(
-            f"{self.__class__.__name__} does not support extensible KV cache."
-        )
-
-    def extensible_kv_cache_unsupported_reason(self) -> str | None:
-        """Return why this worker cannot use the extensible KV cache, or None."""
-        return f"not supported by {self.__class__.__name__}"
-
     def compile_or_warm_up_model(self) -> CompilationTimes:
         """Prepare model for execution through compilation/warmup.
 
         Returns:
-            Compilation times in seconds and warmup memory in bytes.
+            Compilation times (language_model, encoder) in seconds.
         """
         raise NotImplementedError
 
@@ -331,26 +318,11 @@ class WorkerWrapperBase:
             # To make vLLM config available during worker initialization
             self.worker = worker_class(**kwargs)
 
-    def initialize_from_config(
-        self,
-        kv_cache_configs: list[Any],
-        extensible: bool = False,
-    ) -> None:
+    def initialize_from_config(self, kv_cache_configs: list[Any]) -> 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):
-            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 extend_kv_cache(self, kv_cache_configs: list[Any]) -> 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.extend_kv_cache(kv_cache_config)  # type: ignore
+            self.worker.initialize_from_config(kv_cache_config)  # type: ignore
 
     def init_device(self):
         assert self.vllm_config is not None