Use -1 expert ID for padding and non-local tokens in DeepEP v2

Padding rows and non-local expert tokens now get expert_id=-1 instead of rank_expert_offset. This allows DeepGemm's is_computation_valid() to skip computation on these rows, which is both a correctness fix (avoids wasting GEMM compute on junk data) and enables future kernel optimizations for early CTA exit on -1 blocks. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Fix local→global expert ID conversion in DeepEP v2 decode path
2026-05-02 11:51:27 -04:00 · 2026-05-02 02:15:31 -04:00 · 2026-04-30 22:54:59 -04:00 · 2026-04-30 22:19:23 -04:00 · 2026-04-30 21:55:14 -04:00 · 2026-04-30 21:51:16 -04:00
90 changed files with 3226 additions and 3193 deletions
@@ -89,16 +89,6 @@ 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,16 +97,3 @@ 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,7 +16,11 @@ permissions:

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

  pre-commit:
    needs: pre-run-check
-    if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
+    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')
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
@@ -217,7 +217,6 @@ 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,
@@ -226,9 +225,6 @@ 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}
@@ -423,7 +419,6 @@ async def send_turn(
        min_tokens,
        max_tokens,
        req_args.timeout_sec,
-        conversation_id=conv_id,
    )

    if response.valid is False:
@@ -82,73 +82,18 @@ 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;
-    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));
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
+        smem_size);
    if (occupancy < 1) occupancy = 1;

-    // 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 max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
    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="f0acb0c6"
+ARG ROCSHMEM_BRANCH="ba0bf0f3"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-ARG DEEPEP_BRANCH="a9ea9774"
+ARG DEEPEP_BRANCH="5d90af8b"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
 ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
@@ -213,15 +213,18 @@ RUN git clone ${ROCSHMEM_REPO} \
 && git checkout ${ROCSHMEM_BRANCH} \
 && mkdir -p projects/rocshmem/build \
 && cd projects/rocshmem/build \
- && INSTALL_PREFIX=${ROCSHMEM_DIR} \
-    ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF
+ && bash ../scripts/build_configs/all_backends \
+      -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
+      -DROCM_PATH=/opt/rocm \
+      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
+      -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 --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
+ && python3 setup.py --variant rocm --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.
@@ -385,16 +388,6 @@ 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,6 +5,9 @@ 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 \
@@ -23,20 +26,8 @@ RUN apt clean && apt-get update -y && \
    python3.12-dev \
    python3-pip

-# 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/*
+RUN apt update && apt upgrade -y && \
+    apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3

 # Install UMD
 RUN mkdir neo && \
@@ -38,20 +38,8 @@ 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, PluggableLayer=MockPluggableLayer),
-)
+mock_if_no_torch("vllm.model_executor.custom_op", MagicMock(CustomOp=MockCustomOp))
 mock_if_no_torch(
    "vllm.utils.torch_utils", MagicMock(direct_register_custom_op=lambda *a, **k: None)
 )
@@ -1,562 +0,0 @@
-# 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,8 +105,7 @@ plugins:
          - https://docs.aiohttp.org/en/stable/objects.inv
          - https://pillow.readthedocs.io/en/stable/objects.inv
          - https://numpy.org/doc/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
+          - https://pytorch.org/docs/stable/objects.inv
  - redirects:
      redirect_maps:
        features/spec_decode/README.md: features/speculative_decoding/README.md
@@ -61,11 +61,7 @@ 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 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
+terratorch >= 1.2.2 # Required for Prithvi tests
 imagehash # Required for Prithvi tests
 segmentation-models-pytorch > 0.4.0 # Required for Prithvi tests

@@ -1,9 +1,15 @@
 # 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
+    # via
+    #   rouge-score
+    #   tensorboard
 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
@@ -19,14 +25,22 @@ 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
+    # via
+    #   -r requirements/test/cuda.in
+    #   terratorch
 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
@@ -40,10 +54,12 @@ arrow==1.3.0
 attrs==24.2.0
    # via
    #   aiohttp
+    #   fiona
    #   hypothesis
    #   jsonlines
    #   jsonschema
    #   pytest-subtests
+    #   rasterio
    #   referencing
 audioread==3.0.1
    # via librosa
@@ -62,7 +78,9 @@ backoff==2.2.1
    #   -r requirements/test/cuda.in
    #   schemathesis
 bitsandbytes==0.49.2
-    # via -r requirements/test/cuda.in
+    # via
+    #   -r requirements/test/cuda.in
+    #   lightning
 black==24.10.0
    # via datamodel-code-generator
 blobfile==3.0.0
@@ -85,9 +103,15 @@ 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
@@ -101,12 +125,25 @@ 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
@@ -154,6 +191,8 @@ 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
@@ -168,10 +207,14 @@ 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
@@ -201,10 +244,13 @@ 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
@@ -221,6 +267,9 @@ fsspec==2024.12.0
    #   evaluate
    #   fastparquet
    #   huggingface-hub
+    #   lightning
+    #   pytorch-lightning
+    #   tacoreader
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
@@ -228,6 +277,12 @@ 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
@@ -262,6 +317,7 @@ 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
@@ -270,6 +326,8 @@ 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
@@ -285,6 +343,7 @@ httpcore==1.0.6
 httpx==0.27.2
    # via
    #   -r requirements/test/cuda.in
+    #   diffusers
    #   huggingface-hub
    #   perceptron
    #   schemathesis
@@ -292,17 +351,23 @@ 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
@@ -327,7 +392,11 @@ imagehash==4.3.2
 imageio==2.37.0
    # via scikit-image
 importlib-metadata==8.7.0
-    # via opentelemetry-api
+    # via
+    #   diffusers
+    #   opentelemetry-api
+importlib-resources==6.5.2
+    # via typeshed-client
 inflect==5.6.2
    # via datamodel-code-generator
 iniconfig==2.0.0
@@ -357,8 +426,14 @@ 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
@@ -377,6 +452,10 @@ 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
@@ -385,6 +464,21 @@ 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
@@ -396,6 +490,8 @@ 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
@@ -404,7 +500,11 @@ markupsafe==3.0.1
    #   mako
    #   werkzeug
 matplotlib==3.9.2
-    # via -r requirements/test/cuda.in
+    # via
+    #   -r requirements/test/cuda.in
+    #   lightning
+    #   pycocotools
+    #   torchgeo
 mbstrdecoder==1.1.3
    # via
    #   dataproperty
@@ -459,6 +559,7 @@ numpy==2.2.6
    # via
    #   -r requirements/test/cuda.in
    #   accelerate
+    #   albucore
    #   albumentations
    #   bitsandbytes
    #   bm25s
@@ -466,14 +567,19 @@ 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
@@ -485,7 +591,11 @@ numpy==2.2.6
    #   patsy
    #   peft
    #   perceptron
+    #   pycocotools
+    #   pyogrio
    #   pywavelets
+    #   rasterio
+    #   rioxarray
    #   rouge-score
    #   runai-model-streamer
    #   sacrebleu
@@ -493,14 +603,21 @@ 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
@@ -540,6 +657,10 @@ 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
@@ -554,6 +675,7 @@ 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
@@ -588,27 +710,44 @@ 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
@@ -623,20 +762,25 @@ 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
@@ -673,7 +817,10 @@ protobuf==6.33.6
    #   opentelemetry-proto
    #   proto-plus
    #   ray
+    #   tensorboard
+    #   tensorboardx
    #   tensorizer
+    #   wandb
 psutil==6.1.0
    # via
    #   accelerate
@@ -687,12 +834,16 @@ 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
@@ -707,11 +858,13 @@ 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
@@ -720,8 +873,17 @@ pygments==2.18.0
    # via rich
 pyjwt==2.11.0
    # via msal
+pyogrio==0.11.0
+    # via geopandas
 pyparsing==3.2.0
-    # via matplotlib
+    # via
+    #   matplotlib
+    #   rasterio
+pyproj==3.7.1
+    # via
+    #   geopandas
+    #   rioxarray
+    #   torchgeo
 pyrate-limiter==3.7.0
    # via schemathesis
 pystemmer==3.0.0
@@ -758,15 +920,22 @@ 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
@@ -783,16 +952,26 @@ 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
@@ -803,6 +982,7 @@ referencing==0.35.1
    #   jsonschema-specifications
 regex==2026.2.28
    # via
+    #   diffusers
    #   nltk
    #   open-clip-torch
    #   sacrebleu
@@ -814,11 +994,13 @@ 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
@@ -828,7 +1010,9 @@ requests==2.32.3
    #   responses
    #   schemathesis
    #   starlette-testclient
+    #   tacoreader
    #   tiktoken
+    #   wandb
 responses==0.25.3
    # via genai-perf
 rfc3339-validator==0.1.4
@@ -838,9 +1022,13 @@ 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
@@ -849,6 +1037,8 @@ 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
@@ -864,6 +1054,7 @@ sacrebleu==2.4.3
 safetensors==0.4.5
    # via
    #   accelerate
+    #   diffusers
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
@@ -872,7 +1063,9 @@ safetensors==0.4.5
 schemathesis==3.39.15
    # via -r requirements/test/cuda.in
 scikit-image==0.25.2
-    # via albumentations
+    # via
+    #   albumentations
+    #   terratorch
 scikit-learn==1.5.2
    # via
    #   albumentations
@@ -880,6 +1073,7 @@ scikit-learn==1.5.2
    #   lm-eval
    #   mteb
    #   sentence-transformers
+    #   terratorch
 scipy==1.13.1
    # via
    #   albumentations
@@ -893,16 +1087,27 @@ scipy==1.13.1
    #   statsmodels
    #   vocos
 segmentation-models-pytorch==0.5.0
-    # via -r requirements/test/cuda.in
+    # via
+    #   -r requirements/test/cuda.in
+    #   terratorch
+    #   torchgeo
 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
@@ -911,12 +1116,15 @@ 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
@@ -958,6 +1166,8 @@ 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
@@ -967,14 +1177,26 @@ 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
+    # via
+    #   gpt-oss
+    #   terratorch
+terratorch==1.2.2
+    # via -r requirements/test/cuda.in
 threadpoolctl==3.5.0
    # via scikit-learn
 tifffile==2025.3.30
-    # via scikit-image
+    # via
+    #   scikit-image
+    #   terratorch
 tiktoken==0.12.0
    # via
    #   -c requirements/common.txt
@@ -986,6 +1208,8 @@ 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
@@ -1003,14 +1227,21 @@ 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
@@ -1020,18 +1251,31 @@ 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
@@ -1039,8 +1283,11 @@ tqdm==4.67.3
    #   optuna
    #   peft
    #   pqdm
+    #   pytorch-lightning
    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   tacoreader
+    #   terratorch
    #   transformers
 transformers==5.5.3
    # via
@@ -1069,6 +1316,8 @@ 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
@@ -1083,6 +1332,8 @@ typing-extensions==4.15.0
    #   grpcio
    #   huggingface-hub
    #   librosa
+    #   lightning
+    #   lightning-utilities
    #   lm-eval
    #   mistral-common
    #   mteb
@@ -1093,12 +1344,16 @@ 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
@@ -1110,8 +1365,10 @@ urllib3==2.2.3
    #   blobfile
    #   botocore
    #   docker
+    #   lightly
    #   requests
    #   responses
+    #   sentry-sdk
    #   tritonclient
 uvicorn==0.35.0
    # via gpt-oss
@@ -1121,16 +1378,22 @@ 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
+    # via
+    #   schemathesis
+    #   tensorboard
 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,11 +61,7 @@ pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0

 # 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
+terratorch>=1.2.2
 imagehash # Required for Prithvi tests
 segmentation-models-pytorch>0.4.0 # Required for Prithvi tests

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

 # ROCm-specific extras (not in CUDA cuda.in)
 rapidfuzz
