From 8a7cfeee2325dbeacadc8ecb2aa9c441100a66ba Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 1 Oct 2026 12:47:06 +0000 Subject: [PATCH 1/3] fix(speculative): keep the DFlash draft's RoPE frequencies in fp32 on a fresh convert restore_draft_precision() is documented as a no-op on a freshly converted model, and examples/speculative_decoding/main.py calls it unconditionally on that basis. It was not one. modify() has already cast the draft to the base dtype and then built the rotary buffer -- after the cast on purpose, so inv_freq stays fp32 -- and restore_draft_precision() cast the draft again. Module.to() casts floating buffers too, so inv_freq came back in bf16. The rounding hits the frequencies rather than the cos/sin values, so the phase error grows with position: up to ~0.4% per frequency (measured max 3.65e-3), several radians in the highest-frequency pairs by position 4096. vLLM computes the angles from fp32 frequencies and only stores cos/sin in the model dtype, so a draft trained this way meets different RoPE phases at serve time. This affects every DFlash-family draft trained from scratch through main.py since #2483; the from_pretrained resume path was already correct. modify() now records whether it placed the draft, and restore_draft_precision() returns early when it did. On the from_pretrained path (base on meta) it still places the draft and builds the buffer. The two new tests pin both sides: the first fails on the previous code (inv_freq comes back bf16), the second fails if restore_draft_precision() always returns. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- .../torch/speculative/plugins/hf_dflash.py | 7 ++++-- .../speculative/plugins/test_hf_dflash.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index c1c3389e428..cd4103a035d 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -462,7 +462,8 @@ def modify(self, config): # lists. After the cast, because `Module.to` casts float buffers and building first # would round the RoPE frequencies. base_device = self._base_device() - if base_device.type != "meta": + self._draft_placed = base_device.type != "meta" + if self._draft_placed: self.dflash_module.to(device=base_device, dtype=self._base_model.dtype) self.dflash_module._maybe_init_rotary_emb(device=base_device) @@ -492,10 +493,12 @@ def restore_draft_precision(self): placement entirely. A no-op on a freshly converted model. """ base_device = self._base_device() - if base_device.type == "meta": + # Placing an already-placed draft again would round its fp32 rotary buffer to bf16. + if self._draft_placed or base_device.type == "meta": return self.dflash_module.to(device=base_device, dtype=self._base_model.dtype) self.dflash_module._maybe_init_rotary_emb(device=base_device) + self._draft_placed = True # Draft-module entries that legitimately come from the base model rather than the # exported draft checkpoint, so their absence (or presence) is not an error. diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ff8931149e5..4242270a59f 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -423,6 +423,29 @@ def test_rotary_emb_init_is_idempotent(self): assert dflash_mod.rotary_emb is first assert not any(b.is_meta for b in dflash_mod.rotary_emb.buffers()) + def test_restore_draft_precision_keeps_fp32_rope_after_convert(self): + """On a fresh convert the draft is already placed; re-casting would round RoPE to bf16.""" + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", get_dflash_config())]) + inv_freq = model.dflash_module.rotary_emb.inv_freq.clone() + assert inv_freq.dtype == torch.float32 + + model.restore_draft_precision() + + assert model.dflash_module.rotary_emb.inv_freq.dtype == torch.float32 + assert torch.equal(model.dflash_module.rotary_emb.inv_freq, inv_freq) + + def test_restore_draft_precision_places_draft_converted_on_meta(self): + """After a meta-device convert (the from_pretrained path) it still places the draft.""" + model = get_tiny_llama(num_hidden_layers=4).to("meta") + mtsp.convert(model, [("dflash", get_dflash_config())]) + model.to_empty(device="cpu") # stands in for from_pretrained loading the weights + + model.restore_draft_precision() + + assert model.dflash_module.rotary_emb.inv_freq.dtype == torch.float32 + assert all(p.dtype == torch.bfloat16 for p in model.dflash_module.parameters()) + class TestBuildTargetLayerIds: """Test target layer selection.""" From a5132fc1a760c5e94ccf503d5e4e925ef965f611 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 23 Sep 2026 07:03:02 +0000 Subject: [PATCH 2/3] fix(speculative): correct the streaming hidden-state connector on current vLLM Three defects in RdmaHiddenStatesConnector, two of which produce wrong data without producing an error. All were found while bringing streaming drafter training up on vLLM 0.29. **The KV-cache layout moved.** vLLM PR #51718 ("Standardize KV cache layout", in 0.29 and later) changed the per-layer view from [B, N, H, C] to [B, H, N, C] and deleted KVConnectorBase_V1.prefer_cross_layer_blocks, which is how this connector used to ask for the capture layer to get its own block table. Reading the old layout on a new build is silent: the extraction returns feat=(N, C), so the trainer sees one plane where it expects several, peels it off as the base hidden and hands the draft an empty aux tensor -- "mat1 and mat2 shapes cannot be multiplied (3072x0 and 40960x8192)", pointing nowhere near here. planes_dim() now identifies the plane axis by matching the capture count, so both layouts work, and the group is resolved from the static kv-cache spec instead of the deleted property. **Captures were addressed by position in the batch.** save_kv_layer walked slot_mapping with a running offset, assuming newly scheduled prefill requests sit contiguously at its front. vLLM appends new requests to the persistent batch, so any request already decoding occupies earlier positions and shifts that window: each capture then read a *different* request's hidden states. Nothing downstream can catch this -- the shapes are right, the pool slot is right, the recorded token_ids are the request's own, and the reconstructed distribution is a real, sharp one. It simply belongs to another prompt, so the drafter trains on noise. Measured with the serve capped at max_num_seqs=1 (a batch can never hold a second request) the per-sequence hidden/token match is 0.96-1.00; with concurrency it is 0.00-0.02, and the drafter plateaus at bigram-level accuracy. Each request's slots now come from its own block ids. **The capture index tensors were built on the wrong stream.** They are consumed on the connector's side stream, and nothing ordered that stream against the allocation, so the caching allocator could hand their memory out while the copy was still reading. The garbage indices surface as an out-of-bounds assert inside whatever kernel runs next -- minutes into a run, not reproducibly, and never pointing at the connector. The code this replaced was safe only incidentally: it sliced vLLM's long-lived slot_mapping rather than allocating anything. Because two of these are silent, the new checks are deliberately loud: the resolved group and cache view are printed at startup, requests whose block ids cannot cover their tokens or run past the cache raise by name, and the first few requests each verify that every derived slot is one this step actually writes. Targets vLLM >= 0.29. Validated by a full 2-epoch production run and by a Qwen3-8B DFlash2 streaming run (600 steps, loss 17.0 -> 3.4). Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../plugins/rdma_hidden_states_connector.py | 219 ++++++++++++++++-- 1 file changed, 195 insertions(+), 24 deletions(-) diff --git a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py index e68262a97ed..8a09a797c9c 100644 --- a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py +++ b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py @@ -63,10 +63,74 @@ 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] +def planes_dim(kv_cache, n_planes: int) -> int: + """Which axis of a per-layer KV cache view holds the captured planes. + + vLLM changed the logical axis order of the per-layer view. Newer builds hand + out ``[B, H, N, C]`` unconditionally -- `create_kv_cache_views` documents + "one 4D [B, H, N, C] view per layer", and the KVCacheLayout only permutes + strides, not axes. Older builds used ``[B, N, H, C]`` (block, token, head, + content), which is what this connector was written against. + + For the hidden-state pseudo-layer H is the number of captured planes and C is + hidden_size, so the two orders are distinguishable by looking for ``n_planes``. + Getting it wrong is not an error, just wrong data: on a new vLLM the old + reading returns ``feat=(N, C)``, the trainer sees ONE plane instead of six, + peels it off as the base hidden and hands the draft an empty aux tensor + ("mat1 and mat2 shapes cannot be multiplied (3072x0 and 40960x8192)"). + """ + if kv_cache.ndim != 4: + raise RuntimeError( + f"hidden-state cache view has {kv_cache.ndim} dims, expected 4; " + f"shape={tuple(kv_cache.shape)}" + ) + if kv_cache.shape[1] == n_planes and kv_cache.shape[2] != n_planes: + return 1 + if kv_cache.shape[2] == n_planes and kv_cache.shape[1] != n_planes: + return 2 + raise RuntimeError( + f"cannot tell which axis of the hidden-state cache view {tuple(kv_cache.shape)} " + f"holds the {n_planes} captured planes; both or neither of dims 1 and 2 match. " + "Set a block size that differs from the capture count to disambiguate." + ) + + +def extract_from_kv_cache(kv_cache, slot_mapping, num_tokens, pdim: int = 2): + """Gather the first ``num_tokens`` tokens of ``kv_cache`` addressed by ``slot_mapping``. + + Returns ``[num_tokens, n_planes, hidden]``. ``pdim`` is :func:`planes_dim`. + """ + if pdim == 1: # [B, H, N, C] + block_size = kv_cache.shape[2] + # Advanced indices separated by a slice, so the gathered dim comes first: + # result is [n, H, C]. + out = kv_cache[slot_mapping // block_size, :, slot_mapping % block_size, :] + else: # [B, N, H, C] + block_size = kv_cache.shape[1] + out = kv_cache[slot_mapping // block_size, slot_mapping % block_size] + return out[:num_tokens] + + +def _hidden_state_group_index(kv_cache_config) -> int: + """Index of the kv-cache group holding the capture layer. + + ``build_connector_meta`` runs in the SCHEDULER process, which never calls + ``register_kv_caches`` and has no worker forward-context to look layers up in, so the + group is identified from the static spec instead. ``prefer_cross_layer_blocks`` puts + the capture layer in a group of its own, and a request's ``block_ids`` is one list per + group -- picking the wrong one addresses the attention KV instead of the hidden states. + """ + from vllm.v1.kv_cache_interface import HiddenStateCacheSpec + + groups = getattr(kv_cache_config, "kv_cache_groups", None) or [] + matches = [i for i, g in enumerate(groups) if isinstance(g.kv_cache_spec, HiddenStateCacheSpec)] + if len(matches) != 1: + raise RuntimeError( + "RdmaHiddenStatesConnector: expected exactly one HiddenStateCacheSpec group in " + f"the kv-cache config, found {len(matches)} of {len(groups)}. Without it the " + "per-request block ids cannot be resolved." + ) + return matches[0] @dataclass @@ -76,11 +140,20 @@ class ReqMeta: req_id: str token_ids: torch.Tensor slot: int + # The request's OWN blocks in the hidden-state cache. Its slots are derived from + # these rather than from a positional window into the batch's slot_mapping -- see + # save_kv_layer. + block_ids: tuple = () @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=tuple(block_ids), + ) @dataclass @@ -89,9 +162,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 +235,9 @@ 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 + self._kv_cache_config = kv_cache_config + self._hs_group = _hidden_state_group_index(kv_cache_config) + self._cache_block_size = 0 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")) @@ -184,6 +260,9 @@ def __init__(self, vllm_config, role, kv_cache_config): self._slot_gen: dict[int, int] = {} self._accum_finished: set[str] = set() self._oversize_warned = False # one-shot guard against log spam on oversized prompts + self._num_blocks = 0 + self._firstcap_dumped = False + self._slotcheck_left = 8 # bounded deep validation; the cheap bound check is always on self._lock = threading.Lock() # scheduler state self._slot_ctr = 0 @@ -229,9 +308,50 @@ 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:]) + # per-token feature = (n_captured_planes, hidden). Which axes those are + # depends on the vLLM build; see planes_dim(). + spec_cfg = self._vllm_config.speculative_config + n_planes = len( + getattr(spec_cfg.draft_model_config.hf_config, "eagle_aux_hidden_state_layer_ids", []) + ) + if not n_planes: + raise RuntimeError( + "RdmaHiddenStatesConnector: eagle_aux_hidden_state_layer_ids is empty; " + "cannot interpret the hidden-state cache layout." + ) + self._planes_dim = planes_dim(kv, n_planes) + self._cache_block_size = kv.shape[2] if self._planes_dim == 1 else kv.shape[1] + self._num_blocks = kv.shape[0] + # The block ids we address with come from the scheduler, per kv-cache group. If the + # group we picked is not this cache's, its ids index a different (usually much + # larger) pool -- and because every group starts allocating from low ids, the first + # requests can look fine and only later requests run off the end. Dump what we + # resolved so that mismatch is visible at startup rather than as an assert later. + for i, g in enumerate(getattr(self._kv_cache_config, "kv_cache_groups", None) or []): + spec = g.kv_cache_spec + print( + f"[HSCONN] kv-cache group {i}: spec={type(spec).__name__} " + f"layers={len(g.layer_names)} spec_block_size={getattr(spec, 'block_size', '?')}" + f"{' <-- capturing from this one' if i == self._hs_group else ''}", + flush=True, + ) + print( + f"[HSCONN] capture cache view {tuple(kv.shape)} planes_dim={self._planes_dim} " + f"blocks={self._num_blocks} block_size={self._cache_block_size}", + flush=True, + ) + self._feat_shape = ( + (kv.shape[1], kv.shape[3]) if self._planes_dim == 1 else tuple(kv.shape[2:]) + ) + self._per_token_elems = int(prod(self._feat_shape)) + logger.info( + "RdmaHiddenStatesConnector: hidden-state cache view %s -> %d planes x %d " + "hidden (planes on dim %d)", + tuple(kv.shape), + self._feat_shape[0], + self._feat_shape[1], + self._planes_dim, + ) self._dtype = kv.dtype self._slot_elems = self._max_tokens * self._per_token_elems @@ -309,15 +429,24 @@ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs): ready = torch.cuda.Event() ready.record() cs.wait_event(ready) + # Each request's slots come from ITS OWN blocks, not from a positional window + # into the batch's slot_mapping. + # + # The previous version walked slot_mapping with a running offset, assuming the + # newly scheduled (prefill) requests sat contiguously at its front. vLLM appends + # new requests to the persistent batch, so any request already decoding occupies + # earlier positions and shifts that window: every capture then read a DIFFERENT + # request's hidden states. Nothing downstream could catch it -- the shapes are + # right, the pool slot is right, the recorded token_ids are the request's own, and + # the reconstructed distribution is a real, sharp one. It simply belongs to another + # prompt, which merely makes the drafter train on noise. + # + # Measured: with the serve capped at max_num_seqs=1 (so a batch can never hold a + # second request) the per-sequence hidden/token match is 0.96-1.00; with concurrency + # it is 0.00-0.02, and the drafter plateaus at bigram-level accuracy. 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 + block_size = self._cache_block_size + device = slot_mapping.device for req in md.requests: n = req.token_ids.shape[0] nelems = n * self._per_token_elems @@ -337,12 +466,49 @@ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs): self._max_tokens, ) self._oversize_warned = True - offset += n continue + need = (n + block_size - 1) // block_size + if len(req.block_ids) < need: + raise RuntimeError( + f"RdmaHiddenStatesConnector: request {req.req_id} has " + f"{len(req.block_ids)} blocks in the hidden-state cache group but needs " + f"{need} for {n} tokens at block_size {block_size}. Capturing from the " + f"batch's slot_mapping by position instead would silently read another " + f"request's hidden states." + ) + hi = max(req.block_ids) + if hi >= self._num_blocks: + raise RuntimeError( + f"RdmaHiddenStatesConnector: request {req.req_id} references block {hi} " + f"but the capture cache only has {self._num_blocks} blocks " + f"(group {self._hs_group}, block_size {block_size}, n={n}). The block " + f"ids being handed out are not this cache's." + ) with torch.cuda.stream(cs): - rsm = slot_mapping[offset : offset + n] - offset += n - hs_gpu = extract_from_kv_cache(kv_layer, rsm, n) # [n, *feat] + # Built HERE, on the stream that consumes them. Built on the default stream + # instead they are a cross-stream hazard: nothing orders cs against the + # allocation, so the caching allocator can hand their memory to someone else + # while cs is still reading, and the resulting garbage indices surface as an + # out-of-bounds assert inside whatever kernel happens to run next -- minutes + # into a run, not reproducibly, and never pointing at this line. The code + # this replaced was safe only incidentally: it sliced vLLM's long-lived + # slot_mapping rather than allocating anything. + blocks = torch.as_tensor(req.block_ids, dtype=torch.long, device=device) + pos = torch.arange(n, device=device) + rsm = blocks[pos // block_size] * block_size + (pos % block_size) + if self._slotcheck_left > 0: + # Bounded deep check: every derived slot must be one this step actually + # writes. torch.isin over the batch is not free, so only the first few. + self._slotcheck_left -= 1 + if not bool(torch.isin(rsm, slot_mapping).all()): + raise RuntimeError( + f"RdmaHiddenStatesConnector: block-derived slots for " + f"{req.req_id} are not in this step's slot_mapping (n={n}, " + f"block_size={block_size}, blocks={len(req.block_ids)}). The " + f"hidden-state cache group or its slot encoding is not what " + f"this connector assumes." + ) + hs_gpu = extract_from_kv_cache(kv_layer, rsm, n, self._planes_dim) # [n, *feat] # copy into the pre-registered pool slot (flattened) self._pool[slot, :nelems].copy_(hs_gpu.reshape(-1), non_blocking=True) ev = torch.cuda.Event() @@ -390,7 +556,12 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnector ) 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._hs_group], + ) return meta def request_finished(self, request, block_ids): From 25f14f49f9d599aba6ec4fbee58b33a869f6131a Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 23 Sep 2026 07:03:37 +0000 Subject: [PATCH 3/3] fix(speculative): find the capture layer when vLLM merges it into the attention group Streaming hidden-state capture could not start at all on a model whose attention layers are all one type -- Qwen3, Llama, most dense models. The serve died at connector init with "expected exactly one HiddenStateCacheSpec group in the kv-cache config, found 0 of 1". vLLM's get_kv_cache_groups splits hidden-state layers into their own group, with the comment that they "use their own block table and must not be absorbed into a compatible attention bucket" -- but that split sits after two early returns for uniform specs, and HiddenStateCacheSpec subclasses FullAttentionSpec. So when every attention layer has the same spec, UniformTypeKVCacheSpecs.from_specs accepts the capture layer as one more full-attention layer and returns a single merged group, and the split never runs. A model with mixed attention is not uniform, falls through, and does get its own group, which is why this never showed up on the models the connector was developed against. The merged group is still usable -- it keeps a per-layer spec dict and one block table indexing every layer in it -- so the fix is in the lookup, not in the capture. Two shapes have to be recognised because the two processes do not see the same object: the worker's copy of the config carries the real UniformTypeKVCacheSpecs and can be inspected per layer, while the copy the scheduler holds has been flattened to a representative FullAttentionSpec by the time it crosses the process boundary. The scheduler side therefore falls back on the one unambiguous case: a single group has a single block table, which necessarily indexes the capture layer too. Anything else still raises, and now names the group specs and their layer counts. Verified on vLLM 0.29.0 and on a later nightly, both of which order the branches this way, so this is not a transient. Qwen3-8B DFlash2 streaming trains after the change: 600 steps, loss 17.0 -> 3.4, drafter exported. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../plugins/rdma_hidden_states_connector.py | 63 +++++++++++++++---- 1 file changed, 50 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py index 8a09a797c9c..f904b1e2c85 100644 --- a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py +++ b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py @@ -111,26 +111,63 @@ def extract_from_kv_cache(kv_cache, slot_mapping, num_tokens, pdim: int = 2): return out[:num_tokens] +def _group_holds_hidden_state(group_spec) -> bool: + """Whether a kv-cache group's spec covers the hidden-state capture layer. + + Two shapes, because vLLM only sometimes gives the capture layer a group of its own. + ``get_kv_cache_groups`` splits hidden-state layers out -- "they use their own block + table and must not be absorbed into a compatible attention bucket" -- only AFTER two + early returns for uniform specs, and ``HiddenStateCacheSpec`` subclasses + ``FullAttentionSpec``. So on a model whose attention layers are all one type (Qwen3, + Llama, most dense models) ``UniformTypeKVCacheSpecs.from_specs`` accepts the capture + layer as just another full-attention layer and returns a SINGLE merged group before + the split is reached. A model with mixed attention -- some sliding-window layers, say + -- is not uniform, falls through, and does get the separate group. + + Merged is still usable: the group keeps a per-layer spec dict and one block table that + indexes every layer in it, so a request's block ids address the capture layer + correctly. Only the lookup has to know about both shapes. + """ + from vllm.v1.kv_cache_interface import HiddenStateCacheSpec + + if isinstance(group_spec, HiddenStateCacheSpec): + return True + per_layer = getattr(group_spec, "kv_cache_specs", None) + if isinstance(per_layer, dict): + return any(isinstance(spec, HiddenStateCacheSpec) for spec in per_layer.values()) + return False + + def _hidden_state_group_index(kv_cache_config) -> int: """Index of the kv-cache group holding the capture layer. ``build_connector_meta`` runs in the SCHEDULER process, which never calls ``register_kv_caches`` and has no worker forward-context to look layers up in, so the - group is identified from the static spec instead. ``prefer_cross_layer_blocks`` puts - the capture layer in a group of its own, and a request's ``block_ids`` is one list per - group -- picking the wrong one addresses the attention KV instead of the hidden states. + group is identified from the static spec instead. A request's ``block_ids`` is one + list per group -- picking the wrong one addresses the attention KV instead of the + hidden states, which is not an error, just wrong data. """ - from vllm.v1.kv_cache_interface import HiddenStateCacheSpec - groups = getattr(kv_cache_config, "kv_cache_groups", None) or [] - matches = [i for i, g in enumerate(groups) if isinstance(g.kv_cache_spec, HiddenStateCacheSpec)] - if len(matches) != 1: - raise RuntimeError( - "RdmaHiddenStatesConnector: expected exactly one HiddenStateCacheSpec group in " - f"the kv-cache config, found {len(matches)} of {len(groups)}. Without it the " - "per-request block ids cannot be resolved." - ) - return matches[0] + matches = [i for i, g in enumerate(groups) if _group_holds_hidden_state(g.kv_cache_spec)] + if len(matches) == 1: + return matches[0] + if not matches and len(groups) == 1: + # The merged case again, seen from the scheduler. The worker's copy of the config + # carries the real UniformTypeKVCacheSpecs, whose per-layer dict the test above can + # inspect; the copy the scheduler holds has been flattened to a representative + # FullAttentionSpec by the time it crosses the process boundary, so that test finds + # nothing even though the capture layer IS in this group. With exactly one group + # there is exactly one block table and it necessarily indexes every layer, capture + # included. If capture were genuinely absent, register_kv_caches on the worker + # fails first and louder: it cannot find a view whose plane count matches the + # configured capture ids. + return 0 + raise RuntimeError( + "RdmaHiddenStatesConnector: expected exactly one kv-cache group holding a " + f"HiddenStateCacheSpec, found {len(matches)} of {len(groups)}. Without it the " + "per-request block ids cannot be resolved. Groups: " + f"{[(type(g.kv_cache_spec).__name__, len(g.layer_names)) for g in groups]}" + ) @dataclass