From 76dad2b04be96f09cda51a4876bb2ee9f76130f6 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Mon, 14 Sep 2026 05:23:48 -0700 Subject: [PATCH] Fix RDMA hidden-state request addressing Use vLLM LBHNC capture geometry and carry each requests hidden-state cache blocks through connector metadata. This prevents runner reordering from associating a request with another requests captured hidden states. Signed-off-by: Andrii Skliar --- .../plugins/rdma_hidden_states_connector.py | 117 +++++++++---- .../test_rdma_hidden_states_connector.py | 163 ++++++++++++++++++ 2 files changed, 248 insertions(+), 32 deletions(-) create mode 100644 tests/unit/torch/speculative/plugins/test_rdma_hidden_states_connector.py diff --git a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py index e68262a97ed..04b19dcfdc2 100644 --- a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py +++ b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py @@ -27,8 +27,9 @@ "kv_role":"kv_producer", "kv_connector_extra_config":{"sidecar_port":"18999","pool_slots":"64","max_tokens":"512"}}' -Scope: TP>=1 (hidden states are replicated across TP ranks; only rank 0 owns the -pool + sidecar and serves them), host(pinned) memory (container UCX has no CUDA), +Scope: TP>=1 using vLLM's LBHNC cache layout (the connector receives each layer +as a BHNC view). Hidden states are replicated across TP ranks; only rank 0 owns the +pool + sidecar and serves them. Uses host-pinned memory (container UCX has no CUDA), ring-slot reuse (fine when in-flight requests < pool_slots; no credit protocol yet). PP>1 is not supported (the capture layer lives on one PP rank; owner election only covers TP). The sidecar is unauthenticated and binds all interfaces, so it assumes a @@ -54,19 +55,39 @@ KVConnectorMetadata, SupportsHMA, ) -from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput -from .hf_streaming_dataset import nixl_backends_from_env +from modelopt.torch.speculative.plugins.hf_streaming_dataset import nixl_backends_from_env logger = init_logger(__name__) def extract_from_kv_cache(kv_cache, slot_mapping, num_tokens): - """Gather the first ``num_tokens`` rows of ``kv_cache`` addressed by ``slot_mapping``.""" - block_size = kv_cache.shape[1] - return kv_cache[slot_mapping // block_size, slot_mapping % block_size][:num_tokens] + """Gather hidden states from vLLM's per-layer ``[B, H, N, C]`` cache view. + + ``H`` is the number of requested capture planes and ``N`` is the KV-cache + block size. Keeping those axes distinct is essential: treating ``N`` as the + feature axis turns a six-plane capture into sixteen planes at block size 16. + """ + if kv_cache.ndim != 4: + raise ValueError( + "RdmaHiddenStatesConnector expected a [blocks, heads, block_size, head_size] " + f"cache tensor, got shape {tuple(kv_cache.shape)}." + ) + block_size = kv_cache.shape[2] + block_ids = slot_mapping[:num_tokens] // block_size + block_offsets = slot_mapping[:num_tokens] % block_size + return kv_cache[block_ids, :, block_offsets, :] + + +def build_request_slot_mapping(block_ids, num_tokens, block_size, device): + """Map one request's logical token positions to its allocated cache blocks.""" + if len(block_ids) * block_size < num_tokens: + raise ValueError("Hidden-state cache blocks do not cover the prompt") + blocks = torch.as_tensor(block_ids, device=device, dtype=torch.long) + positions = torch.arange(num_tokens, device=device) + return blocks[positions // block_size] * block_size + positions % block_size @dataclass @@ -76,11 +97,17 @@ class ReqMeta: req_id: str token_ids: torch.Tensor slot: int + block_ids: list[int] @staticmethod - def make(req_id, token_ids, slot): + def make(req_id, token_ids, slot, block_ids): """Build a :class:`ReqMeta`, tensorizing ``token_ids``.""" - return ReqMeta(req_id=req_id, token_ids=torch.tensor(token_ids), slot=slot) + return ReqMeta( + req_id=req_id, + token_ids=torch.tensor(token_ids), + slot=slot, + block_ids=list(block_ids), + ) @dataclass @@ -89,9 +116,9 @@ class RdmaConnMeta(KVConnectorMetadata): requests: list = field(default_factory=list) - def add(self, req_id, token_ids, slot): + def add(self, req_id, token_ids, slot, block_ids): """Append one request's capture metadata.""" - self.requests.append(ReqMeta.make(req_id, token_ids, slot)) + self.requests.append(ReqMeta.make(req_id, token_ids, slot, block_ids)) class _Sidecar(BaseHTTPRequestHandler): @@ -162,6 +189,20 @@ def __init__(self, vllm_config, role, kv_cache_config): """Read pool/sidecar config from ``kv_connector_extra_config`` and init state.""" super().__init__(vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config) self._role = role + from vllm.v1.kv_cache_interface import HiddenStateCacheSpec + + groups = kv_cache_config.kv_cache_groups if kv_cache_config is not None else [] + capture_groups = [ + index + for index, group in enumerate(groups) + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ] + if len(capture_groups) == 1: + self._capture_group_id = capture_groups[0] + elif kv_cache_config is None or (len(groups) == 1 and not capture_groups): + self._capture_group_id = 0 + else: + raise ValueError("RDMA capture requires exactly one hidden-state cache group") ex = self._kv_transfer_config.get_from_extra_config self._sidecar_port = int(ex("sidecar_port", "18999")) self._pool_slots = int(ex("pool_slots", "64")) @@ -229,9 +270,22 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ) return kv = kv_caches[self.cache_layers[0]] - # per-token feature = everything past [num_blocks, block_size] - self._per_token_elems = int(prod(kv.shape[2:])) - self._feat_shape = tuple(kv.shape[2:]) + if kv.ndim != 4: + raise ValueError( + "RdmaHiddenStatesConnector expected vLLM's per-layer [B, H, N, C] " + f"cache view, got shape {tuple(kv.shape)}." + ) + hf_config = self._vllm_config.speculative_config.draft_model_config.hf_config + capture_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", ()) + if kv.shape[1] != len(capture_ids): + raise ValueError( + f"RdmaHiddenStatesConnector cache exposes {kv.shape[1]} hidden-state planes " + f"but EAGLE requested {len(capture_ids)} ids: {list(capture_ids)}." + ) + # Per-token feature is [capture_planes, hidden_size]. The cache's third + # dimension is its block size and must not leak into the transferred shape. + self._per_token_elems = int(kv.shape[1] * prod(kv.shape[3:])) + self._feat_shape = (kv.shape[1], *kv.shape[3:]) self._dtype = kv.dtype self._slot_elems = self._max_tokens * self._per_token_elems @@ -309,15 +363,6 @@ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs): ready = torch.cuda.Event() ready.record() cs.wait_event(ready) - slot_mapping = get_forward_context().slot_mapping[layer_name] - # Assumes new-request tokens sit contiguously at the front of slot_mapping; bound - # the offset walk so an unexpected layout fails loud instead of short-slicing. - n_capture = sum(req.token_ids.shape[0] for req in md.requests) - assert n_capture <= slot_mapping.shape[0], ( - f"RdmaHiddenStatesConnector: capturing {n_capture} tokens but slot_mapping has " - f"only {slot_mapping.shape[0]}; unexpected batch layout (chunked prefill?)." - ) - offset = 0 for req in md.requests: n = req.token_ids.shape[0] nelems = n * self._per_token_elems @@ -337,11 +382,13 @@ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs): self._max_tokens, ) self._oversize_warned = True - offset += n continue with torch.cuda.stream(cs): - rsm = slot_mapping[offset : offset + n] - offset += n + # The runner can reorder requests independently of scheduler metadata. + # Address this request's allocated cache blocks, not batch offsets. + rsm = build_request_slot_mapping( + req.block_ids, n, kv_layer.shape[2], kv_layer.device + ) hs_gpu = extract_from_kv_cache(kv_layer, rsm, n) # [n, *feat] # copy into the pre-registered pool slot (flattened) self._pool[slot, :nelems].copy_(hs_gpu.reshape(-1), non_blocking=True) @@ -375,9 +422,8 @@ def update_state_after_alloc(self, request, blocks, num_external_tokens): def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata: """Assign each newly scheduled request a ring-pool slot for this step. - ``save_kv_layer`` walks ``slot_mapping`` with a running offset, so it needs each - prompt computed whole in one step; assert that here (chunked prefill would split - it and silently misalign the capture). + Each request carries its own hidden-state cache blocks, independent of the + runner's batch order. Only complete, uncached prompts are supported. """ meta = RdmaConnMeta() for nr in scheduler_output.scheduled_new_reqs: @@ -388,9 +434,16 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnector f"(disable chunked prefill): req {nr.req_id} scheduled {n_sched} of " f"{len(prompt)} prompt tokens." ) + if nr.num_computed_tokens != 0: + raise ValueError("RDMA capture requires an uncached complete prompt") slot = self._slot_ctr % self._pool_slots self._slot_ctr += 1 - meta.add(nr.req_id, token_ids=prompt, slot=slot) + meta.add( + nr.req_id, + token_ids=prompt, + slot=slot, + block_ids=nr.block_ids[self._capture_group_id], + ) return meta def request_finished(self, request, block_ids): @@ -435,7 +488,7 @@ def get_finished(self, finished_req_ids): @classmethod def get_required_kvcache_layout(cls, vllm_config): - """Require NHD layout so each token's hidden state stays contiguous.""" + """Use vLLM's layer/block/head/token/channel cache layout.""" if cls is KVConnectorBase_V1: raise TypeError("not on base class") - return "NHD" + return "LBHNC" diff --git a/tests/unit/torch/speculative/plugins/test_rdma_hidden_states_connector.py b/tests/unit/torch/speculative/plugins/test_rdma_hidden_states_connector.py new file mode 100644 index 00000000000..dcdd71a98e8 --- /dev/null +++ b/tests/unit/torch/speculative/plugins/test_rdma_hidden_states_connector.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Regressions for vLLM hidden-state cache layout and request addressing.""" + +import threading +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("vllm") + +from modelopt.torch.speculative.plugins.rdma_hidden_states_connector import ( + RdmaConnMeta, + RdmaHiddenStatesConnector, + ReqMeta, + build_request_slot_mapping, + extract_from_kv_cache, +) + + +def test_extract_from_bhnc_cache_preserves_capture_planes(): + """The KV block axis must not be mistaken for the hidden-state plane axis.""" + num_blocks, num_planes, block_size, hidden_size = 10, 6, 4, 3 + cache = torch.arange( + num_blocks * num_planes * block_size * hidden_size, dtype=torch.float32 + ).reshape(num_blocks, num_planes, block_size, hidden_size) + slot_mapping = torch.tensor([1, block_size + 2]) + + actual = extract_from_kv_cache(cache, slot_mapping, num_tokens=2) + expected = torch.stack((cache[0, :, 1, :], cache[1, :, 2, :])) + + assert actual.shape == (2, num_planes, hidden_size) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize( + ("block_ids", "num_tokens"), + [([7, 2], 6), ([4], 3), ([9, 1], 8)], +) +def test_request_addressing_is_independent_of_batch_order(block_ids, num_tokens): + """Each request follows its own noncontiguous blocks, not a batch offset.""" + block_size = 4 + mapping = build_request_slot_mapping(block_ids, num_tokens, block_size, "cpu") + expected = torch.tensor( + [ + block_ids[position // block_size] * block_size + position % block_size + for position in range(num_tokens) + ] + ) + torch.testing.assert_close(mapping, expected) + + +def test_request_addressing_rejects_insufficient_blocks(): + with pytest.raises(ValueError, match="do not cover"): + build_request_slot_mapping([1], 5, 4, "cpu") + + +def test_scheduler_metadata_carries_request_specific_capture_blocks(): + """The worker receives the capture group's blocks keyed to each request.""" + connector = object.__new__(RdmaHiddenStatesConnector) + connector._slot_ctr = 0 + connector._pool_slots = 4 + connector._capture_group_id = 1 + request = SimpleNamespace( + req_id="request-b", + prompt_token_ids=list(range(6)), + num_computed_tokens=0, + block_ids=[[0, 1], [7, 2]], + ) + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[request], + num_scheduled_tokens={"request-b": 6}, + ) + + metadata = connector.build_connector_meta(scheduler_output) + + assert len(metadata.requests) == 1 + assert metadata.requests[0].req_id == "request-b" + assert metadata.requests[0].block_ids == [7, 2] + + +def test_save_kv_layer_captures_each_requests_own_noncontiguous_blocks(monkeypatch): + """Production capture must not associate scheduler order with cache position.""" + from vllm.model_executor.models import extract_hidden_states + + class AttentionMetadata: + pass + + class Event: + def record(self, stream=None): + pass + + class Stream: + def wait_event(self, event): + pass + + class Nixl: + def get_xfer_descs(self, tensors): + return tensors + + def get_serialized_descs(self, descriptors): + return b"descriptor" + + monkeypatch.setattr(extract_hidden_states, "CacheOnlyAttentionMetadata", AttentionMetadata) + monkeypatch.setattr(torch.cuda, "Event", Event) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + + block_size = 4 + kv_layer = torch.arange(8 * 2 * block_size * 3, dtype=torch.float32).reshape( + 8, 2, block_size, 3 + ) + requests = [ + ReqMeta.make("request-b", range(6), slot=1, block_ids=[5, 1]), + ReqMeta.make("request-a", range(5), slot=0, block_ids=[3, 0]), + ] + connector = object.__new__(RdmaHiddenStatesConnector) + connector._owner = True + connector.cache_layers = ["capture"] + connector._get_connector_metadata = lambda: RdmaConnMeta(requests=requests) + connector._cs = lambda: Stream() + connector._per_token_elems = 2 * 3 + connector._feat_shape = (2, 3) + connector._dtype = kv_layer.dtype + connector._slot_elems = 6 * connector._per_token_elems + connector._pool = torch.full((2, connector._slot_elems), -1.0) + connector._nixl = Nixl() + connector._lock = threading.Lock() + connector._slot_gen = {} + connector._bufs = {} + connector._oversize_warned = False + connector._max_tokens = 6 + + connector.save_kv_layer("capture", kv_layer, AttentionMetadata()) + + for request in requests: + mapping = build_request_slot_mapping( + request.block_ids, len(request.token_ids), block_size, "cpu" + ) + expected = extract_from_kv_cache(kv_layer, mapping, len(request.token_ids)).flatten() + captured = connector._pool[request.slot, : expected.numel()] + torch.testing.assert_close(captured, expected) + assert connector._bufs[request.req_id]["slot"] == request.slot + + +def test_extract_from_kv_cache_rejects_unknown_layout(): + """Fail explicitly if a future vLLM release changes the per-layer cache view.""" + with pytest.raises(ValueError, match=r"\[blocks, heads, block_size, head_size\]"): + extract_from_kv_cache(torch.empty(2, 16, 4), torch.tensor([0]), num_tokens=1)