-# torchgeo also pulled in `lightning` transitively; disabled for the same
-# quarantine reason as terratorch above. Restore once the quarantine clears.
-# torchgeo==0.7.0
+torchgeo==0.7.0
 multiprocess==0.70.16
@@ -1,9 +1,15 @@
 # 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
+    # via
+    #   rouge-score
+    #   tensorboard
 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
@@ -19,8 +25,12 @@ 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
+    # via
+    #   -r requirements/test/rocm.in
+    #   terratorch
 alembic==1.18.4
    # via optuna
 annotated-doc==0.0.4
@@ -33,6 +43,10 @@ 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
@@ -53,9 +67,11 @@ astor==0.8.1
 attrs==26.1.0
    # via
    #   aiohttp
+    #   fiona
    #   jsonlines
    #   jsonschema
    #   pytest-subtests
+    #   rasterio
    #   referencing
 audioread==3.0.1
    # via librosa
@@ -74,7 +90,9 @@ backoff==2.2.1
    #   -r requirements/test/rocm.in
    #   schemathesis
 bitsandbytes==0.49.2
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/rocm.in
+    #   lightning
 black==26.3.1
    # via datamodel-code-generator
 blake3==1.0.8
@@ -101,8 +119,13 @@ 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
@@ -120,13 +143,24 @@ 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
@@ -177,6 +211,8 @@ 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
@@ -201,12 +237,16 @@ docker==7.1.0
 docopt==0.6.2
    # via num2words
 docstring-parser==0.17.0
-    # via anthropic
+    # via
+    #   anthropic
+    #   jsonargparse
 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
@@ -243,11 +283,14 @@ 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
@@ -264,6 +307,9 @@ fsspec==2025.3.0
    #   evaluate
    #   fastparquet
    #   huggingface-hub
+    #   lightning
+    #   pytorch-lightning
+    #   tacoreader
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
@@ -271,10 +317,16 @@ 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
@@ -314,6 +366,7 @@ grpcio==1.78.0
    #   grpcio-reflection
    #   opentelemetry-exporter-otlp-proto-grpc
    #   ray
+    #   tensorboard
 grpcio-reflection==1.78.0
    # via
    #   -c requirements/rocm.txt
@@ -324,6 +377,8 @@ 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
@@ -342,6 +397,7 @@ httpx==0.27.2
    # via
    #   -r requirements/test/rocm.in
    #   anthropic
+    #   diffusers
    #   fastapi
    #   fastapi-cloud-cli
    #   huggingface-hub
@@ -356,17 +412,23 @@ 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
@@ -393,7 +455,11 @@ imagehash==4.3.2
 imageio==2.37.3
    # via scikit-image
 importlib-metadata==8.7.1
-    # via opentelemetry-api
+    # via
+    #   diffusers
+    #   opentelemetry-api
+importlib-resources==6.5.2
+    # via typeshed-client
 inflect==7.5.0
    # via datamodel-code-generator
 iniconfig==2.3.0
@@ -431,8 +497,14 @@ 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
@@ -452,6 +524,10 @@ 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
@@ -464,6 +540,21 @@ 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
@@ -489,6 +580,8 @@ 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
@@ -497,7 +590,10 @@ markupsafe==3.0.3
    #   mako
    #   werkzeug
 matplotlib==3.10.8
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/rocm.in
+    #   lightning
+    #   torchgeo
 mbstrdecoder==1.1.4
    # via
    #   dataproperty
@@ -564,11 +660,14 @@ 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
@@ -576,15 +675,20 @@ 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
@@ -596,8 +700,12 @@ numpy==2.2.6
    #   patsy
    #   peft
    #   perceptron
+    #   pycocotools
+    #   pyogrio
    #   pytrec-eval-terrier
    #   pywavelets
+    #   rasterio
+    #   rioxarray
    #   rouge-score
    #   runai-model-streamer
    #   sacrebleu
@@ -606,16 +714,27 @@ 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
@@ -705,29 +824,46 @@ 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
@@ -745,14 +881,18 @@ 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
@@ -760,6 +900,7 @@ platformdirs==4.3.6
    #   pooch
    #   python-discovery
    #   virtualenv
+    #   wandb
 plotly==6.6.0
    # via
    #   -r requirements/test/rocm.in
@@ -805,7 +946,10 @@ protobuf==6.33.6
    #   opentelemetry-proto
    #   proto-plus
    #   ray
+    #   tensorboard
+    #   tensorboardx
    #   tensorizer
+    #   wandb
 psutil==7.2.2
    # via
    #   -r requirements/test/../common.txt
@@ -822,12 +966,16 @@ 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
@@ -846,6 +994,7 @@ pydantic==2.12.5
    #   fastapi
    #   fastapi-cloud-cli
    #   gpt-oss
+    #   lightly
    #   lm-format-enforcer
    #   mcp
    #   mistral-common
@@ -856,6 +1005,7 @@ pydantic==2.12.5
    #   pydantic-extra-types
    #   pydantic-settings
    #   ray
+    #   wandb
    #   xgrammar
 pydantic-core==2.41.5
    # via pydantic
@@ -873,8 +1023,17 @@ pyjwt==2.12.1
    # via
    #   mcp
    #   msal
+pyogrio==0.12.1
+    # via geopandas
 pyparsing==3.3.2
-    # via matplotlib
+    # via
+    #   matplotlib
+    #   rasterio
+pyproj==3.7.2
+    # via
+    #   geopandas
+    #   rioxarray
+    #   torchgeo
 pyrate-limiter==3.9.0
    # via schemathesis
 pystemmer==3.0.0
@@ -911,10 +1070,13 @@ 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
@@ -934,6 +1096,10 @@ 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
@@ -950,9 +1116,13 @@ pyyaml==6.0.3
    #   genai-perf
    #   gguf
    #   huggingface-hub
+    #   jsonargparse
+    #   lightning
    #   lm-format-enforcer
+    #   omegaconf
    #   optuna
    #   peft
+    #   pytorch-lightning
    #   ray
    #   responses
    #   schemathesis
@@ -960,6 +1130,7 @@ pyyaml==6.0.3
    #   transformers
    #   uvicorn
    #   vocos
+    #   wandb
 pyzmq==27.1.0
    # via
    #   -c requirements/common.txt
@@ -968,6 +1139,11 @@ 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
@@ -979,6 +1155,7 @@ referencing==0.37.0
 regex==2026.2.28
    # via
    #   -r requirements/test/../common.txt
+    #   diffusers
    #   nltk
    #   open-clip-torch
    #   sacrebleu
@@ -991,12 +1168,14 @@ 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
@@ -1007,7 +1186,9 @@ requests==2.32.5
    #   responses
    #   schemathesis
    #   starlette-testclient
+    #   tacoreader
    #   tiktoken
+    #   wandb
 responses==0.26.0
    # via genai-perf
 rfc3339-validator==0.1.4
@@ -1017,9 +1198,11 @@ rfc3987==1.3.8
 rich==14.3.3
    # via
    #   genai-perf
+    #   lightning
    #   mteb
    #   perceptron
    #   rich-toolkit
+    #   terratorch
    #   typer
 rich-toolkit==0.19.7
    # via
@@ -1027,12 +1210,16 @@ 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
@@ -1050,6 +1237,7 @@ sacrebleu==2.6.0
 safetensors==0.7.0
    # via
    #   accelerate
+    #   diffusers
    #   open-clip-torch
    #   peft
    #   segmentation-models-pytorch
@@ -1058,7 +1246,9 @@ safetensors==0.7.0
 schemathesis==3.39.15
    # via -r requirements/test/rocm.in
 scikit-image==0.26.0
-    # via albumentations
+    # via
+    #   albumentations
+    #   terratorch
 scikit-learn==1.8.0
    # via
    #   albumentations
@@ -1066,6 +1256,7 @@ scikit-learn==1.8.0
    #   lm-eval
    #   mteb
    #   sentence-transformers
+    #   terratorch
 scipy==1.17.1
    # via
    #   albumentations
@@ -1080,7 +1271,10 @@ scipy==1.17.1
    #   statsmodels
    #   vocos
 segmentation-models-pytorch==0.5.0
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/rocm.in
+    #   terratorch
+    #   torchgeo
 sentence-transformers==5.3.0
    # via
    #   -r requirements/test/rocm.in
@@ -1088,7 +1282,9 @@ sentence-transformers==5.3.0
 sentencepiece==0.2.1
    # via -r requirements/test/../common.txt
 sentry-sdk==2.55.0
-    # via fastapi-cloud-cli
+    # via
+    #   fastapi-cloud-cli
+    #   wandb
 setproctitle==1.3.7
    # via -r requirements/test/../common.txt
 setuptools==79.0.1
@@ -1098,7 +1294,12 @@ 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
@@ -1110,12 +1311,15 @@ 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
@@ -1154,6 +1358,8 @@ 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
@@ -1166,6 +1372,8 @@ 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
@@ -1174,16 +1382,28 @@ 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
+    # via
+    #   gpt-oss
+    #   terratorch
+terratorch==1.2.2
+    # via -r requirements/test/rocm.in
 threadpoolctl==3.6.0
    # via scikit-learn
 tifffile==2026.3.3
-    # via scikit-image
+    # via
+    #   scikit-image
+    #   terratorch
 tiktoken==0.12.0
    # via
    #   -c requirements/common.txt
@@ -1197,6 +1417,8 @@ 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
@@ -1207,6 +1429,16 @@ 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
@@ -1214,6 +1446,8 @@ tqdm==4.67.3
    #   evaluate
    #   gguf
    #   huggingface-hub
+    #   lightly
+    #   lightning
    #   lm-eval
    #   mteb
    #   nltk
@@ -1222,8 +1456,11 @@ tqdm==4.67.3
    #   optuna
    #   peft
    #   pqdm
+    #   pytorch-lightning
    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   tacoreader
+    #   terratorch
    #   transformers
 transformers==5.5.3
    # via
@@ -1255,6 +1492,8 @@ 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
@@ -1272,6 +1511,8 @@ typing-extensions==4.15.0
    #   grpcio
    #   huggingface-hub
    #   librosa
+    #   lightning
+    #   lightning-utilities
    #   lm-eval
    #   mcp
    #   mistral-common
@@ -1286,14 +1527,18 @@ 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
@@ -1310,6 +1555,7 @@ urllib3==2.6.3
    #   blobfile
    #   botocore
    #   docker
+    #   lightly
    #   requests
    #   responses
    #   sentry-sdk
@@ -1329,6 +1575,8 @@ 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
@@ -1340,11 +1588,15 @@ webcolors==25.10.0
 websockets==16.0
    # via uvicorn
 werkzeug==3.1.6
-    # via schemathesis
+    # via
+    #   schemathesis
+    #   tensorboard
 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,11 +150,13 @@ def test_full_graph(
        if is_torch_equal_or_newer("2.9.0.dev")
    ]
    + [
-        # Cover compile_sizes autotune path.
+        # 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+
        (
            CompilationConfig(
                mode=CompilationMode.VLLM_COMPILE,
-                compile_sizes=[1, 2],  # Triggers the autotune path.
+                compile_sizes=[1, 2],  # Triggers autotune which uses get_raw_stream
                cudagraph_mode=CUDAGraphMode.NONE,
            ),
            "facebook/opt-125m",
@@ -24,6 +24,7 @@ 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

@@ -42,6 +43,29 @@ 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)
@@ -19,6 +19,7 @@ from vllm.utils.flashinfer import (
    has_flashinfer_nvlink_one_sided,
    has_flashinfer_nvlink_two_sided,
 )
+from vllm.utils.import_utils import has_deep_ep_v2
 from vllm.utils.network_utils import get_open_port

 from ..utils import init_test_distributed_environment
@@ -194,6 +195,10 @@ requires_ptrace = pytest.mark.skipif(
    not _has_sys_ptrace(),
    reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)",
 )
