remove torch 2.9, 2.10 workround

Signed-off-by: yewentao256 <zhyanwentao@126.com>
[CI/Build] Skip Prithvi/Terratorch model-registry tests when terratorch is missing (#41389 )
2026-04-30 21:10:50 +00:00 · 2026-04-30 12:47:55 -07:00 · 2026-04-30 12:24:09 -07:00 · 2026-04-30 09:57:27 -07:00 · 2026-04-30 16:02:23 +00:00 · 2026-04-30 07:59:07 -07:00
79 changed files with 3182 additions and 1951 deletions
@@ -89,6 +89,16 @@ steps:
  commands:
    - pytest -v -s kernels/mamba

+- label: Kernels KDA Test
+  timeout_in_minutes: 20
+  source_file_dependencies:
+  - vllm/model_executor/layers/fla/ops/kda.py
+  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
+  - vllm/model_executor/layers/fla/ops/l2norm.py
+  - tests/kernels/test_kda.py
+  commands:
+    - pytest -v -s kernels/test_kda.py
+
 - label: Kernels DeepGEMM Test (H100)
  key: kernels-deepgemm-test-h100
  timeout_in_minutes: 45
@@ -97,3 +97,16 @@ steps:
  commands:
    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
+
+- label: Spec Decode MTP hybrid (B200)
+  timeout_in_minutes: 30
+  device: b200
+  optional: true
+  source_file_dependencies:
+    - vllm/v1/spec_decode/
+    - vllm/v1/worker/gpu/spec_decode/
+    - vllm/model_executor/models/qwen3_5.py
+    - vllm/model_executor/models/qwen3_5_mtp.py
+    - tests/v1/e2e/spec_decode/
+  commands:
+    - pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
@@ -16,11 +16,7 @@ permissions:

 jobs:
  pre-run-check:
-    if: >-
-      github.event_name == 'pull_request' &&
-      (github.event.action != 'labeled' ||
-       github.event.label.name == 'ready' ||
-       github.event.label.name == 'verified')
+    if: github.event_name == 'pull_request'
    runs-on: ubuntu-latest
    steps:
    - name: Check PR label and author merge count
@@ -49,12 +45,7 @@ jobs:

  pre-commit:
    needs: pre-run-check
-    if: >-
-      always() &&
-      (github.event.action != 'labeled' ||
-       github.event.label.name == 'ready' ||
-       github.event.label.name == 'verified') &&
-      (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
+    if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
@@ -217,6 +217,7 @@ async def send_request(
    min_tokens: int | None = None,
    max_tokens: int | None = None,
    timeout_sec: int = 120,
+    conversation_id: str | None = None,
 ) -> ServerResponse:
    payload = {
        "model": model,
@@ -225,6 +226,9 @@ async def send_request(
        "temperature": 0.0,
    }

+    if conversation_id is not None:
+        payload["conversation_id"] = conversation_id
+
    if stream:
        payload["stream"] = True
        payload["stream_options"] = {"include_usage": False}
@@ -419,6 +423,7 @@ async def send_turn(
        min_tokens,
        max_tokens,
        req_args.timeout_sec,
+        conversation_id=conv_id,
    )

    if response.valid is False:
@@ -82,18 +82,73 @@ void launch_persistent_topk(const torch::Tensor& logits,
    size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
    if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;

+    // Query occupancy for the instantiation that will actually launch;
+    // overestimating it deadlocks the cooperative barrier.
    int occupancy = 1;
-    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-        &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
-        smem_size);
+    cudaError_t occ_err = cudaSuccess;
+    if (vec_size == 4) {
+      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+          &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
+          smem_size);
+    } else if (vec_size == 2) {
+      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+          &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
+          smem_size);
+    } else {
+      occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+          &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
+          smem_size);
+    }
+    TORCH_CHECK(occ_err == cudaSuccess,
+                "persistent_topk occupancy query failed: ",
+                cudaGetErrorString(occ_err));
    if (occupancy < 1) occupancy = 1;

-    uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
+    // The cooperative spin-wait barrier only runs when at least one row hits
+    // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
+    // early-exit, so oversubscription can't deadlock and headroom is wasted.
+    const bool needs_cooperative =
+        static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
+
+    const uint32_t hw_resident_cap =
+        static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
+    uint32_t max_resident_ctas = hw_resident_cap;
+    if (needs_cooperative) {
+      // Reserve one CTA per SM when occupancy allows; fall back to a single
+      // CTA when occupancy == 1 (the most deadlock-prone case — any straggler
+      // kernel that takes the only slot on one SM hangs the barrier). Never
+      // drop below one full group's worth.
+      uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
+      if (max_resident_ctas >= headroom + ctas_per_group) {
+        max_resident_ctas -= headroom;
+      }
+    }
    uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
                                   static_cast<uint32_t>(num_rows));
    if (num_groups == 0) num_groups = 1;
    uint32_t total_ctas = num_groups * ctas_per_group;

+    // If the cooperative launch wouldn't fit, fall back to FilteredTopK
+    // instead of deadlocking. Only relevant when needs_cooperative.
+    if (needs_cooperative && total_ctas > hw_resident_cap) {
+      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
+                  "persistent_topk would oversubscribe and the FilteredTopK "
+                  "fallback requires >=128KB smem per block (have ",
+                  max_smem_per_block, "). total_ctas=", total_ctas,
+                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
+                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
+                  ", smem=", smem_size, ").");
+      cudaError_t status =
+          vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
+              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
+              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
+              stream);
+      TORCH_CHECK(status == cudaSuccess,
+                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
+      return;
+    }
+
    size_t state_bytes = num_groups * sizeof(P::RadixRowState);
    TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
                "workspace too small, need ", state_bytes, " bytes");
@@ -200,9 +200,9 @@ RUN cd /opt/rixl && \

 # DeepEP build stage
 FROM base AS build_deep
-ARG ROCSHMEM_BRANCH="ba0bf0f3"
+ARG ROCSHMEM_BRANCH="f0acb0c6"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-ARG DEEPEP_BRANCH="5d90af8b"
+ARG DEEPEP_BRANCH="a9ea9774"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
 ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
@@ -213,18 +213,15 @@ RUN git clone ${ROCSHMEM_REPO} \
 && git checkout ${ROCSHMEM_BRANCH} \
 && mkdir -p projects/rocshmem/build \
 && cd projects/rocshmem/build \
- && bash ../scripts/build_configs/all_backends \
-      -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-      -DROCM_PATH=/opt/rocm \
-      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
-      -DUSE_EXTERNAL_MPI=OFF
+ && INSTALL_PREFIX=${ROCSHMEM_DIR} \
+    ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF

 # Build DeepEP wheel.
 # DeepEP looks for rocshmem at ROCSHMEM_DIR.
 RUN git clone ${DEEPEP_REPO} \
 && cd DeepEP \
 && git checkout ${DEEPEP_BRANCH} \
- && python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
+ && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install

 # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
 # not force users to rebuild the long-lived Dockerfile.rocm_base image.
