From 080dfda0f75ce989ffaef464b2fa850d24c604de 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/8] 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) <noreply@anthropic.com>
---
 .../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 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/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."""

From eef4f0e09d0c0f652ec25cd6848768ebefb7e1d4 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/8] 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) <noreply@anthropic.com>
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 7ad46e4cb4224192cf85ed1160163364c68def75 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/8] 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) <noreply@anthropic.com>
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

From 464ffc141437b4db9033bdf45b14524ad86cf40c Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Fri, 30 Oct 2026 11:23:56 +0000
Subject: [PATCH 4/8] fix(speculative): keep the streaming corpus memory-mapped
 instead of copying it per rank

StreamingDataset turned its entries into a list. The trainer passes an
Arrow-backed datasets.Dataset, so list() decoded every row into Python objects
in every rank, and every forked DataLoader worker then copy-on-wrote the pages
it touched. An indexable source is now kept as is; only an iterator is still
materialized.

On a 3.1M-row Muse-Spark corpus the list measured 80 GB per rank, and with 4
ranks and 8 DataLoader workers each it OOM-killed a trainer node with 950 GB of
host memory. On the 3.48M-row corpus of the released Muse-Spark-Open DFlash2
drafter, a trainer node stopped responding 20 steps into training. Kept as
Arrow, the rows are memory-mapped and shared by the forked workers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 CHANGELOG.rst                                        |  1 +
 .../speculative/plugins/hf_streaming_dataset.py      | 12 ++++++++----
 .../speculative/plugins/test_hf_streaming_dataset.py | 10 ++++++++++
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 89fb2a8b8d7..6834c166a02 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -105,6 +105,7 @@ 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.
 
 0.48.0 (2026-10-xx)
 ^^^^^^^^^^^^^^^^^^^
diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
index 5e02097c96a..f226899baaf 100644
--- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
+++ b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
@@ -218,12 +218,16 @@ def __init__(
             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 = 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
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..bdd4517384b 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,16 @@ 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 = 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

From 5a8a8b3ed1ca068343e2923bc0c5acc242e0f3d7 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Fri, 30 Oct 2026 10:22:35 +0000
Subject: [PATCH 5/8] fix(speculative): weigh DFlash2's selector loss like the
 backbone loss

The selector cross-entropy now uses the backbone loss's own per-position
weights (decay or D-PACE) and divides by the backbone's weight sum, misses
included. Before, every covered position counted equally and the term was
divided by the covered positions alone.

A block position's choice only matters if every earlier position was
accepted, which is what the backbone's weighting models, so the selector now
spends its capacity the same way. Dividing by the covered positions alone
also scaled the term up whenever the backbone's top-k recall dropped. This
matches the SGLang/SpecForge DFlash2 reference, and it is the objective the
released Muse-Spark-Open DFlash2 drafter was trained with.

The selector reads the backbone's mask and weights from _DFlashLossTerms
instead of rebuilding the mask, so the two cannot drift apart. The module
docstring still described the gold-token substitution that the strict top-k
replaced; it now describes the current objective.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../torch/speculative/plugins/hf_dflash2.py   | 57 +++++++++++--------
 .../speculative/plugins/test_hf_dflash2.py    | 32 ++++++++++-
 2 files changed, 61 insertions(+), 28 deletions(-)

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/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.
 

From d45a8875c9daecbae5670e30362294fe070f237c Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Sat, 31 Oct 2026 14:12:01 +0000
Subject: [PATCH 6/8] fix(speculative): make the streaming connector's startup
 log visible; test its addressing

Review follow-ups for the streaming hidden-state connector.

The module's logger was named modelopt.torch...; vLLM configures only the
"vllm" logger tree, so every logger.info in this module was dropped (WARNING
and above still reached stderr through Python's last-resort handler). That is
why the startup diagnostics were print() calls. The logger now sits under
vllm.*, and the two [HSCONN] prints are folded into one logger.info that lists
every kv-cache group and names the one being captured.

Also:
- Drop the prefer_cross_layer_blocks override. vLLM #51718 deleted the base
  property, and before that its default was already False, the value the
  override returned, so it changed nothing on any vLLM version.
- Drop the unused _firstcap_dumped attribute.
- Give build_connector_meta's full-prompt assert its current reason: a request
  is captured in the step that first schedules it, and under chunked prefill
  that step computes only the first chunk and its block_ids cover only that.
- Lift the block-to-slot derivation into _request_slots so it can be tested.

New CPU tests sit in tests/gpu_vllm because the module imports vLLM at the top
and only the gpu_vllm job has vLLM. They check that a request reads its own
tokens through its own blocks on both cache layouts, that an ambiguous cache
view is refused, and that the capture group is found whether vLLM gives it a
group of its own, merges it into the attention group, or hands the scheduler a
flattened spec. Each fails with its fix reverted.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../plugins/rdma_hidden_states_connector.py   | 61 ++++++------
 .../test_rdma_hidden_states_connector.py      | 97 +++++++++++++++++++
 2 files changed, 125 insertions(+), 33 deletions(-)
 create mode 100644 tests/gpu_vllm/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 f904b1e2c85..025052ce7af 100644
--- a/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py
+++ b/modelopt/torch/speculative/plugins/rdma_hidden_states_connector.py
@@ -60,7 +60,8 @@
 
 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 planes_dim(kv_cache, n_planes: int) -> int:
@@ -111,6 +112,13 @@ def extract_from_kv_cache(kv_cache, slot_mapping, num_tokens, pdim: int = 2):
     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.
 
@@ -263,11 +271,6 @@ 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)
@@ -298,7 +301,6 @@ def __init__(self, vllm_config, role, kv_cache_config):
         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
@@ -359,35 +361,30 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         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))
+        # 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: hidden-state cache view %s -> %d planes x %d "
-            "hidden (planes on dim %d)",
+            "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
@@ -530,9 +527,7 @@ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
                 # 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)
+                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.
@@ -578,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:
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))

From 771feca941614b76b9349655dc1b331972674675 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Sat, 31 Oct 2026 14:12:14 +0000
Subject: [PATCH 7/8] refactor(speculative): type StreamingDataset's entries as
 the iterable it accepts

Review follow-up. Since the corpus fix, the constructor keeps a sized,
indexable source as is, so list[dict] understated what callers can pass: the
Arrow-backed datasets.Dataset that main.py hands it is not even an
abc.Sequence. Annotate Iterable[dict] on StreamingDataset and
EagleVllmStreamingDataset, and say in the docstring which inputs are kept and
which are materialized. Also note why the test imports datasets inside it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 .../speculative/plugins/hf_streaming_dataset.py      | 12 ++++++++----
 .../speculative/plugins/test_hf_streaming_dataset.py |  1 +
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py
index f226899baaf..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,6 +214,9 @@ 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
@@ -223,7 +227,7 @@ def __init__(
         # 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 = entries
+            self.entries = cast("Sequence[dict]", entries)
         else:
             self.entries = list(entries)
         if len(self.entries) == 0:
@@ -409,7 +413,7 @@ class EagleVllmStreamingDataset(StreamingDataset):
 
     def __init__(
         self,
-        entries: list[dict],
+        entries: Iterable[dict],
         tokenizer,
         config: EagleVllmStreamingConfig,
     ):
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 bdd4517384b..f7a41f42a45 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py
@@ -58,6 +58,7 @@ def test_empty_corpus_raises():
 
 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())

From 3eb6189439ab2b33181fadd33c2af836d41149c9 Mon Sep 17 00:00:00 2001
From: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Date: Sat, 31 Oct 2026 14:12:18 +0000
Subject: [PATCH 8/8] docs: add changelog entries for the streaming connector,
 RoPE and selector fixes

The connector shipped in 0.46.0, and the 0.48.0 release branch already carries
both the RoPE regression (#2483) and DFlash2, so all three fix released code.
The connector's HiddenStateCacheSpec import needs vLLM 0.22 or later; it is
tested on 0.29 and later, hence the backward-breaking entry.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
---
 CHANGELOG.rst | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 6834c166a02..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**
 
@@ -106,6 +107,9 @@ Changelog
 - 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)
 ^^^^^^^^^^^^^^^^^^^