+requires_deep_ep_v2 = pytest.mark.skipif(
+    not has_deep_ep_v2(),
+    reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4",
+)

 # NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have
 # their own @requires_two_sided / @requires_one_sided decorators, and
@@ -772,3 +777,76 @@ def _one_sided_data_worker(rank, world_size):
 def test_one_sided_dispatch_combine(world_size):
    """Test FlashInfer one-sided dispatch/combine with actual data flow."""
    _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)
+
+
+# ---------------------------------------------------------------------------
+# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle
+# ---------------------------------------------------------------------------
+#
+# Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using
+# the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4.
+#
+# Uses EP group because the DeepEP v2 manager is constructed with an
+# EP-scoped communicator in production. With tp=world_size the EP group
+# spans all ranks.
+# ---------------------------------------------------------------------------
+
+
+def _deepep_v2_lifecycle_worker(rank, world_size):
+    from vllm.distributed.device_communicators.all2all import (
+        DeepEPV2All2AllManager,
+    )
+
+    cpu_group = get_ep_group().cpu_group
+    manager = DeepEPV2All2AllManager(cpu_group)
+
+    assert manager.rank == rank
+    assert manager.world_size == world_size
+    assert manager._num_sms is None
+
+    hidden_size = 7168
+    num_experts = world_size * 32
+    num_topk = 8
+    max_tokens = 256
+
+    handle_kwargs = dict(
+        num_max_tokens_per_rank=max_tokens,
+        hidden=hidden_size,
+        num_topk=num_topk,
+        num_experts=num_experts,
+        use_fp8_dispatch=False,
+    )
+
+    handle = manager.get_handle(handle_kwargs)
+    assert handle is not None
+    assert manager._num_sms is not None
+    assert manager._num_sms > 0
+
+    torch.distributed.barrier()
+
+    # get_handle again with same args should return cached handle
+    handle2 = manager.get_handle(dict(handle_kwargs))
+    assert handle2 is handle
+
+    torch.distributed.barrier()
+
+    # Destroy clears the cache
+    manager.destroy()
+    assert len(manager.handle_cache._cache) == 0
+
+    torch.distributed.barrier()
+
+    # Re-create after destroy
+    handle3 = manager.get_handle(dict(handle_kwargs))
+    assert handle3 is not None
+
+    torch.distributed.barrier()
+    manager.destroy()
+
+
+@requires_multi_gpu
+@requires_deep_ep_v2
+@pytest.mark.parametrize("world_size", [2])
+def test_deepep_v2_manager_lifecycle(world_size):
+    """Test DeepEP v2 ElasticBuffer manager init, caching, and destroy."""
+    _spawn_workers(_deepep_v2_lifecycle_worker, world_size)
@@ -1,8 +1,6 @@
 # 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
@@ -12,16 +10,7 @@ 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,8 +167,9 @@ 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,9 +6,7 @@ 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
@@ -62,35 +60,3 @@ 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
@@ -15,7 +15,7 @@ from torch.distributed import ProcessGroup
 from torch.multiprocessing import spawn  # pyright: ignore[reportPrivateImportUsage]
 from typing_extensions import ParamSpec

-from vllm.utils.import_utils import has_deep_ep
+from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2
 from vllm.utils.network_utils import get_open_port

 if has_deep_ep():
@@ -26,6 +26,11 @@ if has_deep_ep():
        DeepEPLLPrepareAndFinalize,
    )

+if has_deep_ep_v2():
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
+        DeepEPV2PrepareAndFinalize,
+    )
+
 ## Parallel Processes Utils

 P = ParamSpec("P")
