[TRTLLM-14874][feat] Refactor advanced-sampling CUDA graph capture

tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py

@@ -24,6 +24,7 @@
 
 if TYPE_CHECKING:
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, SparseAttentionConfig
+    from tensorrt_llm.sampling_params import SamplingParams
 
 ModelConfig = tensorrt_llm.bindings.ModelConfig
 
@@ -124,6 +125,7 @@ def add_dummy_requests(
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
         draft_kv_cache_manager=None,
+        capture_sampling_params: Optional["SamplingParams"] = None,
     ):
         requests = super().add_dummy_requests(
             request_ids=request_ids,
@@ -136,6 +138,7 @@ def add_dummy_requests(
             max_beam_width=max_beam_width,
             encoder_output_lens=encoder_output_lens,
             draft_kv_cache_manager=draft_kv_cache_manager,
+            capture_sampling_params=capture_sampling_params,
         )
         if prepare_resource:
             for req in requests:

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -528,37 +528,6 @@ def maybe_get_cuda_graph(
             graph_spec_metadata = None
         return graph_attn_metadata, graph_spec_metadata, key
 
-    def clear_capture_only_spec_state(self) -> int:
-        """Clear capture-scoped state from every cached graph SpecMetadata.
-
-        ``create_cuda_graph_metadata`` shallow-copies the live SpecMetadata, so a
-        copy made while ``_run_capture_pass(force_non_greedy=True)`` is active
-        inherits ``_force_non_greedy_for_capture=True``. That copy is cached here
-        and reseated as the live spec_metadata on every later replay of its graph,
-        while the capture pass clears the flag on the base object only. Without
-        this cleanup the copies keep the flag forever and
-        ``_scan_one_model_sampling`` rewrites EVERY serving request's sampling
-        params to the synthetic capture values (temperature 0.7 / top_k 50 /
-        top_p 0.9), silently ignoring what the client asked for.
-
-        The flag must NOT be cleared at copy time instead: it is load-bearing
-        *during* capture. It is what makes the pass-2 populate scan non-greedy on
-        parameter-less warmup requests, so that the advanced-sampling branch (not
-        the argmax fast path, and with the top-k/top-p kernels present) is the one
-        recorded into the graph. Clearing it here -- after the pass has captured
-        every graph -- keeps capture correct and serving clean.
-
-        Returns the number of cached metadata objects cleared.
-        """
-        cleared = 0
-        for stored in self.graph_metadata.values():
-            spec_metadata = stored.get("spec_metadata")
-            if spec_metadata is not None and getattr(
-                    spec_metadata, "_force_non_greedy_for_capture", False):
-                spec_metadata._force_non_greedy_for_capture = False
-                cleared += 1
-        return cleared
-
     def needs_capture(self, key: KeyType):
         return self._capture_allowed and key not in self.graph_outputs
 

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -3358,6 +3358,7 @@ def add_dummy_requests(
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
         draft_kv_cache_manager: Optional["BaseResourceManager"] = None,
+        capture_sampling_params: Optional[SamplingParams] = None,
     ):
         _kv_draft = (
             kv_reserve_draft_tokens if kv_reserve_draft_tokens is not None else max_num_draft_tokens
@@ -3381,7 +3382,18 @@ def release_resources(
         for i, req_id in enumerate(request_ids):
             # exact choice of n can be ignored for dummy requests
             sampling_params = SamplingParams(
-                n=beam_width, best_of=beam_width, use_beam_search=beam_width > 1
+                n=beam_width,
+                best_of=beam_width,
+                use_beam_search=beam_width > 1,
+                temperature=capture_sampling_params.temperature
+                if capture_sampling_params is not None
+                else None,
+                top_k=capture_sampling_params.top_k
+                if capture_sampling_params is not None
+                else None,
+                top_p=capture_sampling_params.top_p
+                if capture_sampling_params is not None
+                else None,
             )
             # Here 1+max_num_draft_tokens is used to extend the prompt length to
             # a non-zero number to skip illegal memory access issue in MLA kernel

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -27,6 +27,7 @@
 if TYPE_CHECKING:
     from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
+    from tensorrt_llm.sampling_params import SamplingParams
 
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import (
     BlockReusePolicy, KVCacheManagerV2, Role)
@@ -2502,6 +2503,7 @@ def add_dummy_requests(
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
         draft_kv_cache_manager: Optional[KVCacheManager] = None,
+        capture_sampling_params: Optional["SamplingParams"] = None,
     ) -> List[LlmRequest]:
         requests = super().add_dummy_requests(
             request_ids=request_ids,
@@ -2514,6 +2516,7 @@ def add_dummy_requests(
             max_beam_width=max_beam_width,
             encoder_output_lens=encoder_output_lens,
             draft_kv_cache_manager=draft_kv_cache_manager,
+            capture_sampling_params=capture_sampling_params,
         )
         if requests:
             self.requests.extend(requests)
@@ -3559,6 +3562,7 @@ def add_dummy_requests(
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
         draft_kv_cache_manager: Optional[BaseResourceManager] = None,
+        capture_sampling_params: Optional["SamplingParams"] = None,
     ) -> List[LlmRequest]:
         requests = super().add_dummy_requests(
             request_ids=request_ids,
@@ -3571,6 +3575,7 @@ def add_dummy_requests(
             max_beam_width=max_beam_width,
             encoder_output_lens=encoder_output_lens,
             draft_kv_cache_manager=draft_kv_cache_manager,
+            capture_sampling_params=capture_sampling_params,
         )
         if requests and prepare_resource:
             self._setup_state_indices(requests)

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -45,6 +45,7 @@
                                           TorchCompileConfig, TorchLlmArgs)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import CpType, Mapping
+from tensorrt_llm.sampling_params import SamplingParams
 
 from ..attention_backend.interface import (AttentionMetadata,
                                            AttentionRuntimeFeatures)
@@ -334,6 +335,12 @@ def _filter_cuda_graph_seq_lens(cuda_graph_seq_lens: list[int],
 
 _DEEP_GEMM_PDL_CONFIGURED = False
 
+# Arbitrary non-greedy params used to force the advanced-sampling CUDA graph
+# warmup capture path.
+NON_GREEDY_CAPTURE_SAMPLING_PARAMS = SamplingParams(temperature=0.7,
+                                                    top_k=50,
+                                                    top_p=0.9)
+
 
 def _configure_deep_gemm_pdl() -> None:
     global _DEEP_GEMM_PDL_CONFIGURED
@@ -2373,25 +2380,20 @@ def prepare_cross_batch(batch: ScheduledRequests,
 
         def _run_capture_pass(force_non_greedy: bool, label: str,
                               force_lora_graph: bool) -> None:
-            spec_metadata = self.spec_metadata
             assert self._force_lora_graph_for_capture is None
             self._force_lora_graph_for_capture = force_lora_graph
-            if force_non_greedy and spec_metadata is not None:
-                spec_metadata._force_non_greedy_for_capture = True
-                # maybe_get_cuda_graph reads spec_metadata.is_all_greedy_sample
-                # to build the graph cache key BEFORE populate runs inside
-                # _prepare_inputs. Pre-flip it here so the very first capture
-                # in this pass uses the non-greedy key; populate's override
-                # below will keep it False on every subsequent iteration.
-                spec_metadata.is_all_greedy_sample = False
             try:
                 for bs, draft_len in graphs_to_capture:
                     if bs > self.batch_size:
                         continue
 
                     for max_seq_len in max_seq_len_list:
                         warmup_request = self._create_cuda_graph_warmup_request(
-                            resource_manager, bs, draft_len, max_seq_len)
+                            resource_manager,
+                            bs,
+                            draft_len,
+                            max_seq_len,
+                            force_non_greedy=force_non_greedy)
                         with self._release_batch_context(
                                 warmup_request, resource_manager) as batch:
                             if batch is None:
@@ -2422,25 +2424,6 @@ def _run_capture_pass(force_non_greedy: bool, label: str,
                             torch.cuda.synchronize()
             finally:
                 self._force_lora_graph_for_capture = None
-                if force_non_greedy and spec_metadata is not None:
-                    spec_metadata._force_non_greedy_for_capture = False
-                    # The base object is not the only holder of the flag: every
-                    # graph captured during this pass cached its own SHALLOW COPY
-                    # of spec_metadata (create_cuda_graph_metadata -> copy.copy),
-                    # which inherited the flag. Those copies are reseated as the
-                    # live spec_metadata on every later replay, so leaving the
-                    # flag set there makes _scan_one_model_sampling overwrite
-                    # every serving request's sampling params with the synthetic
-                    # capture values (0.7 / 50 / 0.9). Defer cleanup during
-                    # warmup-only because the real capture pass reuses that
-                    # pass's cached metadata.
-                    if not self.cuda_graph_runner.is_warmup_only:
-                        cleared = (self.cuda_graph_runner.
-                                   clear_capture_only_spec_state())
-                        logger.info(
-                            "Cleared capture-only sampling override from "
-                            f"{cleared} cached CUDA graph spec metadata "
-                            "object(s).")
 
         if self.cuda_graph_lora_manager is None:
             lora_graph_cases = [False]
@@ -2479,11 +2462,11 @@ def _run_capture_pass(force_non_greedy: bool, label: str,
                                   force_lora_graph=use_lora_graph)
         # Set the value back to the original value after cuda graph warmups are complete
         self.enable_spec_decode = self.is_spec_decode
-        # The advanced-sampling capture pass above leaves is_all_greedy_sample
-        # set to False on spec_metadata. Reset it to the default so the first
-        # real iteration's graph-key selection is not seeded with this
-        # capture-only value. (update_is_all_greedy_sample refreshes it every
-        # iteration; this is a defensive guard.)
+        # update_is_all_greedy_sample inside each forward call during the
+        # non-greedy capture pass leaves is_all_greedy_sample=False on
+        # spec_metadata. Reset it so the first real iteration starts clean;
+        # update_is_all_greedy_sample will refresh it on every iteration anyway.
+        # This is a defensive guard.
         if self.spec_metadata is not None:
             self.spec_metadata.is_all_greedy_sample = True
 
@@ -2864,8 +2847,11 @@ def _create_cuda_graph_warmup_request(
         max_seq_len: int = None,
         mixed_context_encoder_output_lens: Optional[Sequence[int]] = None,
         mixed_context_query_len: int = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM,
+        force_non_greedy: bool = False,
     ) -> Optional[ScheduledRequests]:
         """Creates a dummy ScheduledRequests tailored for CUDA graph capture."""
+        capture_sampling_params = (NON_GREEDY_CAPTURE_SAMPLING_PARAMS
+                                   if force_non_greedy else None)
         kv_cache_manager = resource_manager.get_resource_manager(
             self.kv_cache_manager_key)
         spec_resource_manager = resource_manager.get_resource_manager(
@@ -2906,7 +2892,8 @@ def _create_cuda_graph_warmup_request(
                 use_mrope=self.use_mrope,
                 max_beam_width=self.max_beam_width,
                 encoder_output_lens=list(mixed_context_encoder_output_lens),
-                draft_kv_cache_manager=draft_kv_cache_manager)
+                draft_kv_cache_manager=draft_kv_cache_manager,
+                capture_sampling_params=capture_sampling_params)
             if context_requests is None:
                 return None
 
@@ -2925,7 +2912,8 @@ def _create_cuda_graph_warmup_request(
                     max_beam_width=self.max_beam_width,
                     encoder_output_lens=[max_encoder_output_len] *
                     len(generation_request_ids),
-                    draft_kv_cache_manager=draft_kv_cache_manager)
+                    draft_kv_cache_manager=draft_kv_cache_manager,
+                    capture_sampling_params=capture_sampling_params)
                 if generation_requests is None:
                     for request in context_requests:
                         kv_cache_manager.free_resources(request)
@@ -2947,7 +2935,8 @@ def _create_cuda_graph_warmup_request(
                 use_mrope=self.use_mrope,
                 max_beam_width=self.max_beam_width,
                 encoder_output_lens=encoder_output_lens,
-                draft_kv_cache_manager=draft_kv_cache_manager)
+                draft_kv_cache_manager=draft_kv_cache_manager,
+                capture_sampling_params=capture_sampling_params)
             if requests is None:
                 return None
 
@@ -3012,7 +3001,8 @@ def free_warmup_requests() -> None:
             max_beam_width=self.max_beam_width,
             encoder_output_lens=[max_encoder_output_len]
             if is_enc_dec else None,
-            draft_kv_cache_manager=draft_kv_cache_manager)
+            draft_kv_cache_manager=draft_kv_cache_manager,
+            capture_sampling_params=capture_sampling_params)
 
         if max_seq_len_request is None:
             free_warmup_requests()
@@ -6353,16 +6343,6 @@ def previous_seq_slots_device():
                 num_accepted_draft_tokens)]
             if isinstance(spec_metadata, Eagle3SpecMetadata):
                 spec_metadata.request_accepted_path = request_accepted_path
-            # The capture-only sampling override must never be live outside CUDA
-            # graph warmup: it replaces every request's sampling params with
-            # synthetic capture values. It leaked here once already (inherited by
-            # the cached graph metadata shallow copies), so assert rather than
-            # trust the teardown.
-            assert self.is_warmup or not getattr(
-                spec_metadata, '_force_non_greedy_for_capture', False
-            ), ("capture-only sampling override (_force_non_greedy_for_capture) "
-                "is set outside CUDA graph warmup; serving requests would be "
-                "silently decoded with the synthetic capture sampling params")
             # No-op for non 1-model
             spec_metadata.populate_sampling_params_for_one_model(
                 scheduled_requests.all_requests())

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -1141,6 +1141,7 @@ def add_dummy_requests(
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
         draft_kv_cache_manager: Optional[BaseResourceManager] = None,
+        capture_sampling_params: Optional[SamplingParams] = None,
     ):
         _kv_draft = kv_reserve_draft_tokens if kv_reserve_draft_tokens is not None else max_num_draft_tokens
         available_blocks = self.get_num_free_blocks()
@@ -1156,9 +1157,17 @@ def add_dummy_requests(
         draft_batch_llm_requests = []
         for i, req_id in enumerate(request_ids):
             # exact choice of n can be ignored for dummy requests
-            sampling_params = SamplingParams(n=beam_width,
-                                             best_of=beam_width,
-                                             use_beam_search=beam_width > 1)
+            sampling_params = SamplingParams(
+                n=beam_width,
+                best_of=beam_width,
+                use_beam_search=beam_width > 1,
+                temperature=capture_sampling_params.temperature
+                if capture_sampling_params is not None else None,
+                top_k=capture_sampling_params.top_k
+                if capture_sampling_params is not None else None,
+                top_p=capture_sampling_params.top_p
+                if capture_sampling_params is not None else None,
+            )
             # Here 1+max_num_draft_tokens is used to extend the prompt length to
             # a non-zero number to skip illegal memory access issue in MLA kernel
             # during warmup.

tensorrt_llm/_torch/speculative/interface.py

@@ -1035,7 +1035,7 @@ def _scan_one_model_sampling(
         """Single source of truth for one-engine sampling-param detection.
 
         Scans the batch's sampling configs and sets skip_*/is_all_greedy_sample
-        (honoring the warmup capture override). Returns
+        (honoring the group-synchronized value, see below). Returns
         ``(per_request_normalized, per_request_slot_ids)`` for buffer
         population. Does NOT allocate or fill GPU buffers, so it is safe to call
         before the CUDA graph key is built.
@@ -1139,22 +1139,10 @@ def _normalize_request_sampling_params(
         # unset), so this cannot be derived from which filters are in use.
         self.is_all_greedy_sample = not has_non_greedy_requests
 
-        # Warmup-time override: force the advanced-sampling path so the CUDA
-        # graph for the (is_all_greedy_sample=False) key gets captured. Dummy
-        # warmup requests carry no sampling params, so substitute synthetic
-        # non-greedy scalars to populate the GPU buffers.
-        if getattr(self, '_force_non_greedy_for_capture', False):
-            self.is_all_greedy_sample = False
-            per_request_normalized = [
-                (0.7, 50, 0.9, num_tokens)
-                for (_, _, _, num_tokens) in per_request_normalized
-            ]
-
-        # Apply the group-synchronized override last (semantics: see the
-        # ``group_all_greedy_sample`` field comment). Local contract: the
-        # synced value already incorporates any capture override, and rescans
-        # (e.g. populate after the graph key) must converge to it rather than
-        # resurrect the local value.
+        # Apply the group-synchronized value last (semantics: see the
+        # ``group_all_greedy_sample`` field comment). Local contract: rescans
+        # (e.g. populate after the graph key) must converge to the synced
+        # value rather than resurrect the local value.
         if self.group_all_greedy_sample is not None:
             self.is_all_greedy_sample = self.group_all_greedy_sample
         return per_request_normalized, per_request_slot_ids

tests/integration/defs/.test_durations

@@ -1294,7 +1294,6 @@
    "unittest/_torch/sampler/test_torch_sampler.py": 251.46527294117647,
    "unittest/_torch/sampler/test_trtllm_sampler.py": 78.89402733485194,
    "unittest/_torch/speculative/hw_agnostic": 192.419382629108,
-   "unittest/_torch/speculative/test_capture_override_leak.py": 18.178845569620254,
    "unittest/_torch/speculative/test_dspark_cute_dsl_attention.py": 27.91303125,
    "unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py": 20.8966,
    "unittest/_torch/speculative/test_eagle3.py": 556.1181581081081,