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/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py index e68262a97ed..f904b1e2c85 100644 --- a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py +++ b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py @@ -63,10 +63,111 @@ 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 _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. 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. + """ + groups = getattr(kv_cache_config, "kv_cache_groups", None) or [] + 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 @@ -76,11 +177,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 +199,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 +272,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 +297,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 +345,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 +466,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 +503,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 +593,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): 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."""