@@ -55,11 +60,10 @@ def _worker_parallel_launch(
    torch.accelerator.set_device_index(local_rank)
    device = torch.device("cuda", local_rank)
    torch.distributed.init_process_group(
-        backend="cpu:gloo,cuda:nccl",
+        backend="nccl",
        init_method=init_method,
        rank=rank,
        world_size=world_size,
-        device_id=device,
    )
    barrier = torch.tensor([rank], device=device)
    torch.distributed.all_reduce(barrier)
@@ -200,3 +204,40 @@ def make_deepep_a2a(

    assert deepep_ll_args is not None
    return make_deepep_ll_a2a(pg, pgi, deepep_ll_args, q_dtype, block_shape)
+
+
+@dataclasses.dataclass
+class DeepEPV2Args:
+    num_local_experts: int
+    num_experts: int
+    num_topk: int
+    hidden_size: int
+    max_tokens_per_rank: int
+    use_fp8_dispatch: bool
+
+
+def make_deepep_v2_a2a(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    v2_args: DeepEPV2Args,
+):
+    import deep_ep
+
+    buffer = deep_ep.ElasticBuffer(
+        group=pg,
+        num_max_tokens_per_rank=v2_args.max_tokens_per_rank,
+        hidden=v2_args.hidden_size,
+        num_topk=v2_args.num_topk,
+        use_fp8_dispatch=v2_args.use_fp8_dispatch,
+        explicitly_destroy=True,
+    )
+    return DeepEPV2PrepareAndFinalize(
+        buffer=buffer,
+        num_dispatchers=pgi.world_size,
+        dp_size=dp_size,
+        rank_expert_offset=pgi.rank * v2_args.num_local_experts,
+        num_experts=v2_args.num_experts,
+        num_topk=v2_args.num_topk,
+        use_fp8_dispatch=v2_args.use_fp8_dispatch,
+    )
@@ -0,0 +1,537 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Test DeepEP v2 (ElasticBuffer) dispatch-combine logic.
+Compares against a pure-PyTorch reference MoE implementation.
+"""
+
+import dataclasses
+
+import pytest
+import torch.distributed
+from torch.distributed import ProcessGroup
+
+from tests.kernels.moe.utils import make_dummy_moe_config
+from vllm import _custom_ops as ops
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe import TritonExperts
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
+from vllm.utils.import_utils import has_deep_ep_v2
+from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.worker.workspace import init_workspace_manager
+
+from ...utils import multi_gpu_test
+from .parallel_utils import ProcessGroupInfo, parallel_launch
+
+if has_deep_ep_v2():
+    from .parallel_utils import DeepEPV2Args, make_deepep_v2_a2a
+
+requires_deep_ep_v2 = pytest.mark.skipif(
+    not has_deep_ep_v2(),
+    reason="Requires DeepEP v2 (ElasticBuffer)",
+)
+
+
+def make_weights(
+    e, n, k, dtype
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if dtype in [torch.float16, torch.bfloat16]:
+        w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
+        w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
+        return w1, w2, None, None
+
+    assert dtype == torch.float8_e4m3fn
+    w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16)
+    w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16)
+
+    n_b_scales = 2 * n
+    k_b_scales = k
+    w1_q = torch.empty_like(w1, dtype=dtype)
+    w2_q = torch.empty_like(w2, dtype=dtype)
+    w1_scale = torch.empty((e, n_b_scales, 1), device="cuda", dtype=torch.float32)
+    w2_scale = torch.empty((e, k_b_scales, 1), device="cuda", dtype=torch.float32)
+    for expert in range(e):
+        w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
+            w1[expert], use_per_token_if_dynamic=True
+        )
+        w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
+            w2[expert], use_per_token_if_dynamic=True
+        )
+    return w1_q, w2_q, w1_scale, w2_scale
+
+
+@dataclasses.dataclass
+class TestConfig:
+    dtype: torch.dtype
+    topk: int
+    m: int
+    k: int
+    n: int
+    num_experts: int
+
+
+@dataclasses.dataclass
+class TestTensors:
+    rank_tokens: torch.Tensor
+    rank_token_scales: torch.Tensor | None
+    topk: torch.Tensor
+    topk_weights: torch.Tensor
+    config: TestConfig
+
+    @staticmethod
+    def make(config: TestConfig) -> "TestTensors":
+        assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn]
+        token_dtype = (
+            torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype
+        )
+        rank_tokens = (
+            torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10
+        )
+
+        topk = torch.stack(
+            [
+                torch.randperm(config.num_experts, device="cuda")[: config.topk]
+                for _ in range(config.m)
+            ]
+        ).to(dtype=torch.int64)
+        topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda")
+        return TestTensors(
+            rank_tokens=rank_tokens,
+            rank_token_scales=None,
+            topk=topk,
+            topk_weights=topk_weights,
+            config=config,
+        )
+
+
+def make_modular_kernel(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    hidden_size: int,
+    num_experts: int,
+    num_local_experts: int,
+    topk: int,
+    q_dtype: torch.dtype | None,
+    use_fp8_dispatch: bool,
+    quant_config: FusedMoEQuantConfig,
+) -> FusedMoEKernel:
+    v2_args = DeepEPV2Args(
+        num_local_experts=num_local_experts,
+        num_experts=num_experts,
+        num_topk=topk,
+        hidden_size=hidden_size,
+        max_tokens_per_rank=8192,
+        use_fp8_dispatch=use_fp8_dispatch,
+    )
+
+    a2a = make_deepep_v2_a2a(
+        pg=pg,
+        pgi=pgi,
+        dp_size=dp_size,
+        v2_args=v2_args,
+    )
+
+    moe_config = make_dummy_moe_config()
+
+    fused_experts = TritonExperts(
+        moe_config=moe_config,
+        quant_config=quant_config,
+    )
+
+    mk = FusedMoEKernel(
+        prepare_finalize=a2a,
+        fused_experts=fused_experts,
+        inplace=False,
+    )
+    return mk
+
+
+def deepep_v2_moe_impl(
+    pg: ProcessGroup,
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    test_tensors: TestTensors,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    num_experts: int,
+    topk: int,
+    use_fp8_dispatch: bool,
+    per_act_token_quant: bool,
+) -> torch.Tensor:
+    num_local_experts = w1.size(0)
+
+    def build_expert_map():
+        expert_map = torch.full((num_experts,), fill_value=-1, dtype=torch.int32)
+        s = pgi.rank * num_local_experts
+        e = s + num_local_experts
+        expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
+        device = torch.accelerator.current_device_index()
+        return expert_map.to(device=device, dtype=torch.int32)
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    q_dtype = torch.float8_e4m3fn if is_quantized else None
+
+    quant_config = FusedMoEQuantConfig.make(
+        q_dtype,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        per_act_token_quant=per_act_token_quant,
+        a1_scale=test_tensors.rank_token_scales,
+    )
+
+    hidden_size = test_tensors.rank_tokens.size(1)
+
+    mk: FusedMoEKernel = make_modular_kernel(
+        pg,
+        pgi,
+        dp_size,
+        hidden_size,
+        num_experts,
+        num_local_experts,
+        topk,
+        q_dtype,
+        use_fp8_dispatch,
+        quant_config,
+    )
+
+    out = mk.apply(
+        hidden_states=test_tensors.rank_tokens,
+        w1=w1,
+        w2=w2,
+        topk_weights=test_tensors.topk_weights,
+        topk_ids=test_tensors.topk,
+        activation=MoEActivation.SILU,
+        global_num_experts=num_experts,
+        expert_map=build_expert_map(),
+        apply_router_weight_on_input=False,
+    )
+
+    return out
+
+
+def torch_moe_impl(
+    test_tensors: TestTensors,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    per_act_token_quant: bool,
+):
+    a, topk_ids, topk_weights = (
+        test_tensors.rank_tokens,
+        test_tensors.topk,
+        test_tensors.topk_weights,
+    )
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    a_dtype = a.dtype
+    if is_quantized:
+        w1 = w1.to(dtype=torch.float32) * w1_scale
+        w2 = w2.to(dtype=torch.float32) * w2_scale
+        a = a.to(dtype=torch.float32)
+
+    m, _ = a.shape
+    topk = topk_ids.size(1)
+    out = torch.zeros_like(a)
+
+    for i in range(m):
+        a_i = a[i]
+        o_i = out[i]
+        for j in range(topk):
+            e = topk_ids[i][j]
+            e_w = topk_weights[i][j]
+            w1_e = w1[e]
+            w2_e = w2[e]
+            o_i += (
+                SiluAndMul()(a_i @ w1_e.transpose(0, 1)) @ w2_e.transpose(0, 1)
+            ) * e_w
+
+    if is_quantized:
+        out = out.to(dtype=a_dtype)
+
+    return out
+
+
+def _deep_ep_v2_moe(
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    config: TestConfig,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+    use_fp8_dispatch: bool,
+    per_act_token_quant: bool,
+):
+    device = torch.device(f"cuda:{pgi.local_rank}")
+    init_workspace_manager(device)
+
+    is_quantized = w1.dtype == torch.float8_e4m3fn
+    device_idx = torch.accelerator.current_device_index()
+    w1 = w1.to(device=device_idx)
+    w2 = w2.to(device=device_idx)
+    if is_quantized:
+        assert w1_scale is not None and w2_scale is not None
+        w1_scale = w1_scale.to(device=device_idx)
+        w2_scale = w2_scale.to(device=device_idx)
+
+    pg = torch.distributed.new_group(list(range(pgi.world_size)))
+    test_tensors = TestTensors.make(config)
+
+    with set_current_vllm_config(VllmConfig()):
+        # Reference
+        torch_combined = torch_moe_impl(
+            test_tensors,
+            w1,
+            w2,
+            w1_scale,
+            w2_scale,
+            per_act_token_quant,
+        )
+
+        # Splice experts for this rank
+        num_local_experts = config.num_experts // pgi.world_size
+        e_start = num_local_experts * pgi.rank
+        e_end = e_start + num_local_experts
+        w1_ep = w1[e_start:e_end]
+        w2_ep = w2[e_start:e_end]
+
+        w1_scale_ep, w2_scale_ep = None, None
+        if is_quantized:
+            w1_scale_ep = w1_scale[e_start:e_end]  # type: ignore
+            w2_scale_ep = w2_scale[e_start:e_end]  # type: ignore
+
+        deepep_combined = deepep_v2_moe_impl(
+            pg,
+            pgi,
+            dp_size,
+            test_tensors,
+            w1_ep,
+            w2_ep,
+            w1_scale_ep,
+            w2_scale_ep,
+            config.num_experts,
+            config.topk,
+            use_fp8_dispatch,
+            per_act_token_quant,
+        )
+
+    torch.testing.assert_close(
+        torch_combined,
+        deepep_combined,
+        atol=6e-2,
+        rtol=6e-2,
+    )
+
+
+MNKs = [
+    (1, 256, 256),
+    (2, 256, 512),
+    (3, 1024, 2048),
+    (32, 256, 1024),
+    (45, 512, 2048),
+    (64, 1024, 1024),
+    (222, 1024, 2048),
+]
+
+DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("m,n,k", MNKs)
+@pytest.mark.parametrize("num_experts", [32])
+@pytest.mark.parametrize("topk", [6])
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
+@multi_gpu_test(num_gpus=2)
+@requires_deep_ep_v2
+def test_deep_ep_v2_moe(
+    dtype: torch.dtype,
+    m: int,
+    n: int,
+    k: int,
+    num_experts: int,
+    topk: int,
+    world_dp_size: tuple[int, int],
+    workspace_init,
+):
+    per_act_token_quant = False
+    use_fp8_dispatch = False
+
+    set_random_seed(7)
+    world_size, dp_size = world_dp_size
+    config = TestConfig(dtype=dtype, topk=topk, m=m, k=k, n=n, num_experts=num_experts)
+
+    w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype)
+
+    parallel_launch(
+        world_size,
+        _deep_ep_v2_moe,
+        dp_size,
+        config,
+        w1,
+        w2,
+        w1_scale,
+        w2_scale,
+        use_fp8_dispatch,
+        per_act_token_quant,
+    )
+
+
+def _deep_ep_v2_moe_cudagraph(
+    pgi: ProcessGroupInfo,
+    dp_size: int,
+    config: TestConfig,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor | None,
+    w2_scale: torch.Tensor | None,
+):
+    """Worker function: verify DeepEP v2 MoE works under cudagraph capture."""
+    device = torch.device(f"cuda:{pgi.local_rank}")
+    init_workspace_manager(device)
+
+    device_idx = torch.accelerator.current_device_index()
+    w1 = w1.to(device=device_idx)
+    w2 = w2.to(device=device_idx)
+
+    pg = torch.distributed.new_group(list(range(pgi.world_size)))
+    test_tensors = TestTensors.make(config)
+
+    num_local_experts = config.num_experts // pgi.world_size
+    e_start = num_local_experts * pgi.rank
+    e_end = e_start + num_local_experts
+    w1_ep = w1[e_start:e_end]
+    w2_ep = w2[e_start:e_end]
+
+    with set_current_vllm_config(VllmConfig()):
+        # Reference
+        torch_combined = torch_moe_impl(
+            test_tensors, w1, w2, None, None, False,
+        )
+
+        # Build kernel
+        num_local_experts = w1_ep.size(0)
+        hidden_size = test_tensors.rank_tokens.size(1)
+
+        quant_config = FusedMoEQuantConfig.make(None)
+
+        mk_kernel = make_modular_kernel(
+            pg, pgi, dp_size, hidden_size,
+            config.num_experts, num_local_experts,
+            config.topk, None, False, quant_config,
+        )
+
+        def build_expert_map():
+            expert_map = torch.full(
+                (config.num_experts,), fill_value=-1, dtype=torch.int32,
+            )
+            s = pgi.rank * num_local_experts
+            expert_map[s:s + num_local_experts] = torch.tensor(
+                list(range(num_local_experts)),
+            )
+            return expert_map.to(device=device_idx, dtype=torch.int32)
+
+        expert_map = build_expert_map()
+
+        # Warmup (non-captured)
+        out = mk_kernel.apply(
+            hidden_states=test_tensors.rank_tokens,
+            w1=w1_ep, w2=w2_ep,
+            topk_weights=test_tensors.topk_weights,
+            topk_ids=test_tensors.topk,
+            activation=MoEActivation.SILU,
+            global_num_experts=config.num_experts,
+            expert_map=expert_map,
+            apply_router_weight_on_input=False,
+        )
+
+        torch.testing.assert_close(
+            torch_combined, out, atol=6e-2, rtol=6e-2,
+        )
+
+        # Cudagraph capture
+        torch.cuda.synchronize()
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+
+        # Warmup on capture stream
+        with torch.cuda.stream(s):
+            for _ in range(3):
+                mk_kernel.apply(
+                    hidden_states=test_tensors.rank_tokens,
+                    w1=w1_ep, w2=w2_ep,
+                    topk_weights=test_tensors.topk_weights,
+                    topk_ids=test_tensors.topk,
+                    activation=MoEActivation.SILU,
+                    global_num_experts=config.num_experts,
+                    expert_map=expert_map,
+                    apply_router_weight_on_input=False,
+                )
+        torch.cuda.current_stream().wait_stream(s)
+        torch.cuda.synchronize()
+
+        # Capture
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g, stream=s):
+            graph_out = mk_kernel.apply(
+                hidden_states=test_tensors.rank_tokens,
+                w1=w1_ep, w2=w2_ep,
+                topk_weights=test_tensors.topk_weights,
+                topk_ids=test_tensors.topk,
+                activation=MoEActivation.SILU,
+                global_num_experts=config.num_experts,
+                expert_map=expert_map,
+                apply_router_weight_on_input=False,
+            )
+
+        # Replay
+        g.replay()
+        torch.cuda.synchronize()
+
+        torch.testing.assert_close(
+            torch_combined, graph_out, atol=6e-2, rtol=6e-2,
+        )
+
+
+@pytest.mark.parametrize("m,n,k", [(32, 256, 1024)])
+@pytest.mark.parametrize("num_experts", [32])
+@pytest.mark.parametrize("topk", [6])
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
+@multi_gpu_test(num_gpus=2)
+@requires_deep_ep_v2
+def test_deep_ep_v2_moe_cudagraph(
+    m: int,
+    n: int,
+    k: int,
+    num_experts: int,
+    topk: int,
+    world_dp_size: tuple[int, int],
+    workspace_init,
+):
+    set_random_seed(7)
+    world_size, dp_size = world_dp_size
+    config = TestConfig(
+        dtype=torch.bfloat16, topk=topk, m=m, k=k, n=n,
+        num_experts=num_experts,
+    )
+
+    w1, w2, _, _ = make_weights(num_experts, n, k, torch.bfloat16)
+
+    parallel_launch(
+        world_size,
+        _deep_ep_v2_moe_cudagraph,
+        dp_size,
+        config,
+        w1,
+        w2,
+        None,
+        None,
+    )
@@ -1,157 +0,0 @@
-# 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,18 +1,11 @@
 # 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,9 +311,6 @@ def _test_processing_correctness(
            baseline_processor,
            cached_processor,
            batch_idx,
-            hit_rate,
-            num_batches,
-            simplify_rate,
        )


@@ -323,9 +320,6 @@ 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

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

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

        _assert_inputs_equal(
            baseline_text_result,
            baseline_tokenized_result,
            ignore_mm_keys=ignore_mm_keys,
-            msg=(
-                f"Failed ({batch_idx=}, {hit_rate=}, "
-                f"{num_batches=}, {simplify_rate=}, "
-                f"{text_prompt=}, {token_prompt=}, {mm_data=})"
-            ),
+            msg=f"Failed ({batch_idx=}, {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=}, {hit_rate=}, "
-                f"{num_batches=}, {simplify_rate=}, "
-                f"{text_prompt=}, {token_prompt=}, {mm_data=})"
-            ),
+            msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})",
        )


@@ -430,8 +408,6 @@ 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,7 +1396,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
    ),
    # [Encoder-decoder]
    "CohereAsrForConditionalGeneration": _HfExamplesInfo(
-        "CohereLabs/cohere-transcribe-03-2026", trust_remote_code=True
+        "CohereLabs/cohere-transcribe-03-2026",
+        trust_remote_code=True,
+        is_available_online=False,  # TODO (ekagra): revert after asr release
    ),
    "NemotronParseForConditionalGeneration": _HfExamplesInfo(
        "nvidia/NVIDIA-Nemotron-Parse-v1.1", trust_remote_code=True
@@ -109,16 +109,6 @@ 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,17 +36,6 @@ 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,19 +1,12 @@
 # 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,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import importlib.util
 import io

 import imagehash
@@ -12,11 +11,6 @@ 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
@@ -0,0 +1,103 @@
+# 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,52 +432,3 @@ 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,13 +82,6 @@ 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",
-    ),
 ]


@@ -111,14 +104,6 @@ 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,
@@ -129,7 +114,6 @@ 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,6 +1,5 @@
 # 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
@@ -724,13 +723,8 @@ 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", "qwen3_5-hybrid"],
+    ids=["mimo", "deepseek"],
 )
@single_gpu_only
@large_gpu_mark(min_gb=20)
@@ -756,29 +750,13 @@ 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(
@@ -799,7 +777,6 @@ 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.base import (
+from vllm.v1.kv_offload.abstract import (
    OffloadingEvent,
    OffloadingManager,
    ReqContext,
@@ -20,7 +20,7 @@ from vllm.v1.kv_cache_interface import (
    MLAAttentionSpec,
    UniformTypeKVCacheSpecs,
 )
-from vllm.v1.kv_offload.base import (
+from vllm.v1.kv_offload.spec import (
    CanonicalKVCacheRef,
    CanonicalKVCaches,
    OffloadingSpec,
@@ -37,15 +37,15 @@ from vllm.v1.kv_cache_interface import (
    KVCacheConfig,
    KVCacheGroupSpec,
 )
-from vllm.v1.kv_offload.base import (
-    GPULoadStoreSpec,
+from vllm.v1.kv_offload.abstract import (
    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,
@@ -1,915 +0,0 @@
-# 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,18 +8,13 @@ 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,
-    SupportsHMA,
-    supports_hma,
-)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
    MultiConnector,
@@ -88,43 +83,8 @@ class MockConnector(KVConnectorBase_V1):
        pass


-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
+# Register the mock connector
 KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
-KVConnectorFactory.register_connector(
-    "MockHMAConnector", __name__, MockHMAConnector.__name__
-)


@pytest.fixture
@@ -960,133 +920,3 @@ 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,15 +9,14 @@ import torch

 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
-from vllm.v1.kv_offload.base import (
+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 (
    CanonicalKVCacheRef,
    CanonicalKVCaches,
    CanonicalKVCacheTensor,
-    GPULoadStoreSpec,
 )
-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
+from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers

 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.base import (
+from vllm.v1.kv_offload.abstract import (
    LoadStoreSpec,
    OffloadingEvent,
    OffloadKey,
@@ -14,9 +14,9 @@ from vllm.v1.kv_offload.base 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.base import LoadStoreSpec
+from vllm.v1.kv_offload.abstract import LoadStoreSpec
 from vllm.v1.kv_offload.worker.worker import (
    OffloadingHandler,
    OffloadingWorker,
@@ -8,7 +8,7 @@ set -ex
 #   --nvshmem-ver <ver>  NVSHMEM version 

 CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
-DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"73b6ea4"}
+DEEPEP_COMMIT_HASH=${DEEPEP_COMMIT_HASH:-"b306af06af"}
 NVSHMEM_VER=${NVSHMEM_VER:-"3.3.24"}  # Default supports both CUDA 12 and 13
 WORKSPACE=${WORKSPACE:-$(pwd)/ep_kernels_workspace}
 MODE=${MODE:-install}
@@ -203,6 +203,47 @@ 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.
@@ -216,6 +257,8 @@ 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:
@@ -42,6 +42,7 @@ All2AllBackend = Literal[
    "pplx",
    "deepep_high_throughput",
    "deepep_low_latency",
+    "deepep_v2",
    "mori",
    "nixl_ep",
    "allgather_reducescatter",
@@ -1227,12 +1227,14 @@ class VllmConfig:
            assert a2a_backend in [
                "deepep_low_latency",
                "deepep_high_throughput",
+                "deepep_v2",
            ], (
-                "Microbatching currently only supports the deepep_low_latency and "
-                f"deepep_high_throughput all2all backend. {a2a_backend} is not "
-                "supported. To fix use --all2all-backend=deepep_low_latency or "
-                "--all2all-backend=deepep_high_throughput and install the DeepEP"
-                " kernels."
+                "Microbatching currently only supports the deepep_low_latency, "
+                "deepep_high_throughput, and deepep_v2 all2all backends. "
+                f"{a2a_backend} is not supported. To fix use "
+                "--all2all-backend=deepep_low_latency, "
+                "--all2all-backend=deepep_high_throughput, or "
+                "--all2all-backend=deepep_v2 and install the DeepEP kernels."
            )

            if not self.model_config.disable_cascade_attn:
@@ -10,11 +10,12 @@ 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,
 )
-from vllm.utils.import_utils import has_deep_ep, has_mori
+from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori

 from .base_device_communicator import All2AllManagerBase, Cache

@@ -224,8 +225,11 @@ 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):
@@ -299,10 +303,13 @@ 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):
@@ -753,3 +760,63 @@ class MoriAll2AllManager(All2AllManagerBase):
            mori_kwargs, self._make_handle
        )
        return handle
+
+
+class DeepEPV2All2AllManager(All2AllManagerBase):
+    """
+    All2All communication based on DeepEP v2 ElasticBuffer (unified API).
+    Uses NCCL Gin backend with analytical SM calculation.
+    """
+
+    def __init__(self, cpu_group, tcp_store_group=None):
+        assert has_deep_ep_v2(), (
+            "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 "
+            "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4."
+        )
+        super().__init__(cpu_group, tcp_store_group)
+        self.handle_cache = Cache()
+        self._num_sms: int | None = None
+
+    def _make_all2all_kwargs(
+        self,
+        num_max_tokens_per_rank: int,
+        hidden: int,
+        num_topk: int,
+        use_fp8_dispatch: bool,
+    ) -> dict:
+        return dict(
+            group=self.cpu_group,
+            num_max_tokens_per_rank=num_max_tokens_per_rank,
+            hidden=hidden,
+            num_topk=num_topk,
+            use_fp8_dispatch=use_fp8_dispatch,
+            allow_hybrid_mode=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE,
+            prefer_overlap_with_compute=envs.VLLM_DEEPEP_V2_PREFER_OVERLAP,
+            allow_multiple_reduction=(envs.VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION),
+            explicitly_destroy=True,
+        )
+
+    def get_handle(self, kwargs):
+        import deep_ep  # type: ignore[import-not-found]
+
+        num_experts = kwargs.pop("num_experts", 256)
+        buffer_kwargs = self._make_all2all_kwargs(**kwargs)
+        logger.debug("DeepEP v2 all2all args %s", buffer_kwargs)
+        handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create(
+            buffer_kwargs, deep_ep.ElasticBuffer
+        )
+        if self._num_sms is None:
+            self._num_sms = handle.get_theoretical_num_sms(
+                num_experts=num_experts,
+                num_topk=kwargs["num_topk"],
+            )
+        return handle
+
+    def max_sms_used(self) -> int | None:
+        return self._num_sms
+
+    def destroy(self):
+        with self.handle_cache._lock:
+            for _, handle in self.handle_cache._cache.items():
+                handle.destroy()
+            self.handle_cache._cache.clear()
@@ -137,6 +137,12 @@ class CudaCommunicator(DeviceCommunicatorBase):
                from .all2all import MoriAll2AllManager

                self.all2all_manager = MoriAll2AllManager(self.cpu_group)
+            elif self.all2all_backend == "deepep_v2":
+                from .all2all import DeepEPV2All2AllManager
+
+                self.all2all_manager = DeepEPV2All2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
            elif self.all2all_backend == "nixl_ep":
                from .all2all import NixlEPAll2AllManager

@@ -11,7 +11,6 @@ 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,
@@ -48,25 +47,15 @@ class EplbCommunicator(ABC):
    """Abstract EPLB communicator for expert weight transfers."""

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

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

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

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

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

-    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_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
+        self._ops.append(("send", tensor, dst_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 add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
+        self._ops.append(("recv", tensor, src_rank))

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

@@ -242,17 +207,17 @@ class NixlEplbCommunicator(EplbCommunicator):
        self._cuda_stream = cuda_stream
        self._world_size = cpu_group.size()
        self._rank = cpu_group.rank()
-        # 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._send_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {}
+        self._recv_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {}
+        self._dtypes: list[torch.dtype] = []
        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)
@@ -263,12 +228,13 @@ 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]] = {}
+        self._remote_send_meta: dict[int, tuple[int, int, int]] = {}
        self._send_buffer: torch.Tensor = torch.empty(0)
        self._recv_buffer: torch.Tensor = torch.empty(0)
-        self._expert_bytes: int = 0
-
+        self._peer_partition_bytes: int = 0
+        self._dtype_max_bytes: dict[torch.dtype, int] = {}
        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)
@@ -292,33 +258,34 @@ class NixlEplbCommunicator(EplbCommunicator):
        uid = uuid.uuid4().hex[:8]
        return f"eplb-{self._rank}{pp_suffix}-{uid}"

-    def add_send(
+    def _get_peer_buckets(
        self,
-        tensors: list[torch.Tensor],
-        dst_rank: int,
-        expert_id: int,
-    ) -> None:
+        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:
        assert dst_rank != self._rank, (
            "EPLB communicator should not enqueue same-rank sends: "
            f"rank={self._rank}, dst_rank={dst_rank}"
        )
-        # 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
+        self._get_peer_buckets(self._send_tensors, tensor.dtype)[dst_rank].append(
+            tensor
+        )

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

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

    def _init_registered_buffers(self, expert_weights: Sequence[torch.Tensor]) -> None:
-        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
+        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

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

        descs = self._nixl_wrapper.get_reg_descs([self._send_buffer, self._recv_buffer])
@@ -357,11 +336,12 @@ 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] = (
+        local_meta: tuple[int, int, int] = (
            self._send_buffer.data_ptr(),
+            self._peer_partition_bytes,
            self._cuda_device_id,
        )
-        gathered_meta: list[tuple[int, int] | None] = [None] * self._world_size
+        gathered_meta: list[tuple[int, int, int] | None] = [None] * self._world_size
        torch.distributed.all_gather_object(
            gathered_meta, local_meta, group=self._cpu_group
        )
@@ -373,11 +353,14 @@ class NixlEplbCommunicator(EplbCommunicator):

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

    @staticmethod
    def _unpack_recv_buffer(
        recv_buffer: torch.Tensor,
-        out_tensors: list[torch.Tensor],
+        peer_tensors: list[torch.Tensor],
        byte_offset: int,
-    ) -> None:
-        for tensor in out_tensors:
+    ) -> int:
+        """
+        Returns the byte offset after the last read byte.
+        """
+        for tensor in peer_tensors:
            num_bytes = tensor.numel() * tensor.element_size()
            if num_bytes == 0:
                continue
@@ -401,6 +388,19 @@ 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,68 +418,82 @@ class NixlEplbCommunicator(EplbCommunicator):
            if pending:
                time.sleep(0.0005)

-    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.
+    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]

-        Each element in *local_descs* / *remote_descs* is an
-        ``(address, size, device_id)`` tuple.
-
-        Returns ``(local_dlist, remote_dlist, xfer_handle)``.
-        """
+        recv_base = self._recv_buffer.data_ptr()
        local_desc = self._nixl_wrapper.get_xfer_descs(
-            local_descs, self._nixl_memory_type
+            [
+                (
+                    recv_base + recv_offset,
+                    total_bytes,
+                    self._cuda_device_id,
+                )
+            ],
+            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_descs, self._nixl_memory_type
+            [
+                (
+                    remote_base + self._rank * remote_part_bytes,
+                    total_bytes,
+                    remote_dev,
+                )
+            ],
+            self._nixl_memory_type,
        )
        remote_handle = self._nixl_wrapper.prep_xfer_dlist(
-            self._remote_agents[src],
+            agent_name,
            remote_desc,
        )

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

-    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]] = []
+    def execute(self) -> None:
+        xfer_handles: list[int] = []
        try:
