[None][chore] Drop the skip_* sampling flags from one-model spec metadata - #17554
Conversation
…data SpecMetadata carried skip_temperature / skip_top_k / skip_top_p, recomputed each step from whether any request in the batch used that filter, alongside AdvancedSamplingMode which expresses the same "disable this filter" intent as deploy-time configuration. The two were introduced independently and never reconciled. skip_temperature had no reader left: it was a parameter of the sampling op until the FlashInfer sampler refactors removed it from the signature, leaving the assignment behind. skip_top_k / skip_top_p had one reader, in the rejection path, where they selected between passing a tensor and passing None. That choice is made in Python, so CUDA graph capture bakes it in, and unlike is_all_greedy_sample these flags are not part of the graph key -- a replayed graph keeps whichever variant capture happened to see, so the per-batch adaptivity they suggest does not hold once graphs are enabled. AdvancedSamplingMode already covers all four sampling sites and, being fixed for the deployment, has no such gap. Requests that set top_k / top_p now always get the filter applied under the default FULL mode; a deployment that wants the kernels skipped selects NO_TOPK / NO_TOPP / NO_TOPK_NO_TOPP explicitly. The capture-time override no longer has to force these flags: FULL keeps both filters live, so warmup captures the complete variant on its own. is_all_greedy_sample is unaffected -- it is derived from per-request greediness, not from which filters are in use, and stays in the graph key. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
WalkthroughThe speculative sampling metadata no longer contains batch-level temperature, top-k, or top-p skip flags. Request normalization retains normalized sampling values and greediness. Rejection sampling resolves filters directly from sampling mode and per-token tensors. Related capture assertions were removed. ChangesSampling metadata flow
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/speculative/interface.py (1)
824-828: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCorrect the return annotation after removing the skip flags.
_normalize_request_sampling_paramsreturns four values on Lines 820-825, but its annotation on Line 796 still declares a seven-element tuple. Change it totuple[float, int, float, bool]so type checkers and callers see the correct contract.Proposed annotation fix
- ) -> tuple[float, int, float, bool, bool, bool, bool]: + ) -> tuple[float, int, float, bool]:As per coding guidelines, functions must be annotated with precise types.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/speculative/interface.py` around lines 824 - 828, Update the return annotation of _normalize_request_sampling_params to tuple[float, int, float, bool], matching its four returned values and the contract consumed by its callers.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 1723-1729: Update the direct one-model rejection-sampling flow
around resolve_advanced_sampling_filters and compute_probs_from_logits so
skips_top_p disables filtering without passing None or selecting a different
probability implementation; preserve a graph-safe top_p=1.0 tensor or reuse the
existing top-p kernel. Add token-equivalence coverage for NO_TOPP and
NO_TOPK_NO_TOPP.
---
Outside diff comments:
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 824-828: Update the return annotation of
_normalize_request_sampling_params to tuple[float, int, float, bool], matching
its four returned values and the contract consumed by its callers.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 542aa4ec-dfd6-4a86-a607-d4a0d37eafa1
📒 Files selected for processing (2)
tensorrt_llm/_torch/speculative/interface.pytests/unittest/_torch/speculative/test_capture_override_leak.py
💤 Files with no reviewable changes (1)
- tests/unittest/_torch/speculative/test_capture_override_leak.py
|
/bot run |
|
PR_Github #65571 [ run ] triggered by Bot. Commit: |
|
PR_Github #65571 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65742 [ run ] triggered by Bot. Commit: |
|
PR_Github #65742 [ run ] completed with state |
…data (NVIDIA#17554) Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
5 participants
Description
SpecMetadatacarriedskip_temperature/skip_top_k/skip_top_p, recomputed each step from whether any request in the batch used that filter, alongsideAdvancedSamplingMode, which expresses the same "disable this filter" intent as deploy-time configuration. The two were introduced independently and never reconciled. This PR removes theskip_*flags.skip_temperaturehad no reader left. It was a parameter of the sampling op until the FlashInfer sampler refactors ([TRTLLM-13212][refactor] Centralize sampling logic, split backends into isolated modules #15542, [TRTLLM-14198][refactor] Make FlashInfer a hard dependency for the Torch sampler #16160) dropped it from the signature, and the assignment in_scan_one_model_samplingwas left behind.skip_top_k/skip_top_phad one reader, in the rejection path, where they selected between passing a tensor and passingNone. That choice is made in Python, so CUDA graph capture bakes it in, and unlikeis_all_greedy_samplethese flags are not part of the graph key: a replayed graph keeps whichever variant capture happened to see, so the per-batch adaptivity the flags suggest does not survive once graphs are enabled. The capture-time override (_force_non_greedy_for_capture) having to force them toFalseis the same observation from the other side — if they were live at replay, the override would be unnecessary.AdvancedSamplingModecovers the same four sampling sites and, being fixed for the deployment, has no such gap.Under the default
FULLmode a request that setstop_k/top_pnow always gets that filter applied. Previously a batch where no request used a filter would skip its kernel, subject to the capture caveat above; a deployment that wants those kernels skipped selectsNO_TOPK/NO_TOPP/NO_TOPK_NO_TOPPexplicitly. On the eager path this costs two extra kernel launches (~7 us) for batches that use neither filter. Under CUDA graphs there is no change, since the captured graph already contained whatever the capture-time flags selected.is_all_greedy_sampleis unaffected: it is derived from per-request greediness, not from which filters are in use, and stays in the graph key.Not addressed here:
AdvancedSamplingModeis applied at four sampling sites ininterface.py, but three further sampling calls bypass it —eagle3_dynamic_tree.py:775and:883, anddynamic_tree_ops.py:271— so a deployment using dynamic tree does not get the filter kernels skipped there.skip_*never covered those sites either, so this PR neither improves nor regresses them. Worth a follow-up.Test Coverage
Validated on H200 NVL (SM 9.0).
test_capture_override_leak.pypasses (4 cases). It covers the capture-time override path these flags participated in; the threeskip_*assertions are dropped and theis_all_greedy_sampleassertion, the load-bearing one, stays.test_llama_eagle3_rejection_sampling_modespasses (4 cases), coveringuse_cuda_graphxuse_dynamic_treewithtemperature=1.0/top_p=0.9/top_k=50so the non-greedy path is exercised.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.