[TRTLLM-14620][feat] Add DSpark ragged graph and runtime integration

tensorrt_llm/_torch/pyexecutor/engine/metadata.py

@@ -59,6 +59,7 @@ def build_attention_metadata(
     draft_kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None = None,
     enable_context_mla_with_cached_kv: bool | None = None,
     max_num_heads_per_kv: int | None = None,
+    enable_ragged_verification: bool = False,
 ) -> AttentionMetadata:
     """Construct attention metadata and resolve model-derived inputs."""
     pretrained_config = model_config.pretrained_config
@@ -75,6 +76,15 @@ def build_attention_metadata(
         if sparse_attention_config is not None
         else None
     )
+    metadata_kwargs = {}
+    if enable_ragged_verification:
+        from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import (
+            DSAtrtllmAttentionMetadata,
+        )
+
+        if not issubclass(attention_backend.Metadata, DSAtrtllmAttentionMetadata):
+            raise ValueError("Ragged verification requires the DSA target metadata backend")
+        metadata_kwargs["enable_ragged_verification"] = True
     return attention_backend.Metadata(
         max_num_requests=max_batch_size,
         max_num_tokens=max_num_tokens,
@@ -88,4 +98,5 @@ def build_attention_metadata(
         cache_indirection=cache_indirection,
         max_num_heads_per_kv=max_num_heads_per_kv,
         sparse_metadata_params=sparse_metadata_params,
+        **metadata_kwargs,
     )

tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/graphs.py

@@ -58,6 +58,7 @@
 from ..common import get_top_level_model, prepare_multimodal_indices
 from .config import DecoderRunnerConfig
 from .dummy_requests import DummyRequestFactory
+from .inputs import DecoderInputPreparer
 from .metadata import DecoderMetadata
 from .warmup import WarmupAgreement
 
@@ -205,8 +206,10 @@ def __init__(
         torch_compile_bypass: Callable[[], contextlib.AbstractContextManager[None]],
         context_promotion_supported: bool,
         warmup_inputs: torch.Tensor | None = None,
+        inputs: DecoderInputPreparer | None = None,
     ) -> None:
         self.cuda_graph_runner = cuda_graph_runner
+        self._inputs = inputs
         self.breakable_cuda_graph_runner = breakable_cuda_graph_runner
         self._config = config
         self.mapping = mapping
@@ -329,7 +332,7 @@ def captures_model_body(self, key: KeyType | None) -> bool:
 
     def reset_padding_dummies(self) -> None:
         """Drop padding dummies that belong to an earlier KV cache manager."""
-        self.cuda_graph_runner.padding_dummy_requests = {}
+        self.cuda_graph_runner.reset_padding_dummies()
 
     def preallocate_padding_dummies(self, resource_manager: ResourceManager) -> None:
         """Allocate the padding dummies while the KV cache is still empty."""
@@ -602,7 +605,9 @@ def capture_model(model_inputs: dict[str, Any]) -> Any:
                     restore_attn_metadata_after_draft_replay(attn_metadata, saved_draft)
         return outputs
 
-    def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tuple[int, int]]:
+    def _get_graphs_to_capture(
+        self, cuda_graph_batch_sizes: list[int]
+    ) -> list[tuple[int, int] | tuple[int, int, int]]:
         """Determine which (batch_size, draft_len) graphs to capture.
 
         Returns:
@@ -624,6 +629,33 @@ def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tupl
             )
             return graphs
 
+        if self._inputs is not None and self._inputs._dspark_confidence_enabled:
+            max_draft_len = int(self._config.spec_config.max_draft_len)
+            native_graphs = [(bs, max_draft_len) for bs in cuda_graph_batch_sizes]
+            if not self._inputs._dspark_trims_submitted_tokens:
+                return native_graphs
+            from tensorrt_llm._torch.speculative.dspark_planner import ExactSpsCostTable
+
+            exact_table = self._inputs._dspark_sps_cost_table
+            if not isinstance(exact_table, ExactSpsCostTable):
+                raise RuntimeError(
+                    "DSpark confidence graph capture requires an authenticated exact SPS table"
+                )
+            ragged_graphs = [
+                (bs, max_draft_len, verifier_budget)
+                for bs in cuda_graph_batch_sizes
+                for verifier_budget in exact_table.production_candidate_budgets(bs)
+            ]
+            # Full K keeps the ordinary native graph. Compact layouts have
+            # distinct keys; they cannot alias its uniform attention layout.
+            graphs = native_graphs + ragged_graphs
+            logger.info(
+                f"DSpark ragged verification: capturing {len(graphs)} graphs "
+                f"({len(native_graphs)} native static-K{max_draft_len} graphs + "
+                f"{len(ragged_graphs)} exact measured (G,V) cells)."
+            )
+            return graphs
+
         # Case 2: Static draft length
         # Match the runtime_draft_len semantics enforced in resolve_runtime_draft_len:
         # logical K for linear-tree modes, total tree tokens for tree decoding.
@@ -758,7 +790,9 @@ def _run_capture_pass(
             pinned_tier = sample_type or SampleType.FULL
             self.cuda_graph_runner.set_capture_sample_type(pinned_tier)
             try:
-                for bs, draft_len in graphs_to_capture:
+                for entry in graphs_to_capture:
+                    bs, draft_len = entry[:2]
+                    verify_bucket = entry[2] if len(entry) == 3 else None
                     if bs > self._config.max_batch_size:
                         continue
 
@@ -773,6 +807,8 @@ def _run_capture_pass(
                         with self._dummy_requests.release_batch_context(
                             warmup_request, resource_manager
                         ) as batch:
+                            if batch is not None and verify_bucket is not None:
+                                self._set_warmup_ragged_windows(batch, verify_bucket, draft_len)
                             if batch is None:
                                 # No KV cache space for this batch size. During KV
                                 # cache estimation this makes the profiling peak
@@ -814,6 +850,8 @@ def _run_capture_pass(
             finally:
                 self._force_lora_graph_for_capture = None
                 self.cuda_graph_runner.set_capture_sample_type(None)
+                if self._inputs is not None and self._inputs._dspark_confidence_enabled:
+                    self.cuda_graph_runner.agreed_ragged_bucket = None
 
         if self._lora.cuda_graph_manager is None:
             lora_graph_cases = [False]
@@ -1009,3 +1047,27 @@ def release(self, *, release_nccl_window_owners: bool = True) -> None:
             self.breakable_cuda_graph_runner.clear(
                 release_nccl_window_owners=release_nccl_window_owners
             )
+
+    def _set_warmup_ragged_windows(self, batch, verify_bucket: int, draft_len: int) -> None:
+        """Give a capture batch per-request windows summing to ``verify_bucket``.
+
+        The warmup batch must hit the bucket total exactly, or the graph is
+        keyed under a total no runtime batch produces and every ragged step
+        silently falls back to eager. Any split with the right sum captures
+        the right shape (the raggedness lives in the contents of qo_indptr).
+        """
+        requests = batch.generation_requests
+        if not requests:
+            return
+        from tensorrt_llm._torch.speculative.ragged_helpers import RaggedVerifyLayout
+
+        lens = torch.ones(len(requests), dtype=torch.int32)
+        layout = RaggedVerifyLayout.from_verify_lens(
+            lens, graph_num_tokens=int(verify_bucket), total_verify_tokens=len(requests)
+        )
+        filled = layout.fill_bucket(max_verify_len=1 + int(draft_len))
+        for request, tokens in zip(requests, filled.verify_lens.tolist()):
+            request.py_verify_len = int(tokens) - 1
+        # Publish to the runner too: the graph key carries this value, and the
+        # runtime fit never runs during capture.
+        self.cuda_graph_runner.agreed_ragged_bucket = int(verify_bucket)