@@ -388,6 +385,16 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
 # above are not available once that RUN step completes.
 COPY --from=export_vllm /*.whl /opt/vllm-wheels/

+# Update rdma-core to support latest rocshmem
+ARG DEEPEP_NIC
+RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
+    git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
+    cd /tmp/rdma-core && \
+    mkdir -p build && cd build && \
+    cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
+    ninja &&     ninja install &&     ldconfig &&     rm -rf /tmp/rdma-core; \
+fi
+
 # Install RIXL wheel
 RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
    uv pip install --system /rixl_install/*.whl
@@ -5,9 +5,6 @@ WORKDIR /workspace/
 ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"

-RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
-    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
-
 RUN apt clean && apt-get update -y && \
    apt-get install -y --no-install-recommends --fix-missing \
    curl \
@@ -26,8 +23,20 @@ RUN apt clean && apt-get update -y && \
    python3.12-dev \
    python3-pip

-RUN apt update && apt upgrade -y && \
-    apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
+# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
+RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
+    printf '%s\n' \
+    'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
+    'Pin: version 2025.3*' \
+    'Pin-Priority: 1001' \
+    > /etc/apt/preferences.d/oneapi-2025.3.pref && \
+    apt-get update -y && \
+    apt-get install -y --no-install-recommends \
+    intel-oneapi-compiler-dpcpp-cpp-2025.3 \
+    intel-oneapi-mkl-devel-2025.3 \
+    intel-oneapi-dnnl-devel-2025.3 && \
+    rm -rf /var/lib/apt/lists/*

 # Install UMD
 RUN mkdir neo && \
@@ -38,8 +38,20 @@ class MockCustomOp:
        return decorator


+class MockPluggableLayer:
+    @staticmethod
+    def register(name):
+        def decorator(cls):
+            return cls
+
+        return decorator
+
+
 mock_if_no_torch("vllm._C", MagicMock())
-mock_if_no_torch("vllm.model_executor.custom_op", MagicMock(CustomOp=MockCustomOp))
+mock_if_no_torch(
+    "vllm.model_executor.custom_op",
+    MagicMock(CustomOp=MockCustomOp, PluggableLayer=MockPluggableLayer),
+)
 mock_if_no_torch(
    "vllm.utils.torch_utils", MagicMock(direct_register_custom_op=lambda *a, **k: None)
 )
@@ -0,0 +1,562 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer
+
+This proxy sits between clients and a vLLM Prefill/Decode (P/D) deployment,
+routing multi-turn chat requests so that each turn reuses KV cache blocks
+from the previous turn's Decode node via bidirectional KV transfer.
+
+Architecture:
+    Client  ──►  Proxy  ──►  Prefill (P)  ──►  Decode (D)
+                   │              │                 │
+                   │   kv_transfer_params flow:     │
+                   │   D finish ──► proxy caches    │
+                   │   next turn ──► proxy sends    │
+                   │   cached D blocks to P ──►     │
+                   │   P reads D blocks (bidir)     │
+                   │   P sends its blocks to D      │
+
+Per-request flow:
+    1. Client sends chat/completions request to proxy.
+    2. Proxy looks up cached D block info from the previous turn
+       (keyed by conversation_id).
+    3. If cache hit, proxy attaches D's block info to the request
+       so P can read D's KV blocks instead of recomputing.
+    4. Proxy sends request to P (max_tokens=1, non-streaming).
+    5. P returns kv_transfer_params with its own block info.
+    6. Proxy forwards request + P's block info to D (streaming).
+    7. D streams the response. The final chunk includes D's
+       kv_transfer_params, which the proxy caches for the next turn.
+    8. Proxy returns D's response to the client.
+
+Conversation isolation:
+    Each request must include a ``conversation_id`` field (top-level in
+    the JSON body) to scope the KV cache across turns. Without it, the
+    proxy cannot link turns and falls back to no-cache behavior.
+
+Usage:
+    python disagg_proxy_multiturn.py \\
+        --host 0.0.0.0 --port 8000 \\
+        --prefiller-host 10.0.0.1 --prefiller-port 8100 \\
+        --decoder-host 10.0.0.2 --decoder-port 8200
+
+Dependencies:
+    pip install fastapi uvicorn httpx
+"""
+
+from __future__ import annotations
+
+import argparse
+import itertools
+import json
+import logging
+import os
+import time
+import uuid
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import Any
+
+import httpx
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+# Logging
+logging.basicConfig(
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=logging.INFO,
+)
+logger = logging.getLogger("disagg_proxy")
+
+
+# Data structures
+@dataclass
+class CachedKVEntry:
+    """KV transfer parameters cached from D's response for one turn."""
+
+    kv_transfer_params: dict[str, Any]
+    timestamp: float = field(default_factory=time.time)
+
+
+class ConversationKVCache:
+    """Per-conversation KV block cache.
+
+    Each conversation is identified by a ``conversation_id`` supplied by
+    the client. After D finishes a turn, its ``kv_transfer_params`` are
+    stored here. On the next turn, the proxy retrieves them so P can
+    read D's blocks via bidirectional KV transfer.
+    """
+
+    def __init__(self, ttl_seconds: float = 600.0) -> None:
+        self._store: dict[str, CachedKVEntry] = {}
+        self._ttl = ttl_seconds
+
+    def get(self, conversation_id: str) -> dict[str, Any] | None:
+        """Retrieve and consume cached KV params for a conversation.
+
+        Returns a *copy* of the kv_transfer_params dict, or None.
+        The entry is removed after retrieval (single-use).
+        """
+        entry = self._store.pop(conversation_id, None)
+        if entry is None:
+            return None
+        age = time.time() - entry.timestamp
+        if age > self._ttl:
+            logger.info(
+                "conv=%s: stale cache entry (age=%.1fs > ttl=%.1fs), discarding",
+                conversation_id,
+                age,
+                self._ttl,
+            )
+            return None
+        logger.info(
+            "conv=%s: cache HIT (age=%.1fs)",
+            conversation_id,
+            age,
+        )
+        return dict(entry.kv_transfer_params)
+
+    def put(self, conversation_id: str, kv_params: dict[str, Any]) -> None:
+        """Store D's kv_transfer_params for a conversation."""
+        self._store[conversation_id] = CachedKVEntry(
+            kv_transfer_params=dict(kv_params),  # defensive copy
+        )
+        logger.info(
+            "conv=%s: cached D blocks (remote_request_id=%s, blocks=%d)",
+            conversation_id,
+            kv_params.get("remote_request_id", "?"),
+            len(kv_params.get("remote_block_ids", [[]])[0])
+            if kv_params.get("remote_block_ids")
+            else 0,
+        )
+
+    def evict_stale(self) -> int:
+        """Remove entries older than TTL. Returns count of evicted entries."""
+        now = time.time()
+        stale = [
+            cid
+            for cid, entry in self._store.items()
+            if now - entry.timestamp > self._ttl
+        ]
+        for cid in stale:
+            del self._store[cid]
+        return len(stale)
+
+    @property
+    def size(self) -> int:
+        return len(self._store)
+
+
+# Global state
+kv_cache = ConversationKVCache(
+    ttl_seconds=450.0
+)  # Must be < VLLM_NIXL_ABORT_REQUEST_TIMEOUT (480s)
+
+
+# Service client helpers
+@dataclass
+class ServiceClient:
+    """Wrapper around an httpx.AsyncClient for a P or D instance."""
+
+    client: httpx.AsyncClient
+    host: str
+    port: int
+    id: int
+
+
+def _make_headers(request_id: str) -> dict[str, str]:
+    """Build HTTP headers for upstream requests."""
+    headers = {"X-Request-Id": request_id}
+    api_key = os.environ.get("OPENAI_API_KEY")
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+    return headers
+
+
+async def _send_to_prefill(
+    client: ServiceClient,
+    endpoint: str,
+    req_data: dict[str, Any],
+    request_id: str,
+) -> dict[str, Any]:
+    """Send a non-streaming prefill request (max_tokens=1).
+
+    Returns the JSON response from P, which includes kv_transfer_params.
+    """
+    payload = req_data.copy()
+    payload["stream"] = False
+    payload["max_tokens"] = 1
+    payload.pop("max_completion_tokens", None)
+    payload.pop("min_tokens", None)
+    payload.pop("stream_options", None)
+
+    resp = await client.client.post(
+        endpoint,
+        json=payload,
+        headers=_make_headers(request_id),
+    )
+    resp.raise_for_status()
+    return resp.json()
+
+
+async def _stream_from_decode(
+    client: ServiceClient,
+    endpoint: str,
+    req_data: dict[str, Any],
+    request_id: str,
+    conversation_id: str,
+) -> tuple[str, str | None, dict[str, Any] | None, str, str | None, int | None]:
+    """Stream response from D, capturing text and kv_transfer_params.
+
+    Returns (collected_text, finish_reason, kv_params, response_id, created).
+    Also stores kv_params in the conversation cache.
+    """
+    payload = req_data.copy()
+    payload["stream"] = True
+
+    collected_text = ""
+    finish_reason: str | None = None
+    response_id: str | None = None
+    model_name: str | None = None
+    created: int | None = None
+    captured_kv: dict[str, Any] | None = None
+
+    async with client.client.stream(
+        "POST",
+        endpoint,
+        json=payload,
+        headers=_make_headers(request_id),
+    ) as resp:
+        resp.raise_for_status()
+        async for line in resp.aiter_lines():
+            if not line or not line.startswith("data: "):
+                continue
+            if line == "data: [DONE]":
+                break
+            try:
+                chunk = json.loads(line[6:])
+            except json.JSONDecodeError:
+                continue
+
+            if response_id is None:
+                response_id = chunk.get("id")
+                model_name = chunk.get("model")
+                created = chunk.get("created")
+
+            for choice in chunk.get("choices", []):
+                collected_text += choice.get("text", "")
+                delta = choice.get("delta", {})
+                collected_text += delta.get("content", "")
+                if choice.get("finish_reason"):
+                    finish_reason = choice["finish_reason"]
+
+            kv_params = chunk.get("kv_transfer_params")
+            if kv_params:
+                kv_params["remote_host"] = client.host
+                captured_kv = kv_params
+                if conversation_id:
+                    kv_cache.put(conversation_id, kv_params)
+
+    return (
+        collected_text,
+        finish_reason,
+        captured_kv,
+        response_id or request_id,
+        model_name,
+        created,
+    )
+
+
+async def _stream_from_decode_sse(
+    client: ServiceClient,
+    endpoint: str,
+    req_data: dict[str, Any],
+    request_id: str,
+    conversation_id: str,
+):
+    """Yield SSE chunks from D to the client, capturing kv_transfer_params."""
+    payload = req_data.copy()
+    payload["stream"] = True
+
+    async with client.client.stream(
+        "POST",
+        endpoint,
+        json=payload,
+        headers=_make_headers(request_id),
+    ) as resp:
+        resp.raise_for_status()
+        async for line in resp.aiter_lines():
+            if not line:
+                yield "\n"
+                continue
+
+            if line.startswith("data: ") and line != "data: [DONE]":
+                try:
+                    chunk = json.loads(line[6:])
+                    kv_params = chunk.get("kv_transfer_params")
+                    if kv_params and conversation_id:
+                        kv_params["remote_host"] = client.host
+                        kv_cache.put(conversation_id, kv_params)
+                except json.JSONDecodeError:
+                    pass
+
+            yield line + "\n"
+
+
+# FastAPI application
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Initialize HTTP clients for P and D instances."""
+    app.state.prefill_clients: list[ServiceClient] = []
+    app.state.decode_clients: list[ServiceClient] = []
+
+    for i, (host, port) in enumerate(global_args.prefiller_instances):
+        app.state.prefill_clients.append(
+            ServiceClient(
+                client=httpx.AsyncClient(
+                    timeout=None,
+                    base_url=f"http://{host}:{port}/v1",
+                ),
+                host=host,
+                port=port,
+                id=i,
+            )
+        )
+
+    for i, (host, port) in enumerate(global_args.decoder_instances):
+        app.state.decode_clients.append(
+            ServiceClient(
+                client=httpx.AsyncClient(
+                    timeout=None,
+                    base_url=f"http://{host}:{port}/v1",
+                ),
+                host=host,
+                port=port,
+                id=i,
+            )
+        )
+
+    app.state.prefill_iter = itertools.cycle(range(len(app.state.prefill_clients)))
+    app.state.decode_iter = itertools.cycle(range(len(app.state.decode_clients)))
+
+    logger.info(
+        "Ready: %d prefill, %d decode instances",
+        len(app.state.prefill_clients),
+        len(app.state.decode_clients),
+    )
+    yield
+
+    for sc in app.state.prefill_clients + app.state.decode_clients:
+        await sc.client.aclose()
+
+
+app = FastAPI(title="Disaggregated P/D Proxy (Multi-turn)", lifespan=lifespan)
+
+
+def _next_client(app_state, role: str) -> ServiceClient:
+    if role == "prefill":
+        return app_state.prefill_clients[next(app_state.prefill_iter)]
+    return app_state.decode_clients[next(app_state.decode_iter)]
+
+
+# Request handler
+async def _handle_request(api_path: str, request: Request):
+    """Core request handler for both /v1/chat/completions and /v1/completions."""
+    req_data = await request.json()
+    request_id = str(uuid.uuid4())
+    conversation_id: str = req_data.pop("conversation_id", "")
+    client_wants_stream = req_data.get("stream", False)
+
+    if not conversation_id:
+        logger.warning(
+            "[%s] No conversation_id provided — KV cache reuse disabled "
+            "for this request. Add a 'conversation_id' field to enable "
+            "cross-turn KV sharing.",
+            request_id,
+        )
+
+    # Step 1: Look up cached D blocks from the previous turn
+    cached_kv = kv_cache.get(conversation_id) if conversation_id else None
+
+    if cached_kv:
+        # Tell P to read D's blocks (bidirectional transfer)
+        cached_kv["do_remote_decode"] = True
+        cached_kv["do_remote_prefill"] = False
+        req_data["kv_transfer_params"] = cached_kv
+        logger.info(
+            "[%s] conv=%s: sending D's cached blocks to P (remote_request_id=%s)",
+            request_id,
+            conversation_id,
+            cached_kv.get("remote_request_id"),
+        )
+    else:
+        # No cached blocks — P recomputes from scratch
+        req_data["kv_transfer_params"] = {
+            "do_remote_decode": True,
+            "do_remote_prefill": False,
+            "remote_engine_id": None,
+            "remote_block_ids": None,
+            "remote_host": None,
+            "remote_port": None,
+        }
+        logger.info("[%s] conv=%s: cache MISS", request_id, conversation_id)
+
+    # Step 2: Send to Prefill node (non-streaming, max_tokens=1)
+    prefill_client = _next_client(request.app.state, "prefill")
+    t0 = time.time()
+    prefill_resp = await _send_to_prefill(
+        prefill_client,
+        api_path,
+        req_data,
+        request_id,
+    )
+    logger.info(
+        "[%s] Prefill done in %.0fms",
+        request_id,
+        (time.time() - t0) * 1000,
+    )
+
+    # Attach P's kv_transfer_params for D to read P's blocks
+    p_kv_params = prefill_resp.get("kv_transfer_params", {})
+    if p_kv_params:
+        p_kv_params["remote_host"] = prefill_client.host
+        req_data["kv_transfer_params"] = p_kv_params
+
+    # Step 3: Stream from Decode node, capturing kv_transfer_params
+    decode_client = _next_client(request.app.state, "decode")
+
+    if client_wants_stream:
+        return StreamingResponse(
+            _stream_from_decode_sse(
+                decode_client,
+                api_path,
+                req_data,
+                request_id,
+                conversation_id,
+            ),
+            media_type="text/event-stream",
+        )
+
+    text, finish_reason, _, resp_id, model, created = await _stream_from_decode(
+        decode_client,
+        api_path,
+        req_data,
+        request_id,
+        conversation_id,
+    )
+
+    # Build OpenAI-compatible response
+    is_chat = "messages" in req_data
+    if is_chat:
+        body = {
+            "id": resp_id,
+            "object": "chat.completion",
+            "created": created or int(time.time()),
+            "model": model or req_data.get("model", ""),
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": text},
+                    "finish_reason": finish_reason,
+                }
+            ],
+            "usage": None,
+        }
+    else:
+        body = {
+            "id": resp_id,
+            "object": "text_completion",
+            "created": created or int(time.time()),
+            "model": model or req_data.get("model", ""),
+            "choices": [
+                {
+                    "index": 0,
+                    "text": text,
+                    "logprobs": None,
+                    "finish_reason": finish_reason,
+                }
+            ],
+            "usage": None,
+        }
+    return JSONResponse(content=body)
+
+
+# Routes
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request):
+    return await _handle_request("/chat/completions", request)
+
+
+@app.post("/v1/completions")
+async def completions(request: Request):
+    return await _handle_request("/completions", request)
+
+
+@app.get("/health")
+async def health():
+    evicted = kv_cache.evict_stale()
+    return {
+        "status": "ok",
+        "cached_conversations": kv_cache.size,
+        "evicted_stale": evicted,
+    }
+
+
+# CLI
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Disaggregated P/D proxy with bidirectional KV transfer",
+    )
+    p.add_argument("--host", default="0.0.0.0")
+    p.add_argument("--port", type=int, default=8000)
+    p.add_argument(
+        "--prefiller-host",
+        "--prefiller-hosts",
+        dest="prefiller_hosts",
+        nargs="+",
+        default=["localhost"],
+    )
+    p.add_argument(
+        "--prefiller-port",
+        "--prefiller-ports",
+        dest="prefiller_ports",
+        type=int,
+        nargs="+",
+        default=[8100],
+    )
+    p.add_argument(
+        "--decoder-host",
+        "--decoder-hosts",
+        dest="decoder_hosts",
+        nargs="+",
+        default=["localhost"],
+    )
+    p.add_argument(
+        "--decoder-port",
+        "--decoder-ports",
+        dest="decoder_ports",
+        type=int,
+        nargs="+",
+        default=[8200],
+    )
+    args = p.parse_args()
+
+    if len(args.prefiller_hosts) != len(args.prefiller_ports):
+        p.error("Number of prefiller hosts must match ports")
+    if len(args.decoder_hosts) != len(args.decoder_ports):
+        p.error("Number of decoder hosts must match ports")
+
+    args.prefiller_instances = list(zip(args.prefiller_hosts, args.prefiller_ports))
+    args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
+    return args
+
+
+if __name__ == "__main__":
+    global global_args
+    global_args = parse_args()
+
+    import uvicorn
+
+    uvicorn.run(app, host=global_args.host, port=global_args.port)
@@ -105,7 +105,8 @@ plugins:
          - https://docs.aiohttp.org/en/stable/objects.inv
          - https://pillow.readthedocs.io/en/stable/objects.inv
          - https://numpy.org/doc/stable/objects.inv
-          - https://pytorch.org/docs/stable/objects.inv
+          # TODO revert to stable once https://github.com/pytorch/pytorch/issues/182007 is fixed
+          - https://pytorch.org/docs/2.11/objects.inv
  - redirects:
      redirect_maps:
        features/spec_decode/README.md: features/speculative_decoding/README.md
@@ -61,7 +61,11 @@ fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
 instanttensor>=0.1.5
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0; platform_machine == "x86_64"
-terratorch >= 1.2.2 # Required for Prithvi tests
+# terratorch is temporarily disabled while PyPI has the `lightning` package
+# in `quarantined` status (every published terratorch version transitively
+# requires `lightning`, so the resolver fails with "no versions of lightning").
+# Re-enable once PyPI lifts the quarantine. Tracked in #41376.
+# terratorch >= 1.2.2 # Required for Prithvi tests
 imagehash # Required for Prithvi tests
 segmentation-models-pytorch > 0.4.0 # Required for Prithvi tests

@@ -1,15 +1,9 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements/test/cuda.in -c requirements/cuda.txt -o requirements/test/cuda.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_28 --python-version 3.12
 absl-py==2.1.0
-    # via
-    #   rouge-score
-    #   tensorboard
+    # via rouge-score
 accelerate==1.13.0
    # via peft
-aenum==3.1.16
-    # via lightly
-affine==2.4.0
-    # via rasterio
 aiohappyeyeballs==2.6.1
    # via aiohttp
 aiohttp==3.13.3
@@ -25,22 +19,14 @@ aiohttp-cors==0.8.1
    # via ray
 aiosignal==1.4.0
    # via aiohttp
-albucore==0.0.16
-    # via terratorch
 albumentations==1.4.6
-    # via
-    #   -r requirements/test/cuda.in
-    #   terratorch
+    # via -r requirements/test/cuda.in
 alembic==1.16.4
    # via optuna
 annotated-doc==0.0.4
    # via fastapi
 annotated-types==0.7.0
    # via pydantic
-antlr4-python3-runtime==4.9.3
-    # via
-    #   hydra-core
-    #   omegaconf
 anyio==4.6.2.post1
    # via
    #   httpx
@@ -54,12 +40,10 @@ arrow==1.3.0
 attrs==24.2.0
    # via
    #   aiohttp
-    #   fiona
    #   hypothesis
    #   jsonlines
    #   jsonschema
    #   pytest-subtests
-    #   rasterio
    #   referencing
 audioread==3.0.1
    # via librosa
@@ -78,9 +62,7 @@ backoff==2.2.1
    #   -r requirements/test/cuda.in
    #   schemathesis
 bitsandbytes==0.49.2
-    # via
-    #   -r requirements/test/cuda.in
-    #   lightning
+    # via -r requirements/test/cuda.in
 black==24.10.0
    # via datamodel-code-generator
 blobfile==3.0.0
@@ -103,15 +85,9 @@ cachetools==5.5.2
    # via google-auth
 certifi==2024.8.30
    # via
-    #   fiona
    #   httpcore
    #   httpx
-    #   lightly
-    #   pyogrio
-    #   pyproj
-    #   rasterio
    #   requests
-    #   sentry-sdk
 cffi==2.0.0
    # via
    #   cryptography
@@ -125,25 +101,12 @@ chz==0.3.0
 click==8.1.7
    # via
    #   black
-    #   click-plugins
-    #   cligj
-    #   fiona
    #   jiwer
    #   nltk
-    #   rasterio
    #   ray
    #   schemathesis
    #   typer
    #   uvicorn
-    #   wandb
-click-plugins==1.1.1.2
-    # via
-    #   fiona
-    #   rasterio
-cligj==0.7.2
-    # via
-    #   fiona
-    #   rasterio
 colorama==0.4.6
    # via
    #   perceptron
@@ -191,8 +154,6 @@ decorator==5.1.1
    # via librosa
 decord==0.6.0
    # via -r requirements/test/cuda.in
-diffusers==0.36.0
-    # via terratorch
 dill==0.3.8
    # via
    #   datasets
@@ -207,14 +168,10 @@ docker==7.1.0
    # via gpt-oss
 docopt==0.6.2
    # via num2words
-docstring-parser==0.17.0
-    # via jsonargparse
 einops==0.8.1
    # via
    #   -r requirements/test/cuda.in
    #   encodec
-    #   terratorch
-    #   torchgeo
    #   vector-quantize-pytorch
    #   vocos
 einx==0.3.0
@@ -244,13 +201,10 @@ filelock==3.16.1
    #   -c requirements/common.txt
    #   blobfile
    #   datasets
-    #   diffusers
    #   huggingface-hub
    #   ray
    #   torch
    #   virtualenv
-fiona==1.10.1
-    # via torchgeo
 fonttools==4.55.0
    # via matplotlib
 fqdn==1.5.1
@@ -267,9 +221,6 @@ fsspec==2024.12.0
    #   evaluate
    #   fastparquet
    #   huggingface-hub
-    #   lightning
-    #   pytorch-lightning
-    #   tacoreader
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
@@ -277,12 +228,6 @@ genai-perf==0.0.16
    # via -r requirements/test/cuda.in
 genson==1.3.0
    # via datamodel-code-generator
-geopandas==1.0.1
-    # via terratorch
-gitdb==4.0.12
-    # via gitpython
-gitpython==3.1.44
-    # via wandb
 google-api-core==2.24.2
    # via
    #   google-cloud-core
@@ -317,7 +262,6 @@ grpcio==1.78.0
    #   -r requirements/test/cuda.in
    #   grpcio-reflection
    #   ray
-    #   tensorboard
 grpcio-reflection==1.78.0
    # via -r requirements/test/cuda.in
 h11==0.14.0
@@ -326,8 +270,6 @@ h11==0.14.0
    #   uvicorn
 h2==4.3.0
    # via httpx
-h5py==3.13.0
-    # via terratorch
 harfile==0.3.0
    # via schemathesis
 hf-xet==1.4.3
@@ -343,7 +285,6 @@ httpcore==1.0.6
 httpx==0.27.2
    # via
    #   -r requirements/test/cuda.in
-    #   diffusers
    #   huggingface-hub
    #   perceptron
    #   schemathesis
@@ -351,23 +292,17 @@ huggingface-hub==1.10.2
    # via
    #   accelerate
    #   datasets
-    #   diffusers
    #   evaluate
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
    #   sentence-transformers
-    #   terratorch
    #   timm
    #   tokenizers
    #   transformers
    #   vocos
 humanize==4.11.0
    # via runai-model-streamer
-hydra-core==1.3.2
-    # via
-    #   lightly
-    #   lightning
 hyperframe==6.1.0
    # via h2
 hypothesis==6.131.0
@@ -392,11 +327,7 @@ imagehash==4.3.2
 imageio==2.37.0
    # via scikit-image
 importlib-metadata==8.7.0
-    # via
-    #   diffusers
-    #   opentelemetry-api
-importlib-resources==6.5.2
-    # via typeshed-client
+    # via opentelemetry-api
 inflect==5.6.2
    # via datamodel-code-generator
 iniconfig==2.0.0
@@ -426,14 +357,8 @@ joblib==1.4.2
    #   librosa
    #   nltk
    #   scikit-learn
-jsonargparse==4.46.0
-    # via
-    #   lightning
-    #   terratorch
 jsonlines==4.0.0
    # via lm-eval
-jsonnet==0.21.0
-    # via jsonargparse
 jsonpointer==3.0.0
    # via jsonschema
 jsonschema==4.23.0
@@ -452,10 +377,6 @@ kaleido==0.2.1
    # via genai-perf
 kiwisolver==1.4.7
    # via matplotlib
-kornia==0.8.1
-    # via torchgeo
-kornia-rs==0.1.9
-    # via kornia
 lazy-loader==0.4
    # via
    #   librosa
@@ -464,21 +385,6 @@ libnacl==2.1.0
    # via tensorizer
 librosa==0.10.2.post1
    # via -r requirements/test/cuda.in
-lightly==1.5.22
-    # via
-    #   terratorch
-    #   torchgeo
-lightly-utils==0.0.2
-    # via lightly
-lightning==2.6.1
-    # via
-    #   terratorch
-    #   torchgeo
-lightning-utilities==0.14.3
-    # via
-    #   lightning
-    #   pytorch-lightning
-    #   torchmetrics
 llvmlite==0.47.0
    # via numba
 lm-eval==0.4.11
@@ -490,8 +396,6 @@ lxml==5.3.0
    #   sacrebleu
 mako==1.3.10
    # via alembic
-markdown==3.8.2
-    # via tensorboard
 markdown-it-py==3.0.0
    # via rich
 markupsafe==3.0.1
@@ -500,11 +404,7 @@ markupsafe==3.0.1
    #   mako
    #   werkzeug
 matplotlib==3.9.2
-    # via
-    #   -r requirements/test/cuda.in
-    #   lightning
-    #   pycocotools
-    #   torchgeo
+    # via -r requirements/test/cuda.in
 mbstrdecoder==1.1.3
    # via
    #   dataproperty
@@ -559,7 +459,6 @@ numpy==2.2.6
    # via
    #   -r requirements/test/cuda.in
    #   accelerate
-    #   albucore
    #   albumentations
    #   bitsandbytes
    #   bm25s
@@ -567,19 +466,14 @@ numpy==2.2.6
    #   cupy-cuda12x
    #   datasets
    #   decord
-    #   diffusers
    #   einx
    #   encodec
    #   evaluate
    #   fastparquet
    #   genai-perf
-    #   geopandas
-    #   h5py
    #   imagehash
    #   imageio
    #   librosa
-    #   lightly
-    #   lightly-utils
    #   lm-eval
    #   matplotlib
    #   mistral-common
@@ -591,11 +485,7 @@ numpy==2.2.6
    #   patsy
    #   peft
    #   perceptron
-    #   pycocotools
-    #   pyogrio
    #   pywavelets
-    #   rasterio
-    #   rioxarray
    #   rouge-score
    #   runai-model-streamer
    #   sacrebleu
@@ -603,21 +493,14 @@ numpy==2.2.6
    #   scikit-learn
    #   scipy
    #   segmentation-models-pytorch
-    #   shapely
    #   soxr
    #   statsmodels
-    #   tensorboard
-    #   tensorboardx
    #   tensorizer
-    #   terratorch
    #   tifffile
-    #   torchgeo
-    #   torchmetrics
    #   torchvision
    #   transformers
    #   tritonclient
    #   vocos
-    #   xarray
 nvidia-cublas==13.1.0.3
    # via
    #   cuda-toolkit
@@ -657,10 +540,6 @@ nvidia-nvshmem-cu13==3.4.5
    # via torch
 nvidia-nvtx==13.0.85
    # via cuda-toolkit
-omegaconf==2.3.0
-    # via
-    #   hydra-core
-    #   lightning
 open-clip-torch==2.32.0
    # via -r requirements/test/cuda.in
 openai-harmony==0.0.4
@@ -675,7 +554,6 @@ opencv-python-headless==4.13.0.90
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/cuda.in
-    #   albucore
    #   albumentations
    #   mistral-common
 openpyxl==3.1.5
@@ -710,44 +588,27 @@ packaging==24.2
    #   datasets
    #   evaluate
    #   fastparquet
-    #   geopandas
    #   huggingface-hub
-    #   hydra-core
-    #   kornia
    #   lazy-loader
-    #   lightning
-    #   lightning-utilities
    #   matplotlib
    #   optuna
    #   peft
    #   plotly
    #   pooch
-    #   pyogrio
    #   pytest
    #   pytest-rerunfailures
-    #   pytorch-lightning
    #   ray
-    #   rioxarray
    #   scikit-image
    #   statsmodels
-    #   tensorboard
-    #   tensorboardx
-    #   torchmetrics
    #   transformers
    #   typepy
-    #   wandb
-    #   xarray
 pandas==2.2.3
    # via
    #   datasets
    #   evaluate
    #   fastparquet
    #   genai-perf
-    #   geopandas
    #   statsmodels
-    #   tacoreader
-    #   torchgeo
-    #   xarray
 pathspec==0.12.1
    # via black
 pathvalidate==3.2.1
@@ -762,25 +623,20 @@ perf-analyzer==0.1.0
    # via genai-perf
 pillow==10.4.0
    # via
-    #   diffusers
    #   genai-perf
    #   imagehash
    #   imageio
-    #   lightly-utils
    #   matplotlib
    #   mistral-common
    #   perceptron
    #   scikit-image
    #   segmentation-models-pytorch
-    #   tensorboard
-    #   torchgeo
    #   torchvision
 platformdirs==4.3.6
    # via
    #   black
    #   pooch
    #   virtualenv
-    #   wandb
 plotly==5.24.1
    # via
    #   -r requirements/test/cuda.in
@@ -817,10 +673,7 @@ protobuf==6.33.6
    #   opentelemetry-proto
    #   proto-plus
    #   ray
-    #   tensorboard
-    #   tensorboardx
    #   tensorizer
-    #   wandb
 psutil==6.1.0
    # via
    #   accelerate
@@ -834,16 +687,12 @@ pyarrow==23.0.0
    # via
    #   datasets
    #   genai-perf
-    #   tacoreader
-    #   terratorch
 pyasn1==0.6.1
    # via
    #   pyasn1-modules
    #   rsa
 pyasn1-modules==0.4.2
    # via google-auth
-pycocotools==2.0.8
-    # via terratorch
 pycountry==24.6.1
    # via pydantic-extra-types
 pycparser==2.22
@@ -858,13 +707,11 @@ pydantic==2.12.0
    #   datamodel-code-generator
    #   fastapi
    #   gpt-oss
-    #   lightly
    #   mistral-common
    #   mteb
    #   openai-harmony
    #   pydantic-extra-types
    #   ray
-    #   wandb
 pydantic-core==2.41.1
    # via pydantic
 pydantic-extra-types==2.10.5
@@ -873,17 +720,8 @@ pygments==2.18.0
    # via rich
 pyjwt==2.11.0
    # via msal
-pyogrio==0.11.0
-    # via geopandas
 pyparsing==3.2.0
-    # via
-    #   matplotlib
-    #   rasterio
-pyproj==3.7.1
-    # via
-    #   geopandas
-    #   rioxarray
-    #   torchgeo
+    # via matplotlib
 pyrate-limiter==3.7.0
    # via schemathesis
 pystemmer==3.0.0
@@ -920,22 +758,15 @@ pytest-subtests==0.14.1
    # via schemathesis
 pytest-timeout==2.3.1
    # via -r requirements/test/cuda.in
-python-box==7.3.2
-    # via terratorch
 python-dateutil==2.9.0.post0
    # via
    #   arrow
    #   botocore
-    #   lightly
    #   matplotlib
    #   pandas
    #   typepy
 python-rapidjson==1.20
    # via tritonclient
-pytorch-lightning==2.5.2
-    # via
-    #   lightly
-    #   lightning
 pytrec-eval-terrier==0.5.7
    # via mteb
 pytz==2024.2
@@ -952,26 +783,16 @@ pyyaml==6.0.2
    #   datasets
    #   genai-perf
    #   huggingface-hub
-    #   jsonargparse
-    #   lightning
-    #   omegaconf
    #   optuna
    #   peft
-    #   pytorch-lightning
    #   ray
    #   responses
    #   schemathesis
    #   timm
    #   transformers
    #   vocos
-    #   wandb
 rapidfuzz==3.12.1
    # via jiwer
-rasterio==1.4.3
-    # via
-    #   rioxarray
-    #   terratorch
-    #   torchgeo
 ray==2.48.0
    # via -r requirements/test/cuda.in
 redis==5.2.0
@@ -982,7 +803,6 @@ referencing==0.35.1
    #   jsonschema-specifications
 regex==2026.2.28
    # via
-    #   diffusers
    #   nltk
    #   open-clip-torch
    #   sacrebleu
@@ -994,13 +814,11 @@ requests==2.32.3
    #   azure-core
    #   buildkite-test-collector
    #   datasets
-    #   diffusers
    #   docker
    #   evaluate
    #   google-api-core
    #   google-cloud-storage
    #   gpt-oss
-    #   lightly
    #   lm-eval
    #   mistral-common
    #   msal
@@ -1010,9 +828,7 @@ requests==2.32.3
    #   responses
    #   schemathesis
    #   starlette-testclient
-    #   tacoreader
    #   tiktoken
-    #   wandb
 responses==0.25.3
    # via genai-perf
 rfc3339-validator==0.1.4
@@ -1022,13 +838,9 @@ rfc3987==1.3.8
 rich==13.9.4
    # via
    #   genai-perf
-    #   lightning
    #   mteb
    #   perceptron
-    #   terratorch
    #   typer
-rioxarray==0.19.0
-    # via terratorch
 rouge-score==0.1.2
    # via lm-eval
 rpds-py==0.20.1
@@ -1037,8 +849,6 @@ rpds-py==0.20.1
    #   referencing
 rsa==4.9.1
    # via google-auth
-rtree==1.4.0
-    # via torchgeo
 runai-model-streamer==0.15.7
    # via -r requirements/test/cuda.in
 runai-model-streamer-azure==0.15.7
@@ -1054,7 +864,6 @@ sacrebleu==2.4.3
 safetensors==0.4.5
    # via
    #   accelerate
-    #   diffusers
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
@@ -1063,9 +872,7 @@ safetensors==0.4.5
 schemathesis==3.39.15
    # via -r requirements/test/cuda.in
 scikit-image==0.25.2
-    # via
-    #   albumentations
-    #   terratorch
+    # via albumentations
 scikit-learn==1.5.2
    # via
    #   albumentations
@@ -1073,7 +880,6 @@ scikit-learn==1.5.2
    #   lm-eval
    #   mteb
    #   sentence-transformers
-    #   terratorch
 scipy==1.13.1
    # via
    #   albumentations
@@ -1087,27 +893,16 @@ scipy==1.13.1
    #   statsmodels
    #   vocos
 segmentation-models-pytorch==0.5.0
-    # via
-    #   -r requirements/test/cuda.in
-    #   terratorch
-    #   torchgeo
+    # via -r requirements/test/cuda.in
 sentence-transformers==5.2.0
    # via
    #   -r requirements/test/cuda.in
    #   mteb
-sentry-sdk==2.52.0
-    # via wandb
 setuptools==77.0.3
    # via
    #   -c requirements/common.txt
-    #   lightning-utilities
    #   pytablewriter
-    #   tensorboard
    #   torch
-shapely==2.1.1
-    # via
-    #   geopandas
-    #   torchgeo
 shellingham==1.5.4
    # via
    #   perceptron
@@ -1116,15 +911,12 @@ six==1.16.0
    # via
    #   -c requirements/common.txt
    #   junit-xml
-    #   lightly
    #   opencensus
    #   python-dateutil
    #   rfc3339-validator
    #   rouge-score
 smart-open==7.1.0
    # via ray
-smmap==5.0.2
-    # via gitdb
 sniffio==1.3.1
    # via
    #   anyio
@@ -1166,8 +958,6 @@ tabledata==1.3.3
    # via pytablewriter
 tabulate==0.9.0
    # via sacrebleu
-tacoreader==0.5.6
-    # via terratorch
 tblib==3.1.0
    # via -r requirements/test/cuda.in
 tcolorpy==0.1.6
@@ -1177,26 +967,14 @@ tenacity==9.1.2
    #   gpt-oss
    #   lm-eval
    #   plotly
-tensorboard==2.20.0
-    # via terratorch
-tensorboard-data-server==0.7.2
-    # via tensorboard
-tensorboardx==2.6.4
-    # via lightning
 tensorizer==2.10.1
    # via -r requirements/test/cuda.in
 termcolor==3.1.0
-    # via
-    #   gpt-oss
-    #   terratorch
-terratorch==1.2.2
-    # via -r requirements/test/cuda.in
+    # via gpt-oss
 threadpoolctl==3.5.0
    # via scikit-learn
 tifffile==2025.3.30
-    # via
-    #   scikit-image
-    #   terratorch
+    # via scikit-image
 tiktoken==0.12.0
    # via
    #   -c requirements/common.txt
@@ -1208,8 +986,6 @@ timm==1.0.17
    #   -r requirements/test/cuda.in
    #   open-clip-torch
    #   segmentation-models-pytorch
-    #   terratorch
-    #   torchgeo
 tokenizers==0.22.2
    # via
    #   -c requirements/common.txt
@@ -1227,21 +1003,14 @@ torch==2.11.0+cu130
    #   bitsandbytes
    #   encodec
    #   instanttensor
-    #   kornia
-    #   lightly
-    #   lightning
    #   mteb
    #   open-clip-torch
    #   peft
-    #   pytorch-lightning
    #   runai-model-streamer
    #   segmentation-models-pytorch
    #   sentence-transformers
    #   tensorizer
-    #   terratorch
    #   timm
-    #   torchgeo
-    #   torchmetrics
    #   torchvision
    #   vector-quantize-pytorch
    #   vocos
@@ -1251,31 +1020,18 @@ torchaudio==2.11.0+cu130
    #   -r requirements/test/cuda.in
    #   encodec
    #   vocos
-torchgeo==0.7.0
-    # via terratorch
-torchmetrics==1.7.4
-    # via
-    #   lightning
-    #   pytorch-lightning
-    #   terratorch
-    #   torchgeo
 torchvision==0.26.0+cu130
    # via
    #   -c requirements/cuda.txt
    #   -r requirements/test/cuda.in
-    #   lightly
    #   open-clip-torch
    #   segmentation-models-pytorch
-    #   terratorch
    #   timm
-    #   torchgeo
 tqdm==4.67.3
    # via
    #   datasets
    #   evaluate
    #   huggingface-hub
-    #   lightly
-    #   lightning
    #   lm-eval
    #   mteb
    #   nltk
@@ -1283,11 +1039,8 @@ tqdm==4.67.3
    #   optuna
    #   peft
    #   pqdm
-    #   pytorch-lightning
    #   segmentation-models-pytorch
    #   sentence-transformers
-    #   tacoreader
-    #   terratorch
    #   transformers
 transformers==5.5.3
    # via
@@ -1316,8 +1069,6 @@ typer==0.15.2
    #   transformers
 types-python-dateutil==2.9.0.20241206
    # via arrow
-typeshed-client==2.8.2
-    # via jsonargparse
 typing-extensions==4.15.0
    # via
    #   -c requirements/common.txt
@@ -1332,8 +1083,6 @@ typing-extensions==4.15.0
    #   grpcio
    #   huggingface-hub
    #   librosa
-    #   lightning
-    #   lightning-utilities
    #   lm-eval
    #   mistral-common
    #   mteb
@@ -1344,16 +1093,12 @@ typing-extensions==4.15.0
    #   pydantic
    #   pydantic-core
    #   pydantic-extra-types
-    #   pytorch-lightning
    #   sentence-transformers
    #   sqlalchemy
    #   starlette
    #   torch
-    #   torchgeo
    #   typer
-    #   typeshed-client
    #   typing-inspection
-    #   wandb
 typing-inspection==0.4.2
    # via pydantic
 tzdata==2024.2
@@ -1365,10 +1110,8 @@ urllib3==2.2.3
    #   blobfile
    #   botocore
    #   docker
-    #   lightly
    #   requests
    #   responses
-    #   sentry-sdk
    #   tritonclient
 uvicorn==0.35.0
    # via gpt-oss
@@ -1378,22 +1121,16 @@ virtualenv==20.31.2
    # via ray
 vocos==0.1.0
    # via -r requirements/test/cuda.in
-wandb==0.24.2
-    # via terratorch
 wcwidth==0.2.13
    # via ftfy
 webcolors==24.11.1
    # via jsonschema
 werkzeug==3.1.3
-    # via
-    #   schemathesis
-    #   tensorboard
+    # via schemathesis
 word2number==1.1
    # via lm-eval
 wrapt==1.17.2
    # via smart-open
-xarray==2025.7.1
-    # via rioxarray
 xxhash==3.5.0
    # via
    #   datasets
@@ -61,7 +61,11 @@ pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0

 # Prithvi tests
-terratorch>=1.2.2
+# terratorch is temporarily disabled while PyPI has the `lightning` package
+# in `quarantined` status (every published terratorch version transitively
+# requires `lightning`, so the resolver fails with "no versions of lightning").
+# Re-enable once PyPI lifts the quarantine. Tracked in #41376.
+# terratorch>=1.2.2
 imagehash # Required for Prithvi tests
 segmentation-models-pytorch>0.4.0 # Required for Prithvi tests

@@ -79,5 +83,7 @@ plotly # required for perf comparison html report

 # ROCm-specific extras (not in CUDA cuda.in)
 rapidfuzz
-torchgeo==0.7.0
+# torchgeo also pulled in `lightning` transitively; disabled for the same
+# quarantine reason as terratorch above. Restore once the quarantine clears.
+# torchgeo==0.7.0
 multiprocess==0.70.16
@@ -1,15 +1,9 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13
 absl-py==2.4.0
-    # via
-    #   rouge-score
-    #   tensorboard
+    # via rouge-score
 accelerate==1.13.0
    # via peft
-aenum==3.1.17
-    # via lightly
-affine==2.4.0
-    # via rasterio
 aiohappyeyeballs==2.6.1
    # via aiohttp
 aiohttp==3.13.3
@@ -25,12 +19,8 @@ aiohttp-cors==0.8.1
    # via ray
 aiosignal==1.4.0
    # via aiohttp
-albucore==0.1.2
-    # via terratorch
 albumentations==1.4.6
-    # via
-    #   -r requirements/test/rocm.in
-    #   terratorch
+    # via -r requirements/test/rocm.in
 alembic==1.18.4
    # via optuna
 annotated-doc==0.0.4
@@ -43,10 +33,6 @@ anthropic==0.93.0
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
-antlr4-python3-runtime==4.9.3
-    # via
-    #   hydra-core
-    #   omegaconf
 anyio==4.13.0
    # via
    #   anthropic
@@ -67,11 +53,9 @@ astor==0.8.1
 attrs==26.1.0
    # via
    #   aiohttp
-    #   fiona
    #   jsonlines
    #   jsonschema
    #   pytest-subtests
-    #   rasterio
    #   referencing
 audioread==3.0.1
    # via librosa
@@ -90,9 +74,7 @@ backoff==2.2.1
    #   -r requirements/test/rocm.in
    #   schemathesis
 bitsandbytes==0.49.2
-    # via
-    #   -r requirements/test/rocm.in
-    #   lightning
+    # via -r requirements/test/rocm.in
 black==26.3.1
    # via datamodel-code-generator
 blake3==1.0.8
@@ -119,13 +101,8 @@ cbor2==5.9.0
    # via -r requirements/test/../common.txt
 certifi==2026.2.25
    # via
-    #   fiona
    #   httpcore
    #   httpx
-    #   lightly
-    #   pyogrio
-    #   pyproj
-    #   rasterio
    #   requests
    #   sentry-sdk
 cffi==1.17.1
@@ -143,24 +120,13 @@ chz==0.4.0
 click==8.3.1
    # via
    #   black
-    #   click-plugins
-    #   cligj
-    #   fiona
    #   jiwer
    #   nltk
-    #   rasterio
    #   ray
    #   rich-toolkit
    #   schemathesis
    #   typer
    #   uvicorn
-    #   wandb
-click-plugins==1.1.1.2
-    # via fiona
-cligj==0.7.2
-    # via
-    #   fiona
-    #   rasterio
 cloudpickle==3.1.2
    # via -r requirements/test/../common.txt
 colorama==0.4.6
@@ -211,8 +177,6 @@ depyf==0.20.0
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
-diffusers==0.37.0
-    # via terratorch
 dill==0.3.8
    # via
    #   datasets
@@ -237,16 +201,12 @@ docker==7.1.0
 docopt==0.6.2
    # via num2words
 docstring-parser==0.17.0
-    # via
-    #   anthropic
-    #   jsonargparse
+    # via anthropic
 einops==0.8.2
    # via
    #   -r requirements/test/../common.txt
    #   -r requirements/test/rocm.in
    #   encodec
-    #   terratorch
-    #   torchgeo
    #   vector-quantize-pytorch
    #   vocos
 einx==0.4.2
@@ -283,14 +243,11 @@ filelock==3.25.2
    #   -r requirements/test/../common.txt
    #   blobfile
    #   datasets
-    #   diffusers
    #   huggingface-hub
    #   python-discovery
    #   ray
    #   torch
    #   virtualenv
-fiona==1.10.1
-    # via torchgeo
 fonttools==4.62.1
    # via matplotlib
 fqdn==1.5.1
@@ -307,9 +264,6 @@ fsspec==2025.3.0
    #   evaluate
    #   fastparquet
    #   huggingface-hub
-    #   lightning
-    #   pytorch-lightning
-    #   tacoreader
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
@@ -317,16 +271,10 @@ genai-perf==0.0.16
    # via -r requirements/test/rocm.in
 genson==1.3.0
    # via datamodel-code-generator
-geopandas==1.1.3
-    # via terratorch
 gguf==0.18.0
    # via
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
-gitdb==4.0.12
-    # via gitpython
-gitpython==3.1.46
-    # via wandb
 google-api-core==2.30.0
    # via
    #   google-cloud-core
@@ -366,7 +314,6 @@ grpcio==1.78.0
    #   grpcio-reflection
    #   opentelemetry-exporter-otlp-proto-grpc
    #   ray
-    #   tensorboard
 grpcio-reflection==1.78.0
    # via
    #   -c requirements/rocm.txt
@@ -377,8 +324,6 @@ h11==0.16.0
    #   uvicorn
 h2==4.3.0
    # via httpx
-h5py==3.16.0
-    # via terratorch
 harfile==0.4.0
    # via schemathesis
 hf-xet==1.4.3
@@ -397,7 +342,6 @@ httpx==0.27.2
    # via
    #   -r requirements/test/rocm.in
    #   anthropic
-    #   diffusers
    #   fastapi
    #   fastapi-cloud-cli
    #   huggingface-hub
@@ -412,23 +356,17 @@ huggingface-hub==1.10.2
    # via
    #   accelerate
    #   datasets
-    #   diffusers
    #   evaluate
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
    #   sentence-transformers
-    #   terratorch
    #   timm
    #   tokenizers
    #   transformers
    #   vocos
 humanize==4.15.0
    # via runai-model-streamer
-hydra-core==1.3.2
-    # via
-    #   lightly
-    #   lightning
 hyperframe==6.1.0
    # via h2
 hypothesis==6.151.9
@@ -455,11 +393,7 @@ imagehash==4.3.2
 imageio==2.37.3
    # via scikit-image
 importlib-metadata==8.7.1
-    # via
-    #   diffusers
-    #   opentelemetry-api
-importlib-resources==6.5.2
-    # via typeshed-client
+    # via opentelemetry-api
 inflect==7.5.0
    # via datamodel-code-generator
 iniconfig==2.3.0
@@ -497,14 +431,8 @@ joblib==1.5.3
    #   librosa
    #   nltk
    #   scikit-learn
-jsonargparse==4.47.0
-    # via
-    #   lightning
-    #   terratorch
 jsonlines==4.0.0
    # via lm-eval
-jsonnet==0.21.0
-    # via jsonargparse
 jsonpointer==3.1.0
    # via jsonschema
 jsonschema==4.26.0
@@ -524,10 +452,6 @@ kaleido==1.0.0
    # via genai-perf
 kiwisolver==1.5.0
    # via matplotlib
-kornia==0.8.2
-    # via torchgeo
-kornia-rs==0.1.10
-    # via kornia
 lark==1.2.2
    # via
    #   -c requirements/common.txt
@@ -540,21 +464,6 @@ libnacl==2.1.0
    # via tensorizer
 librosa==0.10.2.post1
    # via -r requirements/test/rocm.in
-lightly==1.5.22
-    # via
-    #   terratorch
-    #   torchgeo
-lightly-utils==0.0.2
-    # via lightly
-lightning==2.6.1
-    # via
-    #   terratorch
-    #   torchgeo
-lightning-utilities==0.15.3
-    # via
-    #   lightning
-    #   pytorch-lightning
-    #   torchmetrics
 llguidance==1.3.0
    # via
    #   -c requirements/common.txt
@@ -580,8 +489,6 @@ lxml==6.0.2
    #   sacrebleu
 mako==1.3.10
    # via alembic
-markdown==3.10.2
-    # via tensorboard
 markdown-it-py==4.0.0
    # via rich
 markupsafe==3.0.3
@@ -590,10 +497,7 @@ markupsafe==3.0.3
    #   mako
    #   werkzeug
 matplotlib==3.10.8
-    # via
-    #   -r requirements/test/rocm.in
-    #   lightning
-    #   torchgeo
+    # via -r requirements/test/rocm.in
 mbstrdecoder==1.1.4
    # via
    #   dataproperty
@@ -660,14 +564,11 @@ numba==0.65.0
    #   -c requirements/rocm.txt
    #   -r requirements/test/rocm.in
    #   librosa
-numkong==7.1.1
-    # via albucore
 numpy==2.2.6
    # via
    #   -r requirements/test/../common.txt
    #   -r requirements/test/rocm.in
    #   accelerate
-    #   albucore
    #   albumentations
    #   bitsandbytes
    #   bm25s
@@ -675,20 +576,15 @@ numpy==2.2.6
    #   cupy-cuda12x
    #   datasets
    #   decord
-    #   diffusers
    #   einx
    #   encodec
    #   evaluate
    #   fastparquet
    #   genai-perf
-    #   geopandas
    #   gguf
-    #   h5py
    #   imagehash
    #   imageio
    #   librosa
-    #   lightly
-    #   lightly-utils
    #   lm-eval
    #   matplotlib
    #   mistral-common
@@ -700,12 +596,8 @@ numpy==2.2.6
    #   patsy
    #   peft
    #   perceptron
-    #   pycocotools
-    #   pyogrio
    #   pytrec-eval-terrier
    #   pywavelets
-    #   rasterio
-    #   rioxarray
    #   rouge-score
    #   runai-model-streamer
    #   sacrebleu
@@ -714,27 +606,16 @@ numpy==2.2.6
    #   scipy
    #   segmentation-models-pytorch
    #   sentence-transformers
-    #   shapely
    #   soundfile
    #   soxr
    #   statsmodels
-    #   tensorboard
-    #   tensorboardx
    #   tensorizer
-    #   terratorch
    #   tifffile
-    #   torchgeo
-    #   torchmetrics
    #   torchvision
    #   transformers
    #   tritonclient
    #   vocos
-    #   xarray
    #   xgrammar
-omegaconf==2.3.0
-    # via
-    #   hydra-core
-    #   lightning
 open-clip-torch==2.32.0
    # via -r requirements/test/rocm.in
 openai==2.31.0
@@ -824,46 +705,29 @@ packaging==26.0
    #   datasets
    #   evaluate
    #   fastparquet
-    #   geopandas
    #   huggingface-hub
-    #   hydra-core
    #   kaleido
-    #   kornia
    #   lazy-loader
-    #   lightning
-    #   lightning-utilities
    #   lm-format-enforcer
    #   matplotlib
    #   optuna
    #   peft
    #   plotly
    #   pooch
-    #   pyogrio
    #   pytest
    #   pytest-rerunfailures
-    #   pytorch-lightning
    #   ray
-    #   rioxarray
    #   scikit-image
    #   statsmodels
-    #   tensorboard
-    #   tensorboardx
-    #   torchmetrics
    #   transformers
    #   typepy
-    #   wandb
-    #   xarray
 pandas==3.0.1
    # via
    #   datasets
    #   evaluate
    #   fastparquet
    #   genai-perf
-    #   geopandas
    #   statsmodels
-    #   tacoreader
-    #   torchgeo
-    #   xarray
 partial-json-parser==0.2.1.1.post7
    # via -r requirements/test/../common.txt
 pathspec==1.0.4
@@ -881,18 +745,14 @@ perf-analyzer==0.1.0
 pillow==12.1.1
    # via
    #   -r requirements/test/../common.txt
-    #   diffusers
    #   genai-perf
    #   imagehash
    #   imageio
-    #   lightly-utils
    #   matplotlib
    #   mistral-common
    #   perceptron
    #   scikit-image
    #   segmentation-models-pytorch
-    #   tensorboard
-    #   torchgeo
    #   torchvision
 platformdirs==4.3.6
    # via
@@ -900,7 +760,6 @@ platformdirs==4.3.6
    #   pooch
    #   python-discovery
    #   virtualenv
-    #   wandb
 plotly==6.6.0
    # via
    #   -r requirements/test/rocm.in
@@ -946,10 +805,7 @@ protobuf==6.33.6
    #   opentelemetry-proto
    #   proto-plus
    #   ray
-    #   tensorboard
-    #   tensorboardx
    #   tensorizer
-    #   wandb
 psutil==7.2.2
    # via
    #   -r requirements/test/../common.txt
@@ -966,16 +822,12 @@ pyarrow==23.0.1
    # via
    #   datasets
    #   genai-perf
-    #   tacoreader
-    #   terratorch
 pyasn1==0.6.3
    # via pyasn1-modules
 pyasn1-modules==0.4.2
    # via google-auth
 pybase64==1.4.3
    # via -r requirements/test/../common.txt
-pycocotools==2.0.11
-    # via terratorch
 pycountry==26.2.16
    # via pydantic-extra-types
 pycparser==3.0
@@ -994,7 +846,6 @@ pydantic==2.12.5
    #   fastapi
    #   fastapi-cloud-cli
    #   gpt-oss
-    #   lightly
    #   lm-format-enforcer
    #   mcp
    #   mistral-common
@@ -1005,7 +856,6 @@ pydantic==2.12.5
    #   pydantic-extra-types
    #   pydantic-settings
    #   ray
-    #   wandb
    #   xgrammar
 pydantic-core==2.41.5
    # via pydantic
@@ -1023,17 +873,8 @@ pyjwt==2.12.1
    # via
    #   mcp
    #   msal
-pyogrio==0.12.1
-    # via geopandas
 pyparsing==3.3.2
-    # via
-    #   matplotlib
-    #   rasterio
-pyproj==3.7.2
-    # via
-    #   geopandas
-    #   rioxarray
-    #   torchgeo
+    # via matplotlib
 pyrate-limiter==3.9.0
    # via schemathesis
 pystemmer==3.0.0
@@ -1070,13 +911,10 @@ pytest-subtests==0.14.2
    # via schemathesis
 pytest-timeout==2.3.1
    # via -r requirements/test/rocm.in
-python-box==7.4.1
-    # via terratorch
 python-dateutil==2.9.0.post0
    # via
    #   arrow
    #   botocore
-    #   lightly
    #   matplotlib
    #   pandas
    #   typepy
@@ -1096,10 +934,6 @@ python-rapidjson==1.23
    # via tritonclient
 pytokens==0.4.1
    # via black
-pytorch-lightning==2.6.1
-    # via
-    #   lightly
-    #   lightning
 pytrec-eval-terrier==0.5.10
    # via mteb
 pytz==2026.1.post1
@@ -1116,13 +950,9 @@ pyyaml==6.0.3
    #   genai-perf
    #   gguf
    #   huggingface-hub
-    #   jsonargparse
-    #   lightning
    #   lm-format-enforcer
-    #   omegaconf
    #   optuna
    #   peft
-    #   pytorch-lightning
    #   ray
    #   responses
    #   schemathesis
@@ -1130,7 +960,6 @@ pyyaml==6.0.3
    #   transformers
    #   uvicorn
    #   vocos
-    #   wandb
 pyzmq==27.1.0
    # via
    #   -c requirements/common.txt
@@ -1139,11 +968,6 @@ rapidfuzz==3.12.1
    # via
    #   -r requirements/test/rocm.in
    #   jiwer
-rasterio==1.5.0
-    # via
-    #   rioxarray
-    #   terratorch
-    #   torchgeo
 ray==2.54.0
    # via -r requirements/test/rocm.in
 redis==7.3.0
@@ -1155,7 +979,6 @@ referencing==0.37.0
 regex==2026.2.28
    # via
    #   -r requirements/test/../common.txt
-    #   diffusers
    #   nltk
    #   open-clip-torch
    #   sacrebleu
@@ -1168,14 +991,12 @@ requests==2.32.5
    #   azure-core
    #   buildkite-test-collector
    #   datasets
-    #   diffusers
    #   docker
    #   evaluate
    #   gguf
    #   google-api-core
    #   google-cloud-storage
    #   gpt-oss
-    #   lightly
    #   lm-eval
    #   mistral-common
    #   msal
@@ -1186,9 +1007,7 @@ requests==2.32.5
    #   responses
    #   schemathesis
    #   starlette-testclient
-    #   tacoreader
    #   tiktoken
-    #   wandb
 responses==0.26.0
    # via genai-perf
 rfc3339-validator==0.1.4
@@ -1198,11 +1017,9 @@ rfc3987==1.3.8
 rich==14.3.3
    # via
    #   genai-perf
-    #   lightning
    #   mteb
    #   perceptron
    #   rich-toolkit
-    #   terratorch
    #   typer
 rich-toolkit==0.19.7
    # via
@@ -1210,16 +1027,12 @@ rich-toolkit==0.19.7
    #   fastapi-cloud-cli
 rignore==0.7.6
    # via fastapi-cloud-cli
-rioxarray==0.22.0
-    # via terratorch
 rouge-score==0.1.2
    # via lm-eval
 rpds-py==0.30.0
    # via
    #   jsonschema
    #   referencing
-rtree==1.4.1
-    # via torchgeo
 runai-model-streamer==0.15.7
    # via
    #   -c requirements/rocm.txt
@@ -1237,7 +1050,6 @@ sacrebleu==2.6.0
 safetensors==0.7.0
    # via
    #   accelerate
-    #   diffusers
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
@@ -1246,9 +1058,7 @@ safetensors==0.7.0
 schemathesis==3.39.15
    # via -r requirements/test/rocm.in
 scikit-image==0.26.0
-    # via
-    #   albumentations
-    #   terratorch
+    # via albumentations
 scikit-learn==1.8.0
    # via
    #   albumentations
@@ -1256,7 +1066,6 @@ scikit-learn==1.8.0
    #   lm-eval
    #   mteb
    #   sentence-transformers
-    #   terratorch
 scipy==1.17.1
    # via
    #   albumentations
@@ -1271,10 +1080,7 @@ scipy==1.17.1
    #   statsmodels
    #   vocos
 segmentation-models-pytorch==0.5.0
-    # via
-    #   -r requirements/test/rocm.in
-    #   terratorch
-    #   torchgeo
+    # via -r requirements/test/rocm.in
 sentence-transformers==5.3.0
    # via
    #   -r requirements/test/rocm.in
@@ -1282,9 +1088,7 @@ sentence-transformers==5.3.0
 sentencepiece==0.2.1
    # via -r requirements/test/../common.txt
 sentry-sdk==2.55.0
-    # via
-    #   fastapi-cloud-cli
-    #   wandb
+    # via fastapi-cloud-cli
 setproctitle==1.3.7
    # via -r requirements/test/../common.txt
 setuptools==79.0.1
@@ -1294,12 +1098,7 @@ setuptools==79.0.1
    #   -r requirements/test/../common.txt
    #   model-hosting-container-standards
    #   pytablewriter
-    #   tensorboard
    #   torch
-shapely==2.1.2
-    # via
-    #   geopandas
-    #   torchgeo
 shellingham==1.5.4
    # via
    #   perceptron
@@ -1311,15 +1110,12 @@ six==1.17.0
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
    #   junit-xml
-    #   lightly
    #   opencensus
    #   python-dateutil
    #   rfc3339-validator
    #   rouge-score
 smart-open==7.5.1
    # via ray
-smmap==5.0.3
-    # via gitdb
 sniffio==1.3.1
    # via
    #   anthropic
@@ -1358,8 +1154,6 @@ starlette-testclient==0.4.1
    # via schemathesis
 statsmodels==0.14.6
    # via genai-perf
-stringzilla==4.6.0
-    # via albucore
 structlog==25.5.0
    # via gpt-oss
 supervisor==4.3.0
@@ -1372,8 +1166,6 @@ tabledata==1.3.4
    # via pytablewriter
 tabulate==0.10.0
    # via sacrebleu
-tacoreader==0.5.6
-    # via terratorch
 tblib==3.1.0
    # via -r requirements/test/rocm.in
 tcolorpy==0.1.7
@@ -1382,28 +1174,16 @@ tenacity==9.1.4
    # via
    #   gpt-oss
    #   lm-eval
-tensorboard==2.20.0
-    # via terratorch
-tensorboard-data-server==0.7.2
-    # via tensorboard
-tensorboardx==2.6.4
-    # via lightning
 tensorizer==2.10.1
    # via
    #   -c requirements/rocm.txt
    #   -r requirements/test/rocm.in
 termcolor==3.3.0
-    # via
-    #   gpt-oss
-    #   terratorch
-terratorch==1.2.2
-    # via -r requirements/test/rocm.in
+    # via gpt-oss
 threadpoolctl==3.6.0
    # via scikit-learn
 tifffile==2026.3.3
-    # via
-    #   scikit-image
-    #   terratorch
+    # via scikit-image
 tiktoken==0.12.0
    # via
    #   -c requirements/common.txt
@@ -1417,8 +1197,6 @@ timm==1.0.17
    #   -r requirements/test/rocm.in
    #   open-clip-torch
    #   segmentation-models-pytorch
-    #   terratorch
-    #   torchgeo
 tokenizers==0.22.2
    # via
    #   -c requirements/common.txt
@@ -1429,16 +1207,6 @@ tomli==2.4.0
    # via schemathesis
 tomli-w==1.2.0
    # via schemathesis
-torchgeo==0.7.0
-    # via
-    #   -r requirements/test/rocm.in
-    #   terratorch
-torchmetrics==1.9.0
-    # via
-    #   lightning
-    #   pytorch-lightning
-    #   terratorch
-    #   torchgeo
 tqdm==4.67.3
    # via
    #   -r requirements/test/../common.txt
@@ -1446,8 +1214,6 @@ tqdm==4.67.3
    #   evaluate
    #   gguf
    #   huggingface-hub
-    #   lightly
-    #   lightning
    #   lm-eval
    #   mteb
    #   nltk
@@ -1456,11 +1222,8 @@ tqdm==4.67.3
    #   optuna
    #   peft
    #   pqdm
-    #   pytorch-lightning
    #   segmentation-models-pytorch
    #   sentence-transformers
-    #   tacoreader
-    #   terratorch
    #   transformers
 transformers==5.5.3
    # via
@@ -1492,8 +1255,6 @@ typer==0.24.1
    #   huggingface-hub
    #   perceptron
    #   transformers
-typeshed-client==2.9.0
-    # via jsonargparse
 typing-extensions==4.15.0
    # via
    #   -c requirements/common.txt
@@ -1511,8 +1272,6 @@ typing-extensions==4.15.0
    #   grpcio
    #   huggingface-hub
    #   librosa
-    #   lightning
-    #   lightning-utilities
    #   lm-eval
    #   mcp
    #   mistral-common
@@ -1527,18 +1286,14 @@ typing-extensions==4.15.0
    #   pydantic
    #   pydantic-core
    #   pydantic-extra-types
-    #   pytorch-lightning
    #   referencing
    #   rich-toolkit
    #   sentence-transformers
    #   sqlalchemy
    #   starlette
    #   torch
-    #   torchgeo
    #   typeguard
-    #   typeshed-client
    #   typing-inspection
-    #   wandb
    #   xgrammar
 typing-inspection==0.4.2
    # via
@@ -1555,7 +1310,6 @@ urllib3==2.6.3
    #   blobfile
    #   botocore
    #   docker
-    #   lightly
    #   requests
    #   responses
    #   sentry-sdk
@@ -1575,8 +1329,6 @@ virtualenv==21.2.0
    # via ray
 vocos==0.1.0
    # via -r requirements/test/rocm.in
-wandb==0.25.1
-    # via terratorch
 watchfiles==1.1.1
    # via
    #   -r requirements/test/../common.txt
@@ -1588,15 +1340,11 @@ webcolors==25.10.0
 websockets==16.0
    # via uvicorn
 werkzeug==3.1.6
-    # via
-    #   schemathesis
-    #   tensorboard
+    # via schemathesis
 word2number==1.1
    # via lm-eval
 wrapt==2.1.2
    # via smart-open
-xarray==2026.2.0
-    # via rioxarray
 xgrammar==0.1.33
    # via
    #   -c requirements/common.txt
@@ -150,13 +150,11 @@ def test_full_graph(
        if is_torch_equal_or_newer("2.9.0.dev")
    ]
    + [
-        # Test get_raw_stream patch with compile_sizes
-        # This tests that TorchInductor autotune works correctly with get_raw_stream
-        # patch in torch 2.9 and without patch in torch 2.10+
+        # Cover compile_sizes autotune path.
        (
            CompilationConfig(
                mode=CompilationMode.VLLM_COMPILE,
-                compile_sizes=[1, 2],  # Triggers autotune which uses get_raw_stream
+                compile_sizes=[1, 2],  # Triggers the autotune path.
                cudagraph_mode=CUDAGraphMode.NONE,
            ),
            "facebook/opt-125m",
@@ -24,7 +24,6 @@ from vllm.engine.arg_utils import EngineArgs
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import (
    _is_torch_equal_or_newer,
-    is_torch_equal,
 )
 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

@@ -43,29 +42,6 @@ def test_version():
    assert not _is_torch_equal_or_newer("2.7.1", "2.8.0.dev")


-def test_get_raw_stream_patch():
-    """Test that get_raw_stream patch is applied only for torch 2.9.0 or 2.9.1."""
-    import builtins
-
-    # Check if get_raw_stream exists in builtins
-    has_patch = hasattr(builtins, "get_raw_stream")
-
-    # Import torch to get actual version
-
-    is_torch_2_9 = is_torch_equal("2.9.0") or is_torch_equal("2.9.1")
-
-    if is_torch_2_9:
-        # For torch 2.9.x, the patch should be applied
-        assert has_patch, "get_raw_stream should be patched for torch 2.9.x"
-        # Verify it's callable (it should be the _cuda_getCurrentRawStream function)
-        get_raw_stream = builtins.get_raw_stream  # type: ignore[attr-defined]
-        assert callable(get_raw_stream)
-        # Verify it's the correct function from torch._C
-        from torch._C import _cuda_getCurrentRawStream
-
-        assert get_raw_stream is _cuda_getCurrentRawStream
-
-
 def test_copy_pass():
    vllm_config = VllmConfig()
    inductor_pass = FixFunctionalizationPass(vllm_config)
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import importlib.util
+
 import numpy as np
 import pybase64 as base64
 import pytest
@@ -10,7 +12,16 @@ import torch
 from tests.utils import RemoteOpenAIServer
 from vllm.utils.serial_utils import tensor2base64

+# Prithvi requires terratorch, which is temporarily unavailable while PyPI has
+# `lightning` quarantined (#41376). Skip just the Prithvi case; leave the
+# Qwen3-VL case in the same file untouched.
+_TERRATORCH_AVAILABLE = importlib.util.find_spec("terratorch") is not None

+
+@pytest.mark.skipif(
+    not _TERRATORCH_AVAILABLE,
+    reason="terratorch unavailable while PyPI has `lightning` quarantined; see #41376",
+)
@pytest.mark.parametrize(
    "model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
 )
@@ -167,9 +167,8 @@ def run_evaluation(
    "model_config",
    [
        ("openai/whisper-large-v3", 12.744980),
-        # TODO (ekagra): turn on after asr release
        # CohereASR is used to test the variable encoder length code paths
-        # ("CohereLabs/cohere-transcribe-03-2026", 11.92),
+        ("CohereLabs/cohere-transcribe-03-2026", 11.92),
    ],
 )
 # Original dataset is 20GB+ in size, hence we use a pre-filtered slice.
@@ -6,7 +6,9 @@ from unittest.mock import MagicMock

 import pytest

+import vllm.envs as envs
 from vllm.entrypoints.openai.engine.serving import GenerationError, OpenAIServing
+from vllm.envs import disable_envs_cache


@pytest.mark.asyncio
@@ -60,3 +62,35 @@ async def test_convert_generation_error_to_streaming_response():
    assert isinstance(error_json, str)
    assert "Internal server error" in error_json
    assert "InternalServerError" in error_json
+
+
+def test_is_model_supported_skip_name_validation_env(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """When VLLM_SKIP_MODEL_NAME_VALIDATION is set, accept any model id."""
+    disable_envs_cache()
+    monkeypatch.delenv("VLLM_SKIP_MODEL_NAME_VALIDATION", raising=False)
+
+    mock_engine = MagicMock()
+    mock_engine.model_config = MagicMock()
+    mock_engine.model_config.max_model_len = 100
+    mock_models = MagicMock()
+    mock_models.is_base_model.return_value = False
+
+    serving = OpenAIServing(
+        engine_client=mock_engine,
+        models=mock_models,
+        request_logger=None,
+    )
+
+    assert serving._is_model_supported("not-a-registered-model") is False
+
+    monkeypatch.setenv("VLLM_SKIP_MODEL_NAME_VALIDATION", "1")
+    disable_envs_cache()
+    assert envs.VLLM_SKIP_MODEL_NAME_VALIDATION is True
+    assert serving._is_model_supported("not-a-registered-model") is True
+
+    monkeypatch.setenv("VLLM_SKIP_MODEL_NAME_VALIDATION", "true")
+    disable_envs_cache()
+    assert envs.VLLM_SKIP_MODEL_NAME_VALIDATION is True
+    assert serving._is_model_supported("another-alias") is True
@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Precision tests for vllm's chunk_kda Triton operator.
+
+Compares chunk_kda against a naive recurrent reference (float32).
+Uses torch.rand for q/k/v to match FLA's test pattern.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm.model_executor.layers.fla.ops.kda import chunk_kda
+from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd
+
+DEVICE = "cuda"
+
+
+def naive_recurrent_kda(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Naive recurrent KDA reference, ported from FLA's naive.py."""
+    dtype = v.dtype
+    B, T, H, K = q.shape
+    V = v.shape[-1]
+    if scale is None:
+        scale = K**-0.5
+
+    q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta])
+    q = q * scale
+
+    S = k.new_zeros(B, H, K, V).to(q)
+    if initial_state is not None:
+        S += initial_state
+    o = torch.zeros_like(v)
+    for i in range(T):
+        q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
+        S = S * g_i[..., None].exp()
+        S = S + torch.einsum(
+            "bhk,bhv->bhkv",
+            b_i[..., None] * k_i,
+            v_i - (k_i[..., None] * S).sum(-2),
+        )
+        o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
+    if not output_final_state:
+        S = None
+    return o.to(dtype), S
+
+
+def assert_close(
+    name: str,
+    ref: torch.Tensor,
+    tri: torch.Tensor,
+    ratio: float,
+    err_atol: float = 1e-6,
+):
+    """RMSE-based relative error comparison."""
+    abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item()
+    rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item()
+    rmse_base = ref.detach().flatten().square().mean().sqrt().item()
+    rel_err = rmse_diff / (rmse_base + 1e-8)
+    print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}")
+    if abs_err <= err_atol:
+        return
+    assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref"
+    assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri"
+    assert rel_err < ratio, (
+        f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}"
+    )
+
+
+@pytest.mark.parametrize(
+    ("H", "D", "cu_seqlens", "dtype"),
+    [
+        pytest.param(
+            *test,
+            id="H{}-D{}-cu{}-{}".format(*test),
+        )
+        for test in [
+            (32, 128, [0, 64], torch.float16),
+            (32, 128, [0, 1024], torch.float16),
+            (32, 128, [0, 15], torch.float16),
+            (32, 128, [0, 256, 512, 768, 1024], torch.float16),
+            (32, 128, [0, 15, 100, 300, 1200], torch.float16),
+            (64, 128, [0, 256, 500, 1000], torch.float16),
+            (32, 128, [0, 8192], torch.float16),
+            (32, 128, [0, 256, 500, 1000], torch.bfloat16),
+        ]
+    ],
+)
+@torch.inference_mode()
+def test_chunk_kda(
+    H: int,
+    D: int,
+    cu_seqlens: list[int],
+    dtype: torch.dtype,
+):
+    T = cu_seqlens[-1]
+    torch.manual_seed(42)
+    B = 1
+    cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE)
+    N = len(cu_seqlens) - 1
+
+    q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
+    k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
+    v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE)
+    g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to(
+        dtype
+    )
+    beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid()
+    h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
+
+    # Naive reference with l2norm_fwd (same kernel as chunk_kda)
+    ref_outputs = []
+    ref_states = []
+    for i in range(N):
+        s, e = cu_seqlens[i], cu_seqlens[i + 1]
+        q_i = l2norm_fwd(q[:, s:e].contiguous())
+        k_i = l2norm_fwd(k[:, s:e].contiguous())
+        o_i, ht_i = naive_recurrent_kda(
+            q_i,
+            k_i,
+            v[:, s:e],
+            g[:, s:e],
+            beta[:, s:e],
+            initial_state=h0[i],
+            output_final_state=True,
+        )
+        ref_outputs.append(o_i)
+        ref_states.append(ht_i)
+    ref_o = torch.cat(ref_outputs, dim=1)
+    ref_ht = torch.cat(ref_states, dim=0)
+
+    # h0 transposed to (V, K) layout for the kernel; naive uses (K, V)
+    tri_o, tri_ht = chunk_kda(
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        g=g.clone(),
+        beta=beta.clone(),
+        initial_state=h0.transpose(-1, -2).contiguous().clone(),
+        output_final_state=True,
+        cu_seqlens=cu_seqlens_t,
+        use_qk_l2norm_in_kernel=True,
+    )
+
+    assert not torch.isnan(tri_o).any(), "Triton output o contains NaN"
+    assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
+    assert_close("o", ref_o, tri_o, 0.005)
+    assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
@@ -1,11 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import importlib.util
+
 import pytest
 import torch

 from ....conftest import VllmRunner

+pytestmark = pytest.mark.skipif(
+    importlib.util.find_spec("terratorch") is None,
+    reason="terratorch unavailable while PyPI has `lightning` quarantined; see #41376",
+)
+

 def _run_test(
    vllm_runner: type[VllmRunner],
@@ -311,6 +311,9 @@ def _test_processing_correctness(
            baseline_processor,
            cached_processor,
            batch_idx,
+            hit_rate,
+            num_batches,
+            simplify_rate,
        )


@@ -320,6 +323,9 @@ def _test_processing_correctness_one(
    baseline_processor: BaseMultiModalProcessor,
    cached_processor: BaseMultiModalProcessor,
    batch_idx: int,
+    hit_rate: float,
+    num_batches: int,
+    simplify_rate: float,
 ):
    model_type = model_config.hf_config.model_type

@@ -343,7 +349,11 @@ def _test_processing_correctness_one(
        baseline_tokenized_result,
        cached_tokenized_result,
        ignore_mm_keys=ignore_mm_keys,
-        msg=f"Failed ({batch_idx=}, {token_prompt=}, {mm_data=})",
+        msg=(
+            f"Failed ({batch_idx=}, {hit_rate=}, "
+            f"{num_batches=}, {simplify_rate=}, "
+            f"{text_prompt=}, {token_prompt=}, {mm_data=})"
+        ),
    )

    if text_prompt is not None:
@@ -362,21 +372,33 @@ def _test_processing_correctness_one(
            baseline_text_result,
            cached_text_result,
            ignore_mm_keys=ignore_mm_keys,
-            msg=f"Failed ({batch_idx=}, {text_prompt=}, {mm_data=})",
+            msg=(
+                f"Failed ({batch_idx=}, {hit_rate=}, "
+                f"{num_batches=}, {simplify_rate=}, "
+                f"{text_prompt=}, {token_prompt=}, {mm_data=})"
+            ),
        )

        _assert_inputs_equal(
            baseline_text_result,
            baseline_tokenized_result,
            ignore_mm_keys=ignore_mm_keys,
-            msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})",
+            msg=(
+                f"Failed ({batch_idx=}, {hit_rate=}, "
+                f"{num_batches=}, {simplify_rate=}, "
+                f"{text_prompt=}, {token_prompt=}, {mm_data=})"
+            ),
        )

        _assert_inputs_equal(
            cached_text_result,
            cached_tokenized_result,
            ignore_mm_keys=ignore_mm_keys,
-            msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})",
+            msg=(
+                f"Failed ({batch_idx=}, {hit_rate=}, "
+                f"{num_batches=}, {simplify_rate=}, "
+                f"{text_prompt=}, {token_prompt=}, {mm_data=})"
+            ),
        )


@@ -408,6 +430,8 @@ def test_processing_correctness(
            "correctness test as is. Let's revisit adapting this "
            "test once more realtime models exist."
        )
+    if model_id == "CohereLabs/cohere-transcribe-03-2026":
+        pytest.skip("Fix later")

    _test_processing_correctness(
        model_id,
@@ -1396,9 +1396,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
    ),
    # [Encoder-decoder]
    "CohereAsrForConditionalGeneration": _HfExamplesInfo(
-        "CohereLabs/cohere-transcribe-03-2026",
-        trust_remote_code=True,
-        is_available_online=False,  # TODO (ekagra): revert after asr release
+        "CohereLabs/cohere-transcribe-03-2026", trust_remote_code=True
    ),
    "NemotronParseForConditionalGeneration": _HfExamplesInfo(
        "nvidia/NVIDIA-Nemotron-Parse-v1.1", trust_remote_code=True
@@ -109,6 +109,16 @@ def can_initialize(
            "which is not configured in test environment"
        )

+    if model_arch in ("PrithviGeoSpatialMAE", "Terratorch"):
+        import importlib.util
+
+        if importlib.util.find_spec("terratorch") is None:
+            pytest.skip(
+                "terratorch is not installed; "
+                "temporarily skipped while PyPI has `lightning` quarantined "
+                "(see #41376)"
+            )
+
    if model_arch in ["DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"]:
        from vllm.platforms import current_platform

@@ -36,6 +36,17 @@ def test_registry_imports(model_arch):
        check_max_version=False,
        check_version_reason="vllm",
    )
+
+    if model_arch in ("PrithviGeoSpatialMAE", "Terratorch"):
+        import importlib.util
+
+        if importlib.util.find_spec("terratorch") is None:
+            pytest.skip(
+                "terratorch is not installed; "
+                "temporarily skipped while PyPI has `lightning` quarantined "
+                "(see #41376)"
+            )
+
    # Ensure all model classes can be imported successfully
    model_cls = ModelRegistry._try_load_model_cls(model_arch)
    assert model_cls is not None
@@ -1,12 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import importlib.util
+
 import pytest
 import torch

 from tests.conftest import VllmRunner
 from tests.utils import create_new_process_for_each_test

+pytestmark = pytest.mark.skipif(
+    importlib.util.find_spec("terratorch") is None,
+    reason="terratorch unavailable while PyPI has `lightning` quarantined; see #41376",
+)
+

@create_new_process_for_each_test()  # Hangs otherwise
@pytest.mark.parametrize(
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import importlib.util
 import io

 import imagehash
@@ -11,6 +12,11 @@ from PIL import Image
 from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse

+pytestmark = pytest.mark.skipif(
+    importlib.util.find_spec("terratorch") is None,
+    reason="terratorch unavailable while PyPI has `lightning` quarantined; see #41376",
+)
+
 models_config = {
    "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11": {
        "image_url": "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff",  # noqa: E501
@@ -1,103 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the FxGraphCachePickler.dumps ValueError patch in env_override.py.
-
-Validates that _apply_fxgraphcache_pickle_patch correctly wraps a pickler's
-dumps method to convert ValueError into a bypass exception, without affecting
-other exception types or normal return values.
-"""
-
-import pytest
-
-from vllm.env_override import _apply_fxgraphcache_pickle_patch
-
-
-class _BypassStub(Exception):
-    """Stand-in for BypassFxGraphCache in unit tests."""
-
-
-class TestApplyFxgraphcachePicklePatch:
-    def test_valueerror_converted_to_bypass(self):
-        class Pickler:
-            def dumps(self, obj):
-                raise ValueError("can't serialize blocked layout")
-
-        _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
-
-        with pytest.raises(_BypassStub, match="Failed to pickle cache key"):
-            Pickler().dumps(object())
-
-    def test_original_valueerror_chained(self):
-        class Pickler:
-            def dumps(self, obj):
-                raise ValueError("bad tensor layout")
-
-        _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
-
-        with pytest.raises(_BypassStub) as exc_info:
-            Pickler().dumps(object())
-
-        cause = exc_info.value.__cause__
-        assert isinstance(cause, ValueError)
-        assert str(cause) == "bad tensor layout"
-
-    def test_non_valueerror_propagates(self):
-        class Pickler:
-            def dumps(self, obj):
-                raise TypeError("unexpected type")
-
-        _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
-
-        with pytest.raises(TypeError, match="unexpected type"):
-            Pickler().dumps(object())
-
-    def test_normal_return_preserved(self):
-        sentinel = b"serialized-graph-key"
-
-        class Pickler:
-            def dumps(self, obj):
-                return sentinel
-
-        _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
-
-        assert Pickler().dumps(object()) is sentinel
-
-    def test_idempotent(self):
-        class Pickler:
-            def dumps(self, obj):
-                return b"ok"
-
-        _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
-        first_dumps = Pickler.dumps
-        _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
-
-        assert Pickler.dumps is first_dumps
-
-    def test_sentinel_attribute_set(self):
-        class Pickler:
-            def dumps(self, obj):
-                return b"ok"
-
-        assert not hasattr(Pickler.dumps, "_vllm_patched")
-        assert not getattr(Pickler, "_vllm_fxgraph_dumps_patched", False)
-
-        _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
-
-        assert Pickler.dumps._vllm_patched is True  # type: ignore[attr-defined]
-        assert Pickler._vllm_fxgraph_dumps_patched is True  # type: ignore[attr-defined]
-
-
-def test_patch_applied_in_current_environment():
-    """Integration: verify patch state matches current torch version."""
-    from torch._inductor.codecache import FxGraphCachePickler
-
-    from vllm.utils.torch_utils import is_torch_equal_or_newer
-
-    should_be_patched = is_torch_equal_or_newer(
-        "2.10.0"
-    ) and not is_torch_equal_or_newer("2.11.0")
-
-    assert getattr(FxGraphCachePickler, "_vllm_fxgraph_dumps_patched", False) == (
-        should_be_patched
-    )
-    assert hasattr(FxGraphCachePickler.dumps, "_vllm_patched") == should_be_patched
@@ -432,3 +432,52 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
        )
        == 15
    )
+
+
+def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
+    """In forward steps, `get_num_blocks_to_allocate` must return exactly what
+    `allocate_new_blocks` will pull; otherwise `block_pool.get_new_blocks`
+    raises `ValueError: Cannot get N free blocks from the pool`.
+    """
+    block_size = 2
+    sliding_window = 8  # 4-block live window
+    cap = sliding_window // block_size
+
+    spec = SlidingWindowSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        sliding_window=sliding_window,
+    )
+    block_pool = BlockPool(
+        num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
+    )
+    manager = SlidingWindowManager(
+        spec,
+        block_pool=block_pool,
+        enable_caching=False,
+        kv_cache_group_id=0,
+        max_admission_blocks_per_request=cap,
+    )
+
+    request_id = "req"
+    total_computed = 0
+    # Walk through request forward steps. Check num_blocks returned by
+    # `get_num_blocks_to_allocate` matches what `allocate_new_blocks` pulls
+    for num_tokens in (4, 8, 12, 16):
+        predicted = manager.get_num_blocks_to_allocate(
+            request_id=request_id,
+            num_tokens=num_tokens,
+            new_computed_blocks=[],
+            total_computed_tokens=total_computed,
+            num_tokens_main_model=num_tokens,
+        )
+        new_blocks = manager.allocate_new_blocks(
+            request_id, num_tokens=num_tokens, num_tokens_main_model=num_tokens
+        )
+        assert predicted == len(new_blocks), (
+            f"num_tokens={num_tokens}: predictor returned {predicted} "
+            f"but allocator pulled {len(new_blocks)}"
+        )
+        total_computed = num_tokens
@@ -82,6 +82,13 @@ SPEC_DECODE_CONFIGS = [
        2,
        id="eagle-mla-deepseek",
    ),
+    pytest.param(
+        "Qwen/Qwen3.5-0.8B-Base",
+        "Qwen/Qwen3.5-0.8B-Base",
+        "mtp",
+        1,
+        id="mtp-qwen3_5-hybrid",
+    ),
 ]


@@ -104,6 +111,14 @@ def test_no_sync_with_spec_decode(
    from vllm import LLM, SamplingParams
    from vllm.distributed import cleanup_dist_env_and_memory

+    # Qwen3.5 is a VLM; without this, profile_run runs the ViT warmup
+    # and peaks well above the 18GB MIG slice used by one of the CI lanes.
+    # This test only exercises text generation, so the vision tower is
+    # never needed.
+    extra_kwargs: dict = {}
+    if "Qwen3.5" in model:
+        extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
+
    llm = LLM(
        model=model,
        max_model_len=256,
@@ -114,6 +129,7 @@ def test_no_sync_with_spec_decode(
        },
        enforce_eager=True,
        async_scheduling=True,
+        **extra_kwargs,
    )

    # Assert async scheduling is actually active before running inference.
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
 import random
 from collections.abc import Iterable
 from dataclasses import dataclass
@@ -723,8 +724,13 @@ def test_eagle_correctness_heavy(
    [
        (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5),  # ref: 65%-70%
        (("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1), False, 0.0),  # dummy model
+        (
+            ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1),
+            False,
+            0.20,
+        ),  # hybrid + MTP, ref: ~34%-35%
    ],
-    ids=["mimo", "deepseek"],
+    ids=["mimo", "deepseek", "qwen3_5-hybrid"],
 )
@single_gpu_only
@large_gpu_mark(min_gb=20)
@@ -750,13 +756,29 @@ def test_mtp_correctness(
        method, model_name, tp_size = model_setup
        _skip_if_insufficient_gpus_for_tp(tp_size)

+        if "Qwen3.5" in model_name and os.environ.get("VLLM_USE_V2_MODEL_RUNNER"):
+            pytest.skip(
+                "Model Runner V2 does not yet support hybrid models "
+                "(Qwen3.5 mixes Mamba-style GDN with attention layers)."
+            )
+
        attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
+
+        # Qwen3.5 is a VLM; without this, profile_run runs the ViT warmup
+        # and peaks well above the 18GB MIG slice used by one of the CI
+        # lanes. This test only exercises text generation, so the vision
+        # tower is never needed.
+        extra_kwargs: dict[str, Any] = {}
+        if "Qwen3.5" in model_name:
+            extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
+
        ref_llm = LLM(
            model=model_name,
            max_model_len=2048,
            tensor_parallel_size=tp_size,
            trust_remote_code=True,
            attention_backend=attn_backend,
+            **extra_kwargs,
        )
        ref_outputs = ref_llm.chat(test_prompts, sampling_config)
        evaluate_llm_for_gsm8k(
@@ -777,6 +799,7 @@ def test_mtp_correctness(
            },
            max_model_len=2048,
            attention_backend=attn_backend,
+            **extra_kwargs,
        )
        # MTP supports async scheduling; assert it is active by default.
        assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
@@ -15,7 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
    OffloadingConnectorScheduler,
 )
 from vllm.v1.core.kv_cache_utils import BlockHash
-from vllm.v1.kv_offload.abstract import (
+from vllm.v1.kv_offload.base import (
    OffloadingEvent,
    OffloadingManager,
    ReqContext,
@@ -20,7 +20,7 @@ from vllm.v1.kv_cache_interface import (
    MLAAttentionSpec,
    UniformTypeKVCacheSpecs,
 )
-from vllm.v1.kv_offload.spec import (
+from vllm.v1.kv_offload.base import (
    CanonicalKVCacheRef,
    CanonicalKVCaches,
    OffloadingSpec,
@@ -37,15 +37,15 @@ from vllm.v1.kv_cache_interface import (
    KVCacheConfig,
    KVCacheGroupSpec,
 )
-from vllm.v1.kv_offload.abstract import (
+from vllm.v1.kv_offload.base import (
+    GPULoadStoreSpec,
    LoadStoreSpec,
    OffloadingManager,
+    OffloadingSpec,
    OffloadKey,
    PrepareStoreOutput,
    make_offload_key,
 )
-from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
-from vllm.v1.kv_offload.spec import OffloadingSpec
 from vllm.v1.kv_offload.worker.worker import (
    OffloadingHandler,
    TransferResult,
@@ -0,0 +1,915 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for bi-directional KV cache transfer between P and D nodes.
+
+Tests cover the new behaviors added by the bi-directional KV transfer PR:
+1. P-node scheduler lifecycle: P pulls KV from D using remote_block_ids,
+   eliminating redundant prefill computation in multi-turn conversations.
+2. P-node metadata: NixlConnectorMetadata correctly populates recv metadata
+   when P pulls KV from D (do_remote_decode=True + remote_block_ids).
+3. P-node worker: start_load_kv processes reqs_to_recv for KV pull from D.
+4. D-node request_finished: returns kv_transfer_params with remote_block_ids
+   and remote_num_tokens so P can pull KV in future turns.
+5. Edge cases:
+   - No double read after reschedule (_remote_blocks_processed flag)
+   - remote_num_tokens bounded by block capacity (num_computed_tokens)
+   - kv_recompute_threshold skips small transfers
+   - P-node holds blocks for D after finishing
+   - Cache MISS first turn falls back to local prefill
+   - Partial remote coverage: P pulls partial, computes the rest
+   - _remote_blocks_processed flag persists across reschedules
+
+P-node flags: do_remote_prefill=False (prefill locally),
+do_remote_decode=True (don't decode locally, send KV to D).
+P pulls KV from D when remote_block_ids is not None and
+external tokens > 0.
+"""
+
+import copy
+import time
+from unittest.mock import patch
+
+import pytest
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import (
+    NixlConnector,
+    NixlConnectorMetadata,
+)
+from vllm.forward_context import ForwardContext
+from vllm.v1.outputs import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    KVConnectorOutput,
+)
+from vllm.v1.request import RequestStatus
+
+from .test_nixl_connector import FakeNixlConnectorWorker, FakeNixlWrapper
+from .utils import (
+    assert_scheduler_empty,
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+    make_kv_cache_config,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+# Common extra config for all bi-directional KV transfer tests.
+BIDIR_KV_EXTRA_CONFIG = {"bidirectional_kv_xfer": True, "kv_recompute_threshold": 0}
+
+
+# Helpers
+
+
+def _make_p_node_turn2_request(
+    request_id, block_size, num_tokens, num_remote_blocks=3, remote_num_tokens=None
+):
+    """Create a P-node Turn 2 request with remote_block_ids from D."""
+    request = create_request(
+        request_id=request_id,
+        block_size=block_size,
+        num_tokens=num_tokens,
+        do_remote_decode=True,
+    )
+    if remote_num_tokens is None:
+        remote_num_tokens = num_remote_blocks * block_size
+    request.kv_transfer_params["remote_block_ids"] = [list(range(num_remote_blocks))]
+    request.kv_transfer_params["remote_num_tokens"] = remote_num_tokens
+    request.kv_transfer_params["remote_engine_id"] = "decode-engine"
+    request.kv_transfer_params["remote_request_id"] = f"decode-{request_id}"
+    request.kv_transfer_params["remote_host"] = "decode-host"
+    request.kv_transfer_params["remote_port"] = 5678
+    return request
+
+
+def _make_connector_with_fake_worker(
+    hand_shake_latency=0, cycles_before_done=0, do_handshake=True
+):
+    """Create a NixlConnector with FakeNixlConnectorWorker."""
+    vllm_config = create_vllm_config()
+    kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
+    connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
+    connector.connector_worker = FakeNixlConnectorWorker(
+        vllm_config,
+        connector.engine_id,
+        hand_shake_latency=hand_shake_latency,
+        kv_cache_config=kv_cache_config,
+    )
+    worker = connector.connector_worker
+    assert isinstance(worker.nixl_wrapper, FakeNixlWrapper)
+    worker.nixl_wrapper.set_cycles_before_xfer_done(cycles_before_done)
+    worker.kv_cache_layout = "HND"
+    if do_handshake:
+        remote_agents = worker._nixl_handshake(
+            host="localhost",
+            port=1234,
+            remote_tp_size=1,
+            expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
+        )
+        worker._remote_agents[FakeNixlConnectorWorker.REMOTE_ENGINE_ID] = remote_agents
+    return connector, worker
+
+
+def _make_p_node_recv_metadata(request_id, local_blocks, remote_blocks):
+    """Build NixlConnectorMetadata for P-node pulling KV from D."""
+    meta = NixlConnectorMetadata()
+    meta.add_new_req_to_recv(
+        request_id=request_id,
+        local_block_ids=(local_blocks,),
+        kv_transfer_params={
+            "do_remote_prefill": False,
+            "do_remote_decode": True,
+            "remote_block_ids": (remote_blocks,),
+            "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
+            "remote_request_id": f"decode-{request_id}",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "remote_tp_size": 1,
+        },
+    )
+    return meta
+
+
+def _do_load_kv(connector, metadata):
+    """Bind metadata and call start_load_kv."""
+    connector.bind_connector_metadata(metadata)
+    ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
+    connector.start_load_kv(ctx)
+
+
+# 1. P-node scheduler lifecycle tests
+
+
+def test_multiturn_lifecycle():
+    """Full two-turn lifecycle on the P node:
+    Turn 1: P prefills locally (do_remote_prefill=False), sends KV to D
+    (do_remote_decode=True). Finishes LENGTH_CAPPED with remote_block_ids.
+    Turn 2: P receives remote_block_ids from D. P pulls KV from D because
+    remote_block_ids is not None and external tokens > 0. Computes only
+    new tokens, finishes LENGTH_CAPPED."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+
+    t1 = create_request(
+        request_id=100, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
+    )
+    scheduler.add_request(t1)
+    t1_id = t1.request_id
+    so = scheduler.schedule()
+    mro = create_model_runner_output(reqs=[t1])
+    eco = scheduler.update_from_output(so, mro)
+    assert t1.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    kv = eco[0].outputs[0].kv_transfer_params
+    assert kv and sum(len(g) for g in kv["remote_block_ids"]) > 0
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+
+    t2 = _make_p_node_turn2_request(200, BS, int(BS * 2.5))
+    scheduler.add_request(t2)
+    t2_id = t2.request_id
+    so = scheduler.schedule()
+    assert t2.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_recving={t2_id})
+    scheduler.update_from_output(so, mro)
+    so = scheduler.schedule()
+    mro = create_model_runner_output(reqs=[t2])
+    scheduler.update_from_output(so, mro)
+    assert t2.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={t1_id, t2_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_first_turn_no_remote_blocks():
+    """First turn: P has no remote_block_ids from D yet.
+    Standard local prefill, returns kv_transfer_params for future turns."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=3, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    assert req.status != RequestStatus.WAITING_FOR_REMOTE_KVS
+    mro = create_model_runner_output(reqs=[req])
+    eco = scheduler.update_from_output(so, mro)
+    assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    assert eco[0].outputs[0].kv_transfer_params is not None
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_abort_p_side_during_send():
+    """P-side do_remote_decode=True: blocks held until finished_sending."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=42, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    assert req_id in scheduler.requests
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    assert req_id in scheduler.requests
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_abort_p_side_non_length_capped():
+    """P-side abort with non-LENGTH_CAPPED → immediate block free."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=44, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
+    )
+    req.sampling_params.max_tokens = 100
+    req.max_tokens = 100
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    scheduler.finish_requests([req_id], RequestStatus.FINISHED_ABORTED)
+    conn = scheduler.connector.connector_scheduler
+    assert req_id in conn._reqs_not_processed
+    assert req_id not in scheduler.requests
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    assert_scheduler_empty(scheduler)
+
+
+def test_remote_blocks_exceed_prompt_tokens():
+    """D provides more remote tokens than P's prompt needs.
+    P caps external tokens to prompt length."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    NUM_TOKENS = int(BS * 2.5)
+    req = _make_p_node_turn2_request(
+        300, BS, NUM_TOKENS, num_remote_blocks=5, remote_num_tokens=5 * BS
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert req.num_computed_tokens == NUM_TOKENS
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
+    scheduler.update_from_output(so, mro)
+    so = scheduler.schedule()
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_p_node_pulls_partial_last_block_from_d():
+    """D sends remote_block_ids with partially filled last block.
+    remote_num_tokens < len(remote_block_ids) * block_size.
+    P pulls only remote_num_tokens worth of external tokens."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    num_remote_blocks = 3
+    remote_num_tokens = int(BS * 2.5)
+    assert remote_num_tokens < num_remote_blocks * BS
+    NUM_TOKENS = int(BS * 3.5)
+    req = _make_p_node_turn2_request(
+        400,
+        BS,
+        NUM_TOKENS,
+        num_remote_blocks=num_remote_blocks,
+        remote_num_tokens=remote_num_tokens,
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
+    scheduler.update_from_output(so, mro)
+    so = scheduler.schedule()
+    assert len(scheduler.running) == 1
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+# 2. P-node metadata tests
+
+
+def test_add_new_req_to_recv_populates_remote_meta():
+    """add_new_req_to_recv correctly populates RemoteMeta for P-node
+    bi-directional KV pull from D."""
+    meta = NixlConnectorMetadata()
+    kv_params = {
+        "remote_block_ids": [[0, 1, 2]],
+        "remote_engine_id": "decode-engine",
+        "remote_request_id": "decode-req-123",
+        "remote_host": "decode-host",
+        "remote_port": 5678,
+    }
+    local_block_ids = ([10, 11, 12],)
+    meta.add_new_req_to_recv(
+        request_id="test-req",
+        local_block_ids=local_block_ids,
+        kv_transfer_params=kv_params,
+    )
+    assert "test-req" in meta.reqs_to_recv
+    rm = meta.reqs_to_recv["test-req"]
+    assert rm.remote is not None
+    assert rm.remote.block_ids == kv_params["remote_block_ids"]
+    assert rm.remote.engine_id == "decode-engine"
+    assert rm.remote.request_id == "decode-req-123"
+    assert rm.remote.host == "decode-host"
+    assert rm.remote.port == 5678
+    assert rm.local_block_ids == local_block_ids
+
+
+def test_build_connector_meta_recv_entries():
+    """P-node scheduler: do_remote_decode=True + remote_block_ids →
+    _reqs_need_recv populated, build_connector_meta produces reqs_to_recv."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = _make_p_node_turn2_request(1, BS, int(BS * 2.5))
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    meta = so.kv_connector_metadata
+    assert isinstance(meta, NixlConnectorMetadata)
+    assert req_id in meta.reqs_to_recv
+    rm = meta.reqs_to_recv[req_id]
+    assert rm.remote is not None
+    assert rm.remote.engine_id == "decode-engine"
+
+
+def test_build_connector_meta_clears_reqs_need_recv():
+    """After build_connector_meta, _reqs_need_recv is cleared."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = _make_p_node_turn2_request(2, BS, int(BS * 2.5))
+    scheduler.add_request(req)
+    conn = scheduler.connector.connector_scheduler
+    scheduler.schedule()
+    assert len(conn._reqs_need_recv) == 0
+
+
+def test_build_connector_meta_multiple_requests():
+    """Multiple P-node requests all included in reqs_to_recv."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    reqs = [_make_p_node_turn2_request(10 + i, BS, int(BS * 2.5)) for i in range(3)]
+    for r in reqs:
+        scheduler.add_request(r)
+    so = scheduler.schedule()
+    meta = so.kv_connector_metadata
+    assert isinstance(meta, NixlConnectorMetadata)
+    assert len(meta.reqs_to_recv) == 3
+    for r in reqs:
+        assert r.request_id in meta.reqs_to_recv
+
+
+# 3. P-node worker tests (FakeNixlWrapper)
+
+
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
+    FakeNixlWrapper,
+)
+def test_p_node_pull_kv_from_d(dist_init):
+    """P node pulls KV from D via start_load_kv with reqs_to_recv."""
+    connector, worker = _make_connector_with_fake_worker()
+    meta = _make_p_node_recv_metadata("req-p1", [10, 11, 12], [20, 21, 22])
+    _do_load_kv(connector, meta)
+    assert "req-p1" in worker._recving_metadata
+    _, done_recving = connector.get_finished(finished_req_ids=set())
+    assert "req-p1" in done_recving
+
+
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
+    FakeNixlWrapper,
+)
+def test_p_node_pull_then_send_kv(dist_init):
+    """Full P-node bi-directional: pull KV from D → prefill →
+    send KV back to D via notification."""
+    connector, worker = _make_connector_with_fake_worker()
+    meta = _make_p_node_recv_metadata("req-p2", [10, 11], [20, 21])
+    _do_load_kv(connector, meta)
+    _, done_recving = connector.get_finished(finished_req_ids=set())
+    assert "req-p2" in done_recving
+    worker._reqs_to_send["req-p2"] = time.perf_counter() + 60
+    worker._reqs_to_process.add("req-p2")
+    notif = f"req-p2:{worker.world_size}".encode()
+    orig = worker.nixl_wrapper.get_new_notifs
+    worker.nixl_wrapper.get_new_notifs = lambda: {"agent": [notif]}
+    done_sending, _ = connector.get_finished(finished_req_ids=set())
+    assert "req-p2" in done_sending
+    worker.nixl_wrapper.get_new_notifs = orig
+
+
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
+    FakeNixlWrapper,
+)
+def test_p_node_deferred_pull_on_no_handshake(dist_init):
+    """P defers KV pull when no prior handshake exists."""
+    connector, worker = _make_connector_with_fake_worker(
+        hand_shake_latency=0, do_handshake=False
+    )
+    meta = _make_p_node_recv_metadata("req-p3", [10, 11], [20, 21])
+    _do_load_kv(connector, meta)
+    assert "req-p3" in worker._recving_metadata
+    timeout = 3.0
+    start = time.perf_counter()
+    while time.perf_counter() - start < timeout:
+        connector.bind_connector_metadata(NixlConnectorMetadata())
+        ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
+        connector.start_load_kv(ctx)
+        _, done = connector.get_finished(finished_req_ids=set())
+        if "req-p3" in done:
+            return
+        time.sleep(0.2)
+    raise AssertionError("Transfer did not complete after async handshake")
+
+
+# 4. D-node request_finished returns kv_transfer_params (new behavior)
+
+
+def test_d_node_request_finished_returns_kv_params():
+    """D-node request_finished returns kv_transfer_params with
+    do_remote_decode=True, remote_block_ids, remote_num_tokens
+    for P to pull. These params go directly to P node."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=1, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[], finished_recving={req_id})
+    )
+    so = scheduler.schedule()
+    eco = scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[req], use_eos=True)
+    )
+    assert req.status == RequestStatus.FINISHED_STOPPED
+    kv = eco[0].outputs[0].kv_transfer_params
+    assert kv is not None
+    assert kv["do_remote_decode"] is True
+    assert kv["do_remote_prefill"] is False
+    assert "remote_block_ids" in kv
+    assert "remote_num_tokens" in kv
+    assert kv["remote_num_tokens"] > 0
+
+
+def test_d_node_request_finished_delays_block_free():
+    """D-node holds blocks (delay_free=True) until P reads them."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=2, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[], finished_recving={req_id})
+    )
+    so = scheduler.schedule()
+    scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[req], use_eos=True)
+    )
+    assert req_id in scheduler.requests
+    conn = scheduler.connector.connector_scheduler
+    assert req_id in conn._reqs_need_send
+
+
+def test_d_node_request_finished_remote_num_tokens():
+    """D-node kv_transfer_params includes correct remote_num_tokens."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=3, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[], finished_recving={req_id})
+    )
+    so = scheduler.schedule()
+    eco = scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[req], use_eos=True)
+    )
+    kv = eco[0].outputs[0].kv_transfer_params
+    assert kv["remote_num_tokens"] > 0
+    assert sum(len(g) for g in kv["remote_block_ids"]) > 0
+
+
+def test_d_node_partial_last_block_remote_num_tokens():
+    """D-node: remote_num_tokens < len(remote_block_ids) * block_size
+    when last block is partially filled."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=5, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[], finished_recving={req_id})
+    )
+    so = scheduler.schedule()
+    eco = scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[req], use_eos=True)
+    )
+    kv = eco[0].outputs[0].kv_transfer_params
+    total_blocks = sum(len(g) for g in kv["remote_block_ids"])
+    assert total_blocks == 3
+    assert kv["remote_num_tokens"] < total_blocks * BS
+    assert kv["remote_num_tokens"] > 0
+
+
+# 5. Edge case tests
+
+
+def test_no_double_read_blocks_after_reschedule():
+    """Edge case 1: update_state_after_alloc called twice for the same
+    bidirectional request (once on initial schedule, once after
+    WAITING_FOR_REMOTE_KVS → reschedule). The _remote_blocks_processed
+    flag must prevent the request from being added to _reqs_need_recv
+    twice, which would cause P to read D's blocks twice."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = _make_p_node_turn2_request(500, BS, int(BS * 2.5))
+    scheduler.add_request(req)
+    req_id = req.request_id
+    conn = scheduler.connector.connector_scheduler
+
+    # First schedule: request enters WAITING_FOR_REMOTE_KVS,
+    # _reqs_need_recv populated then cleared by build_connector_meta.
+    so = scheduler.schedule()
+    assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    meta = so.kv_connector_metadata
+    assert isinstance(meta, NixlConnectorMetadata)
+    assert req_id in meta.reqs_to_recv
+    # _reqs_need_recv should be cleared after build_connector_meta
+    assert len(conn._reqs_need_recv) == 0
+
+    # Simulate recv completion
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
+    scheduler.update_from_output(so, mro)
+
+    # Second schedule after recv: update_state_after_alloc called again.
+    # The _remote_blocks_processed flag should prevent re-entry.
+    so = scheduler.schedule()
+    meta2 = so.kv_connector_metadata
+    assert isinstance(meta2, NixlConnectorMetadata)
+    # Must NOT be in reqs_to_recv again
+    assert req_id not in meta2.reqs_to_recv
+
+    # Clean up
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_remote_num_tokens_bounded_by_blocks():
+    """Edge case 2: D-node request_finished must return
+    remote_num_tokens <= len(remote_block_ids) * block_size.
+    request.num_tokens includes the last sampled token which has no KV
+    in the cache, so remote_num_tokens must use num_computed_tokens."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=501,
+        block_size=BS,
+        num_tokens=int(BS * 2.5),
+        do_remote_prefill=True,
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[], finished_recving={req_id})
+    )
+    so = scheduler.schedule()
+    eco = scheduler.update_from_output(
+        so, create_model_runner_output(reqs=[req], use_eos=True)
+    )
+    kv = eco[0].outputs[0].kv_transfer_params
+    assert kv is not None
+    total_blocks = sum(len(g) for g in kv["remote_block_ids"])
+    max_tokens_in_blocks = total_blocks * BS
+    assert kv["remote_num_tokens"] <= max_tokens_in_blocks, (
+        f"remote_num_tokens ({kv['remote_num_tokens']}) exceeds "
+        f"block capacity ({max_tokens_in_blocks})"
+    )
+    assert kv["remote_num_tokens"] > 0
+
+
+def test_kv_recompute_threshold_skips_small_transfer():
+    """Edge case 3: When remote tokens are below kv_recompute_threshold,
+    P should skip the remote pull and compute locally instead of
+    entering WAITING_FOR_REMOTE_KVS."""
+    threshold = 256
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config={
+            "bidirectional_kv_xfer": True,
+            "kv_recompute_threshold": threshold,
+        },
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+
+    # Create request where remote tokens (48) < threshold (256)
+    req = _make_p_node_turn2_request(
+        502,
+        BS,
+        int(BS * 2.5),
+        num_remote_blocks=3,
+        remote_num_tokens=3 * BS,
+    )
+    scheduler.add_request(req)
+    so = scheduler.schedule()
+    # Should NOT enter WAITING_FOR_REMOTE_KVS — threshold not met
+    assert req.status != RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert req.status == RequestStatus.RUNNING
+
+    # Clean up
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req.request_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_p_node_finished_holds_blocks_for_d():
+    """Edge case 4: P-node finishes with FINISHED_LENGTH_CAPPED and
+    do_remote_decode=True. P must hold blocks (delay_free=True) and
+    return kv_transfer_params with do_remote_prefill=True so D can
+    read P's blocks."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=503,
+        block_size=BS,
+        num_tokens=int(BS * 2.5),
+        do_remote_decode=True,
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    mro = create_model_runner_output(reqs=[req])
+    eco = scheduler.update_from_output(so, mro)
+    assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    kv = eco[0].outputs[0].kv_transfer_params
+    assert kv is not None
+    # P-node finished: should tell D to pull (do_remote_prefill=True)
+    assert kv["do_remote_prefill"] is True
+    assert kv["do_remote_decode"] is False
+    assert "remote_block_ids" in kv
+    assert sum(len(g) for g in kv["remote_block_ids"]) > 0
+    # Blocks should be held (request still tracked)
+    assert req_id in scheduler.requests
+
+    # Clean up: simulate D reading and notifying
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_cache_miss_first_turn_no_remote_pull():
+    """Edge case 5: First turn with do_remote_decode=True but no
+    remote_block_ids (cache MISS). P should prefill locally with
+    num_external_tokens=0 and not enter WAITING_FOR_REMOTE_KVS."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=504,
+        block_size=BS,
+        num_tokens=int(BS * 2.5),
+        do_remote_decode=True,
+    )
+    # No remote_block_ids set — this is a cache MISS
+    assert req.kv_transfer_params.get("remote_block_ids") is None
+    scheduler.add_request(req)
+    so = scheduler.schedule()
+    # Should NOT wait for remote KVs
+    assert req.status != RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert req.status == RequestStatus.RUNNING
+
+    # Clean up
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req.request_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_partial_remote_tokens_less_than_prompt():
+    """Edge case 6: D's remote_num_tokens covers only part of P's
+    prompt. P should pull remote_num_tokens worth of external tokens
+    and compute the rest locally."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    NUM_TOKENS = int(BS * 4.5)  # 72 tokens
+    # D provides only 2 blocks (32 tokens) out of 72
+    req = _make_p_node_turn2_request(
+        505,
+        BS,
+        NUM_TOKENS,
+        num_remote_blocks=2,
+        remote_num_tokens=2 * BS,
+    )
+    scheduler.add_request(req)
+    req_id = req.request_id
+    so = scheduler.schedule()
+    assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    # num_computed_tokens should reflect the external tokens pulled
+    # (capped to remote_num_tokens, not full prompt)
+    assert req.num_computed_tokens < NUM_TOKENS
+
+    # Complete the transfer and finish
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
+    scheduler.update_from_output(so, mro)
+    so = scheduler.schedule()
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
+
+
+def test_remote_blocks_processed_flag_persists():
+    """Edge case 7: After recv completes and request is rescheduled,
+    the _remote_blocks_processed flag in kv_transfer_params prevents
+    the bidirectional path from re-entering _reqs_need_recv."""
+    vllm_config = create_vllm_config(
+        kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
+    )
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = _make_p_node_turn2_request(506, BS, int(BS * 2.5))
+    scheduler.add_request(req)
+    req_id = req.request_id
+    conn = scheduler.connector.connector_scheduler
+
+    # First schedule → WAITING_FOR_REMOTE_KVS
+    so = scheduler.schedule()
+    assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+
+    # Recv completes
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
+    scheduler.update_from_output(so, mro)
+
+    # Verify the flag is set
+    assert req.kv_transfer_params.get("_remote_blocks_processed") is True
+
+    # Next schedule: update_state_after_alloc is called again.
+    # _reqs_need_recv must NOT contain this request.
+    so = scheduler.schedule()
+    assert req_id not in conn._reqs_need_recv
+    meta = so.kv_connector_metadata
+    assert isinstance(meta, NixlConnectorMetadata)
+    assert req_id not in meta.reqs_to_recv
+
+    # Clean up
+    mro = create_model_runner_output(reqs=[req])
+    scheduler.update_from_output(so, mro)
+    so = scheduler.schedule()
+    scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
+    so = scheduler.schedule()
+    mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
+    scheduler.update_from_output(so, mro)
+    assert_scheduler_empty(scheduler)
@@ -8,13 +8,18 @@ from typing import Any
 from unittest.mock import MagicMock

 import pytest
+import torch

 from tests.v1.kv_connector.unit.utils import create_vllm_config
 from vllm import LLM, SamplingParams
 from vllm.config import KVTransferConfig
 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
-from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    SupportsHMA,
+    supports_hma,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
    MultiConnector,
@@ -83,8 +88,43 @@ class MockConnector(KVConnectorBase_V1):
        pass


-# Register the mock connector
+class MockHMAConnector(KVConnectorBase_V1, SupportsHMA):
+    """Mock connector that supports HMA for testing."""
+
+    def __new__(cls, *args, **kwargs):
+        mock = MagicMock(spec_set=cls)
+        return mock
+
+    def start_load_kv(self, forward_context, **kwargs):
+        pass
+
+    def wait_for_layer_load(self, layer_name):
+        pass
+
+    def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
+        pass
+
+    def wait_for_save(self):
+        pass
+
+    def build_connector_meta(self, scheduler_output):
+        return None
+
+    def get_num_new_matched_tokens(self, request, num_computed_tokens):
+        return (0, False)
+
+    def update_state_after_alloc(self, request, blocks, num_tokens) -> None:
+        pass
+
+    def request_finished_all_groups(self, request, block_ids):
+        return (False, None)
+
+
+# Register mock connectors
 KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
+KVConnectorFactory.register_connector(
+    "MockHMAConnector", __name__, MockHMAConnector.__name__
+)


@pytest.fixture
@@ -920,3 +960,133 @@ def test_multi_connector_worker_metadata(mc):
    mc.update_connector_output(kv_connector_output)
    assert_update_connector_output_called(mc)
    assert kv_connector_output.kv_connector_worker_meta == mc_worker_meta_01a_01b
+
+
+def _make_multi_connector(connector_names: list[str]) -> MultiConnector:
+    """Build a MultiConnector wrapping the given registered connectors."""
+    vllm_config = create_vllm_config()
+    connectors = [
+        {
+            "kv_connector": name,
+            "kv_role": "kv_both",
+            "kv_connector_module_path": "tests.v1.kv_connector.unit.test_multi_connector",  # noqa: E501
+        }
+        for name in connector_names
+    ]
+    vllm_config.kv_transfer_config = KVTransferConfig(
+        kv_connector="MultiConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={"connectors": connectors},
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=0,
+        kv_cache_tensors=[],
+        kv_cache_groups=[],
+    )
+    return MultiConnector(
+        vllm_config=vllm_config,
+        role=KVConnectorRole.WORKER,
+        kv_cache_config=kv_cache_config,
+    )
+
+
+def test_multi_connector_hma_opt_in():
+    """
+    MultiConnector currently assumes HMA is opt-in: it needs
+    --no-disable-hybrid-kv-cache-manager to be enabled.
+
+    At runtime, _all_support_hma is True only when every sub-connector
+    implements SupportsHMA. Test all combinations of HMA / non-HMA
+    sub-connectors.
+    """
+
+    assert supports_hma(MultiConnector)
+
+    # -- All non-HMA connectors => _all_support_hma is False --
+    mc_none = _make_multi_connector(["MockConnector", "MockConnector"])
+    assert not supports_hma(mc_none._connectors[0])
+    assert not supports_hma(mc_none._connectors[1])
+    assert mc_none._all_support_hma is False
+
+    # -- All HMA connectors => _all_support_hma is True --
+    mc_all = _make_multi_connector(["MockHMAConnector", "MockHMAConnector"])
+    assert supports_hma(mc_all._connectors[0])
+    assert supports_hma(mc_all._connectors[1])
+    assert mc_all._all_support_hma is True
+
+    # -- Mixed: first HMA, second non-HMA => _all_support_hma is False --
+    mc_mixed1 = _make_multi_connector(["MockHMAConnector", "MockConnector"])
+    assert supports_hma(mc_mixed1._connectors[0])
+    assert not supports_hma(mc_mixed1._connectors[1])
+    assert mc_mixed1._all_support_hma is False
+
+    # -- Mixed: first non-HMA, second HMA => _all_support_hma is False --
+    mc_mixed2 = _make_multi_connector(["MockConnector", "MockHMAConnector"])
+    assert not supports_hma(mc_mixed2._connectors[0])
+    assert supports_hma(mc_mixed2._connectors[1])
+    assert mc_mixed2._all_support_hma is False
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="Requires GPU to instantiate LLM"
+)
+def test_multi_connector_mixed_hma_disables_hybrid_kv_cache(monkeypatch):
+    """
+    When MultiConnector wraps a mix of HMA (NixlConnector) and non-HMA
+    (MockConnector) sub-connectors, verify that:
+    1. The scheduler's MultiConnector has _all_support_hma == False.
+    2. vLLM auto-disables the hybrid KV cache manager (no preference expressed by user)
+    """
+    from unittest.mock import patch
+
+    from tests.v1.kv_connector.unit.test_nixl_connector import FakeNixlWrapper
+
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+    kv_transfer_config = KVTransferConfig(
+        kv_connector="MultiConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={
+            "connectors": [
+                {
+                    "kv_connector": "NixlConnector",
+                    "kv_role": "kv_both",
+                },
+                {
+                    "kv_connector": "MockConnector",
+                    "kv_role": "kv_both",
+                    "kv_connector_module_path": (
+                        "tests.v1.kv_connector.unit.test_multi_connector"
+                    ),
+                },
+            ],
+        },
+    )
+
+    with patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
+        FakeNixlWrapper,
+    ):
+        llm = LLM(
+            model="Qwen/Qwen3-0.6B",
+            enforce_eager=True,
+            gpu_memory_utilization=0.3,
+            max_model_len=128,
+            max_num_seqs=1,
+            max_num_batched_tokens=128,
+            kv_transfer_config=kv_transfer_config,
+        )
+        try:
+            # HMA should be auto-disabled when user has not expressed a preference.
+            assert (
+                llm.llm_engine.vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
+                is True
+            )
+            # The scheduler-side MultiConnector should detect the mixed
+            # HMA support among its sub-connectors.
+            scheduler = llm.llm_engine.engine_core.engine_core.scheduler
+            mc = scheduler.connector
+            assert isinstance(mc, MultiConnector)
+            assert mc._all_support_hma is False
+        finally:
+            llm.llm_engine.engine_core.shutdown()
@@ -9,14 +9,15 @@ import torch

 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
-from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
-from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
-from vllm.v1.kv_offload.spec import (
+from vllm.v1.kv_offload.base import (
    CanonicalKVCacheRef,
    CanonicalKVCaches,
    CanonicalKVCacheTensor,
+    GPULoadStoreSpec,
 )
-from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
+from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
+from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
+from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion

 NUM_GPU_BLOCKS = [64]
 NUM_CPU_BLOCKS = [256]
@@ -6,7 +6,7 @@ from dataclasses import dataclass
 import numpy as np
 import pytest

-from vllm.v1.kv_offload.abstract import (
+from vllm.v1.kv_offload.base import (
    LoadStoreSpec,
    OffloadingEvent,
    OffloadKey,
@@ -14,9 +14,9 @@ from vllm.v1.kv_offload.abstract import (
    ReqContext,
    make_offload_key,
 )
+from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
-from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
 from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager


@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from vllm.v1.kv_offload.abstract import LoadStoreSpec
+from vllm.v1.kv_offload.base import LoadStoreSpec
 from vllm.v1.kv_offload.worker.worker import (
    OffloadingHandler,
    OffloadingWorker,
@@ -203,47 +203,6 @@ def is_compile_cache_enabled(
    )


-def _patch_standalone_compile_atomic_save() -> None:
-    """Backport of pytorch/pytorch#162432 for torch < 2.10.0.
-
-    Patches CompiledArtifact.save() to use write_atomic for binary format,
-    preventing corrupt cache files when multiple processes compile
-    concurrently.
-    """
-    from torch._inductor.codecache import write_atomic
-    from torch._inductor.standalone_compile import CompiledArtifact as cls
-
-    if getattr(cls.save, "_vllm_patched", False):
-        return
-
-    original_save = cls.save
-
-    def _save(
-        self: Any, *, path: str, format: Literal["binary", "unpacked"] = "binary"
-    ) -> None:
-        if format != "binary":
-            return original_save(self, path=path, format=format)
-        from torch._dynamo.utils import dynamo_timed
-        from torch._inductor.codecache import torch_key
-        from torch.utils._appending_byte_serializer import BytesWriter
-
-        with dynamo_timed("CompiledArtifact.save"):
-            assert self._artifacts is not None
-            artifact_bytes, cache_info = self._artifacts
-            assert len(cache_info.aot_autograd_artifacts) == 1, cache_info
-            key = cache_info.aot_autograd_artifacts[0]
-            assert not os.path.isdir(path)
-            writer = BytesWriter()
-            writer.write_bytes(torch_key())
-            writer.write_str(key)
-            writer.write_bytes(artifact_bytes)
-            write_atomic(path, writer.to_bytes())
-
-    _save._vllm_patched = True  # type: ignore[attr-defined]
-    cls.save = _save  # type: ignore[assignment]
-    logger.debug("Patched %s.save for atomic writes (torch < 2.10)", cls.__name__)
-
-
 class InductorStandaloneAdaptor(CompilerInterface):
    """
    The adaptor for the Inductor compiler.
@@ -257,8 +216,6 @@ class InductorStandaloneAdaptor(CompilerInterface):
    name = "inductor_standalone"

    def __init__(self, save_format: Literal["binary", "unpacked"]) -> None:
-        if not is_torch_equal_or_newer("2.10.0"):
-            _patch_standalone_compile_atomic_save()
        self.save_format = save_format

    def compute_hash(self, vllm_config: VllmConfig) -> str:
@@ -10,7 +10,6 @@ import vllm.envs as envs
 from vllm.distributed import get_dp_group, get_ep_group
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
-from vllm.platforms import current_platform
 from vllm.utils.flashinfer import (
    has_flashinfer_nvlink_one_sided,
    has_flashinfer_nvlink_two_sided,
@@ -225,11 +224,8 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
            num_rdma_bytes=num_rdma_bytes,
            low_latency_mode=False,
            num_qps_per_rank=num_qps_per_rank,
+            explicitly_destroy=True,
        )
-        if not current_platform.is_rocm():
-            kwargs.update(
-                explicitly_destroy=True,
-            )
        return kwargs

    def get_handle(self, kwargs):
@@ -303,13 +299,10 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
            num_rdma_bytes=num_rdma_bytes,
            low_latency_mode=True,
            num_qps_per_rank=num_qps_per_rank,
+            allow_nvlink_for_low_latency_mode=True,
+            allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
+            explicitly_destroy=True,
        )
-        if not current_platform.is_rocm():
-            kwargs.update(
-                allow_nvlink_for_low_latency_mode=True,
-                allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
-                explicitly_destroy=True,
-            )
        return kwargs

    def get_handle(self, kwargs):
@@ -11,6 +11,7 @@ from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from datetime import timedelta

+import numpy as np
 import torch
 from torch.distributed import (
    P2POp,
@@ -47,15 +48,25 @@ class EplbCommunicator(ABC):
    """Abstract EPLB communicator for expert weight transfers."""

    @abstractmethod
-    def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,
+    ) -> None:
        pass

    @abstractmethod
-    def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,
+    ) -> None:
        pass

    @abstractmethod
-    def execute(self) -> None:
+    def execute(self, old_indices: np.ndarray | None = None) -> None:
        pass

    @property
@@ -85,27 +96,39 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator):
        self._p2p_ops: list[P2POp] = []
        self._log_initialized()

-    def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
-        self._p2p_ops.append(
-            P2POp(
-                torch.distributed.isend,
-                tensor,
-                dst_rank,
-                self._ep_group,
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._p2p_ops.append(
+                P2POp(
+                    torch.distributed.isend,
+                    tensor,
+                    dst_rank,
+                    self._ep_group,
+                )
            )
-        )

-    def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
-        self._p2p_ops.append(
-            P2POp(
-                torch.distributed.irecv,
-                tensor,
-                src_rank,
-                self._ep_group,
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._p2p_ops.append(
+                P2POp(
+                    torch.distributed.irecv,
+                    tensor,
+                    src_rank,
+                    self._ep_group,
+                )
            )
-        )

-    def execute(self) -> None:
+    def execute(self, old_indices: np.ndarray | None = None) -> None:
        if not self._p2p_ops:
            return
        try:
@@ -130,13 +153,25 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):
        self._ops: list[tuple[str, torch.Tensor, int]] = []
        self._log_initialized()

-    def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
-        self._ops.append(("send", tensor, dst_rank))
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._ops.append(("send", tensor, dst_rank))

-    def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
-        self._ops.append(("recv", tensor, src_rank))
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._ops.append(("recv", tensor, src_rank))

-    def execute(self) -> None:
+    def execute(self, old_indices: np.ndarray | None = None) -> None:
        if not self._ops:
            return

@@ -207,17 +242,17 @@ class NixlEplbCommunicator(EplbCommunicator):
        self._cuda_stream = cuda_stream
        self._world_size = cpu_group.size()
        self._rank = cpu_group.rank()
-        self._send_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {}
-        self._recv_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {}
-        self._dtypes: list[torch.dtype] = []
+        # expert_id -> weight tensors to pack into the send buffer.
+        self._expert_send_map: dict[int, list[torch.Tensor]] = {}
+        # src_rank -> expert_id -> weight tensors to unpack after transfer.
+        self._recv_map: dict[int, dict[int, list[torch.Tensor]]] = {}
+        self._num_local_experts: int = expert_weights[0].shape[0]
        self._device = expert_weights[0].device
        for tensor in expert_weights:
            assert tensor.device == self._device, (
                "All local EPLB tensors are expected to be on the same device: "
                f"expected={self._device}, got={tensor.device}"
            )
-            if tensor.dtype not in self._dtypes:
-                self._dtypes.append(tensor.dtype)

        config = (
            nixl_agent_config(capture_telemetry=False)
@@ -228,13 +263,12 @@ class NixlEplbCommunicator(EplbCommunicator):
        self._nixl_memory_type = "VRAM"
        self._registered_desc: object | None = None
        self._remote_agents: dict[int, str] = {}
-        self._remote_send_meta: dict[int, tuple[int, int, int]] = {}
+        self._remote_send_meta: dict[int, tuple[int, int]] = {}
        self._send_buffer: torch.Tensor = torch.empty(0)
        self._recv_buffer: torch.Tensor = torch.empty(0)
-        self._peer_partition_bytes: int = 0
-        self._dtype_max_bytes: dict[torch.dtype, int] = {}
+        self._expert_bytes: int = 0
+
        self._cuda_device_id = int(self._device.index or 0)
-        self._xfer_cache: dict[tuple[int, int, int], tuple[int, int, int]] = {}
        self._init_step("buffers", self._init_registered_buffers, expert_weights)
        self._init_step("agents", self._init_remote_agents)
        self._init_step("send meta", self._exchange_remote_send_meta)
@@ -258,34 +292,33 @@ class NixlEplbCommunicator(EplbCommunicator):
        uid = uuid.uuid4().hex[:8]
        return f"eplb-{self._rank}{pp_suffix}-{uid}"

-    def _get_peer_buckets(
+    def add_send(
        self,
-        bucket_map: dict[torch.dtype, list[list[torch.Tensor]]],
-        dtype: torch.dtype,
-    ) -> list[list[torch.Tensor]]:
-        peer_buckets = bucket_map.get(dtype)
-        if peer_buckets is None:
-            peer_buckets = [[] for _ in range(self._world_size)]
-            bucket_map[dtype] = peer_buckets
-        return peer_buckets
-
-    def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,
+    ) -> None:
        assert dst_rank != self._rank, (
            "EPLB communicator should not enqueue same-rank sends: "
            f"rank={self._rank}, dst_rank={dst_rank}"
        )
-        self._get_peer_buckets(self._send_tensors, tensor.dtype)[dst_rank].append(
-            tensor
-        )
+        # An expert sent to multiple peers is packed only once; skip duplicates.
+        if expert_id not in self._expert_send_map:
+            self._expert_send_map[expert_id] = tensors

-    def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,
+    ) -> None:
        assert src_rank != self._rank, (
            "EPLB communicator should not enqueue same-rank recvs: "
            f"rank={self._rank}, src_rank={src_rank}"
        )
-        self._get_peer_buckets(self._recv_tensors, tensor.dtype)[src_rank].append(
-            tensor
-        )
+        recv_experts = self._recv_map.setdefault(src_rank, {})
+        if expert_id not in recv_experts:
+            recv_experts[expert_id] = tensors

    def _init_remote_agents(self) -> None:
        local_metadata = self._nixl_wrapper.get_agent_metadata()
@@ -303,30 +336,18 @@ class NixlEplbCommunicator(EplbCommunicator):
            )

    def _init_registered_buffers(self, expert_weights: Sequence[torch.Tensor]) -> None:
-        total_max_bytes = 0
-        for dtype in self._dtypes:
-            max_numel = max(
-                sum(t.numel() for t in expert_weights if t.dtype == dtype), 1
-            )
-            max_bytes = max_numel * dtype.itemsize
-            self._dtype_max_bytes[dtype] = max_bytes
-            total_max_bytes += max_bytes
-
-        self._peer_partition_bytes = total_max_bytes
-
-        # The send buffer needs world_size partitions because remote peers
-        # READ from fixed offsets (rank * partition_bytes).
-        # This allocates world_size * partition_bytes
-        # which can cause OOM on large models.
-        # TODO(ilmarkov): shrink to const * partition_bytes and execute
-        # communication in multiple steps dealing with the worst case.
-        send_total_bytes = self._peer_partition_bytes * self._world_size
+        total_bytes = max(sum(t.nbytes for t in expert_weights), 1)
+        assert total_bytes % self._num_local_experts == 0, (
+            f"Number of bytes in moe layer {total_bytes} is not divisible "
+            f"by number of local experts {self._num_local_experts}"
+        )
+        self._expert_bytes = total_bytes // self._num_local_experts

        self._send_buffer = torch.empty(
-            send_total_bytes, device=self._device, dtype=torch.uint8
+            total_bytes, device=self._device, dtype=torch.uint8
        )
        self._recv_buffer = torch.empty(
-            self._peer_partition_bytes, device=self._device, dtype=torch.uint8
+            total_bytes, device=self._device, dtype=torch.uint8
        )

        descs = self._nixl_wrapper.get_reg_descs([self._send_buffer, self._recv_buffer])
@@ -336,12 +357,11 @@ class NixlEplbCommunicator(EplbCommunicator):
    def _exchange_remote_send_meta(self) -> None:
        """Exchange send-buffer metadata so each rank can build dynamic
        descriptors at execute time."""
-        local_meta: tuple[int, int, int] = (
+        local_meta: tuple[int, int] = (
            self._send_buffer.data_ptr(),
-            self._peer_partition_bytes,
            self._cuda_device_id,
        )
-        gathered_meta: list[tuple[int, int, int] | None] = [None] * self._world_size
+        gathered_meta: list[tuple[int, int] | None] = [None] * self._world_size
        torch.distributed.all_gather_object(
            gathered_meta, local_meta, group=self._cpu_group
        )
@@ -353,14 +373,11 @@ class NixlEplbCommunicator(EplbCommunicator):

    @staticmethod
    def _pack_send_buffer(
-        peer_tensors: list[torch.Tensor],
+        in_tensors: list[torch.Tensor],
        send_buffer: torch.Tensor,
        byte_offset: int,
-    ) -> int:
-        """
-        Returns the byte offset after the last written byte.
-        """
-        for tensor in peer_tensors:
+    ) -> None:
+        for tensor in in_tensors:
            raw = tensor.reshape(-1).view(torch.uint8)
            if raw.numel() == 0:
                continue
@@ -368,18 +385,14 @@ class NixlEplbCommunicator(EplbCommunicator):
                raw, non_blocking=True
            )
            byte_offset += raw.numel()
-        return byte_offset

    @staticmethod
    def _unpack_recv_buffer(
        recv_buffer: torch.Tensor,
-        peer_tensors: list[torch.Tensor],
+        out_tensors: list[torch.Tensor],
        byte_offset: int,
-    ) -> int:
-        """
-        Returns the byte offset after the last read byte.
-        """
-        for tensor in peer_tensors:
+    ) -> None:
+        for tensor in out_tensors:
            num_bytes = tensor.numel() * tensor.element_size()
            if num_bytes == 0:
                continue
@@ -388,19 +401,6 @@ class NixlEplbCommunicator(EplbCommunicator):
                non_blocking=True,
            )
            byte_offset += num_bytes
-        return byte_offset
-
-    def _release_all_cached_handles(self) -> None:
-        """Best-effort release of every cached dlist and xfer handle."""
-        for local_dlist, remote_dlist, xfer in self._xfer_cache.values():
-            for release_fn, handle in (
-                (self._nixl_wrapper.release_xfer_handle, xfer),
-                (self._nixl_wrapper.release_dlist_handle, local_dlist),
-                (self._nixl_wrapper.release_dlist_handle, remote_dlist),
-            ):
-                with contextlib.suppress(Exception):
-                    release_fn(handle)
-        self._xfer_cache.clear()

    def _wait_for_all_transfers(self, handles: list[int]) -> None:
        pending = set(handles)
@@ -418,82 +418,68 @@ class NixlEplbCommunicator(EplbCommunicator):
            if pending:
                time.sleep(0.0005)

-    def _get_or_create_xfer(self, src: int, total_bytes: int, recv_offset: int) -> int:
-        """Return a cached xfer handle or create and cache a new one."""
-        key = (src, total_bytes, recv_offset)
-        cached = self._xfer_cache.get(key)
-        if cached is not None:
-            return cached[2]
+    def _create_peer_xfer(
+        self,
+        src: int,
+        local_descs: list[tuple[int, int, int]],
+        remote_descs: list[tuple[int, int, int]],
+    ) -> tuple[int, int, int]:
+        """Create a batched xfer for multiple descriptors from one peer.

-        recv_base = self._recv_buffer.data_ptr()
+        Each element in *local_descs* / *remote_descs* is an
+        ``(address, size, device_id)`` tuple.
+
+        Returns ``(local_dlist, remote_dlist, xfer_handle)``.
+        """
        local_desc = self._nixl_wrapper.get_xfer_descs(
-            [
-                (
-                    recv_base + recv_offset,
-                    total_bytes,
-                    self._cuda_device_id,
-                )
-            ],
-            self._nixl_memory_type,
+            local_descs, self._nixl_memory_type
        )
        local_handle = self._nixl_wrapper.prep_xfer_dlist(
            "NIXL_INIT_AGENT",
            local_desc,
        )

-        remote_base, remote_part_bytes, remote_dev = self._remote_send_meta[src]
-        agent_name = self._remote_agents[src]
        remote_desc = self._nixl_wrapper.get_xfer_descs(
-            [
-                (
-                    remote_base + self._rank * remote_part_bytes,
-                    total_bytes,
-                    remote_dev,
-                )
-            ],
-            self._nixl_memory_type,
+            remote_descs, self._nixl_memory_type
        )
        remote_handle = self._nixl_wrapper.prep_xfer_dlist(
-            agent_name,
+            self._remote_agents[src],
            remote_desc,
        )

+        indices = list(range(len(local_descs)))
        xfer_handle = self._nixl_wrapper.make_prepped_xfer(
            "READ",
            local_handle,
-            [0],
+            indices,
            remote_handle,
-            [0],
+            indices,
        )
-        self._xfer_cache[key] = (local_handle, remote_handle, xfer_handle)
-        return xfer_handle
+        return (local_handle, remote_handle, xfer_handle)

-    def execute(self) -> None:
-        xfer_handles: list[int] = []
+    def execute(self, old_indices: np.ndarray | None = None) -> None:
+        assert old_indices is not None, (
+            "NixlEplbCommunicator.execute requires old_indices"
+        )
+
+        xfer_entries: list[tuple[int, int, int]] = []
        try:
-            # Phase 1: pack send buffers.
+            n = self._num_local_experts
+            rank_experts = old_indices[: self._world_size * n].reshape(
+                self._world_size, n
+            )
+            # Build expert_id -> send slot mapping per rank.
+            expert_to_send_slot: list[dict[int, int]] = [
+                {int(eid): i for i, eid in enumerate(row) if eid != -1}
+                for row in rank_experts
+            ]
+
+            # Phase 1: pack each expert at its slot offset in the send buffer.
            with torch.cuda.stream(self._cuda_stream):
-                for dst in range(self._world_size):
-                    byte_offset = dst * self._peer_partition_bytes
-                    for dtype in self._dtypes:
-                        peer_tensors = self._send_tensors.get(
-                            dtype, [[] for _ in range(self._world_size)]
-                        )[dst]
-                        actual_bytes = sum(
-                            t.numel() * t.element_size() for t in peer_tensors
-                        )
-                        if actual_bytes > self._dtype_max_bytes[dtype]:
-                            raise RuntimeError(
-                                "NIXL EPLB send overflow for dtype "
-                                f"{dtype}: peer={dst}, "
-                                f"required={actual_bytes}, "
-                                f"capacity={self._dtype_max_bytes[dtype]}"
-                            )
-                        byte_offset = self._pack_send_buffer(
-                            peer_tensors,
-                            self._send_buffer,
-                            byte_offset,
-                        )
+                for expert_id, tensors in self._expert_send_map.items():
+                    slot = expert_to_send_slot[self._rank][expert_id]
+                    byte_offset = slot * self._expert_bytes
+                    self._pack_send_buffer(tensors, self._send_buffer, byte_offset)

            # Ensure all packed data is visible in device memory before pulls.
            if self._cuda_stream is not None:
@@ -508,58 +494,65 @@ class NixlEplbCommunicator(EplbCommunicator):
                timeout=timedelta(minutes=5),
            )

-            # Phase 2: look up or create descriptors and issue all READs.
-            # Data from all peers is packed sequentially into the single
-            # partition-sized recv buffer at running offsets.
-            recv_offsets: dict[int, int] = {}
+            # Phase 2: issue one batched READ per peer.
+            recv_offsets: dict[tuple[int, int], int] = {}
            recv_offset = 0
+            recv_base = self._recv_buffer.data_ptr()
            for src in range(self._world_size):
                if src == self._rank:
                    continue
-                actual_total_bytes = 0
-                for dtype in self._dtypes:
-                    peer_tensors = self._recv_tensors.get(
-                        dtype, [[] for _ in range(self._world_size)]
-                    )[src]
-                    actual_total_bytes += sum(
-                        t.numel() * t.element_size() for t in peer_tensors
-                    )
-                if actual_total_bytes == 0:
+                recv_experts = self._recv_map.get(src)
+                if not recv_experts:
                    continue
-
-                recv_offsets[src] = recv_offset
-                xfer_handle = self._get_or_create_xfer(
-                    src, actual_total_bytes, recv_offset
+                expert_ids = list(recv_experts.keys())
+                remote_base, remote_dev = self._remote_send_meta[src]
+                local_descs: list[tuple[int, int, int]] = []
+                remote_descs: list[tuple[int, int, int]] = []
+                for expert_id in expert_ids:
+                    slot = expert_to_send_slot[src][expert_id]
+                    remote_off = slot * self._expert_bytes
+                    recv_offsets[(src, expert_id)] = recv_offset
+                    local_descs.append(
+                        (
+                            recv_base + recv_offset,
+                            self._expert_bytes,
+                            self._cuda_device_id,
+                        )
+                    )
+                    remote_descs.append(
+                        (remote_base + remote_off, self._expert_bytes, remote_dev)
+                    )
+                    recv_offset += self._expert_bytes
+                    assert recv_offset <= self._recv_buffer.nbytes
+                local_h, remote_h, xfer_h = self._create_peer_xfer(
+                    src, local_descs, remote_descs
                )
-                self._nixl_wrapper.transfer(xfer_handle)
-                xfer_handles.append(xfer_handle)
-                recv_offset += actual_total_bytes
+                self._nixl_wrapper.transfer(xfer_h)
+                xfer_entries.append((local_h, remote_h, xfer_h))

-            # Phase 3: single wait for all in-flight transfers, then unpack.
-            self._wait_for_all_transfers(xfer_handles)
+            # Phase 3: wait for all in-flight transfers, then unpack.
+            self._wait_for_all_transfers([x[2] for x in xfer_entries])

            with torch.cuda.stream(self._cuda_stream):
-                for src, offset in recv_offsets.items():
-                    byte_offset = offset
-                    for dtype in self._dtypes:
-                        peer_tensors = self._recv_tensors.get(
-                            dtype, [[] for _ in range(self._world_size)]
-                        )[src]
-                        byte_offset = self._unpack_recv_buffer(
-                            self._recv_buffer,
-                            peer_tensors,
-                            byte_offset,
-                        )
-        except Exception:
-            self._release_all_cached_handles()
-            raise
+                for (src, expert_id), offset in recv_offsets.items():
+                    self._unpack_recv_buffer(
+                        self._recv_buffer,
+                        self._recv_map[src][expert_id],
+                        offset,
+                    )
        finally:
-            self._send_tensors.clear()
-            self._recv_tensors.clear()
+            for local_h, remote_h, xfer_h in xfer_entries:
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_xfer_handle(xfer_h)
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_dlist_handle(local_h)
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_dlist_handle(remote_h)
+            self._expert_send_map.clear()
+            self._recv_map.clear()

    def __del__(self) -> None:
        try:
-            self._release_all_cached_handles()
            if self._registered_desc is not None:
                self._nixl_wrapper.deregister_memory(self._registered_desc)
                self._registered_desc = None
@@ -588,15 +581,27 @@ class PyNcclEplbCommunicator(EplbCommunicator):
            self._pynccl_comm.group_start()
            self._group_started = True

-    def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
        self._ensure_group_started()
-        self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream)
+        for tensor in tensors:
+            self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream)

-    def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
        self._ensure_group_started()
-        self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)
+        for tensor in tensors:
+            self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)

-    def execute(self) -> None:
+    def execute(self, old_indices: np.ndarray | None = None) -> None:
        if self._group_started:
            self._pynccl_comm.group_end()
            self._group_started = False
@@ -294,9 +294,9 @@ def move_to_buffer(
            recver_pos = remainder_start + sender_pos
            if recver_pos < len(ranks_to_recv):
                recv_ranks.append(ranks_to_recv[recver_pos])
+            expert_tensors = [w[src] for w in expert_weights]
            for dst in recv_ranks:
-                for w in expert_weights:
-                    communicator.add_send(w[src], dst)
+                communicator.add_send(expert_tensors, dst, expert_id=int(expert))

    # 3. Post recvs
    if recv_count > 0:
@@ -325,11 +325,14 @@ def move_to_buffer(
                src = ranks_to_send[recver_pos // num_dst_per_sender]
            else:
                src = ranks_to_send[recver_pos - remainder_start]
-            for b in expert_weights_buffers:
-                communicator.add_recv(b[dst], src)
+            communicator.add_recv(
+                [b[dst] for b in expert_weights_buffers],
+                src,
+                expert_id=int(expert),
+            )

    # 4. Execute the P2P operations. The real communication happens here.
-    communicator.execute()
+    communicator.execute(old_indices=old_indices)
    # wait for the communication to finish
    return TransferMetadata(
        is_unchanged=is_unchanged,
@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import copy
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast

 import torch

@@ -18,6 +18,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
    KVConnectorMetadata,
    KVConnectorRole,
    KVConnectorWorkerMetadata,
+    SupportsHMA,
+    supports_hma,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
    KVConnectorPromMetrics,
@@ -123,7 +125,7 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):
            self._prom_metrics[connector_id].observe(stats_data["data"], engine_idx)


-class MultiConnector(KVConnectorBase_V1):
+class MultiConnector(KVConnectorBase_V1, SupportsHMA):
    """
    A wrapper for using multiple KVConnectors at the same time.

@@ -166,6 +168,12 @@ class MultiConnector(KVConnectorBase_V1):
            self._connectors.append(connector_cls(temp_config, role, kv_cache_config))
            self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)

+        self._all_support_hma = all(supports_hma(c) for c in self._connectors)
+        assert (
+            vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
+            or self._all_support_hma
+        ), "HMA should not be enabled unless all sub-connectors support it"
+
        # A mapping from request id to the index of the connector chosen to
        # load the request from (if any).
        self._requests_to_connector: dict[str, int] = {}
@@ -436,15 +444,17 @@ class MultiConnector(KVConnectorBase_V1):
        for c in self._connectors:
            c.set_xfer_handshake_metadata(metadata)

-    def request_finished(
+    def _aggregate_request_finished(
        self,
        request: "Request",
-        blocks: list[int],
+        per_connector_fn: Callable[
+            [KVConnectorBase_V1], tuple[bool, dict[str, Any] | None]
+        ],
    ) -> tuple[bool, dict[str, Any] | None]:
        async_saves = 0
        kv_txfer_params = None
        for c in self._connectors:
-            async_save, txfer_params = c.request_finished(request, blocks)
+            async_save, txfer_params = per_connector_fn(c)
            if async_save:
                async_saves += 1
            if txfer_params is not None:
@@ -458,11 +468,39 @@ class MultiConnector(KVConnectorBase_V1):
        if async_saves > 1:
            self._extra_async_saves[request.request_id] = async_saves - 1

-        # Clean up other state for this request.
        self._requests_to_connector.pop(request.request_id, None)

        return async_saves > 0, kv_txfer_params

+    def request_finished(
+        self,
+        request: "Request",
+        blocks: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return self._aggregate_request_finished(
+            request,
+            lambda c: c.request_finished(request, blocks),
+        )
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        if not self._all_support_hma:
+            assert len(block_ids) == 1, (
+                "HMA with multiple kv_cache_groups requires all "
+                "sub-connectors to support HMA"
+            )
+            return self.request_finished(request, block_ids[0])
+
+        return self._aggregate_request_finished(
+            request,
+            lambda c: cast(SupportsHMA, c).request_finished_all_groups(
+                request, block_ids
+            ),
+        )
+
    def take_events(self) -> Iterable["KVCacheEvent"]:
        for c in self._connectors:
            yield from c.take_events()
@@ -119,6 +119,30 @@ class NixlConnectorScheduler:
            for n_tokens, block_size in sw_sizes_tokens
        ]

+        # Threshold to decide whether to compute kv cache locally
+        # or pull from a remote node: minimum number of remote
+        # tokens to amortize the xfer latencies
+        self.kv_recompute_threshold: int = int(
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "kv_recompute_threshold", 64
+            )
+        )
+
+        # Bi-directional KV transfer feature supports KV block
+        # transfers from D node to P node
+        self.is_bidirectional_kv_xfer_enabled = (
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "bidirectional_kv_xfer", False
+            )
+        )
+
+        if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0:
+            logger.info(
+                "Bidirectional KV transfer is enabled and the kv "
+                "recompute threshold is set to %d tokens",
+                self.kv_recompute_threshold,
+            )
+
    def shutdown(self):
        self._stop_event.set()
        if self._nixl_handshake_listener_t is not None:
@@ -298,6 +322,44 @@ class NixlConnectorScheduler:
        if params is not None and params.get("do_remote_decode") and self._has_mamba:
            self._truncate_mamba_request_for_prefill(request)

+        if (
+            params is not None
+            and params.get("do_remote_decode")
+            and params.get("remote_block_ids")
+            and all(
+                p in params
+                for p in (
+                    "remote_engine_id",
+                    "remote_request_id",
+                    "remote_host",
+                    "remote_port",
+                )
+            )
+        ):
+            # Decode node has kv blocks for part of prefill request, so, provide them
+            # as an external token count to scheduler.
+            # The tokens will be loaded if not already present
+            # in the prefill node local cache
+            remote_num_tokens = params.get("remote_num_tokens") or 0
+            count = (
+                min(remote_num_tokens, request.num_prompt_tokens) - num_computed_tokens
+            )
+            if count > 0:
+                # Check kv_recompute_threshold: skip pull if
+                # remote tokens are below the threshold.
+                if (
+                    self.kv_recompute_threshold > 0
+                    and count < self.kv_recompute_threshold
+                ):
+                    logger.debug(
+                        "Skipping remote pull for %s: %d remote tokens < threshold %d",
+                        request.request_id,
+                        count,
+                        self.kv_recompute_threshold,
+                    )
+                    return 0, False
+                return count, True
+
        # No remote prefill for this request.
        return 0, False

@@ -315,13 +377,19 @@ class NixlConnectorScheduler:
        if not params:
            return

-        if params.get("do_remote_decode"):
+        if params.get("do_remote_decode") or (
+            params.get("do_remote_prefill") and self.is_bidirectional_kv_xfer_enabled
+        ):
            self._reqs_in_batch.add(request.request_id)
        if self.use_host_buffer and params.get("do_remote_decode"):
            # NOTE: when accelerator is not directly supported by Nixl,
            # prefilled blocks need to be saved to host memory before transfer.
            self._reqs_need_save[request.request_id] = request
-        elif params.get("do_remote_prefill"):
+        elif params.get("do_remote_prefill") or (
+            params.get("do_remote_decode")
+            and self.is_bidirectional_kv_xfer_enabled
+            and not params.get("_remote_blocks_processed")
+        ):
            if params.get("remote_block_ids"):
                if all(
                    p in params
@@ -333,8 +401,8 @@ class NixlConnectorScheduler:
                    )
                ):
                    # If remote_blocks and num_external_tokens = 0, we have
-                    # a full prefix cache hit on the D worker. We need to call
-                    # send_notif in _read_blocks to free the memory on the P.
+                    # a full prefix cache hit on the local node. We need to call
+                    # send_notif in _read_blocks to free the memory on the remote node.

                    unhashed_local_block_ids: BlockIds = (
                        blocks.get_unhashed_block_ids_all_groups()
@@ -362,6 +430,7 @@ class NixlConnectorScheduler:
                assert num_external_tokens == 0
            # Only trigger 1 KV transfer per request.
            params["do_remote_prefill"] = False
+            params["_remote_blocks_processed"] = True

    def _build_save_meta(
        self,
@@ -450,6 +519,9 @@ class NixlConnectorScheduler:
        if not params:
            return False, None

+        is_p_node = bool(params.get("do_remote_decode"))
+        is_d_node = not is_p_node
+
        if params.get("do_remote_prefill"):
            # If do_remote_prefill is still True when the request is finished,
            # update_state_after_alloc must not have been called (the request
@@ -461,9 +533,13 @@ class NixlConnectorScheduler:
            params["do_remote_prefill"] = False
            return False, None

-        if not params.get("do_remote_decode"):
+        if is_d_node and not self.is_bidirectional_kv_xfer_enabled:
            return False, None
-        if request.status != RequestStatus.FINISHED_LENGTH_CAPPED:
+
+        if request.status not in (
+            RequestStatus.FINISHED_LENGTH_CAPPED,
+            RequestStatus.FINISHED_STOPPED,
+        ):
            # Also include the case of a P/D Prefill request with immediate
            # block free (eg abort). Stop tracking this request.
            self._reqs_not_processed.add(request.request_id)
@@ -474,7 +550,7 @@ class NixlConnectorScheduler:
        # TODO: check whether block_ids actually ever be 0. If not we could
        # remove the conditional below
        delay_free_blocks = any(len(group) > 0 for group in block_ids)
-
+        remote_num_tokens = 0
        if delay_free_blocks:
            # Prefill request on remote. It will be read from D upon completion
            logger.debug(
@@ -492,13 +568,16 @@ class NixlConnectorScheduler:
            # Here we "unpad" blocks to send the actual remote blocks to be read.
            block_ids = self.get_sw_clipped_blocks(block_ids)

+            remote_num_tokens = request.num_computed_tokens
+
        return delay_free_blocks, dict(
-            do_remote_prefill=True,
-            do_remote_decode=False,
+            do_remote_prefill=is_p_node,
+            do_remote_decode=is_d_node,
            remote_block_ids=block_ids,
            remote_engine_id=self.engine_id,
            remote_request_id=request.request_id,
            remote_host=self.side_channel_host,
            remote_port=self.side_channel_port,
            tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
+            remote_num_tokens=remote_num_tokens,
        )
@@ -18,15 +18,15 @@ from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.kv_offload.abstract import (
+from vllm.v1.kv_offload.base import (
+    GPULoadStoreSpec,
    OffloadingManager,
+    OffloadingSpec,
    OffloadKey,
    ReqContext,
    get_offload_block_hash,
    make_offload_key,
 )
-from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
-from vllm.v1.kv_offload.spec import OffloadingSpec
 from vllm.v1.outputs import KVConnectorOutput
 from vllm.v1.request import Request

@@ -25,7 +25,7 @@ from vllm.v1.kv_cache_interface import (
    MambaSpec,
    UniformTypeKVCacheSpecs,
 )
-from vllm.v1.kv_offload.spec import (
+from vllm.v1.kv_offload.base import (
    CanonicalKVCacheRef,
    CanonicalKVCaches,
    CanonicalKVCacheTensor,
@@ -755,6 +755,8 @@ class OpenAIServing:
    def _is_model_supported(self, model_name: str | None) -> bool:
        if not model_name:
            return True
+        if envs.VLLM_SKIP_MODEL_NAME_VALIDATION:
+            return True
        return self.models.is_base_model(model_name)


@@ -87,7 +87,7 @@ _maybe_set_cuda_compatibility_path()
 import torch

 from vllm.logger import init_logger
-from vllm.utils.torch_utils import is_torch_equal, is_torch_equal_or_newer
+from vllm.utils.torch_utils import is_torch_equal_or_newer

 logger = init_logger(__name__)

@@ -112,384 +112,6 @@ os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
 # in the environment.
 os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")

-# ===================================================
-# torch 2.9 Inductor PythonWrapperCodegen monkeypatch
-# ===================================================
-# This change monkeypatches memory_plan_reuse in pytorch 2.9.0 to work around
-# a test failure for test_multi_graph_piecewise_compile_outputs_equal.
-# For more context, see https://github.com/pytorch/pytorch/pull/165514.
-
-
-def memory_plan_reuse_patched(self):
-    import torch._inductor.ir as ir
-    from torch._inductor.codegen.wrapper import (
-        EnterSubgraphLine,
-        ExitSubgraphLine,
-        MemoryPlanningLine,
-        MemoryPlanningState,
-        SubgraphPythonWrapperCodegen,
-    )
-    from torch._inductor.virtualized import V
-
-    def get_output_names(graph_outputs) -> list[str]:
-        import itertools
-
-        names = []
-        shape_counter = itertools.count(0)
-        none_counter = itertools.count(0)
-        for node in graph_outputs:
-            if isinstance(node, ir.NoneAsConstantBuffer):
-                names.append(f"{V.graph.name}_none{next(none_counter)}")
-            elif isinstance(node, ir.ShapeAsConstantBuffer):
-                names.append(f"{V.graph.name}_shape{next(shape_counter)}")
-            else:
-                names.append(node.get_name())
-        return names
-
-    if (
-        isinstance(V.graph.wrapper_code, SubgraphPythonWrapperCodegen)
-        and V.graph.wrapper_code.partition_signatures is not None
-    ):
-        out_names = get_output_names(
-            V.graph.wrapper_code.partition_signatures.output_nodes
-        )
-    else:
-        out_names = V.graph.get_output_names()
-
-    while (
-        self.lines
-        and isinstance(self.lines[-1], MemoryPlanningLine)
-        and self.lines[-1].node.name not in out_names  # type: ignore[attr-defined]
-    ):
-        # these lines will be pointless
-        self.lines.pop()
-
-    # codegen allocations in two passes
-    planning_states = [MemoryPlanningState()]
-    past_planning_states = []
-    for i in range(len(self.lines)):
-        line = self.lines[i]
-        if isinstance(line, MemoryPlanningLine):
-            self.lines[i] = line.plan(planning_states[-1])
-        elif isinstance(line, EnterSubgraphLine):
-            planning_states.append(MemoryPlanningState())
-        elif isinstance(line, ExitSubgraphLine):
-            past_planning_states.append(planning_states.pop())
-    past_planning_states.append(planning_states.pop())
-    assert len(planning_states) == 0
-
-
-# ===================================================
-# torch 2.9 Inductor get_graph_partition_signature monkeypatch
-# ===================================================
-# This change monkeypatches get_graph_partition_signature in pytorch 2.9.0 to
-# fix inductor partition + attention-nvfp4 quant fusion, tested in
-# `tests/compile/test_fusion_attn.py::test_attn_quant`.
-# For more context, see https://github.com/pytorch/pytorch/pull/165815.
-
-
-def get_graph_partition_signature_patched(
-    self, partitions, skip_cudagraphs: list[bool]
-):
-    """
-    Gets signature for each graph partition, including input nodes, output nodes, and
-    whether deallocating an input within graph partition.
-    """
-    from torch._inductor import dependencies
-    from torch._inductor.ir import GraphPartitionSignature, MutationOutput, NoneLayout
-    from torch._inductor.virtualized import V
-    from torch.utils._ordered_set import OrderedSet
-
-    signatures = []
-
-    unmet_output_names = OrderedSet(V.graph.get_output_names())
-    name_to_node = self.get_name_to_nodes()
-
-    def is_none_layout(buf_name: str) -> bool:
-        """
-        Checks if buf_name is NoneLayout. Buffers with NoneLayout is not allocated
-        so graph partition should not take it as inputs or outputs.
-        """
-        buf = self.name_to_buf.get(buf_name, None)
-
-        if buf is None:
-            return False
-
-        if isinstance(buf.node.layout, NoneLayout):
-            if isinstance(buf.node, MutationOutput) and (
-                real_name := self.mutation_real_name.get(buf_name, None)
-            ):
-                return is_none_layout(real_name)
-
-            return True
-
-        return False
-
-    for partition, skip_cudagraph in zip(
-        reversed(partitions), reversed(skip_cudagraphs)
-    ):
-        output_names: OrderedSet[str] = OrderedSet()
-
-        for node in partition:
-            output_names.update(node.outputs_by_name.keys())
-
-        returned_output_names = output_names.intersection(unmet_output_names)
-
-        # all reads/writes are partition inputs except those generated
-        # within the partition and tensor constants
-        read_writes = dependencies.ReadWrites.merge_list(
-            [node.read_writes for node in partition]
-        )
-
-        # WeakDep is fake dependency on unused buffer. It should not appear
-        # in partition_input_names for inputs that are actually read or written.
-        partition_input_names = (
-            OrderedSet(
-                [
-                    x.name
-                    for x in read_writes.reads | read_writes.writes
-                    if not is_none_layout(x.name)
-                ]
-            )
-            - output_names
-        )
-
-        partition_input_names = OrderedSet(
-            self.mutation_real_name.get(name, name) for name in partition_input_names
-        )
-
-        buffer_names_to_free: OrderedSet[str] = OrderedSet()
-        for node in partition:
-            buffer_names_to_free.update(node.last_usage)
-
-        # buffer_names_to_free may contain buffers allocated in previous
-        # graph partitions. These buffers should also be a partition
-        # input.
-        extra_input_names = [
-            name
-            for name in (buffer_names_to_free - output_names)
-            if name in name_to_node
-        ]
-        partition_input_names.update(extra_input_names)
-
-        input_nodes = {
-            name: name_to_node[name]
-            for name in partition_input_names
-            if name in name_to_node
-        }
-        input_deallocation = {
-            name: name in buffer_names_to_free
-            for name in partition_input_names
-            if name in name_to_node
-        }
-
-        # if an input tensor is not freed in the partition function, it should
-        # also be returned as an output. This brings benefits to cudagraph
-        # since the returned output tensor is a cudagraph managed tensor with
-        # a static tensor address.
-        extra_output_names = [
-            name
-            for name in partition_input_names
-            if name in name_to_node and name not in buffer_names_to_free
-        ]
-
-        returned_output_names.update(extra_output_names)
-
-        returned_output_names = OrderedSet(
-            self.mutation_real_name.get(name, name) for name in returned_output_names
-        )
-
-        output_nodes = [
-            name_to_node[name]
-            for name in returned_output_names
-            if not is_none_layout(name)
-        ]
-
-        constant_names = [
-            name for name in partition_input_names if name in V.graph.constants
-        ]
-
-        symbol_inputs = self.get_graph_partition_symbol_inputs(partition, input_nodes)
-
-        partition_signature = GraphPartitionSignature(
-            symbol_inputs,
-            input_nodes,
-            output_nodes,
-            input_deallocation,
-            skip_cudagraph,
-            constant_names,
-        )
-
-        signatures.append(partition_signature)
-
-        unmet_output_names = partition_input_names.union(
-            unmet_output_names - returned_output_names
-        )
-
-    return signatures[::-1]
-
-
-# ========================================
-# torch 2.9 Inductor Scheduler monkeypatch
-# ========================================
-# This change monkeypatches a function in Inductor to work around the following
-# bug: https://github.com/vllm-project/vllm/issues/26678
-#
-# The bug occurs when `use_inductor_graph_partition` is turned on and there
-# exists operators inside of `splitting_ops` that have an in-place mutation. In
-# vllm, this specifically occurs on the operator
-# vllm.unified_attention_with_output. In this case, inductor does not populate
-# the inductor IR's `origin_node` field, causing an assertion error when trying
-# to access the node's `origin_node` field.
-#
-# So, we will monkeypatch torch._inductor.scheduler.Scheduler.should_partition
-# so that it does not access the inductor IR node's `origin_node` field and just
-# returns True if a node is registered as having a custom partition function.
-# This is ok for now since vllm's implementation of the custom partition
-# functions just return True.
-# ========================================
-
-
-def should_partition_patched(self, node, should_log: bool = False) -> bool:
-    # This is a patched version of
-    # torch._inductor.scheduler.Scheduler.should_partition that modifies
-    # the following piece of code so that we always return True:
-    # https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/_inductor/scheduler.py#L4712-L4724
-    """Return True if we should partition the inductor graph on this node"""
-
-    import torch._inductor.ir as ir
-    from torch._inductor.scheduler import (
-        BaseSchedulerNode,
-        FusedSchedulerNode,
-    )
-    from torch._inductor.utils import (
-        _unstable_customized_partition_wrapper,
-        is_cudagraph_unsafe_op,
-        maybe_log_cudagraph_partition,
-    )
-
-    # Allow users to manually specify if a node should be partitioned
-    # Can only do this for FallbackKernels
-    ir_node = node.node
-    if isinstance(ir_node, torch._inductor.ir.FallbackKernel) and (
-        op := ir_node.op_overload
-    ):
-        op_overload_packet_name = op.name()
-        op_overload_name = (
-            f"{op_overload_packet_name}.{op._overloadname}"
-            if isinstance(op, torch._ops.OpOverload)
-            else op_overload_packet_name
-        )
-        if (
-            op_overload_packet_name
-            in torch._inductor.config.custom_should_partition_ops
-            or op_overload_name in torch._inductor.config.custom_should_partition_ops
-        ):
-            assert isinstance(op, torch._ops.OpOverload)
-            return True
-
-    # When not using cudagraphs, keep all kernels in the `call` function
-    # instead of graph partition functions, since graph partition only brings
-    # benefit to cudagraph
-    if (
-        not torch._inductor.config.triton.cudagraphs
-        and _unstable_customized_partition_wrapper.wrapper is None
-    ):
-        return True
-
-    # avoid duplicating logs when should_partition is called multiple times
-    # on the same node
-    def noop_log(msg: str, node: BaseSchedulerNode | None) -> None:
-        return
-
-    log_partition_reason = maybe_log_cudagraph_partition if should_log else noop_log
-
-    if isinstance(node, FusedSchedulerNode):
-        return any(self.should_partition(snode) for snode in node.snodes)
-
-    assert node.node is not None
-
-    if not node.is_gpu():
-        log_partition_reason("non gpu ops", node=node)
-
-        return True
-
-    if isinstance(node.node, ir.DeviceCopy):
-        log_partition_reason("DeviceCopy ops", node=node)
-        return True
-
-    if isinstance(node.node, ir.Conditional):
-        log_partition_reason("Conditional ops", node=node)
-        return True
-
-    if getattr(node.node, "unbacked_bindings", None):
-        log_partition_reason("unbacked binding ops", node=node)
-        return True
-
-    if is_cudagraph_unsafe_op(node.node):
-        log_partition_reason("CUDAGraph-unsafe custom ops", node=node)
-        return True
-
-    return False
-
-
-def _update_scheduler_patched(self) -> None:
-    # Copied from torch._inductor.graph.GrahLowering._update_scheduler. Patches
-    # this method so that we can patch Scheduler.should_partition with the
-    # function above
-    """
-    (Re)initializes the scheduler member.  When initializing the scheduler, no CUBIN
-    files should be generated (to avoid biasing any benchmarks and pessimizing
-    fusion decisions).
-    """
-    import torch._inductor.config as config
-    from torch._inductor.scheduler import Scheduler
-
-    Scheduler.should_partition = should_partition_patched
-    Scheduler.get_graph_partition_signature = get_graph_partition_signature_patched
-
-    with config.patch("triton.store_cubin", False):
-        self.scheduler = Scheduler(self.operations)
-
-
-# ===================================================
-# torch 2.9 Inductor get_raw_stream workaround
-# ===================================================
-# Workaround for TorchInductor autotune using get_raw_stream() without defining it.
-# This occurs when compile_sizes > 1 in compilation_config.
-# For more context, see https://github.com/vllm-project/vllm/issues/30905.
-def _patch_get_raw_stream_if_needed():
-    """Workaround for TorchInductor autotune get_raw_stream() bug."""
-    from vllm.utils.torch_utils import is_torch_equal
-
-    # Only apply the patch for torch 2.9.0 or 2.9.1
-    if is_torch_equal("2.9.0") or is_torch_equal("2.9.1"):
-        import builtins
-
-        # Check if CUDA functionality is available without initializing CUDA
-        # _cuda_getCurrentRawStream only exists in CUDA builds of PyTorch
-        if hasattr(torch._C, "_cuda_getCurrentRawStream"):
-            from torch._C import _cuda_getCurrentRawStream as _get_raw_stream
-
-            builtins.get_raw_stream = _get_raw_stream  # type: ignore[attr-defined]
-
-
-_patch_get_raw_stream_if_needed()
-
-if is_torch_equal("2.9.0"):
-    from torch._inductor.codegen.wrapper import PythonWrapperCodegen
-    from torch._inductor.graph import GraphLowering
-    from torch.utils._config_module import _Config, _ConfigEntry
-
-    # `custom_should_partition_ops` is a new config after 2.9.0. So this would
-    # not overwrite any user configs.
-    torch._inductor.config._config["custom_should_partition_ops"] = _ConfigEntry(
-        _Config(default=[])
-    )
-
-    PythonWrapperCodegen.memory_plan_reuse = memory_plan_reuse_patched
-    GraphLowering._update_scheduler = _update_scheduler_patched
-
 # ===================================================
 # torch <2.12 GraphCaptureOutput.get_runtime_env monkeypatch
 # ===================================================
@@ -586,55 +208,6 @@ if is_torch_equal_or_newer("2.10.0") and not is_torch_equal_or_newer("2.12.0.dev

    GraphCaptureOutput.get_runtime_env = _patched_get_runtime_env

-# ===================================================
-# torch 2.10 FxGraphCachePickler.dumps ValueError fix
-# ===================================================
-# PyTorch 2.10's FxGraphCachePickler.dumps() doesn't catch ValueError,
-# causing torch.compile cache failures when tensors with non-standard
-# layouts (e.g. blocked-layout prepacked weights) are serialized.
-# PyTorch mainline fixed this in pytorch/pytorch#176557 (merged 2026-03-04).
-# This is a thin backport for 2.10 users; remove once 2.10 is dropped.
-
-
-def _apply_fxgraphcache_pickle_patch(pickler_cls, bypass_cls):
-    """Wrap pickler_cls.dumps to convert ValueError into bypass_cls.
-
-    Idempotent: sets `_vllm_fxgraph_dumps_patched` on the class after the
-    first apply to prevent re-application. The wrapper function is also
-    marked with `_vllm_patched` as an additional safeguard.
-    """
-    if getattr(pickler_cls, "_vllm_fxgraph_dumps_patched", False):
-        return
-
-    original_dumps = pickler_cls.dumps
-    if hasattr(original_dumps, "_vllm_patched"):
-        return
-
-    def patched_dumps(self, obj):
-        try:
-            return original_dumps(self, obj)
-        except ValueError as e:
-            raise bypass_cls("Failed to pickle cache key") from e
-
-    patched_dumps._vllm_patched = True  # type: ignore[attr-defined]
-    pickler_cls.dumps = patched_dumps
-    pickler_cls._vllm_fxgraph_dumps_patched = True  # type: ignore[attr-defined]
-
-
-def _patch_fxgraphcache_pickle_if_needed():
-    """Apply FxGraphCachePickler.dumps ValueError backport when on torch 2.10.x."""
-    from vllm.utils.torch_utils import is_torch_equal_or_newer
-
-    if not is_torch_equal_or_newer("2.10.0") or is_torch_equal_or_newer("2.11.0"):
-        return
-
-    from torch._inductor.codecache import BypassFxGraphCache, FxGraphCachePickler
-
-    _apply_fxgraphcache_pickle_patch(FxGraphCachePickler, BypassFxGraphCache)
-
-
-_patch_fxgraphcache_pickle_if_needed()
-
 # ===================================================
 # torch 2.11 Inductor cpp codegen indirect_assert scalar-mask fix
 # ===================================================
@@ -255,6 +255,10 @@ if TYPE_CHECKING:
    VLLM_LORA_DISABLE_PDL: bool = False
    VLLM_ENABLE_CUDA_COMPATIBILITY: bool = False
    VLLM_CUDA_COMPATIBILITY_PATH: str | None = None
+    VLLM_SKIP_MODEL_NAME_VALIDATION: bool = False
+    """If set, vLLM will skip model name validation in API requests.
+    This allows any model name to be accepted in the 'model' field of requests,
+    making the server model-name agnostic. Useful for proxy/gateway scenarios."""
    VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False
    VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False
    VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True
@@ -1711,6 +1715,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_CUDA_COMPATIBILITY_PATH": lambda: os.environ.get(
        "VLLM_CUDA_COMPATIBILITY_PATH", None
    ),
+    # Skip model name validation in OpenAI API requests.
+    # When set to 1, any model name will be accepted in the 'model' field
+    # of API requests. This is useful for proxy/gateway scenarios where
+    # the actual model is served but different names may be used in requests.
+    "VLLM_SKIP_MODEL_NAME_VALIDATION": lambda: (
+        os.getenv("VLLM_SKIP_MODEL_NAME_VALIDATION", "0").strip().lower()
+        in ("1", "true")
+    ),
    # Whether it is a scale up launch engine for elastic EP,
    # Should only be set by EngineCoreClient.
    "VLLM_ELASTIC_EP_SCALE_UP_LAUNCH": lambda: bool(
@@ -1886,6 +1898,7 @@ def compile_factors() -> dict[str, object]:
        "VLLM_TEST_FORCE_LOAD_FORMAT",
        "VLLM_ENABLE_CUDA_COMPATIBILITY",
        "VLLM_CUDA_COMPATIBILITY_PATH",
+        "VLLM_SKIP_MODEL_NAME_VALIDATION",
        "LOCAL_RANK",
        "CUDA_VISIBLE_DEVICES",
        "NO_COLOR",
@@ -1076,10 +1076,10 @@ def chunk_gla_fwd_kernel_o(
        )
        p_h = tl.make_block_ptr(
            h + (i_tg * H + i_h) * K * V,
-            (K, V),
-            (V, 1),
-            (i_k * BK, i_v * BV),
-            (BK, BV),
+            (V, K),
+            (K, 1),
+            (i_v * BV, i_k * BK),
+            (BV, BK),
            (1, 0),
        )

@@ -1090,12 +1090,11 @@ def chunk_gla_fwd_kernel_o(
        b_g = tl.load(p_g, boundary_check=(0, 1))
        # [BT, BK]
        b_qg = (b_q * exp(b_g)).to(b_q.dtype)
-        # [BK, BV]
+        # [BV, BK]
        b_h = tl.load(p_h, boundary_check=(0, 1))
-        # works but dkw, owing to divine benevolence
        # [BT, BV]
        if i_k >= 0:
-            b_o += tl.dot(b_qg, b_h.to(b_qg.dtype))
+            b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
    p_v = tl.make_block_ptr(
        v + (bos * H + i_h) * V,
        (T, V),
@@ -1487,15 +1487,19 @@ class FusedMoE(PluggableLayer):
            "w2_input_scale",
        }

+        # Parameters of non-expert submodules that live inside runner (MoERunner).
+        # These must be excluded from EPLB weight rearrangement.
+        NON_EXPERT_PREFIXES = (
+            "runner._shared_experts.",
+            "runner.gate.",
+            "runner.routed_input_transform.",
+            "runner.routed_output_transform.",
+        )
+
        assert all(
            weight.is_contiguous()
            for name, weight in weights
-            if not (
-                name.startswith("_shared_experts.")
-                or name.startswith("_gate.")
-                or name.startswith("_routed_input_transform.")
-                or name.startswith("_routed_output_transform.")
-            )
+            if not name.startswith(NON_EXPERT_PREFIXES)
            and name not in NON_EXPERT_WEIGHTS
        )

@@ -1504,12 +1508,7 @@ class FusedMoE(PluggableLayer):
            for name, weight in weights
            if name not in NON_EXPERT_WEIGHTS
            and weight.shape != torch.Size([])
-            and not name.startswith("_shared_experts.")
-            # exclude parameters from non-expert submodules,
-            # e.g. gate/shared/transforms.
-            and not name.startswith("_gate.")
-            and not name.startswith("_routed_input_transform.")
-            and not name.startswith("_routed_output_transform.")
+            and not name.startswith(NON_EXPERT_PREFIXES)
        ]

    def set_eplb_state(
@@ -16,7 +16,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
    moe_kernel_quantize_input,
    normalize_batched_scales_shape,
 )
-from vllm.platforms import current_platform
 from vllm.v1.worker.ubatching import (
    dbo_current_ubatch_id,
    dbo_enabled,
@@ -290,46 +289,29 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

        # Dispatch
        dispatch_topk_ids = self._map_global_to_physical_ids(topk_ids)
-        if current_platform.is_rocm():
-            (
-                expert_x,
-                expert_num_tokens,
-                handle,
-                _,
-                hook,
-            ) = self.buffer.low_latency_dispatch(
-                a1,
-                dispatch_topk_ids,
-                self.max_tokens_per_rank,
-                num_experts,
-                use_fp8=self.use_fp8_dispatch,
-                async_finish=False,
-                return_recv_hook=True,
-            )
-        else:
-            (
-                expert_x,
-                expert_num_tokens,
-                handle,
-                _,
-                hook,
-            ) = self.buffer.low_latency_dispatch(
-                a1,
-                dispatch_topk_ids,
-                self.max_tokens_per_rank,
-                num_experts,
-                use_fp8=self.use_fp8_dispatch,
-                round_scale=self.use_ue8m0_dispatch,
-                use_ue8m0=self.use_ue8m0_dispatch,
-                **(dict(use_nvfp4=True) if use_nvfp4 else dict()),
-                **(
-                    dict(x_global_scale=qc_a1_gscale_or_scale)
-                    if qc_a1_gscale_or_scale is not None and nvfp4_dispatch
-                    else dict()
-                ),
-                async_finish=False,
-                return_recv_hook=True,
-            )
+        (
+            expert_x,
+            expert_num_tokens,
+            handle,
+            _,
+            hook,
+        ) = self.buffer.low_latency_dispatch(
+            a1,
+            dispatch_topk_ids,
+            self.max_tokens_per_rank,
+            num_experts,
+            use_fp8=self.use_fp8_dispatch,
+            round_scale=self.use_ue8m0_dispatch,
+            use_ue8m0=self.use_ue8m0_dispatch,
+            **(dict(use_nvfp4=True) if use_nvfp4 else dict()),
+            **(
+                dict(x_global_scale=qc_a1_gscale_or_scale)
+                if qc_a1_gscale_or_scale is not None and nvfp4_dispatch
+                else dict()
+            ),
+            async_finish=False,
+            return_recv_hook=True,
+        )
        self.handles[a2a_idx] = handle

        return (
@@ -220,7 +220,7 @@ class MoERunner(MoERunnerInterface):
        self.routed_output_transform = routed_output_transform
        self.routed_scaling_factor = routed_scaling_factor
        self.gate = gate
-        self.quant_method = quant_method
+        self._quant_method = quant_method
        self.enable_dbo = enable_dbo

        self._shared_experts: SharedExperts | None = None
@@ -263,7 +263,7 @@ class MoERunner(MoERunnerInterface):
    def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
        if self._shared_experts is not None:
            self._shared_experts._quant_method = quant_method
-        self.quant_method = quant_method
+        self._quant_method = quant_method

    def is_internal_router(self) -> bool:
        return self.gate is not None
@@ -330,8 +330,8 @@ class MoERunner(MoERunnerInterface):
    @property
    def _fused_output_is_reduced(self) -> bool:
        return (
-            self.quant_method.moe_kernel is not None
-            and self.quant_method.moe_kernel.output_is_reduced()
+            self._quant_method.moe_kernel is not None
+            and self._quant_method.moe_kernel.output_is_reduced()
        )

    def _maybe_reduce_shared_expert_output(
@@ -407,7 +407,7 @@ class MoERunner(MoERunnerInterface):
        )
        transformed_hidden_dim = hidden_states.shape[-1]
        if (
-            not self.quant_method.skip_forward_padding
+            not self._quant_method.skip_forward_padding
            and self.moe_config.hidden_dim != transformed_hidden_dim
        ):
            hidden_states = F.pad(
@@ -451,8 +451,8 @@ class MoERunner(MoERunnerInterface):
            shared_experts_input, SharedExpertsOrder.NO_OVERLAP
        )

-        if self.quant_method.is_monolithic:
-            fused_out = self.quant_method.apply_monolithic(
+        if self._quant_method.is_monolithic:
+            fused_out = self._quant_method.apply_monolithic(
                layer=layer,
                x=hidden_states,
                router_logits=router_logits,
@@ -467,7 +467,7 @@ class MoERunner(MoERunnerInterface):

            # Passing shared_experts_input in case SharedExpertsOrder is
            # MK_INTERNAL_OVERLAPPED.
-            fused_out = self.quant_method.apply(
+            fused_out = self._quant_method.apply(
                layer=layer,
                x=hidden_states,
                topk_weights=topk_weights,
@@ -618,7 +618,7 @@ class MoERunner(MoERunnerInterface):
    @property
    def do_naive_dispatch_combine(self) -> bool:
        return (
-            self.moe_config.dp_size > 1 and not self.quant_method.supports_internal_mk
+            self.moe_config.dp_size > 1 and not self._quant_method.supports_internal_mk
        )

    def _maybe_dispatch(
@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod

 import torch

+from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
    FusedMoEMethodBase,
 )
@@ -12,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
 )


-class MoERunnerInterface(ABC):
+class MoERunnerInterface(PluggableLayer, ABC):
    """
    Abstract base class for Mixture of Experts (MoE) runners.

@@ -3,6 +3,7 @@

 import math
 from collections.abc import Iterable, Mapping, Sequence
+from typing import Any, ClassVar

 import torch
 import torch.nn.functional as F
@@ -14,7 +15,7 @@ from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
 from vllm.config.speech_to_text import SpeechToTextParams
 from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.inputs import MultiModalDataDict, PromptType, TextPrompt
+from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.attention import (
@@ -48,6 +49,7 @@ from vllm.multimodal.processing import (
    PromptUpdate,
 )
 from vllm.renderers import TokenizeParams
+from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.transformers_utils.processors.cohere_asr import (
    INF_VAL,
    CohereASRFeatureExtractor,
@@ -2008,6 +2010,9 @@ class CohereAsrForConditionalGeneration(
    supported_languages = ISO639_1_SUPPORTED_LANGS
    skip_warmup_audio_preprocessing = True
    no_space_languages = {"ja", "zh"}
+    _default_prompt_token_ids_cache: ClassVar[
+        dict[tuple[str | None, str | None, str], tuple[int, ...]]
+    ] = {}

    @classmethod
    def validate_language(cls, language: str | None) -> str | None:
@@ -2025,30 +2030,81 @@ class CohereAsrForConditionalGeneration(
        audio = stt_params.audio
        stt_config = stt_params.stt_config
        language = stt_params.language
-        request_prompt = stt_params.request_prompt
+        model_config = stt_params.model_config

        if language is None:
            raise ValueError(
                "Language must be specified when creating the CohereASR prompt"
            )

-        # NOTE: this function is used only by online inference and not offline inference
-        # CohereASR doesnt have encoder prompt
-        language_tag = f"<|{language}|><|{language}|>"
-        pnc = True  # TODO(ekagra): make this configurable later
-        pnc_tag = "<|pnc|>" if pnc else "<|nopnc|>"
-        default_prompt = (
-            f"<|startofcontext|><|startoftranscript|>"
-            f"<|emo:undefined|>{language_tag}{pnc_tag}"
-            f"<|noitn|><|notimestamp|><|nodiarize|>"
-        )
-        prompt_text = request_prompt if request_prompt else default_prompt
+        tokenizer = cached_tokenizer_from_config(model_config)

-        return TextPrompt(
+        # prompt_text is None because CoherASR uses fast implementation of
+        # sentencepiece tokenizer which needs "▁" as the first token
+        # (which is different from "_") and encode("▁ABC") ignores the first token
+        # so the prompt_text is unreliable. However, prompt_token_ids can be used
+        # to get prompt_text but it wont have the first token "▁".
+        prompt_text = None
+        prompt_token_ids = cls._get_default_prompt_token_ids(
+            tokenizer,
+            model_config,
+            language,
+        )
+
+        return TokensPrompt(
            prompt=prompt_text,
+            prompt_token_ids=prompt_token_ids,
            multi_modal_data={"audio": (audio, stt_config.sample_rate)},
        )

+    @classmethod
+    def _get_default_prompt_tokens(cls, language: str) -> tuple[str, ...]:
+        # Use token-level control tags so fast tokenizers do not have to parse
+        # the raw string form of the decoder prefix.
+        return (
+            "▁",
+            "<|startofcontext|>",
+            "<|startoftranscript|>",
+            "<|emo:undefined|>",
+            f"<|{language}|>",
+            f"<|{language}|>",
+            "<|pnc|>",
+            "<|noitn|>",
+            "<|notimestamp|>",
+            "<|nodiarize|>",
+        )
+
+    @classmethod
+    def _get_default_prompt_token_ids(
+        cls,
+        tokenizer: Any,
+        model_config: ModelConfig,
+        language: str,
+    ) -> list[int]:
+        cache_key = (
+            getattr(model_config, "tokenizer", None),
+            getattr(model_config, "tokenizer_revision", None),
+            language,
+        )
+        prompt_token_ids = cls._default_prompt_token_ids_cache.get(cache_key)
+        if prompt_token_ids is None:
+            prompt_tokens = list(cls._get_default_prompt_tokens(language))
+            token_ids = tokenizer.convert_tokens_to_ids(prompt_tokens)
+            if not isinstance(token_ids, list):
+                token_ids = [token_ids]
+            unk_token_id = getattr(tokenizer, "unk_token_id", None)
+            if unk_token_id is not None and any(
+                token_id == unk_token_id for token_id in token_ids
+            ):
+                raise ValueError(
+                    "Failed to resolve the CohereASR decoder control tokens "
+                    "with the configured tokenizer."
+                )
+            prompt_token_ids = tuple(int(token_id) for token_id in token_ids)
+            cls._default_prompt_token_ids_cache[cache_key] = prompt_token_ids
+
+        return list(prompt_token_ids)
+
    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        # Required as part of SupportsMultiModal interface.
@@ -854,10 +854,9 @@ class DeepseekV4MoE(nn.Module):
    def forward(
        self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
    ) -> torch.Tensor:
-        if self.gate.tid2eid is not None:
-            if input_ids is None:
-                raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
-            input_ids = input_ids.to(dtype=self.hash_indices_dtype)
+        if self.gate.tid2eid is not None and input_ids is None:
+            raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
+
        if not self.use_mega_moe:
            return self._forward_fused_moe(hidden_states, input_ids)

@@ -1225,7 +1224,12 @@ class DeepseekV4Model(nn.Module):
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        self.config = config
-
+        if vllm_config.parallel_config.enable_expert_parallel:
+            self.use_mega_moe = (
+                vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
+            )
+        else:
+            self.use_mega_moe = False
        self.vocab_size = config.vocab_size
        self.hc_eps = config.hc_eps
        self.hc_mult = config.hc_mult
@@ -1309,7 +1313,8 @@ class DeepseekV4Model(nn.Module):
    ) -> torch.Tensor | IntermediateTensors:
        hidden_states = self.embed_input_ids(input_ids)
        hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
-
+        if self.use_mega_moe:
+            input_ids = input_ids.to(torch.int64)
        for layer in islice(self.layers, self.start_layer, self.end_layer):
            hidden_states = layer(
                hidden_states,
@@ -84,6 +84,10 @@ from .utils import (
 logger = init_logger(__name__)


+def _remap_gemma4_expert_weight_name(name: str) -> str:
+    return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
+
+
@triton.jit
 def _gemma4_routing_kernel(
    gating_ptr,
@@ -1650,7 +1654,7 @@ class Gemma4ForCausalLM(
                # Remap individual 2D expert weights:
                # .experts.{id}.{proj} → .moe.experts.{id}.{proj}
                # (This handles per-expert 2D quantized weights)
-                name = re.sub(r"\.experts\.(\d+)\.", r".moe.experts.\1.", name)
+                name = _remap_gemma4_expert_weight_name(name)

                # MoE expert weights: checkpoint stores as 3D packed
                # tensors.  Explode into per-expert 2D weights for
@@ -351,6 +351,9 @@ class CohereAsrModelArchConfigConvertor(ModelArchConfigConvertorBase):
        )
        return enc_num_kv_heads

+    def is_mm_prefix_lm(self) -> bool:
+        return False
+

 class MambaModelArchConfigConvertor(ModelArchConfigConvertorBase):
    def get_head_size(self) -> int:
@@ -85,6 +85,7 @@ class KVCacheCoordinator(ABC):
        num_encoder_tokens: int,
        total_computed_tokens: int,
        num_tokens_main_model: int,
+        apply_admission_cap: bool = False,
    ) -> int:
        """
        Get the number of blocks needed to be allocated for the request.
@@ -101,6 +102,10 @@ class KVCacheCoordinator(ABC):
            num_tokens_main_model: The number of tokens for the main model (aka target
                model in spec decode). w/o spec decode, it is num_tokens;
                with spec decode, it is num_tokens - num_lookahead_tokens.
+            apply_admission_cap: If True, apply the recycling-aware
+                per-request admission cap (SWA / chunked-local). Set only by
+                the full-sequence admission gate; per-step allocation must
+                leave it False so the predictor matches `allocate_new_blocks`.

        Returns:
            The number of blocks to allocate.
@@ -111,7 +116,12 @@ class KVCacheCoordinator(ABC):
                # For cross-attention, we issue a single static allocation
                # of blocks based on the number of encoder input tokens.
                num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
-                    request_id, num_encoder_tokens, [], 0, num_encoder_tokens
+                    request_id,
+                    num_encoder_tokens,
+                    [],
+                    0,
+                    num_encoder_tokens,
+                    apply_admission_cap=apply_admission_cap,
                )
            else:
                num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
@@ -120,6 +130,7 @@ class KVCacheCoordinator(ABC):
                    new_computed_blocks[i],
                    total_computed_tokens,
                    num_tokens_main_model,
+                    apply_admission_cap=apply_admission_cap,
                )
        return num_blocks_to_allocate

@@ -257,6 +257,7 @@ class KVCacheManager:
            num_encoder_tokens=num_encoder_tokens,
            total_computed_tokens=total_computed_tokens,
            num_tokens_main_model=full_num_tokens,
+            apply_admission_cap=True,
        )

        return num_blocks_to_allocate <= self.block_pool.get_num_free_blocks()
@@ -92,6 +92,7 @@ class SingleTypeKVCacheManager(ABC):
        new_computed_blocks: Sequence[KVCacheBlock],
        total_computed_tokens: int,
        num_tokens_main_model: int,
+        apply_admission_cap: bool = False,
    ) -> int:
        """
        Get the number of blocks needed to be allocated for the request.
@@ -107,13 +108,16 @@ class SingleTypeKVCacheManager(ABC):
            num_tokens_main_model: The number of tokens for the main model (aka target
                model in spec decode). w/o spec decode, it is num_tokens;
                with spec decode, it is num_tokens - num_lookahead_tokens.
+            apply_admission_cap: If True, clamp by `num_required_blocks` by
+                `_max_admission_blocks_per_request`for recycling-aware specs
+                (SWA, chunked-local).

        Returns:
            The number of blocks to allocate.
        """

        num_required_blocks = cdiv(num_tokens, self.block_size)
-        if self._max_admission_blocks_per_request is not None:
+        if apply_admission_cap and self._max_admission_blocks_per_request is not None:
            # Recycling-aware specs (SWA, chunked-local) cap the per-request
            # reservation here so admission matches the startup pool sizer
            # (`SlidingWindowSpec.max_admission_blocks_per_request` / its
@@ -893,6 +897,7 @@ class MambaManager(SingleTypeKVCacheManager):
        new_computed_blocks: Sequence[KVCacheBlock],
        total_computed_tokens: int,
        num_tokens_main_model: int,
+        apply_admission_cap: bool = False,
    ) -> int:
        assert isinstance(self.kv_cache_spec, MambaSpec)
        if (
@@ -917,6 +922,7 @@ class MambaManager(SingleTypeKVCacheManager):
                new_computed_blocks,
                total_computed_tokens,
                num_tokens_main_model,
+                apply_admission_cap=apply_admission_cap,
            )
        else:
            # We don't allocate blocks for lookahead tokens in align mode, because if
@@ -1,197 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-OffloadingManager class for managing KV data offloading in vLLM v1
-
-This class runs in the scheduler, tracks which blocks are offloaded
-and their address.
-
-The class provides the following primitives:
-    lookup() - check whether a single block is offloaded and ready.
-    prepare_load() - prepare given blocks to be read.
-        The given blocks will be protected from eviction.
-        This function returns a LoadSpec which encapsulates
-        information required for performing the load.
-    touch() - marks the give blocks as recently used. Can be used
-        to track block's LRU. This function is separated from the
-        prepare_load function to allow setting block recency even
-        for blocks which do not need reading from the cache, such as
-        blocks that are cached by the GPU prefix cache.
-    complete_load() - mark blocks which were previously prepared to be
-        loaded as done loading. This is to re-allow their eviction.
-    prepare_store() - prepare the given blocks to be written.
-        Returns a StoreSpec encapsulating offloading information,
-        as well as a list of blocks that were evicted as a result.
-    complete_store() - marks a previous store as completed.
-        Following this call, the given blocks will become loadable.
-"""
-
-from abc import ABC, abstractmethod
-from collections.abc import Iterable, Sequence
-from dataclasses import dataclass
-from typing import Any, NewType
-
-# `OffloadKey` identifies an offloaded block. It combines a block hash with
-# its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.
-# Use the helper functions below to construct / decompose keys.
-OffloadKey = NewType("OffloadKey", bytes)
-
-
-def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey:
-    """Pack a block hash and group index into an `OffloadKey`."""
-    return OffloadKey(block_hash + group_idx.to_bytes(4, "big", signed=False))
-
-
-def get_offload_block_hash(key: OffloadKey) -> bytes:
-    """Extract the block hash from an `OffloadKey`."""
-    return key[:-4]
-
-
-def get_offload_group_idx(key: OffloadKey) -> int:
-    """Extract the group index from an `OffloadKey`."""
-    return int.from_bytes(key[-4:], "big", signed=False)
-
-
-@dataclass
-class ReqContext:
-    kv_transfer_params: dict[str, Any] | None = None
-
-
-class LoadStoreSpec(ABC):
-    """
-    Abstract metadata that encapsulates information allowing a worker
-    to load, and optionally also to store, blocks of KV data.
-    """
-
-    @staticmethod
-    @abstractmethod
-    def medium() -> str:
-        """
-        Returns a string representation of the medium type
-        this store/load targets.
-        """
-        pass
-
-
-@dataclass
-class PrepareStoreOutput:
-    keys_to_store: list[OffloadKey]
-    store_spec: LoadStoreSpec
-    evicted_keys: list[OffloadKey]
-
-
-@dataclass
-class OffloadingEvent:
-    keys: list[OffloadKey]
-    medium: str
-    # True if blocks are removed, False if stored
-    removed: bool
-
-
-class OffloadingManager(ABC):
-    @abstractmethod
-    def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
-        """
-        Checks whether a single block is offloaded and ready to be read.
-
-        Args:
-            key: the key identifying the block to lookup.
-            req_context: per-request context (e.g. kv_transfer_params).
-
-        Returns:
-            True if the block is offloaded and ready, False if not,
-            or None if the lookup should be retried later.
-            Returning None will delay the request handling by the vLLM
-            scheduler.
-        """
-        pass
-
-    @abstractmethod
-    def prepare_load(
-        self,
-        keys: Sequence[OffloadKey],
-        req_context: ReqContext,
-    ) -> LoadStoreSpec:
-        """
-        Prepare the given blocks to be read.
-        The given blocks will be protected from eviction until
-        complete_load is called.
-        It assumes all given blocks are offloaded.
-
-        Args:
-            keys: the keys identifying the blocks.
-            req_context: per-request context (e.g. kv_transfer_params).
-
-        Returns:
-            A LoadStoreSpec that can be used by a worker to locate and load
-            the actual offloaded KV data.
-        """
-        pass
-
-    def touch(self, keys: Sequence[OffloadKey]):
-        """
-        Mark the given blocks as recently used.
-        This could in practice mean moving them to the end of an LRU list.
-
-        Args:
-            keys: the keys identifying the blocks.
-        """
-        return
-
-    def complete_load(self, keys: Iterable[OffloadKey]):
-        """
-        Marks previous blocks that were prepared to load as done loading.
-
-        Args:
-            keys: the keys identifying the blocks.
-        """
-        return
-
-    @abstractmethod
-    def prepare_store(
-        self,
-        keys: Sequence[OffloadKey],
-        req_context: ReqContext,
-    ) -> PrepareStoreOutput | None:
-        """
-        Prepare the given blocks to be offloaded.
-        The given blocks will be protected from eviction until
-        complete_store is called.
-
-        Args:
-            keys: the keys identifying the blocks.
-            req_context: per-request context (e.g. kv_transfer_params).
-
-        Returns:
-            A PrepareStoreOutput indicating which blocks need storing,
-            where to store them (LoadStoreSpec), and list of blocks that
-            were evicted as a result.
-            None is returned if the blocks cannot be stored.
-        """
-        pass
-
-    def complete_store(self, keys: Iterable[OffloadKey], success: bool = True):
-        """
-        Marks blocks which were previously prepared to be stored, as stored.
-        Following this call, the blocks become loadable.
-        If if_success is False, blocks that were not marked as stored will be
-        removed.
-
-        Args:
-            keys: the keys identifying the blocks.
-            success: whether the blocks were stored successfully.
-        """
-        return
-
-    def take_events(self) -> Iterable[OffloadingEvent]:
-        """
-        Take the offloading events from the manager.
-
-        Yields:
-            New OffloadingEvents collected since the last call.
-        """
-        return ()
-
-    def shutdown(self) -> None:
-        """Shutdown the manager and release any resources."""
-        return
@@ -0,0 +1,390 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Core abstractions for KV cache offloading in vLLM v1.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Iterator, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, NewType
+
+import numpy as np
+import torch
+
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.kv_offload.worker.worker import OffloadingHandler
+
+# `OffloadKey` identifies an offloaded block. It combines a block hash with
+# its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.
+# Use the helper functions below to construct / decompose keys.
+OffloadKey = NewType("OffloadKey", bytes)
+
+logger = init_logger(__name__)
+
+
+def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey:
+    """Pack a block hash and group index into an `OffloadKey`."""
+    return OffloadKey(block_hash + group_idx.to_bytes(4, "big", signed=False))
+
+
+def get_offload_block_hash(key: OffloadKey) -> bytes:
+    """Extract the block hash from an `OffloadKey`."""
+    return key[:-4]
+
+
+def get_offload_group_idx(key: OffloadKey) -> int:
+    """Extract the group index from an `OffloadKey`."""
+    return int.from_bytes(key[-4:], "big", signed=False)
+
+
+@dataclass
+class ReqContext:
+    kv_transfer_params: dict[str, Any] | None = None
+
+
+class LoadStoreSpec(ABC):
+    """
+    Abstract metadata that encapsulates information allowing a worker
+    to load, and optionally also to store, blocks of KV data.
+    """
+
+    @staticmethod
+    @abstractmethod
+    def medium() -> str:
+        """
+        Returns a string representation of the medium type
+        this store/load targets.
+        """
+        pass
+
+
+@dataclass
+class PrepareStoreOutput:
+    keys_to_store: list[OffloadKey]
+    store_spec: LoadStoreSpec
+    evicted_keys: list[OffloadKey]
+
+
+@dataclass
+class OffloadingEvent:
+    keys: list[OffloadKey]
+    medium: str
+    # True if blocks are removed, False if stored
+    removed: bool
+
+
+"""
+OffloadingManager class for managing KV data offloading in vLLM v1
+
+This class runs in the scheduler, tracks which blocks are offloaded
+and their address.
+
+The class provides the following primitives:
+    lookup() - check whether a single block is offloaded and ready.
+    prepare_load() - prepare given blocks to be read.
+        The given blocks will be protected from eviction.
+        This function returns a LoadSpec which encapsulates
+        information required for performing the load.
+    touch() - marks the give blocks as recently used. Can be used
+        to track block's LRU. This function is separated from the
+        prepare_load function to allow setting block recency even
+        for blocks which do not need reading from the cache, such as
+        blocks that are cached by the GPU prefix cache.
+    complete_load() - mark blocks which were previously prepared to be
+        loaded as done loading. This is to re-allow their eviction.
+    prepare_store() - prepare the given blocks to be written.
+        Returns a StoreSpec encapsulating offloading information,
+        as well as a list of blocks that were evicted as a result.
+    complete_store() - marks a previous store as completed.
+        Following this call, the given blocks will become loadable.
+"""
+
+
+class OffloadingManager(ABC):
+    @abstractmethod
+    def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
+        """
+        Checks whether a single block is offloaded and ready to be read.
+
+        Args:
+            key: the key identifying the block to lookup.
+            req_context: per-request context (e.g. kv_transfer_params).
+
+        Returns:
+            True if the block is offloaded and ready, False if not,
+            or None if the lookup should be retried later.
+            Returning None will delay the request handling by the vLLM
+            scheduler.
+        """
+        pass
+
+    @abstractmethod
+    def prepare_load(
+        self,
+        keys: Sequence[OffloadKey],
+        req_context: ReqContext,
+    ) -> LoadStoreSpec:
+        """
+        Prepare the given blocks to be read.
+        The given blocks will be protected from eviction until
+        complete_load is called.
+        It assumes all given blocks are offloaded.
+
+        Args:
+            keys: the keys identifying the blocks.
+            req_context: per-request context (e.g. kv_transfer_params).
+
+        Returns:
+            A LoadStoreSpec that can be used by a worker to locate and load
+            the actual offloaded KV data.
+        """
+        pass
+
+    def touch(self, keys: Sequence[OffloadKey]):
+        """
+        Mark the given blocks as recently used.
+        This could in practice mean moving them to the end of an LRU list.
+
+        Args:
+            keys: the keys identifying the blocks.
+        """
+        return
+
+    def complete_load(self, keys: Iterable[OffloadKey]):
+        """
+        Marks previous blocks that were prepared to load as done loading.
+
+        Args:
+            keys: the keys identifying the blocks.
+        """
+        return
+
+    @abstractmethod
+    def prepare_store(
+        self,
+        keys: Sequence[OffloadKey],
+        req_context: ReqContext,
+    ) -> PrepareStoreOutput | None:
+        """
+        Prepare the given blocks to be offloaded.
+        The given blocks will be protected from eviction until
+        complete_store is called.
+
+        Args:
+            keys: the keys identifying the blocks.
+            req_context: per-request context (e.g. kv_transfer_params).
+
+        Returns:
+            A PrepareStoreOutput indicating which blocks need storing,
+            where to store them (LoadStoreSpec), and list of blocks that
+            were evicted as a result.
+            None is returned if the blocks cannot be stored.
+        """
+        pass
+
+    def complete_store(self, keys: Iterable[OffloadKey], success: bool = True):
+        """
+        Marks blocks which were previously prepared to be stored, as stored.
+        Following this call, the blocks become loadable.
+        If if_success is False, blocks that were not marked as stored will be
+        removed.
+
+        Args:
+            keys: the keys identifying the blocks.
+            success: whether the blocks were stored successfully.
+        """
+        return
+
+    def take_events(self) -> Iterable[OffloadingEvent]:
+        """
+        Take the offloading events from the manager.
+
+        Yields:
+            New OffloadingEvents collected since the last call.
+        """
+        return ()
+
+    def shutdown(self) -> None:
+        """Shutdown the manager and release any resources."""
+        return
+
+
+class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
+    """
+    Spec for loading/storing KV blocks from given block numbers.
+    """
+
+    def __init__(self, block_ids: list[int]):
+        self.block_ids = np.array(block_ids, dtype=np.int64)
+
+    def __repr__(self) -> str:
+        return repr(self.block_ids)
+
+
+class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
+    """
+    Spec for loading/storing a KV block to GPU memory.
+
+    If there are multiple KV groups, the blocks are expected to be
+    ordered by the group index.
+    In that case, group_sizes[i] determines the number of blocks
+    per the i-th KV group, and thus sum(group_sizes) == len(block_ids).
+    group_sizes=None indicates a single KV group.
+
+    If block_indices is given, each group (determined by group_sizes) of block IDs
+    will correspond to logically contiguous blocks, e.g. blocks 5-10 of a some request.
+    block_indices[i] will represent the block index of the first block in group #i.
+    Thus, len(block_indices) == len(group_sizes) = number of KV cache groups.
+    This information is required in order to support off/loading from offloaded blocks
+    which are larger than GPU blocks.
+    In such cases, the first GPU block per each group may be unaligned to the offloaded
+    block size, and so knowing block_indices[i] allows the worker to correctly
+    skip part of the first matching offloaded block.
+    """
+
+    def __init__(
+        self,
+        block_ids: list[int],
+        group_sizes: Sequence[int],
+        block_indices: Sequence[int],
+    ):
+        super().__init__(block_ids)
+        assert sum(group_sizes) == len(block_ids)
+        assert len(block_indices) == len(group_sizes)
+        self.group_sizes: Sequence[int] = group_sizes
+        self.block_indices: Sequence[int] = block_indices
+
+    @staticmethod
+    def medium() -> str:
+        return "GPU"
+
+
+@dataclass
+class CanonicalKVCacheTensor:
+    """
+    A canonicalized KV cache tensor whose first dimension is num_blocks.
+
+    For attention backends where the raw tensor has num_blocks at a
+    non-leading physical dimension (e.g. FlashAttention's
+    (2, num_blocks, ...) layout), the tensor is split so that each
+    resulting CanonicalKVCacheTensor starts with (num_blocks, ...).
+    """
+
+    # The KV cache tensor with shape (num_blocks, ...)
+    tensor: torch.Tensor
+    # The (possibly padded) page size per block in bytes
+    page_size_bytes: int
+
+
+@dataclass
+class CanonicalKVCacheRef:
+    """
+    Per-layer (or group of layers) reference to a specific (by index)
+    CanonicalKVCacheTensor and records the un-padded page size used by that layer.
+    """
+
+    # Index into the list of CanonicalKVCacheTensor objects
+    tensor_idx: int
+    # The un-padded page size per block in bytes
+    page_size_bytes: int
+
+
+@dataclass
+class CanonicalKVCaches:
+    """
+    Canonicalized block-level representation of the KV caches.
+
+    Composed of:
+        - Unique list of KV cache data tensors,
+          each with shape (num_blocks, page_size_in_bytes) and int8 dtype.
+        - Per-group data references of the tensors.
+          i.e. how each KV cache group maps to the tensors.
+    """
+
+    # Ordered list of unique block tensors, each with shape
+    # (num_blocks, ...).
+    tensors: list[CanonicalKVCacheTensor]
+    # Per-KV-cache-group list of data references that map each layer
+    # in the group to the appropriate entry in the tensors list.
+    group_data_refs: list[list[CanonicalKVCacheRef]]
+
+
+class OffloadingSpec(ABC):
+    """Spec for an offloading connector"""
+
+    def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
+        logger.warning(
+            "Initializing OffloadingSpec. This API is experimental and "
+            "subject to change in the future as we iterate the design."
+        )
+        self.vllm_config = vllm_config
+        self.kv_cache_config = kv_cache_config
+
+        kv_transfer_config = vllm_config.kv_transfer_config
+        assert kv_transfer_config is not None
+        self.extra_config = kv_transfer_config.kv_connector_extra_config
+
+        # block size used by vLLM for hashing request tokens for the sake
+        # of enabling prefix caching
+        self.hash_block_size = vllm_config.cache_config.block_size
+        # gpu block size per group
+        self.gpu_block_size: tuple[int, ...] = tuple(
+            kv_cache_group.kv_cache_spec.block_size
+            for kv_cache_group in kv_cache_config.kv_cache_groups
+        )
+
+        for block_size in self.gpu_block_size:
+            assert block_size % self.hash_block_size == 0, (
+                f"gpu_block_size={block_size} not divisible by "
+                f"hash_block_size={self.hash_block_size}. "
+                f"Hybrid models (e.g. Mamba+Attention) need "
+                f"--enable-prefix-caching to align block sizes."
+            )
+
+        # offloaded_block_size / gpu_block_size
+        self.block_size_factor: int = 1
+
+        offloaded_block_size = self.extra_config.get("block_size")
+        if offloaded_block_size is not None:
+            offloaded_block_size_int = int(offloaded_block_size)
+            gpu_block_sizes = set(self.gpu_block_size)
+            assert len(gpu_block_sizes) == 1, (
+                "If 'block_size' is specified in kv_connector_extra_config, "
+                "there must be at least one KV cache group, "
+                "and all groups must have the same block size."
+            )
+            gpu_block_size = gpu_block_sizes.pop()
+
+            assert offloaded_block_size_int % gpu_block_size == 0
+            self.block_size_factor = offloaded_block_size_int // gpu_block_size
+
+    @abstractmethod
+    def get_manager(self) -> OffloadingManager:
+        """
+        Get an OffloadingManager that will be used
+        by the scheduler-side offloading connector to track
+        offloaded blocks and manage evictions.
+        """
+        pass
+
+    @abstractmethod
+    def get_handlers(
+        self, kv_caches: CanonicalKVCaches
+    ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
+        """
+        Get offloading handlers along with their respective src and dst types.
+
+        Args:
+            kv_caches: Canonicalized KV caches.
+
+        Yields:
+            Tuples of (src_type, dst_type, offloading_handler).
+        """
+        pass
@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec
+
+
+class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
+    """
+    Spec for loading/storing a KV block to CPU memory.
+    """
+
+    @staticmethod
+    def medium() -> str:
+        return "CPU"
@@ -11,9 +11,13 @@ from vllm import _custom_ops as ops
 from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
 from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.v1.kv_offload.base import (
+    BlockIDsLoadStoreSpec,
+    CanonicalKVCacheRef,
+    CanonicalKVCaches,
+    GPULoadStoreSpec,
+)
 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
-from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec, GPULoadStoreSpec
-from vllm.v1.kv_offload.spec import CanonicalKVCacheRef, CanonicalKVCaches
 from vllm.v1.kv_offload.worker.worker import (
    OffloadingHandler,
    TransferResult,
@@ -3,7 +3,7 @@
 from collections.abc import Iterable, Sequence
 from typing import Literal

-from vllm.v1.kv_offload.abstract import (
+from vllm.v1.kv_offload.base import (
    LoadStoreSpec,
    OffloadingEvent,
    OffloadingManager,
@@ -11,10 +11,10 @@ from vllm.v1.kv_offload.abstract import (
    PrepareStoreOutput,
    ReqContext,
 )
-from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
+from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
+from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
 from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
-from vllm.v1.kv_offload.mediums import CPULoadStoreSpec

 _CACHE_POLICIES: dict[str, type[CachePolicy]] = {
    "lru": LRUCachePolicy,
@@ -3,8 +3,8 @@
 from collections import OrderedDict
 from collections.abc import Iterable

-from vllm.v1.kv_offload.abstract import OffloadKey
-from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
+from vllm.v1.kv_offload.base import OffloadKey
+from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy


 class ARCCachePolicy(CachePolicy):
@@ -4,7 +4,7 @@ import ctypes
 from abc import ABC, abstractmethod
 from collections.abc import Iterable

-from vllm.v1.kv_offload.abstract import OffloadKey
+from vllm.v1.kv_offload.base import OffloadKey


 class BlockStatus(ctypes.Structure):
@@ -3,8 +3,8 @@
 from collections import OrderedDict
 from collections.abc import Iterable

-from vllm.v1.kv_offload.abstract import OffloadKey
-from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
+from vllm.v1.kv_offload.base import OffloadKey
+from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy


 class LRUCachePolicy(CachePolicy):
@@ -5,12 +5,17 @@ from collections.abc import Iterator
 from vllm.config import VllmConfig
 from vllm.platforms import current_platform
 from vllm.v1.kv_cache_interface import KVCacheConfig
-from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
+from vllm.v1.kv_offload.base import (
+    CanonicalKVCaches,
+    GPULoadStoreSpec,
+    LoadStoreSpec,
+    OffloadingManager,
+    OffloadingSpec,
+)
+from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
+from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
-from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
 from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
-from vllm.v1.kv_offload.spec import CanonicalKVCaches, OffloadingSpec
-from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
 from vllm.v1.kv_offload.worker.worker import OffloadingHandler


@@ -5,7 +5,7 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING

 from vllm.logger import init_logger
-from vllm.v1.kv_offload.spec import OffloadingSpec
+from vllm.v1.kv_offload.base import OffloadingSpec

 if TYPE_CHECKING:
    from vllm.config import VllmConfig
@@ -1,68 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from abc import ABC
-from collections.abc import Sequence
-
-import numpy as np
-
-from vllm.v1.kv_offload.abstract import LoadStoreSpec
-
-
-class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
-    """
-    Spec for loading/storing KV blocks from given block numbers.
-    """
-
-    def __init__(self, block_ids: list[int]):
-        self.block_ids = np.array(block_ids, dtype=np.int64)
-
-    def __repr__(self) -> str:
-        return repr(self.block_ids)
-
-
-class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
-    """
-    Spec for loading/storing a KV block to GPU memory.
-
-    If there are multiple KV groups, the blocks are expected to be
-    ordered by the group index.
-    In that case, group_sizes[i] determines the number of blocks
-    per the i-th KV group, and thus sum(group_sizes) == len(block_ids).
-    group_sizes=None indicates a single KV group.
-
-    If block_indices is given, each group (determined by group_sizes) of block IDs
-    will correspond to logically contiguous blocks, e.g. blocks 5-10 of a some request.
-    block_indices[i] will represent the block index of the first block in group #i.
-    Thus, len(block_indices) == len(group_sizes) = number of KV cache groups.
-    This information is required in order to support off/loading from offloaded blocks
-    which are larger than GPU blocks.
-    In such cases, the first GPU block per each group may be unaligned to the offloaded
-    block size, and so knowing block_indices[i] allows the worker to correctly
-    skip part of the first matching offloaded block.
-    """
-
-    def __init__(
-        self,
-        block_ids: list[int],
-        group_sizes: Sequence[int],
-        block_indices: Sequence[int],
-    ):
-        super().__init__(block_ids)
-        assert sum(group_sizes) == len(block_ids)
-        assert len(block_indices) == len(group_sizes)
-        self.group_sizes: Sequence[int] = group_sizes
-        self.block_indices: Sequence[int] = block_indices
-
-    @staticmethod
-    def medium() -> str:
-        return "GPU"
-
-
-class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
-    """
-    Spec for loading/storing a KV block to CPU memory.
-    """
-
-    @staticmethod
-    def medium() -> str:
-        return "CPU"
@@ -10,7 +10,7 @@ FilterReusedOffloadingManager — OffloadingManager decorator that skips
 from collections import OrderedDict
 from collections.abc import Iterable, Sequence

-from vllm.v1.kv_offload.abstract import (
+from vllm.v1.kv_offload.base import (
    LoadStoreSpec,
    OffloadingEvent,
    OffloadingManager,
@@ -1,142 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from abc import ABC, abstractmethod
-from collections.abc import Iterator
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
-
-import torch
-
-from vllm.logger import init_logger
-from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
-from vllm.v1.kv_offload.worker.worker import OffloadingHandler
-
-if TYPE_CHECKING:
-    from vllm.config import VllmConfig
-    from vllm.v1.kv_cache_interface import KVCacheConfig
-
-logger = init_logger(__name__)
-
-
-@dataclass
-class CanonicalKVCacheTensor:
-    """
-    A canonicalized KV cache tensor whose first dimension is num_blocks.
-
-    For attention backends where the raw tensor has num_blocks at a
-    non-leading physical dimension (e.g. FlashAttention's
-    (2, num_blocks, ...) layout), the tensor is split so that each
-    resulting CanonicalKVCacheTensor starts with (num_blocks, ...).
-    """
-
-    # The KV cache tensor with shape (num_blocks, ...)
-    tensor: torch.Tensor
-    # The (possibly padded) page size per block in bytes
-    page_size_bytes: int
-
-
-@dataclass
-class CanonicalKVCacheRef:
-    """
-    Per-layer (or group of layers) reference to a specific (by index)
-    CanonicalKVCacheTensor and records the un-padded page size used by that layer.
-    """
-
-    # Index into the list of CanonicalKVCacheTensor objects
-    tensor_idx: int
-    # The un-padded page size per block in bytes
-    page_size_bytes: int
-
-
-@dataclass
-class CanonicalKVCaches:
-    """
-    Canonicalized block-level representation of the KV caches.
-
-    Composed of:
-        - Unique list of KV cache data tensors,
-          each with shape (num_blocks, page_size_in_bytes) and int8 dtype.
-        - Per-group data references of the tensors.
-          i.e. how each KV cache group maps to the tensors.
-    """
-
-    # Ordered list of unique block tensors, each with shape
-    # (num_blocks, ...).
-    tensors: list[CanonicalKVCacheTensor]
-    # Per-KV-cache-group list of data references that map each layer
-    # in the group to the appropriate entry in the tensors list.
-    group_data_refs: list[list[CanonicalKVCacheRef]]
-
-
-class OffloadingSpec(ABC):
-    """Spec for an offloading connector"""
-
-    def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
-        logger.warning(
-            "Initializing OffloadingSpec. This API is experimental and "
-            "subject to change in the future as we iterate the design."
-        )
-        self.vllm_config = vllm_config
-        self.kv_cache_config = kv_cache_config
-
-        kv_transfer_config = vllm_config.kv_transfer_config
-        assert kv_transfer_config is not None
-        self.extra_config = kv_transfer_config.kv_connector_extra_config
-
-        # block size used by vLLM for hashing request tokens for the sake
-        # of enabling prefix caching
-        self.hash_block_size = vllm_config.cache_config.block_size
-        # gpu block size per group
-        self.gpu_block_size: tuple[int, ...] = tuple(
-            kv_cache_group.kv_cache_spec.block_size
-            for kv_cache_group in kv_cache_config.kv_cache_groups
-        )
-
-        for block_size in self.gpu_block_size:
-            assert block_size % self.hash_block_size == 0, (
-                f"gpu_block_size={block_size} not divisible by "
-                f"hash_block_size={self.hash_block_size}. "
-                f"Hybrid models (e.g. Mamba+Attention) need "
-                f"--enable-prefix-caching to align block sizes."
-            )
-
-        # offloaded_block_size / gpu_block_size
-        self.block_size_factor: int = 1
-
-        offloaded_block_size = self.extra_config.get("block_size")
-        if offloaded_block_size is not None:
-            offloaded_block_size_int = int(offloaded_block_size)
-            gpu_block_sizes = set(self.gpu_block_size)
-            assert len(gpu_block_sizes) == 1, (
-                "If 'block_size' is specified in kv_connector_extra_config, "
-                "there must be at least one KV cache group, "
-                "and all groups must have the same block size."
-            )
-            gpu_block_size = gpu_block_sizes.pop()
-
-            assert offloaded_block_size_int % gpu_block_size == 0
-            self.block_size_factor = offloaded_block_size_int // gpu_block_size
-
-    @abstractmethod
-    def get_manager(self) -> OffloadingManager:
-        """
-        Get an OffloadingManager that will be used
-        by the scheduler-side offloading connector to track
-        offloaded blocks and manage evictions.
-        """
-        pass
-
-    @abstractmethod
-    def get_handlers(
-        self, kv_caches: CanonicalKVCaches
-    ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
-        """
-        Get offloading handlers along with their respective src and dst types.
-
-        Args:
-            kv_caches: Canonicalized KV caches.
-
-        Yields:
-            Tuples of (src_type, dst_type, offloading_handler).
-        """
-        pass
@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass

 from vllm.logger import init_logger
-from vllm.v1.kv_offload.abstract import LoadStoreSpec
+from vllm.v1.kv_offload.base import LoadStoreSpec

 # a single transfer spec (src_blocks_spec, dst_blocks_spec)
 TransferSpec = tuple[LoadStoreSpec, LoadStoreSpec]
@@ -5,12 +5,13 @@
 import gc
 import os
 from collections.abc import Callable
-from contextlib import AbstractContextManager, nullcontext
+from contextlib import AbstractContextManager, contextmanager, nullcontext
 from datetime import timedelta
 from types import NoneType
 from typing import TYPE_CHECKING, Any

 import numpy as np
+import regex as re
 import torch
 import torch.nn as nn

@@ -208,6 +209,30 @@ class Worker(WorkerBase):
            )
        return allocator.use_memory_pool(tag=tag)

+    @contextmanager
+    def _scoped_allocator_max_split(self, max_split_size_mb: int):
+        """Temporarily set max_split_size_mb to reduce allocator fragmentation at the
+        cost of more cudaMalloc calls (negligible in practice). Restores the original
+        value on exit."""
+        if not current_platform.is_cuda():
+            yield
+            return
+
+        conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
+        match = re.search(r"max_split_size_mb:(\d+)", conf)
+        original_value = match.group(1) if match else None
+
+        torch._C._accelerator_setAllocatorSettings(
+            f"max_split_size_mb:{max_split_size_mb}"
+        )
+        try:
+            yield
+        finally:
+            # PyTorch defaults to SIZE_MAX (no limit).
+            _SIZE_MAX_MB = (2**64 - 1) // (1024 * 1024)
+            restore = original_value if original_value else str(_SIZE_MAX_MB)
+            torch._C._accelerator_setAllocatorSettings(f"max_split_size_mb:{restore}")
+
    @instrument(span_name="Init device")
    def init_device(self):
        if self.device_config.device_type == "cuda":
@@ -312,6 +337,8 @@ class Worker(WorkerBase):
        with (
            self._maybe_get_memory_pool_context(tag="weights"),
            set_current_vllm_config(self.vllm_config),
+            # 20 MiB is the minimum PyTorch allows for max_split_size_mb.
+            self._scoped_allocator_max_split(max_split_size_mb=20),
        ):
            self.model_runner.load_model(load_dummy_weights=load_dummy_weights)
Author	SHA1	Message	Date
yewentao256	77230471c0	remove torch 2.9, 2.10 workround Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-04-30 21:10:50 +00:00
Stefano CastagnettaandGitHub	efb4cdf2b8	[CI/Build] Skip Prithvi/Terratorch model-registry tests when terratorch is missing (#41389 ) Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>	2026-04-30 12:47:55 -07:00
Stefano Castagnetta GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	92a7c121b6	[CI] Add MTP coverage: Qwen3.5 correctness + no-sync spec decode (#40472 ) Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-04-30 12:24:09 -07:00
Jee Jee LiandGitHub	307b17ce33	[DSV4] Avoid redundant dtype conversion. (#41374 ) Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>	2026-04-30 09:57:27 -07:00
wenjun liu GitHub Kunshang Ji	3ca6ca210f	xpu docker: pin oneAPI to 2025.3 and avoid unintended 2026 upgrade (#41380 ) Signed-off-by: wendyliu235 <wenjun.liu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-04-30 16:02:23 +00:00
Stefano CastagnettaandGitHub	10558f5f46	[CI/Build] Skip terratorch + torchgeo while PyPI has lightning quarantined (#41377 ) Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>	2026-04-30 07:59:07 -07:00
tej GitHub TJian HAIAI	121dbe7a22	[ROCm] ROCm DeepEP API updated to latest (#39721 ) Signed-off-by: Tej Kiran <vpolamre@amd.com> Signed-off-by: tej <37236721+itej89@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: HAIAI <39548240+HAIAI@users.noreply.github.com>	2026-04-30 07:46:59 -07:00
Matthew BonanniandGitHub	f03d82efdd	[UX][Bugfix] Fix OOM by setting PyTorch `max_split_size_mb` during model loading (#41268 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>	2026-04-30 07:46:54 -07:00
Ilya Markov GitHub Markov Ilya Nicolò Lucchesi	a7fb008510	[EPLB] Optimize memory overhead in Nixl communicator (#40013 ) Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>	2026-04-30 07:46:49 -07:00
Harry MellorandGitHub	ff449b6426	Stop mergify labelling from skipping pre-commit (#41362 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-04-30 05:48:38 -07:00
Stefano Castagnetta GitHub Harry Mellor	3527229517	[Doc] Fix RTD build: pytorch.org/docs/stable/objects.inv returns 404 (#41353 ) Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-04-30 05:06:44 -07:00
Xiaoshuang Wang GitHub Cyrus Leungmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	b55b26520c	[MoE] Make MoERunnerInterface a PluggableLayer for OOT support (#35178 ) Signed-off-by: wxsIcey <1790571317@qq.com> Signed-off-by: Icey <1790571317@qq.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-04-30 03:31:08 -07:00
snadampalandGitHub	3179e53135	[P/D] Prefill compute optimizations with bi-directional KV cache transfers between P and D nodes (#32553 ) Signed-off-by: Sunita Nadampalli <nadampal@amazon.com>	2026-04-30 10:14:20 +00:00
Nicolò LucchesiandGitHub	efdc95674d	[KVConnector] MultiConnector SupportsHMA (#39571 ) Signed-off-by: NickLucche <nlucches@redhat.com>	2026-04-30 02:10:50 -07:00
Chenxi Qian GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	54146a9bf9	[Bugfix] correct h matrix layout in chunk_kda output kernel (#40956 ) Signed-off-by: ChenxiQian <chenxi.qian.cq@outlook.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-04-30 16:22:41 +08:00
Baekpica GitHub sunghoon.baek OpenAI Codex	ca97f7b9bb	Fix Gemma4 MoE expert weight remapping (#41206 ) Signed-off-by: sunghoon.baek <sunghoon.baek@connectfy.cloud> Co-authored-by: sunghoon.baek <sunghoon.baek@connectfy.cloud> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-04-30 00:12:42 -07:00
Ekagra RanjanandGitHub	a04e0cf3b8	Fix Cohere ASR after HF upgrade (#40582 ) Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com>	2026-04-29 23:39:04 -07:00
Dhruv Singal GitHub Claude Opus 4.6 Dhruv Singal Your Name OpenCode Simon Mo	cb1b02d0e8	[Frontend] Add VLLM_SKIP_MODEL_NAME_VALIDATION environment variable (#34676 ) Signed-off-by: Dhruv Singal <dhruvsingalabc@gmail.com> Signed-off-by: Dhruv Singal <dsingal@Dhruvs-MacBook-Pro.local> Signed-off-by: Your Name <you@example.com> Signed-off-by: vLLM Assistant <assistant@vllm.ai> Signed-off-by: Simon Mo <simon.mo@hey.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Dhruv Singal <dsingal@Dhruvs-MacBook-Pro.local> Co-authored-by: Your Name <you@example.com> Co-authored-by: OpenCode <noreply@openai.com> Co-authored-by: Simon Mo <simon.mo@hey.com>	2026-04-29 23:19:09 -07:00
Yongye Zhu GitHub Claude Opus 4.7	a749a33d8d	[Bugfix] Fix persistent_topk cooperative deadlock at TopK=1024 (#41189 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-04-29 21:03:45 -07:00
Martin Hickey GitHub Or Ozeri	c42981d034	[Refactor][kv_offload] KV Offloading maintainability improvements (#40538 ) Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>	2026-04-30 05:55:31 +03:00
Wei ZhaoandGitHub	0ff1bf9bb1	[Bugfix] Fix failure to allocate KV blocks error (#41282 ) Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>	2026-04-29 18:44:07 -07:00