[TRTLLM-14620][feat] Add ragged DSpark verification semantics

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -1717,3 +1717,22 @@ def get_draft_token_length(request: LlmRequest) -> int:
     if request.py_draft_tokens is not None:
         return len(request.py_draft_tokens)
     return 0
+
+
+def get_request_tokens_per_gen_step(request: LlmRequest,
+                                    runtime_tokens_per_gen_step: int) -> int:
+    """How many tokens this request contributes to one generation step.
+
+    Uniform speculation gives every request the batch-wide
+    ``runtime_tokens_per_gen_step`` (1 accepted token + ``runtime_draft_len``
+    draft positions). Ragged verification -- where a scheduler hands each
+    request its own verify window -- records the per-request draft length in
+    ``py_verify_len``.
+
+    ``py_verify_len`` is set only by the ragged path, so every other
+    speculation mode gets the batch-wide value back unchanged.
+    """
+    verify_len = getattr(request, "py_verify_len", None)
+    if verify_len is None:
+        return runtime_tokens_per_gen_step
+    return 1 + int(verify_len)

tensorrt_llm/_torch/pyexecutor/sampler/penalties.py

@@ -672,6 +672,12 @@ class SpecMetadataLike(Protocol):
     vocab_size: int
     penalty_state: Optional["PenaltyState"]
     batch_slot_ids: Optional[torch.Tensor]
+    verify_lens: Optional[torch.Tensor]
+    qo_indptr: Optional[torch.Tensor]
+    total_verify_tokens: Optional[int]
+
+    @property
+    def is_ragged_verify(self) -> bool: ...
 
 
 @dataclass(kw_only=True)
