[None][fix] Restore DSpark disaggregated decoding accuracy

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -293,11 +293,39 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice:
             block_ids = adapter.get_block_ids(req, idx, lg)
             # Limit to prompt_len blocks, matching C++ cacheFormatter behavior.
             total_blocks = (req.prompt_len + tpb - 1) // tpb
-            if block_ids.size > total_blocks:
-                block_ids = block_ids[:total_blocks]
             window_size = lg.sliding_window_size
 
             if window_size is not None:
+                allocated_blocks = (
+                    req.prompt_len + self._kv_cache_manager.num_extra_kv_tokens + tpb - 1
+                ) // tpb
+                beam0_block_ids, tail_block_ids = self._split_packed_beam_block_ids(
+                    block_ids,
+                    req.py_beam_width,
+                    allocated_blocks,
+                )
+                if beam0_block_ids.size > allocated_blocks:
+                    beam0_block_ids = beam0_block_ids[:allocated_blocks]
+                    block_ids = (
+                        np.concatenate([beam0_block_ids, tail_block_ids])
+                        if tail_block_ids.size > 0
+                        else beam0_block_ids
+                    )
+                # Current PyExecutor cache managers disable KV-cache token sinks,
+                # so SWA block lists contain an evictable prompt prefix followed
+                # by the speculative scratch tail. If token sinks are enabled,
+                # this must use block-ordinal metadata to preserve the sink prefix.
+                # Remove scratch before trimming stale prompt blocks; otherwise a
+                # boundary-crossing allocation can displace initialized prompt KV.
+                scratch_blocks = max(0, allocated_blocks - total_blocks)
+                if scratch_blocks > 0:
+                    if req.py_beam_width != 1:
+                        raise ValueError("speculative scratch blocks require beam_width == 1")
+                    block_ids = (
+                        block_ids[:-scratch_blocks]
+                        if scratch_blocks < block_ids.size
+                        else np.array([], dtype=np.int64)
+                    )
                 # Drop stale blocks the manager may still expose (V1 pre-eviction).
                 stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb)
                 expected_valid = max(0, total_blocks - stale_end)
