diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 89fb2a8b8d7..3d0112d84d6 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -59,6 +59,7 @@ Changelog - ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning. - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing. +- Streaming speculative-decoding training's vLLM connector (``modelopt.torch.speculative.plugins.rdma_hidden_states_connector``) now requires vLLM 0.22 or later. It is tested with vLLM 0.29 and later. **Deprecations** @@ -105,6 +106,10 @@ Changelog - Fix ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``), which failed with an ``AttributeError`` on every rank. - AWQ calibration (``awq_lite``, ``awq_full``) of an FSDP2 model sharded over more than one rank is now rejected with ``NotImplementedError`` by ``mtq.quantize``, ``mtq.calibrate`` and ``mtq.auto_quantize`` (before its search starts, its default formats included), and by ``examples/hf_ptq --use_fsdp2`` before loading the model: each rank would fold a different smoothing scale into its weight shard. Calibrate in a single process, or use a non-AWQ algorithm. - ``mtq.compress`` now rejects W4A8 weights with ``NotImplementedError``: a compressed W4A8 model could neither run nor be reloaded. +- Fix streaming speculative-decoding training copying the whole training corpus into every rank and every DataLoader worker, which ran trainer nodes out of host memory on corpora of a few million samples. The corpus now stays memory-mapped. +- Fix streaming speculative-decoding training capturing another request's hidden states when the vLLM serve ran requests concurrently, and failing on vLLM 0.29 and later. Retrain drafters trained in streaming mode against a serve that ran requests concurrently. +- Fix DFlash-family drafters (DFlash, DFlash2, DSpark, Domino, LiLiCorr) trained from scratch with ``examples/speculative_decoding/main.py`` having their RoPE frequencies rounded to BF16, so they trained under different rotary phases than vLLM serves them with. +- Fix DFlash2's candidate-selector loss weighing every block position equally and dividing by its own count of positions, unlike the backbone loss and the SGLang/SpecForge DFlash2 reference. Logged DFlash2 losses change. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 5d3fc293efc..1e20d5b3099 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -562,7 +562,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) @@ -592,10 +593,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/hf_dflash2.py b/modelopt/torch/speculative/plugins/hf_dflash2.py index 6e4171c9210..88c1e3e38f7 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash2.py +++ b/modelopt/torch/speculative/plugins/hf_dflash2.py @@ -29,9 +29,11 @@ - Take the backbone's top-k candidates per block position. - Score each candidate against its *teacher-forced* predecessor token, so the positions train in parallel exactly as the backbone does. -- When the gold token is missing from the top-k, substitute it into the last - candidate slot. Without this the selector sees no positive class on the hard - positions and never learns those edges. +- Supervise only the strict top-k, the candidate set serving builds. A gold token + outside it is a backbone recall miss, not a selector example, and carries no + selector gradient. +- Weight each position as the backbone loss does (decay or D-PACE) and divide by the + backbone's weight sum, misses included. """ import torch @@ -39,7 +41,7 @@ from transformers import PreTrainedModel from ..dflash.conversion import DFlash2DMRegistry -from .hf_dflash import HFDFlashModel, _block_targets, _block_tokens +from .hf_dflash import HFDFlashModel, _block_tokens from .modeling_dflash2 import DFlash2Module __all__ = ["HFDFlash2Model"] @@ -106,7 +108,9 @@ def get_exporter(self): return DFlash2Exporter(self) - def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mask): + def _selector_loss( + self, logits, target_ids, hidden, predecessor_ids, supervised_mask, loss_weights, weight_sum + ): """Cross-entropy over the selector's candidate set, and its top-1 accuracy. Args: @@ -114,7 +118,10 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas target_ids: Gold token ids ``[B, N, block_size]``. hidden: Backbone hidden states ``[B, N, block_size, H]``. predecessor_ids: Teacher-forced predecessor ids ``[B, N, block_size]``. - weight_mask: Per-position loss weights ``[B, N, block_size]``. + supervised_mask: The positions the backbone loss supervises, unweighted + ``[B, N, block_size]``. + loss_weights: The backbone loss's per-position weights ``[B, N, block_size]``. + weight_sum: The backbone loss's denominator, the sum of ``loss_weights``. Returns: ``(loss, accuracy, coverage)`` — coverage is the fraction of supervised @@ -130,7 +137,7 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas # token the backbone did not propose is a backbone recall failure, not a # selector classification example: substituting it in would teach the selector # to override the unary ranking on a set it will never be shown. Those - # positions carry no selector gradient and leave the denominator instead. + # positions carry no selector gradient. gold_matches = candidate_ids == target_ids.unsqueeze(-1) gold_in_topk = gold_matches.any(dim=-1) gold_slot = gold_matches.long().argmax(dim=-1) @@ -139,26 +146,27 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas candidate_ids, unary_logits, hidden, predecessor_ids ) - covered = weight_mask * gold_in_topk.to(weight_mask.dtype) - flat_weights = covered.reshape(-1) - denominator = flat_weights.sum() + 1e-6 + in_topk = gold_in_topk.to(loss_weights.dtype) per_token = F.cross_entropy( selector_logits.float().reshape(-1, top_k), gold_slot.reshape(-1), reduction="none", ) - loss = (per_token * flat_weights).sum() / denominator + # The backbone's denominator, misses included: dividing by the covered positions + # alone would scale the term up whenever backbone recall drops. + loss = (per_token * (loss_weights * in_topk).reshape(-1)).sum() / weight_sum with torch.no_grad(): + covered = (supervised_mask * in_topk).reshape(-1) chosen = selector_logits.argmax(dim=-1).reshape(-1) - accuracy = ( - (chosen == gold_slot.reshape(-1)).float() * flat_weights - ).sum() / denominator + accuracy = ((chosen == gold_slot.reshape(-1)).float() * covered).sum() / ( + covered.sum() + 1e-6 + ) # Coverage keeps the full supervised mask as its denominator: it measures # how often the selector was given a solvable problem at all. # clamp, not an epsilon: a fully covered batch must read exactly 1.0. - supervised = weight_mask.reshape(-1).sum() - coverage = flat_weights.sum() / supervised.clamp(min=1.0) + supervised = supervised_mask.reshape(-1).sum() + coverage = covered.sum() / supervised.clamp(min=1.0) # Detached tensors, not Python scalars: .item() would force a CPU-GPU sync on # every training step. The trainer converts them at the logging boundary. return loss, accuracy.detach(), coverage.detach() @@ -235,15 +243,10 @@ def _compute_loss( block_size = self.dflash_block_size n_blocks = anchor_positions.shape[1] - # The backbone's supervision mask, without its position weighting (decay/D-PACE): - # that shapes *where* the backbone spends capacity, while the selector should learn - # every position's transition equally. - targets = _block_targets( - input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift - ) + target_ids = _block_tokens(input_ids, anchor_positions, block_size, self._label_shift) # Teacher-forced predecessor of block position k is the real token at anchor+k-1. - # Only offsets 1..block_size-1 are supervised (the mask zeroes slot 0), so the + # Only offsets 1..block_size-1 are supervised (the backbone's mask zeroes slot 0), so the # first supervised position, k=1, has the anchor's own token as its predecessor. # Slot 0's entry here resolves to anchor-1 and never contributes. # @@ -254,12 +257,16 @@ def _compute_loss( input_ids, anchor_positions, block_size, self._label_shift - 1 ) + # The backbone's own mask and position weights: a block position's choice only + # matters if every earlier one was accepted, which decay / D-PACE model. selector_loss, selector_accuracy, selector_coverage = self._selector_loss( logits.reshape(bsz, n_blocks, block_size, -1), - targets.token_ids, + target_ids, draft_hidden.reshape(bsz, n_blocks, block_size, -1), predecessor_ids, - targets.mask, + terms.supervised_mask.view_as(target_ids), + terms.weights.view_as(target_ids), + terms.weight_sum, ) self._selector_metrics = { "selector_accuracy": selector_accuracy, diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py index 5e02097c96a..62f17402d8e 100644 --- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py +++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py @@ -41,7 +41,8 @@ import os import re import time -from typing import TypedDict +from collections.abc import Iterable, Sequence +from typing import TypedDict, cast import httpx import torch @@ -199,7 +200,7 @@ class StreamingDataset(Dataset): def __init__( self, - entries: list[dict], + entries: Iterable[dict], tokenizer, config: StreamingConfig | None = None, ): @@ -213,17 +214,24 @@ def __init__( Args: entries: Untokenized per-sample dicts from the input jsonl. Schema is subclass-defined (see :meth:`_tokenize_entry`); passed to :meth:`_fetch`. + A sized, indexable source (e.g. an Arrow-backed ``datasets.Dataset``) is + kept as is, so its rows stay memory-mapped; any other iterable is + materialized into a list. tokenizer: HF tokenizer; used for client-side tokenization and the server/client loss-mask alignment in :meth:`_fetch`. config: Tuning knobs (timeout, answer_only_loss, ...); defaults to ``self.config_cls()``. See :class:`StreamingConfig`. """ - if not entries: - raise ValueError("entries is empty") self.tokenizer = tokenizer self.config = config if config is not None else self.config_cls() - # Materialize to a plain list so DataLoader worker processes fork it cheaply. - self.entries = list(entries) + # Keep an indexable source (e.g. an Arrow-backed datasets.Dataset) as is: list() + # would copy every row into each rank, and each forked worker would copy it again. + if hasattr(entries, "__len__") and hasattr(entries, "__getitem__"): + self.entries = cast("Sequence[dict]", entries) + else: + self.entries = list(entries) + if len(self.entries) == 0: + raise ValueError("entries is empty") # Per-process consecutive-failure counter for the circuit breaker. Reset to 0 # on every successful fetch; tripped only by fetch failures (not unfit entries). self._consecutive_fail = 0 @@ -405,7 +413,7 @@ class EagleVllmStreamingDataset(StreamingDataset): def __init__( self, - entries: list[dict], + entries: Iterable[dict], tokenizer, config: EagleVllmStreamingConfig, ): diff --git a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py index e68262a97ed..025052ce7af 100644 --- a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py +++ b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py @@ -60,13 +60,122 @@ from .hf_streaming_dataset import nixl_backends_from_env -logger = init_logger(__name__) +# Under vllm.*: vLLM configures only its own logger tree, so a modelopt.* logger's INFO is dropped. +logger = init_logger(f"vllm.{__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 _request_slots(block_ids, num_tokens: int, block_size: int, device=None): + """Cache slots of a request's first ``num_tokens`` tokens, addressed through its own blocks.""" + blocks = torch.as_tensor(block_ids, dtype=torch.long, device=device) + pos = torch.arange(num_tokens, device=device) + return blocks[pos // block_size] * block_size + pos % block_size + + +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 +185,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 +207,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): @@ -153,15 +271,13 @@ class RdmaHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): See the module docstring for the pooled design and scope. """ - @property - def prefer_cross_layer_blocks(self) -> bool: - """Keep the capture layer's blocks separate from the drafter's KV cache.""" - return False - 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 +300,8 @@ 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._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 +347,45 @@ 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] + 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)) + # 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. Log what we + # resolved so that mismatch is visible at startup rather than as an assert later. + groups = getattr(self._kv_cache_config, "kv_cache_groups", None) or [] + logger.info( + "RdmaHiddenStatesConnector: capturing from kv-cache group %d of %s; cache view " + "%s -> %d planes x %d hidden (planes on dim %d), %d blocks of %d tokens", + self._hs_group, + [ + (type(g.kv_cache_spec).__name__, len(g.layer_names), g.kv_cache_spec.block_size) + for g in groups + ], + tuple(kv.shape), + self._feat_shape[0], + self._feat_shape[1], + self._planes_dim, + self._num_blocks, + self._cache_block_size, + ) self._dtype = kv.dtype self._slot_elems = self._max_tokens * self._per_token_elems @@ -309,15 +463,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 +500,47 @@ 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. + rsm = _request_slots(req.block_ids, n, block_size, device) + 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() @@ -375,9 +573,9 @@ 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). + A request is captured once, in the step that first schedules it, from the blocks it + holds then. Under chunked prefill that step computes only the first chunk and + ``block_ids`` covers only that chunk, so require the whole prompt here. """ meta = RdmaConnMeta() for nr in scheduler_output.scheduled_new_reqs: @@ -390,7 +588,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/gpu_vllm/torch/speculative/plugins/test_rdma_hidden_states_connector.py b/tests/gpu_vllm/torch/speculative/plugins/test_rdma_hidden_states_connector.py new file mode 100644 index 00000000000..e89a7751cf7 --- /dev/null +++ b/tests/gpu_vllm/torch/speculative/plugins/test_rdma_hidden_states_connector.py @@ -0,0 +1,97 @@ +# 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. + +"""How the streaming hidden-state connector finds a request's hidden states in vLLM's cache. + +CPU-only, but here rather than in tests/unit because the connector imports vLLM at module +level. Every mistake these tests guard against is silent: the draft trains on another +token's or another request's hidden states, and nothing raises. +""" + +from types import SimpleNamespace + +import pytest +import torch +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + HiddenStateCacheSpec, + UniformTypeKVCacheSpecs, +) + +from modelopt.torch.speculative.plugins.rdma_hidden_states_connector import ( + _hidden_state_group_index, + _request_slots, + extract_from_kv_cache, + planes_dim, +) + +NUM_BLOCKS, BLOCK_SIZE, N_PLANES, HIDDEN = 5, 4, 3, 8 + +ATTENTION = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16) +CAPTURE = HiddenStateCacheSpec( + block_size=1, num_kv_heads=N_PLANES, head_size=HIDDEN, dtype=torch.bfloat16 +) + + +@pytest.mark.parametrize("layout", ["BNHC", "BHNC"]) +def test_a_request_reads_its_own_tokens_from_either_layout(layout): + """vLLM #51718 (0.29) moved the captured planes from dim 2 of the cache view to dim 1.""" + tokens = torch.randn(NUM_BLOCKS * BLOCK_SIZE, N_PLANES, HIDDEN) # row i = slot i + cache = tokens.view(NUM_BLOCKS, BLOCK_SIZE, N_PLANES, HIDDEN) + if layout == "BHNC": + cache = cache.transpose(1, 2).contiguous() + # Holding blocks 3 then 1, the request's 6 tokens sit in slots 12-15 and 4-5. + slots = _request_slots((3, 1), 6, BLOCK_SIZE) + got = extract_from_kv_cache(cache, slots, 6, planes_dim(cache, N_PLANES)) + assert torch.equal(got, tokens[[12, 13, 14, 15, 4, 5]]) + + +@pytest.mark.parametrize("dims", [(N_PLANES, N_PLANES), (BLOCK_SIZE, BLOCK_SIZE)]) +def test_a_view_that_does_not_show_where_the_planes_are_is_refused(dims): + with pytest.raises(RuntimeError): + planes_dim(torch.empty(NUM_BLOCKS, *dims, HIDDEN), N_PLANES) + + +def _groups(*specs): + return SimpleNamespace( + kv_cache_groups=[SimpleNamespace(kv_cache_spec=s, layer_names=["layer"]) for s in specs] + ) + + +@pytest.mark.parametrize( + ("config", "index"), + [ + # Mixed attention types: the capture layer gets a group of its own. + (_groups(ATTENTION, CAPTURE), 1), + # One attention type, worker side: the capture layer is merged into its group. + ( + _groups( + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs={"attn": ATTENTION, "capture": CAPTURE} + ) + ), + 0, + ), + # The same group as the scheduler sees it, flattened to one representative spec. + (_groups(ATTENTION), 0), + ], +) +def test_the_capture_group_is_found_however_vllm_groups_it(config, index): + assert _hidden_state_group_index(config) == index + + +def test_no_capture_group_among_several_is_refused(): + with pytest.raises(RuntimeError): + _hidden_state_group_index(_groups(ATTENTION, ATTENTION)) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index 55d2978a7e2..d3bf4162cf7 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -428,6 +428,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.""" diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py index c43677818c0..e1ea8c10dc0 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py @@ -397,9 +397,11 @@ def _selector(self, vocab=16, rank=4, top_k=3, hidden=8): torch.nn.init.normal_(selector.successor_codebook, std=0.5) return selector - def _selector_terms(self, selector, target_ids, candidate_ids): + def _selector_terms(self, selector, target_ids, candidate_ids, loss_weights=None): """Drive HFDFlash2Model._selector_loss with a hand-built candidate set.""" b, length, k = candidate_ids.shape + supervised = torch.ones(b, length, dtype=torch.double) + weights = supervised if loss_weights is None else loss_weights model = HFDFlash2Model.__new__(HFDFlash2Model) model.dflash_module = type("_M", (), {"candidate_selector": selector})() logits = torch.zeros(b, length, 32, dtype=torch.double) @@ -410,9 +412,13 @@ def _selector_terms(self, selector, target_ids, candidate_ids): model, logits, target_ids, - torch.randn(b, length, 8, dtype=torch.double), + torch.randn( + b, length, 8, generator=torch.Generator().manual_seed(1), dtype=torch.double + ), torch.zeros(b, length, dtype=torch.long), - torch.ones(b, length, dtype=torch.double), + supervised, + weights, + weights.sum(), ) def test_a_miss_carries_no_selector_gradient(self): @@ -453,6 +459,26 @@ def test_a_covered_position_does_train_the_selector(self): loss.backward() assert selector.successor_codebook.grad.abs().sum() > 0.0 + def test_positions_weigh_like_the_backbone_and_misses_stay_in_the_denominator(self): + """The term is the backbone-weighted mean over every supervised position. + + A miss adds nothing to the numerator but keeps its weight in the denominator, so + the term does not grow as backbone recall drops. + """ + selector = self._selector(vocab=32, top_k=2) + candidates = torch.tensor([[[5, 6], [5, 6], [5, 6]]]) + targets = torch.tensor([[5, 6, 9]]) # two covered positions, one miss + weights = torch.tensor([[1.0, 3.0, 2.0]], dtype=torch.double) + + loss, _, _ = self._selector_terms(selector, targets, candidates, weights) + per_position = [ + self._selector_terms(selector, targets, candidates, torch.eye(3)[i : i + 1].double())[0] + for i in range(3) + ] + assert float(per_position[2]) == 0.0 + expected = sum(w * term for w, term in zip(weights[0], per_position)) / weights.sum() + torch.testing.assert_close(loss, expected) + def test_greedy_path_position_zero_is_seeded_by_the_anchor(self): """``greedy_path`` position 0 scores against the anchor, i.e. block offset 1. diff --git a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py index 00b5a1f6ab5..f7a41f42a45 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py +++ b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py @@ -56,6 +56,17 @@ def test_empty_corpus_raises(): StreamingDataset([], tokenizer=MagicMock(), config=StreamingConfig()) +def test_an_indexable_corpus_is_kept_and_an_iterator_materialized(): + """An Arrow-backed corpus stays memory-mapped instead of being copied into every rank.""" + # datasets is an optional dependency; skip only this test without it. + datasets = pytest.importorskip("datasets") + corpus = datasets.Dataset.from_list(_entries(5)) + ds = StreamingDataset(corpus, tokenizer=MagicMock(), config=StreamingConfig()) + assert ds.entries is corpus + ds = StreamingDataset(iter(_entries(5)), tokenizer=MagicMock(), config=StreamingConfig()) + assert ds.entries == _entries(5) + + def test_len_matches_corpus(): ds = StreamingDataset(_entries(37), tokenizer=MagicMock(), config=StreamingConfig()) assert len(ds) == 37