diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py index 454164115ff9..a30a0774f5a9 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py +++ b/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: diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 32ea096b621f..382119e108e4 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 43f7d9773876..c15c64a7e235 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 409352b25af1..9ff887cc46d9 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index a4a3129a064d..bc9ebb526f79 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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,17 +2380,8 @@ 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: @@ -2391,7 +2389,11 @@ def _run_capture_pass(force_non_greedy: bool, label: str, 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()) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 81ba164f9693..4c4c89a702bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/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. diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 79498baab45f..313a081b87de 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/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 diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 416f7f65463c..4c495a4815c7 100644 --- a/tests/integration/defs/.test_durations +++ b/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, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ac609a96d45c..66de85bc40c7 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -55,8 +55,7 @@ # Keep helper definitions below imports so new imports do not need E402 # suppressions in this legacy test file. def patch_mpi_pool_session_for_env(mocker, env_vars: dict): - """ - Patch MpiPoolSession._start_mpi_pool to propagate environment variables to MPI child processes. + """Patch MpiPoolSession._start_mpi_pool to propagate environment variables to MPI child processes. Uses MPIPoolExecutor's built-in `env` parameter instead of `initializer` to avoid segfault issues during process cleanup (UCX memory cache conflicts with PyTorch @@ -116,6 +115,51 @@ def _compute_acceptance_length(llm) -> float: return (accepted + reqs) / reqs +def _assert_non_greedy_cuda_graph_matches_eager(build_llm_kwargs, + graph_cuda_graph_config=None): + """TRTLLM-14874 regression: non-greedy output must match with and without CUDA graphs. + + Capture-only sampling state must not leak into cached graph metadata. + `build_llm_kwargs(cuda_graph_config)` returns the kwargs for constructing + the `LLM` under test for either the graph run (`cuda_graph_config` set) or + the eager run (`cuda_graph_config=None`). + """ + raw_prompts = [ + "Write a short story about a robot who discovers it can dream.", + "Explain quantum entanglement as if you were a pirate.", + "Compose a haiku about the last byte of memory before a crash.", + ] + # Not the synthetic capture params (0.7/50/0.9). max_tokens kept + # short to avoid unrelated graph-vs-eager fp drift over a long run. + sampling_params = SamplingParams(temperature=0.3, + top_k=10, + top_p=0.8, + seed=42, + max_tokens=32) + + def run(cuda_graph_config): + with LLM(**build_llm_kwargs(cuda_graph_config)) as llm: + prompts = [ + llm.tokenizer.apply_chat_template( + [{ + "role": "user", + "content": p + }], + tokenize=False, + add_generation_prompt=True, + ) for p in raw_prompts + ] + return [ + output.outputs[0].token_ids + for output in llm.generate(prompts, + sampling_params=sampling_params) + ] + + graph_token_ids = run(graph_cuda_graph_config or CudaGraphConfig()) + eager_token_ids = run(None) + assert graph_token_ids == eager_token_ids + + def _run_multinode_accuracy(model_path, model_name, *, @@ -847,6 +891,44 @@ def test_suffix_automaton_dynamic_draft_len(self): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @skip_pre_hopper + @parametrize_with_ids("spec_dec_algo", ["eagle3", "suffix_automaton"]) + # Builds two LLM() instances back-to-back; the MPI-pool-reuse test layer + # would otherwise hand the second one the first's just-used worker pool. + @pytest.mark.private_mpi_session + def test_one_engine_non_greedy_cuda_graph_matches_eager( + self, spec_dec_algo): + """TRTLLM-14874 regression: capture-only sampling state must not leak into cached graph metadata, non-greedy output must match with and without CUDA graphs.""" + max_bs = 4 + target_model_dir = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct" + + if spec_dec_algo == "eagle3": + eagle_model_dir = f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B" + spec_config = Eagle3DecodingConfig( + max_draft_len=4, + speculative_model=eagle_model_dir, + eagle3_one_model=True) + kv_cache_config = KvCacheConfig(enable_block_reuse=True, + free_gpu_memory_fraction=0.8) + else: + spec_config = SADecodingConfig( + max_draft_len=4, + max_matching_ngram_size=-1, # longest match via suffix automaton + ) + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.8) + + def build_llm_kwargs(cuda_graph_config): + return dict(model=target_model_dir, + disable_overlap_scheduler=True, + cuda_graph_config=cuda_graph_config, + kv_cache_config=kv_cache_config, + speculative_config=spec_config, + max_batch_size=max_bs) + + _assert_non_greedy_cuda_graph_matches_eager( + build_llm_kwargs, CudaGraphConfig(max_batch_size=max_bs)) + @skip_pre_blackwell def test_draft_target_dynamic_draft_len(self): draft_len_schedule = {50: 4, 200: 3, 350: 2} @@ -1751,6 +1833,28 @@ def test_bfloat16_2_model_mtp(self): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip_less_device_memory(60000) + # Builds two LLM() instances back-to-back; the MPI-pool-reuse test layer + # would otherwise hand the second one the first's just-used worker pool. + @pytest.mark.private_mpi_session + def test_mtp_non_greedy_cuda_graph_matches_eager(self): + """TRTLLM-14874 regression on the MTP capture path (see test_one_engine_non_greedy_cuda_graph_matches_eager in TestLlama3_1_8BInstruct for the EAGLE3/SA equivalents).""" + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.3) + mtp_config = MTPDecodingConfig(max_draft_len=3, + mtp_eagle_one_model=False, + speculative_model=self.MODEL_PATH) + + def build_llm_kwargs(cuda_graph_config): + return dict(model=self.MODEL_PATH, + kv_cache_config=kv_cache_config, + enable_chunked_prefill=False, + max_num_tokens=8192, + disable_overlap_scheduler=True, + cuda_graph_config=cuda_graph_config, + speculative_config=mtp_config) + + _assert_non_greedy_cuda_graph_matches_eager(build_llm_kwargs) + @pytest.mark.skip_less_device_memory(60000) def test_bfloat16_mtp_sa(self): """Accuracy test for MTP + Suffix Automaton (MTP+SA) speculative decoding.""" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 6f0b5c320cc9..0b6ac4f22168 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -317,6 +317,7 @@ l0_b200: - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_t2i_4step_example TIMEOUT (30) - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_i2v_4step_example TIMEOUT (45) - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_edge_i2v_example TIMEOUT (30) + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_mtp_non_greedy_cuda_graph_matches_eager - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 52fc2b1168c3..17021a5d1415 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -58,9 +58,9 @@ l0_h100: - unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor" - unittest/_torch/speculative/test_eagle3.py - unittest/_torch/speculative/test_rejection_buffers_guard.py - - unittest/_torch/speculative/test_capture_override_leak.py - unittest/_torch/speculative/test_sa_hybrid_state_promotion.py - unittest/_torch/speculative/hw_agnostic + - unittest/_torch/speculative/test_capture_sampling_params.py - unittest/_torch/thop/parallel - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/thop/parallel_hw_agnostic @@ -144,6 +144,8 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=True-attn_backend=TRTLLM-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=True-overlap_scheduler=True] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-eagle3_one_model=True-overlap_scheduler=True] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_one_engine_non_greedy_cuda_graph_matches_eager[spec_dec_algo=eagle3] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_one_engine_non_greedy_cuda_graph_matches_eager[spec_dec_algo=suffix_automaton] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_pard[overlap_scheduler=True] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_pard[overlap_scheduler=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dflash diff --git a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py index 2efe84988c27..da0173067604 100644 --- a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py +++ b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py @@ -18,6 +18,7 @@ RocketVanillaAttentionMetadata, ) from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import get_size_in_bytes from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, RocketSparseAttentionConfig @@ -749,6 +750,23 @@ def test_rocketkv_kt_cache_sized_by_kt_dtype_not_kv_dtype(): kv_cache_manager.shutdown() +def test_rocket_add_dummy_requests_forwards_capture_sampling_params(mocker): + """RocketKVCacheManager overrides add_dummy_requests with an explicit + signature that must stay in sync with KVCacheManager.add_dummy_requests's + capture_sampling_params kwarg, or CUDA graph warmup breaks with a + TypeError for sparse-attention models. + """ + mgr = object.__new__(RocketKVCacheManager) + base_add_dummy_requests = mocker.patch.object( + KVCacheManager, "add_dummy_requests", return_value=[] + ) + sampling_params = SamplingParams(temperature=0.9, top_p=0.95) + + mgr.add_dummy_requests([123], capture_sampling_params=sampling_params) + + assert base_add_dummy_requests.call_args.kwargs["capture_sampling_params"] is sampling_params + + if __name__ == "__main__": # RocketKV e2e tests print("=== Testing RocketKV E2E tests ===") diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 4f9c6f9ea440..e0cc636cbe26 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -79,6 +79,7 @@ _introspection, ) from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager +from tensorrt_llm.sampling_params import SamplingParams skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -1946,6 +1947,31 @@ def test_v2_hybrid_add_dummy_requests_forwards_encoder_output_lens(mocker): assert base_add_dummy_requests.call_args.kwargs["encoder_output_lens"] == [17] +@pytest.mark.parametrize( + "manager_cls,base_cls", + [ + (MambaHybridCacheManagerV2, KVCacheManagerV2), + (CppMambaHybridCacheManager, KVCacheManager), + ], + ids=["v2", "cpp"], +) +def test_add_dummy_requests_forwards_capture_sampling_params(mocker, manager_cls, base_cls): + """Regression test: model_engine.py's CUDA graph warmup path always + passes capture_sampling_params (None on the greedy pass, a real + SamplingParams on non-greedy passes) to whatever concrete + kv_cache_manager is active. An override with an explicit signature + that doesn't accept/forward this kwarg breaks warmup for that + manager type with a TypeError. + """ + mgr = object.__new__(manager_cls) + base_add_dummy_requests = mocker.patch.object(base_cls, "add_dummy_requests", return_value=[]) + sampling_params = SamplingParams(temperature=0.9, top_p=0.95) + + mgr.add_dummy_requests([123], capture_sampling_params=sampling_params) + + assert base_add_dummy_requests.call_args.kwargs["capture_sampling_params"] is sampling_params + + @pytest.mark.parametrize( "manager_cls", [ diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index ddec7eb5a456..1861e3adce4a 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1422,9 +1422,7 @@ def test_promoted_context_precedes_speculative_overlap_generation( max_num_tokens=32, kv_cache_manager=kv_cache_manager) attn_metadata.is_cuda_graph = False - # A bare Mock auto-vivifies every attribute, so the capture-only - # override has to be pinned off or _prepare_tp_inputs reads it as live. - spec_metadata = Mock(_force_non_greedy_for_capture=False) + spec_metadata = Mock() context = _create_request_with_tokens([11, 22, 33, 44], 1) context.context_current_position = 3 diff --git a/tests/unittest/_torch/speculative/test_capture_override_leak.py b/tests/unittest/_torch/speculative/test_capture_override_leak.py deleted file mode 100644 index 497e8c78b638..000000000000 --- a/tests/unittest/_torch/speculative/test_capture_override_leak.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""CPU unit tests for the lifetime of the capture-only sampling override. - -The advanced-sampling CUDA graph capture pass sets -``_force_non_greedy_for_capture=True`` on the live ``SpecMetadata`` so that -parameter-less warmup requests scan as non-greedy and the advanced-sampling -branch is the one recorded into the graph. - -``create_cuda_graph_metadata`` shallow-copies the live metadata, so every graph -captured during that pass caches a copy that inherited the flag, and those -copies are reseated as the live spec_metadata on every later replay. Clearing -the flag on the base object alone therefore leaves it set forever on the copies, -and ``_scan_one_model_sampling`` then rewrites EVERY serving request's sampling -params to the synthetic capture values. - -These tests use a real (base) ``SpecMetadata`` -- its ``__post_init__`` is a -no-op and none of the fields exercised here are tensors -- plus an unbound call -of ``CUDAGraphRunner.clear_capture_only_spec_state`` on a stand-in holding only -``graph_metadata``, mirroring test_group_all_greedy_sync.py. No GPU, no runner -construction, and no model forward is needed. -""" - -import types - -import torch - -from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.speculative.interface import SpecMetadata - -# The synthetic params the capture override substitutes for the request's own, -# and the sentinel that means "top-k disabled" (see _scan_one_model_sampling). -CAPTURE_TEMPERATURE, CAPTURE_TOP_K, CAPTURE_TOP_P = 0.7, 50, 0.9 -DISABLE_TOPK_VAL = torch.iinfo(torch.int32).max - - -def _base_meta(): - """A live (non-graph) SpecMetadata, as the model engine holds it.""" - return SpecMetadata(max_num_requests=8, max_draft_len=1, max_total_draft_tokens=1) - - -def _graph_copy(meta, batch_size=8): - """The shallow copy maybe_get_cuda_graph caches for one captured graph.""" - graph_meta = meta.create_cuda_graph_metadata(batch_size) - assert graph_meta is not meta - return graph_meta - - -def _clear(graph_metadata): - """CUDAGraphRunner.clear_capture_only_spec_state, called unbound.""" - return CUDAGraphRunner.clear_capture_only_spec_state( - types.SimpleNamespace(graph_metadata=graph_metadata) - ) - - -def _request(temperature=None, top_k=None, top_p=None, slot=0): - return types.SimpleNamespace( - sampling_config=types.SimpleNamespace( - temperature=[temperature] if temperature is not None else None, - top_k=[top_k] if top_k is not None else None, - top_p=[top_p] if top_p is not None else None, - ), - state=LlmRequestState.GENERATION_IN_PROGRESS, - py_seq_slot=slot, - ) - - -def _scan(meta, requests): - normalized, _ = SpecMetadata._scan_one_model_sampling(meta, requests) - # Drop the trailing num_tokens; only the sampling params matter here. - return [entry[:3] for entry in normalized] - - -def test_graph_copy_inherits_flag_and_base_teardown_does_not_reach_it(): - # The mechanism the bug rests on: copy.copy carries the flag over, and the - # copies are independent objects, so clearing the base misses them. - meta = _base_meta() - meta._force_non_greedy_for_capture = True - copies = [_graph_copy(meta, bs) for bs in (1, 2, 4)] - assert all(copy._force_non_greedy_for_capture for copy in copies) - - meta._force_non_greedy_for_capture = False - assert all(copy._force_non_greedy_for_capture for copy in copies) - - -def test_clear_capture_only_spec_state_clears_every_cached_copy(): - meta = _base_meta() - meta._force_non_greedy_for_capture = True - advanced = [_graph_copy(meta, bs) for bs in (1, 2, 4)] - # Graphs captured by the greedy pass never had the flag, and non-spec - # graphs cache no spec_metadata at all; both must be left alone. - meta._force_non_greedy_for_capture = False - greedy = _graph_copy(meta, 8) - - graph_metadata = {("greedy", 8): {"spec_metadata": greedy}} - graph_metadata[("no_spec", 1)] = {"spec_metadata": None} - for i, copy in enumerate(advanced): - graph_metadata[("advanced", i)] = {"spec_metadata": copy} - - assert _clear(graph_metadata) == len(advanced) - assert not any(copy._force_non_greedy_for_capture for copy in advanced) - assert greedy._force_non_greedy_for_capture is False - # Idempotent: a second teardown finds nothing left to clear. - assert _clear(graph_metadata) == 0 - - -def test_serving_scan_honors_client_params_after_capture_teardown(): - # End-to-end property of the fix, and the case that fails without it: with - # only the base-object teardown this scan returns (0.7, 50, 0.9). - meta = _base_meta() - meta._force_non_greedy_for_capture = True - graph_meta = _graph_copy(meta) - - meta._force_non_greedy_for_capture = False # base-object teardown - _clear({("advanced", 8): {"spec_metadata": graph_meta}}) # the fix - - # Replay reseats the cached copy as the live spec_metadata, so the serving - # scan runs on it, not on the base object. - assert _scan(graph_meta, [_request(temperature=1.0, top_p=1.0)]) == [ - (1.0, DISABLE_TOPK_VAL, 1.0) - ] - assert graph_meta.is_all_greedy_sample is False - - -def test_override_stays_live_while_the_flag_is_set(): - # Anti-regression for the rejected "clear at copy time" fix: the flag is - # load-bearing *during* capture. Cleared any earlier, the pass-2 populate - # would scan these parameter-less warmup requests as greedy and bake the - # argmax fast path -- with no top-k/top-p kernels -- into the graph keyed - # as the advanced-sampling variant. - meta = _base_meta() - meta._force_non_greedy_for_capture = True - graph_meta = _graph_copy(meta) - - warmup_requests = [_request(slot=None), _request(slot=None)] - assert _scan(graph_meta, warmup_requests) == [ - (CAPTURE_TEMPERATURE, CAPTURE_TOP_K, CAPTURE_TOP_P) - ] * len(warmup_requests) - assert graph_meta.is_all_greedy_sample is False diff --git a/tests/unittest/_torch/speculative/test_capture_sampling_params.py b/tests/unittest/_torch/speculative/test_capture_sampling_params.py new file mode 100644 index 000000000000..d5b39e7f0c2f --- /dev/null +++ b/tests/unittest/_torch/speculative/test_capture_sampling_params.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the refactored non-greedy CUDA graph capture path +(TRTLLM-14874). + +The advanced-sampling CUDA graph capture pass used to force the non-greedy +branch by mutating a ``_force_non_greedy_for_capture`` flag on the live +``SpecMetadata``. ``create_cuda_graph_metadata`` shallow-copied that flag into +every cached graph entry, and those copies were later reseated as the live +``spec_metadata`` on replay -- leaking the flag into serving and silently +rewriting every request's sampling params to the synthetic capture values +(temperature=0.7, top_k=50, top_p=0.9). + +The fix drives capture with *real* non-greedy warmup ``SamplingParams`` +(``KVCacheManager.add_dummy_requests(capture_sampling_params=...)``) instead +of a metadata flag: warmup requests carry genuine sampling params, so +``SpecMetadata._scan_one_model_sampling`` classifies them as non-greedy the +same way it would classify any real client request. There is no capture-only +mutable state left on ``SpecMetadata`` to leak, so these tests assert that +property rather than a specific teardown step. +""" + +import types +import unittest + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.model_engine import NON_GREEDY_CAPTURE_SAMPLING_PARAMS +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm._torch.speculative.interface import SpecMetadata +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping + +# The synthetic params the refactored capture path warms up with. +CAPTURE_TEMPERATURE = NON_GREEDY_CAPTURE_SAMPLING_PARAMS.temperature +CAPTURE_TOP_K = NON_GREEDY_CAPTURE_SAMPLING_PARAMS.top_k +CAPTURE_TOP_P = NON_GREEDY_CAPTURE_SAMPLING_PARAMS.top_p + + +class TestAddDummyRequestsCaptureSamplingParams(unittest.TestCase): + """`KVCacheManager.add_dummy_requests(capture_sampling_params=...)` must + stamp the synthetic non-greedy values onto the dummy requests it builds, + and must leave requests greedy when no capture params are supplied.""" + + def _kv_cache_manager(self): + return KVCacheManager( + kv_cache_config=KvCacheConfig(max_tokens=256, enable_block_reuse=False), + kv_cache_type=tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=2, + num_kv_heads=2, + head_dim=128, + tokens_per_block=8, + max_seq_len=64, + max_batch_size=1, + mapping=Mapping(), + ) + + def test_capture_sampling_params_stamped_onto_dummy_requests(self): + kv_cache_manager = self._kv_cache_manager() + try: + requests = kv_cache_manager.add_dummy_requests( + [0], + token_nums=[8], + capture_sampling_params=NON_GREEDY_CAPTURE_SAMPLING_PARAMS, + ) + self.assertEqual(len(requests), 1) + sampling_config = requests[0].sampling_config + # Sampling params round-trip through a C++ binding as float32, so + # compare with tolerance rather than exact equality. + self.assertAlmostEqual(sampling_config.temperature[0], CAPTURE_TEMPERATURE, places=6) + self.assertEqual(sampling_config.top_k, [CAPTURE_TOP_K]) + self.assertAlmostEqual(sampling_config.top_p[0], CAPTURE_TOP_P, places=6) + finally: + kv_cache_manager.shutdown() + + def test_no_capture_sampling_params_stays_greedy(self): + kv_cache_manager = self._kv_cache_manager() + try: + requests = kv_cache_manager.add_dummy_requests([0], token_nums=[8]) + self.assertEqual(len(requests), 1) + sampling_config = requests[0].sampling_config + self.assertIsNone(sampling_config.temperature) + self.assertIsNone(sampling_config.top_k) + self.assertIsNone(sampling_config.top_p) + finally: + kv_cache_manager.shutdown() + + +def _request(temperature=None, top_k=None, top_p=None, slot=0): + return types.SimpleNamespace( + sampling_config=types.SimpleNamespace( + temperature=[temperature] if temperature is not None else None, + top_k=[top_k] if top_k is not None else None, + top_p=[top_p] if top_p is not None else None, + ), + state=LlmRequestState.GENERATION_IN_PROGRESS, + py_seq_slot=slot, + ) + + +def _fake_meta(): + return types.SimpleNamespace( + runtime_draft_len=1, dummy_slot_row=0, group_all_greedy_sample=None + ) + + +def _scan(meta, requests): + normalized, _ = SpecMetadata._scan_one_model_sampling(meta, requests) + # Drop the trailing num_tokens; only the sampling params matter here. + return [entry[:3] for entry in normalized] + + +class TestScanOneModelSamplingHonorsRealCaptureParams(unittest.TestCase): + """`_scan_one_model_sampling` needs no capture-only override: warmup + requests carrying real non-greedy params classify as non-greedy on their + own, the same way any client request would.""" + + def test_warmup_requests_with_real_capture_params_scan_non_greedy(self): + meta = _fake_meta() + warmup_requests = [ + _request( + temperature=CAPTURE_TEMPERATURE, top_k=CAPTURE_TOP_K, top_p=CAPTURE_TOP_P, slot=None + ), + _request( + temperature=CAPTURE_TEMPERATURE, top_k=CAPTURE_TOP_K, top_p=CAPTURE_TOP_P, slot=None + ), + ] + normalized = _scan(meta, warmup_requests) + self.assertEqual( + normalized, [(CAPTURE_TEMPERATURE, CAPTURE_TOP_K, CAPTURE_TOP_P)] * len(warmup_requests) + ) + self.assertFalse(meta.is_all_greedy_sample) + + def test_parameterless_warmup_requests_still_scan_greedy(self): + # No SpecMetadata flag exists to force these non-greedy anymore: a + # warmup pass that (mis-)builds dummy requests without capture params + # must fall back to the ordinary greedy classification, not silently + # capture the advanced-sampling graph. + meta = _fake_meta() + normalized = _scan(meta, [_request(slot=None), _request(slot=None)]) + self.assertTrue(all(temp != CAPTURE_TEMPERATURE for temp, _, _ in normalized)) + self.assertTrue(meta.is_all_greedy_sample) + + def test_capture_time_values_do_not_leak_into_a_later_serving_scan(self): + # Regression for the original bug class: a SpecMetadata object (e.g. + # a graph copy reseated as the live spec_metadata on replay) that was + # previously scanned with capture-time params must not retain them -- + # every scan is a pure function of the requests passed in, so a + # later scan with the client's real params must reflect only those. + meta = _fake_meta() + _scan( + meta, + [ + _request( + temperature=CAPTURE_TEMPERATURE, + top_k=CAPTURE_TOP_K, + top_p=CAPTURE_TOP_P, + slot=None, + ) + ], + ) + self.assertFalse(meta.is_all_greedy_sample) + + serving_normalized = _scan(meta, [_request(temperature=1.0, top_p=1.0, slot=1)]) + temp, top_k, top_p = serving_normalized[0] + self.assertEqual(temp, 1.0) + self.assertNotEqual(top_k, CAPTURE_TOP_K) + self.assertEqual(top_p, 1.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py index c532bfb246f4..da5b5bdc203f 100644 --- a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py +++ b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py @@ -32,13 +32,12 @@ def _fake_request(temperature=None, top_k=None, top_p=None, slot=0): ) -def _fake_meta(group_all_greedy_sample=None, force_capture=False): +def _fake_meta(group_all_greedy_sample=None): return types.SimpleNamespace( runtime_draft_len=2, dummy_slot_row=0, spec_dec_mode=types.SimpleNamespace(use_one_engine=lambda: True), group_all_greedy_sample=group_all_greedy_sample, - _force_non_greedy_for_capture=force_capture, ) @@ -84,15 +83,6 @@ def test_group_override_true_keeps_greedy(): assert meta.is_all_greedy_sample is True -def test_capture_override_composes_with_group_sync(): - # Warmup forces the advanced variant to capture its CUDA graph; the group - # value is derived from capture-forced locals (all False), so the final - # flag stays False regardless of composition order. - meta = _fake_meta(group_all_greedy_sample=False, force_capture=True) - _scan(meta, [_fake_request()]) - assert meta.is_all_greedy_sample is False - - def test_refresh_discards_warmup_group_value_for_real_request(): meta = _fake_meta()