@@ -307,7 +335,8 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice:
                 # stale region (those blocks were already pruned, no extra skip).
                 cache_skip = max(0, cached_per_lg[idx] // tpb - stale_end)
             else:
-                total_blocks = (req.prompt_len + tpb - 1) // tpb
+                if block_ids.size > total_blocks:
+                    block_ids = block_ids[:total_blocks]
                 expected_valid = total_blocks
                 cache_skip = cached_per_lg[idx] // tpb
 

tensorrt_llm/_torch/models/dspark/attention.py

@@ -165,6 +165,7 @@ def get_dspark_topk_idxs_batched(
     window_size: int,
     block_size: int,
     start_pos: torch.Tensor,
+    valid_len: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``.
 
@@ -173,29 +174,33 @@ def get_dspark_topk_idxs_batched(
     ``start_pos``), this always returns the **fixed** width ``window_size +
     block_size`` and masks the unfilled context slots with ``-1``. The masked
     slots are excluded by :func:`dspark_sparse_attn` exactly as if they were
-    absent, so the result is numerically identical to gathering only the
-    ``min(window_size, start_pos+1)`` valid context positions — but the shape no
-    longer depends on the data, which is what CUDA-graph capture requires.
+    absent while the shape remains CUDA-graph safe.
 
-    Every query position attends to the same set: context window slots
-    ``0..window_size-1`` (slot ``c`` valid iff ``c <= start_pos[g]``, i.e. it has
-    been written) followed by the ``block_size`` block positions at offset
-    ``window_size`` (always valid).
+    Every query attends to the actually written circular-window suffix, followed
+    by the current-block positions. Without ``valid_len`` this preserves the
+    legacy ``start_pos``-only behavior.
 
     Args:
         window_size: sliding-window length of the captured-context KV cache.
         block_size: number of draft positions per request.
         start_pos: ``[G]`` int tensor of per-request absolute decode positions.
+        valid_len: optional ``[G]`` count of actually written rolling-window
+            entries. When omitted, preserve the legacy ``start_pos`` mask.
 
     Returns:
         int32 tensor ``[G, block_size, window_size + block_size]``.
     """
     device = start_pos.device
     g = start_pos.shape[0]
     ctx_cols = torch.arange(window_size, device=device)  # [win]
-    # Context slot c holds a written key iff c <= start_pos (slots 0..start_pos
-    # filled; for start_pos >= window_size-1 the whole rolling window is filled).
-    valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1)  # [G, win]
+    if valid_len is None:
+        valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1)  # [G, win]
+    else:
+        # The valid entries are the contiguous logical suffix ending at
+        # start_pos, but their physical slots wrap modulo window_size.
+        valid_len = valid_len.clamp(min=0, max=window_size)
+        age = torch.remainder(start_pos.unsqueeze(1) - ctx_cols.unsqueeze(0), window_size)
+        valid = age < valid_len.unsqueeze(1)
     ctx_idx = torch.where(
         valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long)
     )
@@ -371,6 +376,7 @@ def dspark_attention_forward_batched(
     start_pos: torch.Tensor,
     kv_cache: torch.Tensor,
     slots: torch.Tensor,
+    valid_len: torch.Tensor | None = None,
     *,
     wq_a: torch.Tensor,
     q_norm_w: torch.Tensor,
@@ -415,6 +421,9 @@ def dspark_attention_forward_batched(
         kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows
             (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers).
         slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row.
+        valid_len: optional ``[G]`` count of actually written context entries;
+            masks holes left when absolute positions are bootstrapped without
+            receiving the corresponding DSpark rolling-window state.
         freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table;
             must satisfy ``maxlen > start_pos.max() + block_size``.
 
@@ -451,7 +460,7 @@ def dspark_attention_forward_batched(
     write_target[slots, slot_pos] = main_kv.squeeze(1).to(write_target.dtype)
     cache_rows = write_target[slots]  # [G, window, head_dim]
     kv_full = torch.cat([cache_rows, kv], dim=1)  # [G, window + block, head_dim]
-    topk = get_dspark_topk_idxs_batched(window_size, block, start_pos)
+    topk = get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len)
     o = dspark_sparse_attn(q, kv_full, attn_sink, topk, softmax_scale)  # [G, block, h, head_dim]
     o = _rope_last_dims_batched(o, rd, blk_freqs, inverse=True)
 

tensorrt_llm/_torch/models/modeling_dspark.py

@@ -824,6 +824,7 @@ def _forward_stage(
         moe_input_ids: torch.Tensor,
         stage_window: Optional[torch.Tensor] = None,
         slots: Optional[torch.Tensor] = None,
+        valid_len: Optional[torch.Tensor] = None,
         all_rank_num_tokens: Optional[List[int]] = None,
     ) -> torch.Tensor:
         """One DSpark stage = reference ``Block.forward`` with captured-context attn.
@@ -874,6 +875,7 @@ def _forward_stage(
                 start_pos,
                 kv_cache,
                 slots,
+                valid_len=valid_len,
                 freqs_cis=freqs_cis,
                 persist=True,
                 **stage._dspark_attn,
@@ -1015,6 +1017,7 @@ def forward_batched(
         *,
         kv_windows: torch.Tensor,
         slots: torch.Tensor,
+        valid_len: Optional[torch.Tensor] = None,
         temperature: float = 0.0,
         confidence_threshold: float = 0.0,
         return_logits: bool = False,
@@ -1039,6 +1042,7 @@ def forward_batched(
             kv_windows: ``[N, num_stages, window_size, head_dim]`` persistent rolling
                 windows; written in place through ``slots``.
             slots: ``[G]`` int tensor mapping each request to its ``kv_windows`` row.
+            valid_len: ``[G]`` count of actually written rolling-window entries.
         Returns:
             ``(draft_tokens [G, block], num_proposed [G])`` from ``forward_head``.
         """
@@ -1064,6 +1068,7 @@ def forward_batched(
                 moe_input_ids,
                 stage_window,
                 slots,
+                valid_len,
                 all_rank_num_tokens=all_rank_num_tokens,
             )
 

tensorrt_llm/_torch/speculative/dspark.py

@@ -111,6 +111,8 @@ def prepare(self):
                 if rid not in current:
                     slot = worker._req_to_slot.pop(rid)
                     worker._ctx_len[slot] = 0
+                    worker._valid_len[slot] = 0
+                    worker._position_initialized[slot] = False
                     worker._kv_windows[slot].zero_()
                     worker._free_slots.append(slot)
             # Assign a persistent rolling-window slot to every real generation
@@ -225,6 +227,8 @@ def __init__(
         self._win_inited = False
         self._kv_windows: Optional[torch.Tensor] = None  # [max_batch, num_stages, win, hd]
         self._ctx_len: Optional[torch.Tensor] = None  # [max_batch] abs decode position
+        self._valid_len: Optional[torch.Tensor] = None  # [max_batch] written window entries
+        self._position_initialized: Optional[torch.Tensor] = None  # [max_batch] bool
         self._win = 0
 
         # Slot management. ``_req_to_slot`` (python dict) + ``_free_slots`` are the
@@ -298,6 +302,8 @@ def _lazy_init(self, draft_model, spec_metadata) -> None:
             device="cuda",
         )
         self._ctx_len = torch.zeros(num_rows, dtype=torch.long, device="cuda")
+        self._valid_len = torch.zeros(num_rows, dtype=torch.long, device="cuda")
+        self._position_initialized = torch.zeros(num_rows, dtype=torch.bool, device="cuda")
         self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda")
         self._free_slots = deque(range(max_batch))
         self._req_to_slot = {}
@@ -313,6 +319,8 @@ def _assign_slot(self, req_id: int, reset: bool) -> int:
         if reset and req_id in self._req_to_slot:
             old = self._req_to_slot.pop(req_id)
             self._ctx_len[old] = 0
+            self._valid_len[old] = 0
+            self._position_initialized[old] = False
             self._kv_windows[old].zero_()
             self._free_slots.append(old)
         if req_id not in self._req_to_slot:
@@ -324,6 +332,8 @@ def _assign_slot(self, req_id: int, reset: bool) -> int:
             slot = self._free_slots.popleft()
             self._req_to_slot[req_id] = slot
             self._ctx_len[slot] = 0
+            self._valid_len[slot] = 0
+            self._position_initialized[slot] = False
             self._kv_windows[slot].zero_()
         return self._req_to_slot[req_id]
 
@@ -355,8 +365,12 @@ def _seed_context_windows(
             first_position = int(chunk_positions[0].item())
             slot = self._assign_slot(req_id, reset=first_position == 0)
             self._ctx_len[slot] = chunk_positions[-1] + 1
+            self._position_initialized[slot] = True
 
             if captured is not None:
+                self._valid_len[slot] = torch.clamp(
+                    self._valid_len[slot] + chunk_len, max=self._win
+                )
                 keep = min(self._win, chunk_len)
                 hidden = captured[context_offset + chunk_len - keep : context_offset + chunk_len]
                 # A prompt token at absolute position p is stored in frame p+1,
@@ -365,6 +379,22 @@ def _seed_context_windows(
                 draft_model.write_context_windows(hidden, window_positions, self._kv_windows[slot])
             context_offset += chunk_len
 
+    def _advance_generation_state(
+        self,
+        slots: torch.Tensor,
+        num_accepted_tokens: torch.Tensor,
+        input_positions: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Bootstrap and advance per-slot decode state without host synchronization."""
+        old = torch.where(self._position_initialized[slots], self._ctx_len[slots], input_positions)
+        start_pos = old + num_accepted_tokens
+        self._ctx_len[slots] = start_pos
+        self._valid_len[slots] = torch.clamp(
+            self._valid_len[slots] + num_accepted_tokens, max=self._win
+        )
+        self._position_initialized[slots] = torch.ones_like(slots, dtype=torch.bool)
+        return old, start_pos
+
     def _draft_gen_block_batched(
         self,
         draft_model,
@@ -375,6 +405,7 @@ def _draft_gen_block_batched(
         num_contexts: int,
         batch_size: int,
         total_target_tokens: int,
+        position_ids: torch.Tensor,
         all_rank_num_tokens: Optional[List[int]] = None,
     ) -> torch.Tensor:
         """CUDA-graph-safe batched gen draft (all gen requests in one forward).
@@ -390,7 +421,6 @@ def _draft_gen_block_batched(
         """
         num_gens = batch_size - num_contexts
         K = self.max_draft_len
-        Kp1 = K + 1
         device = accepted_tokens.device
 
         if num_gens == 0:
@@ -411,16 +441,23 @@ def _draft_gen_block_batched(
             accepted_tokens[num_contexts:batch_size].gather(1, gidx.unsqueeze(1)).squeeze(1).long()
         )  # [G]
 
-        # Captured target hidden at the bonus position within each request's Kp1
-        # processed tokens.
+        # Bootstrap iterations can process one target token per request, while
+        # normal speculative verification processes K+1. Use the actual accepted
+        # row width to index both captured hidden states and position IDs.
+        target_width = accepted_tokens.shape[1]
         arange_g = torch.arange(num_gens, device=device)
-        base = gen_start + arange_g * Kp1  # [G]
+        base = gen_start + arange_g * target_width  # [G]
         main_hidden = captured[base + gidx]  # [G, ncap*hidden]
 
         # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens
         # (everything but the bonus) into the rolling window — same frames as the
         # eager path (old+1 .. old+nacc-1), with j >= nacc-1 masked out.
-        old = self._ctx_len[slots]  # [G] pre-increment decode position
+        # A disaggregated generation worker never sees prompt prefill, so a new
+        # slot has no absolute decode position. Bootstrap it once from the first
+        # target input position; locally-prefilled and existing slots keep their
+        # monotonically advanced position.
+        input_positions = position_ids.reshape(-1)[base].long()
+        old, start_pos = self._advance_generation_state(slots, nacc, input_positions)
         j = torch.arange(K, device=device)  # [K]
         interim_valid = j.unsqueeze(0) < (nacc.unsqueeze(1) - 1)  # [G, K]
         interim_pos = old.unsqueeze(1) + 1 + j.unsqueeze(0)  # [G, K]
@@ -432,11 +469,6 @@ def _draft_gen_block_batched(
             interim_hidden, interim_pos, slots, interim_valid, self._kv_windows
         )
 
-        # Advance the decode position by the accepted count; start_pos (= post-
-        # increment ctx_len) matches the eager path's frame value.
-        start_pos = old + nacc  # [G]
-        self._ctx_len[slots] = start_pos
-
         # Surface the per-position corrected block logits ([num_gens, K, vocab])
         # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection)
         # sampling + TP gather + draft_probs scatter, rather than argmaxing here.
@@ -446,6 +478,7 @@ def _draft_gen_block_batched(
             start_pos,
             kv_windows=self._kv_windows,
             slots=slots,
+            valid_len=self._valid_len[slots],
             temperature=0.0,
             confidence_threshold=0.0,
             return_logits=True,
@@ -538,6 +571,8 @@ def _forward_impl(
         )
         if is_warmup:
             saved_ctx_len = self._ctx_len.clone()
+            saved_valid_len = self._valid_len.clone()
+            saved_position_initialized = self._position_initialized.clone()
             saved_windows = self._kv_windows.clone()
 
         # Assign / reset window slots for context (prefill) requests and seed each
@@ -592,6 +627,7 @@ def _forward_impl(
                 num_contexts,
                 batch_size,
                 total_target_tokens,
+                position_ids,
                 all_rank_num_tokens=all_rank_draft_tokens,
             )
             if gen_logits is not None:
@@ -649,6 +685,8 @@ def _forward_impl(
 
         if is_warmup:
             self._ctx_len.copy_(saved_ctx_len)
+            self._valid_len.copy_(saved_valid_len)
+            self._position_initialized.copy_(saved_position_initialized)
             self._kv_windows.copy_(saved_windows)
 
         return {

tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py

@@ -195,6 +195,23 @@ def test_batched_topk_matches_scalar(start_positions):
             torch.testing.assert_close(valid, scalar[m].to(valid.dtype))
 
 
+def test_batched_topk_respects_partial_window_valid_len():
+    """Only actually written circular slots are visible for partial seeds."""
+    window, block = 8, 3
+    start_pos = torch.tensor([3, 10, 20], dtype=torch.long)
+    valid_len = torch.tensor([0, 3, 8], dtype=torch.long)
+    batched = get_dspark_topk_idxs_batched(window, block, start_pos, valid_len)
+    assert tuple(batched.shape) == (3, block, window + block)
+
+    for i, (pos, count) in enumerate(zip(start_pos.tolist(), valid_len.tolist())):
+        expected_context = sorted((pos - age) % window for age in range(count))
+        row = batched[i, 0]
+        context = row[(row >= 0) & (row < window)].tolist()
+        assert context == expected_context
+        block_indices = row[row >= window].tolist()
+        assert block_indices == list(range(window, window + block))
+
+
 @pytest.mark.parametrize("ndim", [3, 4])
 def test_batched_rotary_matches_scalar_per_row(ndim):
     """apply_dspark_rotary_batched (per-row freqs) == scalar applied row by row."""