-            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.
+            # Phase 1: pack send buffers.
            with torch.cuda.stream(self._cuda_stream):
-                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)
+                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,
+                        )

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

-            # Phase 2: issue one batched READ per peer.
-            recv_offsets: dict[tuple[int, int], int] = {}
+            # 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] = {}
            recv_offset = 0
-            recv_base = self._recv_buffer.data_ptr()
            for src in range(self._world_size):
                if src == self._rank:
                    continue
-                recv_experts = self._recv_map.get(src)
-                if not recv_experts:
+                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:
                    continue
-                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_h)
-                xfer_entries.append((local_h, remote_h, xfer_h))

-            # Phase 3: wait for all in-flight transfers, then unpack.
-            self._wait_for_all_transfers([x[2] for x in xfer_entries])
+                recv_offsets[src] = recv_offset
+                xfer_handle = self._get_or_create_xfer(
+                    src, actual_total_bytes, recv_offset
+                )
+                self._nixl_wrapper.transfer(xfer_handle)
+                xfer_handles.append(xfer_handle)
+                recv_offset += actual_total_bytes
+
+            # Phase 3: single wait for all in-flight transfers, then unpack.
+            self._wait_for_all_transfers(xfer_handles)

            with torch.cuda.stream(self._cuda_stream):
-                for (src, expert_id), offset in recv_offsets.items():
-                    self._unpack_recv_buffer(
-                        self._recv_buffer,
-                        self._recv_map[src][expert_id],
-                        offset,
-                    )
+                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
        finally:
