[None][feat] Add cuda_graph_config.sampling_variants to choose the captured sampling variants

docs/source/features/speculative-decoding.md

@@ -352,6 +352,36 @@ The field name `speculative_model_dir` can also be used as an alias for `specula
       speculative_model_dir: /path/to/draft/model
 ```
 
+### CUDA graph sampling variants
+
+One-engine speculative decoding captures its decode CUDA graphs once per sampling variant: `greedy` for
+batches whose requests all decode greedily, `advanced` for batches with a sampled request (one whose
+temperature, top-k, top-p or min-p does not resolve to greedy decoding; a seed or a penalty alone does not),
+and `temperature_only` for sampled batches whose sampled requests set only a temperature. `temperature_only`
+needs a model whose MTP-Eagle one-model decode samples vocabulary-sharded logits (the spec-decode fast paths;
+today Qwen3.5 checkpoints on SM 10.0/10.3 GPUs), tensor parallelism, and no attention DP, relaxed acceptance,
+rejection sampling or penalties; on other models the first CUDA graph capture rejects it.
+`cuda_graph_config.sampling_variants` (prototype) captures only the listed variants, which shortens startup
+and saves graph memory:
+
+```yaml
+cuda_graph_config:
+  enable_padding: true
+  sampling_variants: [advanced]
+```
+
+- Without `temperature_only`, temperature-only batches replay the `advanced` graphs.
+- Without `greedy`, greedy batches replay the sampled graphs (`temperature_only` if listed, else `advanced`)
+  and still decode greedily, at a cost: they add sampling-parameter and RNG-state uploads, take the sampled
+  path in eager steps too, and on the `advanced` graphs exact ties can break differently than on the `greedy`
+  graphs. They also lose the argmax fast paths: with rejection sampling, argmax drafting and its LM-head-TP
+  path; and the vocabulary-sharded target argmax where the model has one. A model with the spec-decode fast
+  paths keeps that argmax by listing `temperature_only`, unless it uses rejection sampling; DFlash with tensor
+  parallelism keeps its shard selection when it runs without rejection sampling and with
+  `advanced_sampling_mode: no_topk_no_topp`.
+- Without `advanced`, a sampled batch that `temperature_only` does not serve runs eagerly, and the first one
+  logs a warning. Under attention DP, every rank then runs that step eagerly.
+
 ### DSpark attention projection precision
 
 `TRTLLM_DSPARK_LINEAR_BACKEND=bf16|fp8` selects the attention projection

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -153,6 +153,14 @@ class KeyType(NamedTuple):
     ragged_verify_bucket: Optional[int] = None
 
 
+def sampling_variant(spec_metadata: Optional[SpecMetadata]) -> str:
+    """The one-engine sampling variant of a batch, as cuda_graph_config.sampling_variants names it."""
+    if getattr(spec_metadata, "is_all_greedy_sample", True):
+        return "greedy"
+    return ("temperature_only" if getattr(spec_metadata, "temperature_only",
+                                          False) else "advanced")
+
+
 def _save_spec_decode_capture_state(
         attn_metadata: Any, enable_spec_decode: bool) -> Optional[torch.Tensor]:
     if not enable_spec_decode or not hasattr(attn_metadata, 'kv_lens_cuda'):
@@ -232,6 +240,8 @@ class CUDAGraphRunnerConfig:
     private_metadata_buffers: bool = False
     enable_encoder_decoder_mixed_cuda_graph: bool = False
     enable_in_graph_sampling: bool = False