@@ -886,13 +892,16 @@ def build_row_mapping(
     draft_len: int,
     draft_tokens: torch.Tensor,
     device: torch.device,
+    num_logit_rows: Optional[int] = None,
 ) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
     """Derive the per-row slot map and intra-step prefix for a packed logits batch.
 
-    Assumes the layout every linear one-model mode produces once
-    ``_reshape_logits_for_accept`` has run: ``[ctx (1 row each), gen (draft_len + 1
-    rows each)]``, request-major. PARD emits ``2 * draft_len`` rows per gen request
-    but that hook already narrows them to ``draft_len + 1``, so both share this code.
+    Uniform linear modes produce ``[ctx (1 row each), gen (draft_len + 1 rows
+    each)]`` once ``_reshape_logits_for_accept`` has run. Ragged verification
+    instead packs each generation request's ``verify_lens[r]`` token rows back to
+    back on the flat token axis. In either layout, rows for one request are
+    contiguous and ordered by speculative position. PARD emits ``2 * draft_len``
+    rows per gen request but its reshape hook narrows them before this function.
 
     Tree modes do NOT fit: their rows are tree nodes, so a row's prefix is its
     root path rather than the rows before it. They are rejected at admission.
@@ -909,16 +918,56 @@ def build_row_mapping(
         return None
 
     num_gens = batch_size - num_contexts
-    rows_per_gen = draft_len + 1
-    total_rows = num_contexts + num_gens * rows_per_gen
-
-    batch_slots = slot_ids[:batch_size].to(torch.int64)
-    # Context requests own one row; gen requests own rows_per_gen consecutive rows.
+    batch_slots = slot_ids[:batch_size].to(device=device, dtype=torch.int64)
+
+    if spec_metadata.is_ragged_verify:
+        from tensorrt_llm._torch.speculative.ragged_helpers import row_ids_from_lens
+
+        verify_lens = spec_metadata.verify_lens
+        qo_indptr = spec_metadata.qo_indptr
+        total_verify_tokens = spec_metadata.total_verify_tokens
+        if (
+            verify_lens is None
+            or qo_indptr is None
+            or verify_lens.dim() != 1
+            or qo_indptr.dim() != 1
+            or verify_lens.numel() != num_gens
+            or qo_indptr.numel() != num_gens + 1
+        ):
+            return None
+        if total_verify_tokens is None:
+            if num_logit_rows is None:
+                return None
+            total_verify_tokens = num_logit_rows - num_contexts
+        total_gen_rows = int(total_verify_tokens)
+        if total_gen_rows < 0:
+            return None
+        total_rows = num_contexts + total_gen_rows
+        if num_logit_rows is not None and total_rows != num_logit_rows:
+            return None
+
+        lens = verify_lens.to(device=device, dtype=torch.int64)
+        gen_request_ids = row_ids_from_lens(lens, total=total_gen_rows)
+        packed_positions = torch.arange(total_gen_rows, device=device)
+        starts = qo_indptr[:num_gens].to(device=device, dtype=torch.int64)
+        position_in_request = packed_positions - starts[gen_request_ids]
+    else:
+        rows_per_gen = draft_len + 1
+        total_gen_rows = num_gens * rows_per_gen
+        total_rows = num_contexts + total_gen_rows
+        if num_logit_rows is not None and total_rows != num_logit_rows:
+            return None
+        gen_rows = torch.arange(total_gen_rows, device=device)
+        gen_request_ids = gen_rows // rows_per_gen
+        position_in_request = gen_rows % rows_per_gen
+
+    # Context requests own one row. Generation rows are request-contiguous in
+    # both the uniform rectangle and the ragged packed token axis.
     row_slots = (
         torch.cat(
             [
                 batch_slots[:num_contexts],
-                batch_slots[num_contexts:batch_size].repeat_interleave(rows_per_gen),
+                batch_slots[num_contexts:batch_size][gen_request_ids],
             ]
         )
         if num_gens > 0
@@ -935,15 +984,12 @@ def build_row_mapping(
     intra_valid = torch.zeros((total_rows, max(draft_len, 0)), dtype=torch.bool, device=device)
     if num_gens > 0 and draft_len > 0 and draft_tokens.numel() > 0:
         drafts = draft_tokens.reshape(num_gens, -1)[:, :draft_len].to(torch.int64)
-        gen_rows = torch.arange(num_gens * rows_per_gen, device=device)
-        g_of_row = gen_rows // rows_per_gen  # which request owns the row
-        k_of_row = gen_rows % rows_per_gen  # the row's speculative position
         # Every gen row carries its own request's drafts; validity alone selects the
         # strictly-earlier prefix. Column j counts for row k iff j < k.
-        intra_tokens[num_contexts:] = drafts[g_of_row]
+        intra_tokens[num_contexts:] = drafts[gen_request_ids]
         intra_valid[num_contexts:] = torch.arange(draft_len, device=device).unsqueeze(
             0
-        ) < k_of_row.unsqueeze(1)
+        ) < position_in_request.unsqueeze(1)
     return row_slots, intra_tokens, intra_valid
 
 

tensorrt_llm/_torch/speculative/dspark.py

@@ -53,6 +53,59 @@ def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int)
     return int(max_ctx) + int(max_draft_len) + int(block_size) + 3
 
 
+def _publish_policy_window_output(
+    outputs: dict[str, torch.Tensor | bool],
+    verify_lens: Optional[torch.Tensor],
+    *,
+    num_contexts: int,
+    num_generations: int,
+    output_verify_lens: Optional[torch.Tensor],
+) -> None:
+    """Publish executed token windows in full output-row order.
+
+    ``verify_lens`` is generation-ordered, including graph-padding generations.
+    ``output_verify_lens`` is persistent int32 device storage sized at warmup.
+    Context rows have one anchor and no verified drafts. The sampler removes
+    skipped context rows and snapshots the remaining view through its D2H path.
+    """
+    batch_size = num_contexts + num_generations
+    if num_contexts < 0 or num_generations < 0:
+        raise ValueError("verification output row counts must be nonnegative")
+    for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens", "next_new_tokens"):
+        if outputs[name].shape[0] != batch_size:
+            raise ValueError("verification output row counts must match every worker output")
+
+    if verify_lens is None:
+        outputs.pop("verify_lens", None)
+        outputs.pop("verify_lens_in_output_order", None)
+        outputs.pop("host_policy_windows_snapshot", None)
+        outputs["native_uniform_verify"] = True
+        return
+    if (
+        verify_lens.dim() != 1
+        or verify_lens.shape[0] != num_generations
+        or verify_lens.dtype != torch.int32
+        or verify_lens.device != outputs["new_tokens_lens"].device
+    ):
+        raise ValueError("verify_lens must contain exactly one int32 window per generation row")
+    if (
+        output_verify_lens is None
+        or output_verify_lens.dim() != 1
+        or output_verify_lens.shape[0] < batch_size
+        or output_verify_lens.dtype != torch.int32
+        or output_verify_lens.device != verify_lens.device
+    ):
+        raise ValueError("executed-window output storage must cover the full batch on device")
+
+    output_view = output_verify_lens[:batch_size]
+    output_view[:num_contexts].fill_(1)
+    output_view[num_contexts:].copy_(verify_lens)
+    outputs.pop("native_uniform_verify", None)
+    outputs.pop("host_policy_windows_snapshot", None)
+    outputs["verify_lens"] = output_view
+    outputs["verify_lens_in_output_order"] = True
+
+
 @dataclass
 class DSparkSpecMetadata(SpecMetadata):
     """Metadata for DSpark speculative decoding.
@@ -267,6 +320,7 @@ def __init__(
         # Per-slot rolling captured-context KV windows, built lazily on the
         # first forward (fixed-size for slot-indexed reads/writes).
         self._win_inited = False
+        self._output_verify_lens: Optional[torch.Tensor] = None
         self._attention_warmup_attempted = 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
@@ -354,6 +408,7 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None:
 
         if not self._win_inited:
             max_batch = spec_metadata.max_num_requests
+            self._output_verify_lens = torch.empty(max_batch, dtype=torch.int32, device="cuda")
             num_stages = draft_model.num_stages
             self._win = int(draft_model._attn_params["window_size"])
             head_dim = int(draft_model._attn_params["head_dim"])
@@ -819,13 +874,21 @@ def _forward_impl(
             self._position_initialized.copy_(saved_position_initialized)
             self._kv_windows.copy_(saved_windows)
 
-        return {
+        outputs = {
             "logits": raw_logits,
             "new_tokens": accepted_tokens,
             "new_tokens_lens": num_accepted_tokens,
             "next_draft_tokens": next_draft_tokens,
             "next_new_tokens": next_new_tokens,
         }
+        _publish_policy_window_output(
+            outputs,
+            spec_metadata.verify_lens,
+            num_contexts=num_contexts,
+            num_generations=num_gens,
+            output_verify_lens=self._output_verify_lens,
+        )
+        return outputs
 
 
 class DSparkWorker(DFlashWorker):