-            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()
+            self._send_tensors.clear()
+            self._recv_tensors.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
@@ -581,27 +588,15 @@ class PyNcclEplbCommunicator(EplbCommunicator):
            self._pynccl_comm.group_start()
            self._group_started = True

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

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

-    def execute(self, old_indices: np.ndarray | None = None) -> None:
+    def execute(self) -> 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:
-                communicator.add_send(expert_tensors, dst, expert_id=int(expert))
+                for w in expert_weights:
+                    communicator.add_send(w[src], dst)

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

    # 4. Execute the P2P operations. The real communication happens here.
-    communicator.execute(old_indices=old_indices)
+    communicator.execute()
    # 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 Callable, Iterable
+from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any

 import torch

@@ -18,8 +18,6 @@ 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,
@@ -125,7 +123,7 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):
            self._prom_metrics[connector_id].observe(stats_data["data"], engine_idx)


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

@@ -168,12 +166,6 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
            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] = {}
@@ -444,17 +436,15 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
        for c in self._connectors:
            c.set_xfer_handshake_metadata(metadata)

-    def _aggregate_request_finished(
+    def request_finished(
        self,
        request: "Request",
-        per_connector_fn: Callable[
-            [KVConnectorBase_V1], tuple[bool, dict[str, Any] | None]
-        ],
+        blocks: list[int],
    ) -> tuple[bool, dict[str, Any] | None]:
        async_saves = 0
        kv_txfer_params = None
        for c in self._connectors:
-            async_save, txfer_params = per_connector_fn(c)
+            async_save, txfer_params = c.request_finished(request, blocks)
            if async_save:
                async_saves += 1
            if txfer_params is not None:
@@ -468,39 +458,11 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
        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,30 +119,6 @@ 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:
@@ -322,44 +298,6 @@ 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

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

-        if params.get("do_remote_decode") or (
-            params.get("do_remote_prefill") and self.is_bidirectional_kv_xfer_enabled
-        ):
+        if params.get("do_remote_decode"):
            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") or (
-            params.get("do_remote_decode")
-            and self.is_bidirectional_kv_xfer_enabled
-            and not params.get("_remote_blocks_processed")
-        ):
+        elif params.get("do_remote_prefill"):
            if params.get("remote_block_ids"):
                if all(
                    p in params
@@ -401,8 +333,8 @@ class NixlConnectorScheduler:
                    )
                ):
                    # If remote_blocks and num_external_tokens = 0, we have