+    # One-engine sampling variants that have graphs (cuda_graph_config.sampling_variants); None for all.
+    sampling_variants: Optional[frozenset[str]] = None
     static_input_ids: Optional[torch.Tensor] = None
     """The decoder runner's own int32 input-id buffer. When given (with
     ``static_position_ids``), the graphs' static inputs are views of the
@@ -511,11 +521,11 @@ def get_graph_key(
         # both variants and dispatch at replay based on actual batch state.
         # Default to True (greedy fast-path) when the metadata doesn't carry
         # this field (non-one-engine paths or non-spec batches).
-        is_all_greedy_sample = bool(
-            getattr(spec_metadata, "is_all_greedy_sample", True))
+        # maybe_get_cuda_graph checks the same variant against the captured ones.
+        variant = sampling_variant(spec_metadata)
+        is_all_greedy_sample = variant == "greedy"
         # The temperature-only variant of the advanced-sampling graphs.
-        temperature_only = (not is_all_greedy_sample and bool(
-            getattr(spec_metadata, "temperature_only", False)))
+        temperature_only = variant == "temperature_only"
 
         # Sampling tier this graph must contain. Only meaningful when the fast
         # sampler is enabled; otherwise it stays FULL so every batch shares one
@@ -720,7 +730,12 @@ def maybe_get_cuda_graph(
             return None, None, None
 
         is_mixed_encoder_decoder = self._is_mixed_encoder_decoder_batch(batch)
-        can_run_cuda_graph = self._can_run_cuda_graph_batch(batch)
+        # A batch of an uncaptured sampling variant runs eagerly; under attention
+        # DP the agreement below then takes every rank's step eager with it.
+        can_run_cuda_graph = (
+            self._can_run_cuda_graph_batch(batch) and
+            (spec_metadata is None or self.config.sampling_variants is None or
+             sampling_variant(spec_metadata) in self.config.sampling_variants))
         batch_size = batch.batch_size
         shapes_agree, forced_sample_type = self._agree_graph_shape(
             batch, can_run_cuda_graph, promoted_context_request_ids)

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

@@ -25,7 +25,11 @@
 )
 from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs
 from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner
-from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner, KeyType
+from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import (
+    CUDAGraphRunner,
+    KeyType,
+    sampling_variant,
+)
 from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length
 from tensorrt_llm._torch.pyexecutor.resource_manager import (
@@ -56,7 +60,12 @@
 from .config import DecoderRunnerConfig
 from .dummy_requests import DummyRequestFactory
 from .inputs import DecoderInputPreparer
-from .metadata import DecoderMetadata, resolve_mamba_metadata_cls, supports_advanced_draft_sampling
+from .metadata import (
+    DecoderMetadata,
+    resolve_mamba_metadata_cls,
+    supports_advanced_draft_sampling,
+    supports_temperature_only_graphs,
+)
 from .warmup import WarmupAgreement
 
 # Context requests per warmup batch whose breakable prefill graph breaks get
@@ -230,6 +239,8 @@ def __init__(
         # Promoted context rows carry no per-sequence extra model inputs.
         self._context_promotion_supported = context_promotion_supported
         self._force_lora_graph_for_capture: bool | None = None
+        # The one-engine sampling variants captured when cuda_graph_config.sampling_variants lists them.
+        self._sampling_variants: frozenset[str] | None = None
 
     def forward_discarded_warmup(
         self,
@@ -923,10 +934,12 @@ def _capture_variant(
             #
             # Dummy warmup requests carry no sampling params, so the greedy pass
             # needs no override while the advanced one has to force the flag.
-            _capture_variant("greedy")
-            if supports_advanced_draft_sampling(self._config.spec_config):
+            variants = self._sampling_variants_to_capture(temperature_only_graphs)
+            if "greedy" in variants:
+                _capture_variant("greedy")
+            if "advanced" in variants:
                 _capture_variant("advanced sampling", force_non_greedy=True)
-            if temperature_only_graphs:
+            if "temperature_only" in variants:
                 # Temperature-only batches replay the vocabulary-sharded Gumbel-max
                 # target sampling (the advanced pass records the stock sampler: its
                 # params enable top-k / top-p / min-p).
@@ -957,6 +970,43 @@ def _capture_variant(
             self._metadata.spec_metadata.is_all_greedy_sample = True
             self._metadata.spec_metadata.temperature_only = False
 
+    def _sampling_variants_to_capture(self, temperature_only_graphs: bool) -> frozenset[str]:
+        """One-engine sampling variants to capture; cuda_graph_config.sampling_variants narrows them."""
+        listed = getattr(self._config.cuda_graph_config, "sampling_variants", None)
+        if listed is not None:
+            # Judged from the config: the speculative metadata may not exist before the first pass.
+            # Batches never select temperature_only with rejection sampling or penalties.
+            spec_config = self._config.spec_config
+            temperature_only_graphs = supports_temperature_only_graphs(
+                self._config, self.mapping
+            ) and not (spec_config.use_rejection_sampling or spec_config.enable_penalty)
+        supported = {"greedy"}
+        if supports_advanced_draft_sampling(self._config.spec_config):
+            supported.add("advanced")
+        if temperature_only_graphs:
+            supported.add("temperature_only")
+        if listed is None:
+            return frozenset(supported)
+        if not supported.issuperset(listed):
+            raise ValueError(
+                f"cuda_graph_config.sampling_variants lists {sorted(set(listed) - supported)}, "
+                f"but this engine can capture only {sorted(supported)}"
+            )
+        self._sampling_variants = frozenset(listed)
+        return self._sampling_variants
+
+    def warn_uncaptured_sampling_variant(self, spec_metadata: SpecMetadata) -> None:
+        """Log once per variant that a decode batch runs eagerly because its variant is not listed."""
+        if self._sampling_variants is None:
+            return
+        variant = sampling_variant(spec_metadata)
+        if variant not in self._sampling_variants:
+            logger.warning_once(
+                f"A decode batch with {variant} sampling runs eagerly: "
+                f"cuda_graph_config.sampling_variants captured {sorted(self._sampling_variants)}",
+                key=f"uncaptured_sampling_variant_{variant}",
+            )
+
     def _capture_break_graphs(
         self,
         resource_manager: ResourceManager,

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

@@ -316,14 +316,31 @@ def get_spec_managers(
     return spec_resource_manager, getattr(spec_resource_manager, "spec_tree_manager", None)
 
 
+def supports_temperature_only_graphs(config: DecoderRunnerConfig, mapping: Mapping | None) -> bool:
+    """Whether the decode CUDA graphs can capture the temperature-only sampling variant."""
+    return (
+        config.cuda_graph_config is not None
+        and config.spec_decode_fast_paths
+        and supports_sharded_target_sampling(config.spec_config, mapping)
+    )
+
+
+def captures_temperature_only_graphs(config: DecoderRunnerConfig, mapping: Mapping | None) -> bool:
+    """Whether temperature-only batches get graphs of their own rather than the advanced ones."""
+    listed = getattr(config.cuda_graph_config, "sampling_variants", None)
+    return supports_temperature_only_graphs(config, mapping) and (
+        listed is None or "temperature_only" in listed
+    )
+
+
 def create_spec_metadata(
     config: DecoderRunnerConfig,
     pretrained_config: object,
     spec_resource_manager: BaseResourceManager | None,
     mapping: Mapping | None = None,
 ) -> SpecMetadata | None:
     """Create speculative metadata sized by the runner capacities."""
-    return get_spec_metadata(
+    spec_metadata = get_spec_metadata(
         config.spec_config,
         pretrained_config,
         config.max_batch_size,
@@ -332,10 +349,12 @@ def create_spec_metadata(
         max_seq_len=config.max_seq_len,
         num_seq_slots=config.num_seq_slots,
         # The decode CUDA graphs then capture the temperature-only variant.
-        temperature_only_graphs=config.cuda_graph_config is not None
-        and config.spec_decode_fast_paths
-        and supports_sharded_target_sampling(config.spec_config, mapping),
+        temperature_only_graphs=captures_temperature_only_graphs(config, mapping),
     )
+    if spec_metadata is not None:
+        listed = getattr(config.cuda_graph_config, "sampling_variants", None)
+        spec_metadata.greedy_graphs = listed is None or "greedy" in listed
+    return spec_metadata
 
 
 def update_spec_metadata(

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

@@ -490,6 +490,7 @@ def _initialize_cuda_graph_runner(self) -> CUDAGraphRunner:
         return cuda_graph_runner
 
     def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig:
+        sampling_variants = getattr(self._config.cuda_graph_config, "sampling_variants", None)
         return CUDAGraphRunnerConfig(
             use_cuda_graph=self._config.cuda_graph_config is not None,
             cuda_graph_padding_enabled=self._config.cuda_graph_padding_enabled,
@@ -510,6 +511,7 @@ def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig:
             short_long_seq_len_threshold=self._short_long_seq_len_threshold(),
             private_metadata_buffers=self.runtime_hooks.skips_unchanged_metadata_uploads,
             enable_in_graph_sampling=self._config.enable_in_graph_sampling,
+            sampling_variants=None if sampling_variants is None else frozenset(sampling_variants),
             # Speculative steps rewrite every input token and position, so
             # captured graphs can read the persistent buffers directly.
             static_input_ids=(
@@ -700,6 +702,13 @@ def _forward_decoder(
             can_run_graph = key is not None and not (
                 ced_enabled and scheduled_requests.num_context_requests > 0
             )
+            if (
+                key is None
+                and not is_dummy
+                and spec_metadata is not None
+                and scheduled_requests.num_context_requests == 0
+            ):
+                self._graphs.warn_uncaptured_sampling_variant(spec_metadata)
             if can_run_graph:
                 attn_metadata = selection.attn_metadata
                 spec_metadata = selection.spec_metadata

tensorrt_llm/_torch/speculative/interface.py

@@ -641,6 +641,8 @@ def is_ragged_verify(self) -> bool:
     # for the current batch (see _temperature_only).
     temperature_only_graphs: bool = False
     temperature_only: bool = False
+    # Whether all-greedy batches have CUDA graphs of their own; without them they take the sampled path.
+    greedy_graphs: bool = True
     # Skip the Philox seed and offset uploads on all-greedy batches (deploy-time;
     # from the model's RuntimeHooks). The argmax path never reads them, and the
     # next sampled batch rewrites every row it reads.
@@ -1434,7 +1436,8 @@ def _normalize_request_sampling_params(
         # All-greedy iff EVERY request is greedy -- note a non-greedy request
         # may still enable no filter (e.g. temperature=1.0 with top_k/top_p
         # unset), so this cannot be derived from which filters are in use.
-        self.is_all_greedy_sample = not has_non_greedy_requests
+        # Without greedy graphs, greedy rows take the sampled path and still decode greedily.
+        self.is_all_greedy_sample = not has_non_greedy_requests and self.greedy_graphs
 
         # Apply the group-synchronized value last (semantics: see the
         # ``group_all_greedy_sample`` field comment). Local contract: rescans