-                    # 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.
+                    # 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.

                    unhashed_local_block_ids: BlockIds = (
                        blocks.get_unhashed_block_ids_all_groups()
@@ -430,7 +362,6 @@ 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,
@@ -519,9 +450,6 @@ 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
@@ -533,13 +461,9 @@ class NixlConnectorScheduler:
            params["do_remote_prefill"] = False
            return False, None

-        if is_d_node and not self.is_bidirectional_kv_xfer_enabled:
+        if not params.get("do_remote_decode"):
            return False, None
-
-        if request.status not in (
-            RequestStatus.FINISHED_LENGTH_CAPPED,
-            RequestStatus.FINISHED_STOPPED,
-        ):
+        if request.status != RequestStatus.FINISHED_LENGTH_CAPPED:
            # 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)
@@ -550,7 +474,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(
@@ -568,16 +492,13 @@ 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=is_p_node,
-            do_remote_decode=is_d_node,
+            do_remote_prefill=True,
+            do_remote_decode=False,
            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.base import (
-    GPULoadStoreSpec,
+from vllm.v1.kv_offload.abstract import (
    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.base import (
+from vllm.v1.kv_offload.spec import (
    CanonicalKVCacheRef,
    CanonicalKVCaches,
    CanonicalKVCacheTensor,
@@ -755,8 +755,6 @@ 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_or_newer
+from vllm.utils.torch_utils import is_torch_equal, is_torch_equal_or_newer

 logger = init_logger(__name__)

@@ -112,6 +112,384 @@ 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
 # ===================================================
@@ -208,6 +586,55 @@ 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
 # ===================================================
@@ -233,6 +233,9 @@ if TYPE_CHECKING:
    VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024
    VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: bool = False
    VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL: bool = False
+    VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE: bool = True
+    VLLM_DEEPEP_V2_PREFER_OVERLAP: bool = False
+    VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION: bool = False
    VLLM_DBO_COMM_SMS: int = 20
    VLLM_PATTERN_MATCH_DEBUG: str | None = None
    VLLM_DEBUG_DUMP_PATH: str | None = None
@@ -255,10 +258,6 @@ 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
@@ -1627,6 +1626,18 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL": lambda: bool(
        int(os.getenv("VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL", "0"))
    ),
+    # DeepEP v2 ElasticBuffer: enable two-tier NVLink+RDMA hybrid mode
+    "VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE": lambda: bool(
+        int(os.getenv("VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE", "1"))
+    ),
+    # DeepEP v2 ElasticBuffer: use fewer SMs at slight throughput cost
+    "VLLM_DEEPEP_V2_PREFER_OVERLAP": lambda: bool(
+        int(os.getenv("VLLM_DEEPEP_V2_PREFER_OVERLAP", "0"))
+    ),
+    # DeepEP v2 ElasticBuffer: trade precision for transfer size in combine
+    "VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION": lambda: bool(
+        int(os.getenv("VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION", "0"))
+    ),
    # The number of SMs to allocate for communication kernels when running DBO
    # the rest of the SMs on the device will be allocated to compute
    "VLLM_DBO_COMM_SMS": lambda: int(os.getenv("VLLM_DBO_COMM_SMS", "20")),
@@ -1715,14 +1726,6 @@ 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(
@@ -1898,7 +1901,6 @@ 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,
-            (V, K),
-            (K, 1),
-            (i_v * BV, i_k * BK),
-            (BV, BK),
+            (K, V),
+            (V, 1),
+            (i_k * BK, i_v * BV),
+            (BK, BV),
            (1, 0),
        )

@@ -1090,11 +1090,12 @@ 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)
-        # [BV, BK]
+        # [BK, BV]
        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, tl.trans(b_h).to(b_qg.dtype))
+            b_o += tl.dot(b_qg, b_h.to(b_qg.dtype))
    p_v = tl.make_block_ptr(
        v + (bos * H + i_h) * V,
        (T, V),
@@ -29,7 +29,12 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two
    FlashInferNVLinkTwoSidedPrepareAndFinalize,
 )
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep
+from vllm.utils.import_utils import (
+    has_deep_ep,
+    has_deep_ep_v2,
+    has_mori,
+    has_nixl_ep,
+)

 logger = init_logger(__name__)

@@ -40,6 +45,8 @@ if current_platform.is_cuda_alike():
            DEEPEP_QUANT_BLOCK_SHAPE,
            DeepEPLLPrepareAndFinalize,
        )
+    if has_deep_ep_v2():
+        from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize
    if has_mori():
        from .prepare_finalize.mori import MoriPrepareAndFinalize
    if has_nixl_ep():
@@ -78,6 +85,11 @@ def maybe_roundup_layer_hidden_size(
            hidden_size
        )

+    if moe_parallel_config.use_deepep_v2_kernels:
+        hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size(
+            hidden_size, act_dtype
+        )
+
    if moe_parallel_config.use_nixl_ep_kernels:
        hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size(
            hidden_size
@@ -181,6 +193,36 @@ def maybe_make_prepare_finalize(
            physical_to_global=physical_to_global,
            local_expert_global_ids=local_expert_global_ids,
        )
+    elif moe.use_deepep_v2_kernels:
+        assert moe.dp_size == all2all_manager.dp_world_size
+
+        use_fp8_dispatch = (
+            quant_config is not None
+            and quant_config.quant_dtype == current_platform.fp8_dtype()
+            and quant_config.is_block_quantized
+        )
+        all_to_all_args = dict(
+            num_max_tokens_per_rank=moe.max_num_tokens,
+            hidden=moe.hidden_dim,
+            num_topk=moe.experts_per_token,
+            num_experts=moe.num_experts,
+            use_fp8_dispatch=use_fp8_dispatch,
+        )
+        handle = all2all_manager.get_handle(all_to_all_args)
+        vllm_config = get_current_vllm_config()
+        use_cudagraph = not vllm_config.model_config.enforce_eager
+
+        prepare_finalize = DeepEPV2PrepareAndFinalize(
+            buffer=handle,
+            num_dispatchers=all2all_manager.world_size,
+            dp_size=all2all_manager.dp_world_size,
+            rank_expert_offset=all2all_manager.rank * moe.num_local_experts,
+            num_experts=moe.num_experts,
+            num_topk=moe.experts_per_token,
+            use_fp8_dispatch=use_fp8_dispatch,
+            use_cudagraph=use_cudagraph,
+        )
+
    elif moe.use_mori_kernels:
        assert quant_config is not None

@@ -1063,6 +1063,10 @@ class FusedMoEParallelConfig:
    def use_nixl_ep_kernels(self):
        return self.use_all2all_kernels and self.all2all_backend == "nixl_ep"

+    @property
+    def use_deepep_v2_kernels(self):
+        return self.use_all2all_kernels and self.all2all_backend == "deepep_v2"
+
    @staticmethod
    def flatten_tp_across_dp_and_pcp(
        tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
@@ -1350,6 +1354,10 @@ class FusedMoEConfig:
    def use_nixl_ep_kernels(self):
        return self.moe_parallel_config.use_nixl_ep_kernels

+    @property
+    def use_deepep_v2_kernels(self):
+        return self.moe_parallel_config.use_deepep_v2_kernels
+
    @property
    def needs_round_robin_routing_tables(self):
        return self.moe_parallel_config.needs_round_robin_routing_tables
@@ -1487,19 +1487,15 @@ 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(NON_EXPERT_PREFIXES)
+            if not (
+                name.startswith("_shared_experts.")
+                or name.startswith("_gate.")
+                or name.startswith("_routed_input_transform.")
+                or name.startswith("_routed_output_transform.")
+            )
            and name not in NON_EXPERT_WEIGHTS
        )

@@ -1508,7 +1504,12 @@ class FusedMoE(PluggableLayer):
            for name, weight in weights
            if name not in NON_EXPERT_WEIGHTS
            and weight.shape != torch.Size([])
-            and not name.startswith(NON_EXPERT_PREFIXES)
+            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.")
        ]

    def set_eplb_state(
@@ -16,6 +16,7 @@ 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,
@@ -289,29 +290,46 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

        # Dispatch
        dispatch_topk_ids = self._map_global_to_physical_ids(topk_ids)
-        (
-            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,
-        )
+        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,
+            )
        self.handles[a2a_idx] = handle

        return (
@@ -0,0 +1,419 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
+
+import deep_ep
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceContiguous,
+    TopKWeightAndReduceDelegate,
+)
+from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
+from vllm.utils.math_utils import round_up
+from vllm.v1.worker.ubatching import (
+    dbo_current_ubatch_id,
+)
+
+
+class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
+    """
+    Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API).
+
+    Supports two modes controlled by the `use_cudagraph` constructor arg:
+
+    **Decode mode (use_cudagraph=True):**
+      - do_expand=False, do_cpu_sync=False
+      - Tokens returned in original order with recv_topk_idx (global IDs)
+      - Worst-case tensor allocation; padding rows zeroed via
+        handle.psum_num_recv_tokens_per_scaleup_rank
+      - Fully cudagraph-capturable
+      - Expert kernel sorts internally (expert_tokens_meta=None)
+
+    **Prefill mode (use_cudagraph=False):**
+      - do_expand=True, do_cpu_sync=True
+      - Per-expert-contiguous layout; exact memory allocation
+      - Saves GPU memory (no worst-case allocation)
+      - Not cudagraph-capturable (CPU polling), but prefill doesn't
+        use cudagraphs anyway
+      - Provides expert_tokens_meta for efficient batched expert kernels
+
+    Both modes use async_with_compute_stream=False (synchronous from
+    caller's perspective). The ElasticBuffer handles comm internally.
+    """
+
+    @staticmethod
+    def maybe_roundup_layer_hidden_size(hidden_size: int, dtype: torch.dtype) -> int:
+        hidden_size_bytes = hidden_size * dtype.itemsize
+        xfer_atom_size = 512  # 32 * 16 (size(int4))
+        if hidden_size_bytes % xfer_atom_size == 0:
+            return hidden_size
+
+        hidden_size_bytes = round_up(hidden_size_bytes, xfer_atom_size)
+        return hidden_size_bytes // dtype.itemsize
+
+    def __init__(
+        self,
+        buffer: deep_ep.ElasticBuffer,
+        num_dispatchers: int,
+        dp_size: int,
+        rank_expert_offset: int,
+        num_experts: int,
+        num_topk: int,
+        use_fp8_dispatch: bool = False,
+        use_cudagraph: bool = False,
+    ):
+        super().__init__()
+        self.buffer = buffer
+        self.num_dispatchers_ = num_dispatchers
+        self.dp_size = dp_size
+        self.rank_expert_offset = rank_expert_offset
+        self.num_experts = num_experts
+        self.num_topk = num_topk
+        self.use_fp8_dispatch = use_fp8_dispatch
+        self.use_cudagraph = use_cudagraph
+
+        # DBO microbatching: one handle slot per micro-batch.
+        self.handles: list[deep_ep.EPHandle | None] = [None, None]
+
+    def num_dispatchers(self) -> int:
+        return self.num_dispatchers_
+
+    def output_is_reduced(self) -> bool:
+        return True
+
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def max_num_tokens_per_rank(self) -> int | None:
+        return None
+
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        return torch.int64
+
+    def _do_dispatch(
+        self,
+        tokens: torch.Tensor,
+        token_scales: torch.Tensor | None,
+        rank_topk_ids: torch.Tensor,
+        rank_topk_weights: torch.Tensor,
+        num_experts: int,
+        a1_scale: torch.Tensor | None,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool,
+    ) -> Callable:
+        has_scales = token_scales is not None
+
+        token_data = tokens
+        if has_scales:
+            token_data = (tokens, token_scales)
+
+        # Decode: do_expand=False + do_cpu_sync=False (cudagraph-safe)
+        # Prefill: do_expand=True + do_cpu_sync=True (memory-efficient)
+        do_expand = not self.use_cudagraph
+        do_cpu_sync = not self.use_cudagraph
+
+        (
+            recv_x,
+            recv_topk_idx,
+            recv_topk_weights,
+            handle,
+            event,
+        ) = self.buffer.dispatch(
+            x=token_data,
+            topk_idx=rank_topk_ids,
+            topk_weights=rank_topk_weights,
+            num_experts=num_experts,
+            do_expand=do_expand,
+            do_cpu_sync=do_cpu_sync,
+            async_with_compute_stream=False,
+        )
+
+        a2a_idx = dbo_current_ubatch_id()
+        self.handles[a2a_idx] = handle
+
+        return lambda: self._receiver(
+            event,
+            has_scales,
+            recv_x,
+            recv_topk_idx,
+            num_experts,
+            handle.num_recv_tokens_per_expert_list,
+            recv_topk_weights,
+            handle.psum_num_recv_tokens_per_scaleup_rank,
+            a1_scale,
+            quant_config,
+            defer_input_quant=defer_input_quant,
+        )
+
+    def _receiver(
+        self,
+        event: deep_ep.EventOverlap,
+        has_scales: bool,
+        recv_x: tuple[torch.Tensor, torch.Tensor] | torch.Tensor,
+        recv_topk_idx: torch.Tensor | None,
+        num_experts: int,
+        recv_expert_num_tokens: list[int],
+        recv_topk_weights: torch.Tensor | None,
+        psum_recv_per_rank: torch.Tensor,
+        a1_scale: torch.Tensor | None,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool,
+    ) -> mk.PrepareResultType:
+        if event.event is not None:
+            event.current_stream_wait()
+
+        if isinstance(recv_x, tuple):
+            expert_x, expert_x_scale = recv_x
+        else:
+            expert_x, expert_x_scale = recv_x, None
+
+        if recv_topk_idx is None:
+            # do_expand=True (prefill mode): build topk_ids from
+            # per-expert token counts.
+            parts = [
+                torch.full(
+                    (count,),
+                    i + self.rank_expert_offset,
+                    dtype=torch.int64,
+                    device=expert_x.device,
+                )
+                for i, count in enumerate(recv_expert_num_tokens)
+                if count > 0
+            ]
+            recv_topk_idx = (
+                torch.cat(parts)
+                if parts
+                else torch.empty(
+                    0, dtype=torch.int64, device=expert_x.device,
+                )
+            )
+            recv_topk_idx = recv_topk_idx.unsqueeze(1)
+        else:
+            # do_expand=False (decode/cudagraph mode): recv_topk_idx has
+            # LOCAL expert IDs (-1 for non-local). Sanitize padding rows
+            # from worst-case allocation, then convert to global IDs.
+            total_rows = expert_x.shape[0]
+            num_real = psum_recv_per_rank[-1]  # GPU scalar tensor
+            row_indices = torch.arange(
+                total_rows, device=expert_x.device, dtype=num_real.dtype,
+            )
+            is_padding = (row_indices >= num_real).unsqueeze(1)
+
+            expert_x = torch.where(
+                is_padding, torch.zeros_like(expert_x), expert_x)
+            if expert_x_scale is not None:
+                expert_x_scale = torch.where(
+                    is_padding, torch.ones_like(expert_x_scale),
+                    expert_x_scale)
+            if recv_topk_weights is not None:
+                recv_topk_weights = torch.where(
+                    is_padding, torch.zeros_like(recv_topk_weights),
+                    recv_topk_weights)
+
+            # dispatch(do_expand=False) returns LOCAL expert IDs (-1 for
+            # non-local, -1 for padding after the where above).
+            # Convert valid local IDs to global. Keep -1 as -1 so
+            # DeepGemm's is_computation_valid skips those rows.
+            is_invalid = (recv_topk_idx < 0) | is_padding
+            recv_topk_idx = torch.where(
+                is_invalid,
+                -1,
+                recv_topk_idx + self.rank_expert_offset,
+            )
+            if recv_topk_weights is not None:
+                recv_topk_weights = torch.where(
+                    is_invalid,
+                    torch.zeros_like(recv_topk_weights),
+                    recv_topk_weights,
+                )
+
+        # Reshape recv_topk_weights to match recv_topk_idx shape [N, 1]
+        if recv_topk_weights is not None and recv_topk_weights.ndim == 1:
+            recv_topk_weights = recv_topk_weights.unsqueeze(1)
+
+        if recv_topk_idx is not None and not self.use_cudagraph:
+            # do_expand=True: we have exact per-expert counts
+            expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
+                recv_expert_num_tokens, device=expert_x.device,
+            )
+        else:
+            expert_tokens_meta = None
+
+        if not quant_config.is_block_quantized and not defer_input_quant:
+            expert_x_scale = None
+            if expert_x.numel() != 0:
+                expert_x, expert_x_scale = moe_kernel_quantize_input(
+                    expert_x,
+                    a1_scale,
+                    quant_dtype=quant_config.quant_dtype,
+                    per_act_token_quant=False,
+                    block_shape=quant_config.block_shape,
+                    is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled,
+                )
+
+        return (
+            expert_x,
+            expert_x_scale,
+            expert_tokens_meta,
+            recv_topk_idx,
+            recv_topk_weights,
+        )
+
+    def supports_async(self) -> bool:
+        return True
+
+    def prepare_async(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.ReceiverType:
+        if apply_router_weight_on_input:
+            topk = topk_ids.size(1)
+            assert topk == 1, (
+                "apply_router_weight_on_input is only implemented for topk=1"
+            )
+            a1 = a1 * topk_weights.to(a1.dtype)
+
+        if quant_config.is_block_quantized and not defer_input_quant:
+            a1q, a1q_scale = moe_kernel_quantize_input(
+                a1,
+                quant_config.a1_scale,
+                quant_dtype=quant_config.quant_dtype,
+                per_act_token_quant=quant_config.per_act_token_quant,
+                block_shape=quant_config.block_shape,
+            )
+            if a1q_scale is not None and a1q_scale.numel() == 1:
+                a1q_scale = a1q_scale.view(1, 1)
+            a1_post_scale = None
+        else:
+            a1q = a1
+            a1q_scale = None
+            a1_post_scale = (
+                quant_config.a1_gscale
+                if quant_config.quant_dtype == "nvfp4"
+                else quant_config.a1_scale
+            )
+
+        return self._do_dispatch(
+            tokens=a1q,
+            token_scales=a1q_scale,
+            rank_topk_ids=topk_ids,
+            rank_topk_weights=topk_weights,
+            num_experts=num_experts,
+            a1_scale=a1_post_scale,
+            quant_config=quant_config,
+            defer_input_quant=defer_input_quant,
+        )
+
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.PrepareResultType:
+        receiver = self.prepare_async(
+            a1,
+            topk_weights,
+            topk_ids,
+            num_experts,
+            expert_map,
+            apply_router_weight_on_input,
+            quant_config,
+            defer_input_quant,
+        )
+        return receiver()
+
+    def _finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+        do_async: bool,
+    ) -> Callable | None:
+        a2a_idx = dbo_current_ubatch_id()
+        handle = self.handles[a2a_idx]
+        assert handle is not None
+
+        if fused_expert_output.numel() != 0:
+            if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
+                weight_and_reduce_impl = TopKWeightAndReduceContiguous()
+            fused_expert_output = weight_and_reduce_impl.apply(
+                output=None,
+                fused_expert_output=fused_expert_output,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                apply_router_weight_on_input=apply_router_weight_on_input,
+            )
+
+        if fused_expert_output.dtype != torch.bfloat16:
+            raise ValueError(
+                f"DeepEP v2 combine requires bfloat16 input, "
+                f"got {fused_expert_output.dtype}"
+            )
+
+        combined_x, _, event = self.buffer.combine(
+            x=fused_expert_output,
+            handle=handle,
+            topk_weights=None,
+            async_with_compute_stream=False,
+        )
+
+        output.copy_(combined_x, non_blocking=True)
+        return None
+
+    def finalize_async(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> Callable:
+        self._finalize(
+            output,
+            fused_expert_output,
+            topk_weights,
+            topk_ids,
+            apply_router_weight_on_input,
+            weight_and_reduce_impl,
+            False,
+        )
+        return lambda: None
+
+    def finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> None:
+        self._finalize(
+            output,
+            fused_expert_output,
+            topk_weights,
+            topk_ids,
+            apply_router_weight_on_input,
+            weight_and_reduce_impl,
+            False,
+        )
@@ -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,7 +4,6 @@ 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,
 )
@@ -13,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
 )


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

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

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

 import torch
 import torch.nn.functional as F
@@ -15,7 +14,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, TokensPrompt
+from vllm.inputs import MultiModalDataDict, PromptType, TextPrompt
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.attention import (
@@ -49,7 +48,6 @@ 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,
@@ -2010,9 +2008,6 @@ 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:
@@ -2030,81 +2025,30 @@ class CohereAsrForConditionalGeneration(
        audio = stt_params.audio
        stt_config = stt_params.stt_config
        language = stt_params.language
-        model_config = stt_params.model_config
+        request_prompt = stt_params.request_prompt

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

-        tokenizer = cached_tokenizer_from_config(model_config)
-
-        # 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,
+        # 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

-        return TokensPrompt(
+        return TextPrompt(
            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,9 +854,10 @@ 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 and input_ids is None:
-            raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
-
+        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 not self.use_mega_moe:
            return self._forward_fused_moe(hidden_states, input_ids)

@@ -1224,12 +1225,7 @@ 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
@@ -1313,8 +1309,7 @@ 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,10 +84,6 @@ 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,
@@ -1654,7 +1650,7 @@ class Gemma4ForCausalLM(
                # Remap individual 2D expert weights:
                # .experts.{id}.{proj} → .moe.experts.{id}.{proj}
                # (This handles per-expert 2D quantized weights)
-                name = _remap_gemma4_expert_weight_name(name)
+                name = re.sub(r"\.experts\.(\d+)\.", r".moe.experts.\1.", name)

                # MoE expert weights: checkpoint stores as 3D packed
                # tensors.  Explode into per-expert 2D weights for
@@ -351,9 +351,6 @@ 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:
@@ -405,6 +405,61 @@ def has_deep_ep() -> bool:
    return _has_module("deep_ep")


+DEEPEP_V2_MIN_NCCL_VERSION_RAW = 23004  # 2.30.4
+
+
+def _get_runtime_nccl_version() -> int | None:
+    """Get the runtime NCCL version by loading the actual library.
+
+    Returns the raw version int (e.g. 23004 for 2.30.4), or None on failure.
+    torch.cuda.nccl.version() is a compile-time constant from the PyTorch
+    wheel and does not reflect a separately installed NCCL.
+    """
+    import ctypes
+
+    try:
+        from vllm.utils.nccl import find_nccl_library
+
+        lib = ctypes.CDLL(find_nccl_library())
+        version = ctypes.c_int()
+        lib.ncclGetVersion(ctypes.byref(version))
+        return version.value
+    except Exception:
+        return None
+
+
+def _format_nccl_raw_version(raw: int) -> str:
+    s = str(raw)
+    return f"{s[0]}.{s[1:3].lstrip('0') or '0'}.{s[3:].lstrip('0') or '0'}"
+
+
+def has_deep_ep_v2() -> bool:
+    """Whether deep_ep with ElasticBuffer (v2 API) is available.
+
+    Requires both the ElasticBuffer class in the deep_ep module and
+    NCCL >= 2.30.4 (GIN backend), checked against the runtime library.
+    """
+    if not _has_module("deep_ep"):
+        return False
+    import deep_ep  # type: ignore[import-not-found]
+
+    if not hasattr(deep_ep, "ElasticBuffer"):
+        return False
+    try:
+        nccl_ver = _get_runtime_nccl_version()
+        if nccl_ver is None or nccl_ver < DEEPEP_V2_MIN_NCCL_VERSION_RAW:
+            logger.info_once(
+                "DeepEP v2 requires NCCL >= %s but found %s. "
+                "deepep_v2 backend will not be available.",
+                _format_nccl_raw_version(DEEPEP_V2_MIN_NCCL_VERSION_RAW),
+                _format_nccl_raw_version(nccl_ver) if nccl_ver else "unknown",
+            )
+            return False
+    except Exception:
+        return False
+    return True
+
+
 def has_deep_gemm() -> bool:
    """Whether the optional `deep_gemm` package is available.

@@ -85,7 +85,6 @@ 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.
@@ -102,10 +101,6 @@ 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.
@@ -116,12 +111,7 @@ 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,
-                    apply_admission_cap=apply_admission_cap,
+                    request_id, num_encoder_tokens, [], 0, num_encoder_tokens
                )
            else:
                num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
@@ -130,7 +120,6 @@ 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,7 +257,6 @@ 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,7 +92,6 @@ 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.
@@ -108,16 +107,13 @@ 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 apply_admission_cap and self._max_admission_blocks_per_request is not None:
+        if 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
@@ -897,7 +893,6 @@ 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 (
@@ -922,7 +917,6 @@ 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
@@ -0,0 +1,197 @@
+# 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
@@ -1,390 +0,0 @@
-# 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
@@ -1,13 +0,0 @@
-# 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"
@@ -3,7 +3,7 @@
 from collections.abc import Iterable, Sequence
 from typing import Literal

-from vllm.v1.kv_offload.base import (
+from vllm.v1.kv_offload.abstract import (
    LoadStoreSpec,
    OffloadingEvent,
    OffloadingManager,
@@ -11,10 +11,10 @@ from vllm.v1.kv_offload.base import (
    PrepareStoreOutput,
    ReqContext,
 )
-from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
+from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
 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,
@@ -4,7 +4,7 @@ import ctypes
 from abc import ABC, abstractmethod
 from collections.abc import Iterable

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


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

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


 class ARCCachePolicy(CachePolicy):
@@ -3,8 +3,8 @@
 from collections import OrderedDict
 from collections.abc import Iterable

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


 class LRUCachePolicy(CachePolicy):
@@ -5,17 +5,12 @@ 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.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.abstract import LoadStoreSpec, OffloadingManager
 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.base import OffloadingSpec
+from vllm.v1.kv_offload.spec import OffloadingSpec

 if TYPE_CHECKING:
    from vllm.config import VllmConfig
@@ -0,0 +1,68 @@
+# 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.base import (
+from vllm.v1.kv_offload.abstract import (
    LoadStoreSpec,
    OffloadingEvent,
    OffloadingManager,
@@ -0,0 +1,142 @@
+# 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
@@ -11,13 +11,9 @@ 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,
@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass

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

 # a single transfer spec (src_blocks_spec, dst_blocks_spec)
 TransferSpec = tuple[LoadStoreSpec, LoadStoreSpec]
@@ -5,13 +5,12 @@
 import gc
 import os
 from collections.abc import Callable
-from contextlib import AbstractContextManager, contextmanager, nullcontext
+from contextlib import AbstractContextManager, 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

@@ -209,30 +208,6 @@ 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":
@@ -337,8 +312,6 @@ 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
Tyler Michael SmithandClaude Opus 4.6	ae8c445789	Use -1 expert ID for padding and non-local tokens in DeepEP v2 Padding rows and non-local expert tokens now get expert_id=-1 instead of rank_expert_offset. This allows DeepGemm's is_computation_valid() to skip computation on these rows, which is both a correctness fix (avoids wasting GEMM compute on junk data) and enables future kernel optimizations for early CTA exit on -1 blocks. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-05-02 11:51:27 -04:00
Tyler Michael SmithandClaude Opus 4.6	45a0b6d72d	Fix local→global expert ID conversion in DeepEP v2 decode path dispatch(do_expand=False) returns LOCAL expert IDs (-1 for non-local). The previous code only replaced -1 with rank_expert_offset but did not convert valid local IDs to global (missing + rank_expert_offset). This produced correct results only on rank 0 where local == global. Also: - Zero weights for non-local experts (-1) so they don't contribute - Use ones (not zeros) for padding scales to avoid DeepGemm issues - Set padding expert IDs to 0 before local→global conversion to avoid double-offset (0 + offset = offset, not offset + offset) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-05-02 02:15:31 -04:00
Tyler Michael SmithandClaude Opus 4.6	0f94cbef14	Fix FP8 padding: use torch.where instead of masked_fill_ masked_fill_ and indexing assignment are not implemented for float8_e4m3fn. Use torch.where which supports FP8 tensors. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-30 22:54:59 -04:00
Tyler Michael SmithandClaude Opus 4.6	ed29278a33	Split DeepEP v2 into prefill/decode modes for memory vs cudagraph Prefill (use_cudagraph=False): do_expand=True + do_cpu_sync=True — exact memory allocation, per-expert-contiguous layout. Saves GPU memory for large batches. Decode (use_cudagraph=True): do_expand=False + do_cpu_sync=False — worst-case allocation, scattered layout. Fully cudagraph-capturable. Mode selected based on enforce_eager config. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-30 22:19:23 -04:00
Tyler Michael SmithandClaude Opus 4.6	56a8767d94	Update DeepEP v2 docstring with cudagraph design rationale Document the four key design decisions (do_expand=False, do_cpu_sync=False, async_with_compute_stream=False, expert_tokens_meta=None) and why each is necessary for cudagraph + DBO compatibility. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-30 21:55:14 -04:00
Tyler Michael SmithandClaude Opus 4.6	5a766bf3de	Refactor DeepEP v2 for cudagraph compatibility and simplified DBO - Remove explicit two-stream DBO switching (dbo_yield_and_switch_*), use synchronous dispatch/combine (async_with_compute_stream=False). The ElasticBuffer handles comm internally on its comm_stream. - Switch from do_expand=True to do_expand=False for cudagraph compat. do_expand=True requires do_cpu_sync=True (CPU polling loop) which can't be captured in a cudagraph. do_expand=False with do_cpu_sync=False is fully capturable. - Handle worst-case padding from do_cpu_sync=False: use handle.psum_num_recv_tokens_per_scaleup_rank to get real token count, zero out padding rows in recv_x, recv_topk_weights, and expert_x_scale. - Add explicitly_destroy=True to ElasticBuffer creation in all2all.py. - Add cudagraph capture/replay unit test (test_deep_ep_v2_moe_cudagraph). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-30 21:51:16 -04:00
Tyler Michael SmithandClaude Opus 4.6	75149ae1b9	Check runtime NCCL version instead of PyTorch compile-time constant torch.cuda.nccl.version() returns the compile-time NCCL version baked into the PyTorch wheel, not the runtime library. Use ctypes to load the actual libnccl.so and call ncclGetVersion() directly, which respects VLLM_NCCL_SO_PATH and LD_LIBRARY_PATH. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 22:58:10 -04:00
Tyler Michael SmithandClaude Opus 4.6	67d26f78d2	Fix NCCL version check: 2.30.4, not 4.30.4 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 22:58:10 -04:00
Tyler Michael SmithandClaude Opus 4.6	019a41f00a	Add DeepEP v2 lifecycle test to test_mnnvl_alltoall Test DeepEPV2All2AllManager init, ElasticBuffer handle creation and caching, SM calculation, and destroy/re-create cycle. Skipped when DeepEP v2 or NCCL >= 4.30.4 is not available. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 22:58:10 -04:00
Tyler Michael SmithandClaude Opus 4.6	515e36da11	Fix FP8 dispatch test to match production behavior use_fp8_dispatch requires the ElasticBuffer to receive FP8 input. In production, this is ensured by pre-quantizing via moe_kernel_quantize_input when is_block_quantized=True. The test was parametrizing use_fp8_dispatch independently of dtype, allowing bf16 input with use_fp8_dispatch=True which triggers a buffer size assertion in DeepEP v2. Fix: - Derive use_fp8_dispatch from dtype (True only for FP8 weights) - Add block_shape=[128, 128] to quant config for FP8 to enable the block quantization path that pre-quantizes input Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 22:56:16 -04:00
Tyler Michael SmithandClaude Opus 4.6	a2a4b00f83	Add DeepEP v2 (ElasticBuffer) all2all backend for MoE EP Add a new `deepep_v2` all2all backend that uses the DeepEP v2 ElasticBuffer API (NCCL GIN backend). This provides a unified dispatch/combine interface that works for both intra-node and inter-node expert parallelism with analytical SM calculation. Key changes: - New DeepEPV2PrepareAndFinalize class using do_expand=True for per-expert-contiguous layout with weighted reduction in combine - DeepEPV2All2AllManager with ElasticBuffer handle caching and theoretical SM calculation via get_theoretical_num_sms() - NCCL >= 4.30.4 version gating in has_deep_ep_v2() since the GIN backend requires a newer NCCL than PyTorch typically bundles - FP8 block-quantized dispatch support - DBO (micro-batching) support with async prepare/finalize - Environment variables: VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE, VLLM_DEEPEP_V2_PREFER_OVERLAP, VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION - Update DeepEP install script to pin v2.0 release (b306af06af) - Comprehensive multi-process test suite Usage: --all2all-backend=deepep_v2 --enable-expert-parallel Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-04-29 21:02:26 -04:00