[None][feat] MTP one-model advanced_sampling_mode: skip redundant top-k / top-p filter kernels with additional config enum - #16561
Conversation
0e79ccf to
4ef515a
Compare
advanced_sampling_mode: skip redundant top-k / top-p filter kernels with additional config enum
b96248c to
68a00f4
Compare
68a00f4 to
6c7f810
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds deploy-time advanced sampling modes for MTP one-model speculative decoding, propagates the selected mode through metadata and sampling paths, conditionally skips top-k/top-p filtering, and adds validation tests and documentation. ChangesAdvanced sampling modes
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant MTPDecodingConfig
participant SpecMetadata
participant SpeculativeInterface
participant sampling_utils
participant FlashInfer
MTPDecodingConfig->>SpecMetadata: configure advanced_sampling_mode
SpecMetadata->>SpeculativeInterface: provide mode for speculative sampling
SpeculativeInterface->>sampling_utils: resolve effective top_k and top_p
sampling_utils-->>SpeculativeInterface: return filters or None
SpeculativeInterface->>sampling_utils: sample draft tokens or compute rejection probabilities
sampling_utils->>FlashInfer: apply enabled sampling kernels
FlashInfer-->>sampling_utils: return probabilities or token IDs
sampling_utils-->>SpeculativeInterface: return sampling results
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py (2)
39-53: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise rejection-disabled behavior explicitly.
Lines 41 and 49 rely on the default for
use_rejection_sampling; set it toFalseso the tests directly prove both sides of the contract.Proposed test adjustment
- cfg = MTPDecodingConfig(max_draft_len=1, advanced_sampling_mode="no_topk") + cfg = MTPDecodingConfig( + max_draft_len=1, + advanced_sampling_mode="no_topk", + use_rejection_sampling=False, + ) ... - MTPDecodingConfig(max_draft_len=1, advanced_sampling_mode=mode) + MTPDecodingConfig( + max_draft_len=1, + advanced_sampling_mode=mode, + use_rejection_sampling=False, + )As per path instructions, coverage in
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.pyneeds direct validation of the configured behavior.🤖 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 `@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` around lines 39 - 53, Update both MTPDecodingConfig constructions in test_no_topk_does_not_require_rejection and test_topp_disabling_modes_require_rejection to explicitly pass use_rejection_sampling=False for the rejection-disabled cases, while keeping the existing rejection-enabled configuration unchanged.Source: Path instructions
90-125: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover NO_TOPK rejection of greedy requests.
This only tests temperature
0.7. Add a case at or belowGREEDY_TEMPERATURE_THRESHOLDthat asserts the NO_TOPK request-validation path rejects it, plus a just-above-threshold success case.Based on PR objectives, NO_TOPK disallows greedy requests. As per path instructions, coverage in
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.pyis insufficient for that contract.🤖 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 `@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` around lines 90 - 125, Extend test_no_topk_matches_full_bit_for_bit or add a focused test covering NO_TOPK temperature validation. Use temperatures at or below GREEDY_TEMPERATURE_THRESHOLD to assert the NO_TOPK request is rejected, and just above the threshold to assert it succeeds. Reuse sampling_batch_spec_dec_one_model and the existing test setup symbols, while preserving the current bit-for-bit comparison coverage.Source: Path instructions
tensorrt_llm/llmapi/llm_args.py (1)
2484-2504: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the validator return annotation.
_log_advanced_sampling_modeshould declare-> "MTPDecodingConfig".🤖 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/llmapi/llm_args.py` around lines 2484 - 2504, The _log_advanced_sampling_mode validator is missing its return annotation. Update that method signature to declare -> "MTPDecodingConfig", preserving its existing validation logic and return self behavior.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/pyexecutor/sampler/sampling_utils.py`:
- Around line 890-897: Update the advanced_sampling_mode branch in the direct
sampling path to use compute_probs_from_logits with the selected mode, then
sample via sampling_from_probs_op instead of always calling
top_p_sampling_from_probs_op. Ensure both NO_TOPP and NO_TOPK_NO_TOPP bypass
top_p filtering while preserving the existing temperature, seed, and offset
behavior.
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 830-846: Initialize and update self.has_greedy_requests while
scanning per-request sampling parameters, ensuring it reflects any greedy row in
mixed batches. Apply the existing mixed-batch rejection to every non-FULL
advanced sampling mode, including no_topk and no_topk_no_topp, while preserving
the all-greedy and warmup-dummy exemptions.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 2475-2482: Validate advanced_sampling_mode after the final
spec-decoding mode is determined, accepting non-FULL values only for the
MTP-Eagle one-model branch. Reject unsupported non-FULL modes for vanilla and
two-model MTP instead of silently ignoring them, while preserving FULL as the
default and existing metadata behavior.
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py`:
- Line 30: Add complete type annotations to every function in this test module,
including all test functions and the spy helper: annotate each parameter with
its precise type and use -> None for functions that do not return a value.
Update the functions represented by test_advanced_sampling_mode_enum_values and
the other referenced test/spy definitions without changing their behavior.
---
Nitpick comments:
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 2484-2504: The _log_advanced_sampling_mode validator is missing
its return annotation. Update that method signature to declare ->
"MTPDecodingConfig", preserving its existing validation logic and return self
behavior.
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py`:
- Around line 39-53: Update both MTPDecodingConfig constructions in
test_no_topk_does_not_require_rejection and
test_topp_disabling_modes_require_rejection to explicitly pass
use_rejection_sampling=False for the rejection-disabled cases, while keeping the
existing rejection-enabled configuration unchanged.
- Around line 90-125: Extend test_no_topk_matches_full_bit_for_bit or add a
focused test covering NO_TOPK temperature validation. Use temperatures at or
below GREEDY_TEMPERATURE_THRESHOLD to assert the NO_TOPK request is rejected,
and just above the threshold to assert it succeeds. Reuse
sampling_batch_spec_dec_one_model and the existing test setup symbols, while
preserving the current bit-for-bit comparison coverage.
🪄 Autofix (Beta)
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: 87e95728-f8e1-4129-8afa-df75e9e6ff80
📒 Files selected for processing (5)
tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py
6c7f810 to
c8eb118
Compare
|
Hi @mikeiovine Please take a look at this PR — it adds advanced_sampling_mode to skip the top-k / top-p filter kernels for temperature-only deployments. Any comments on this implementation approach? |
mikeiovine
left a comment
There was a problem hiding this comment.
I think the optimization makes sense but I'm a bit worried about the API. As you point out in the description, we could make it more usable by capturing more graphs, but that's going to make the warmup too long and is not sustainable.
At least the default is FULL. This should be considered an optimization for advanced use cases IMO
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
f2479d0 to
729831c
Compare
|
Since this adds a new field to DecodingBaseConfig, per the repo policy for LLM args / nested-config changes, please run python3 scripts/generate_llm_args_golden_manifest.py and commit the updated tensorrt_llm/usage/llm_args_golden_manifest.json — the manifest currently doesn't include advanced_sampling_mode, which will fail CI. Also note that new config fields require approval from the telemetry/privacy CODEOWNERs. |
|
@NVIDIA/trt-llm-doc-owners Please help to review this PR, thanks~ |
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
729831c to
8265a20
Compare
Address PR NVIDIA#16561 reviewer feedback (zhaoyangwang-nvidia, mikeiovine): - Add sample_from_logits_op (mirrors compute_probs_from_logits_op) and resolve advanced_sampling_mode once on the speculative side via resolve_advanced_sampling_filters; remove the sampling_batch_spec_dec_one_model wrapper layer. Callers pass effective (top_k/top_p) tensors, None when the mode disables that filter, so the ops stay mode-agnostic. - Handle greedy rows natively via the sentinel temperature (temperature-scaled softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches. Remove the mixed-batch ValueError guard and the has_greedy_requests tracking in _scan_one_model_sampling. - Move advanced_sampling_mode + its validator from MTPDecodingConfig to DecodingBaseConfig (available to any speculative config). Add AdvancedSamplingMode.skips_top_k / .skips_top_p as the single source of truth for the filter-skip predicate (replaces duplicated string literals). - Delete the MTP-one-model spec-mode gate: non-FULL modes now construct on any spec path and fall back to FULL behavior where unsupported, rather than raising. - Migrate eagle3_dynamic_tree to the unified sample_from_logits_op. - Document advanced_sampling_mode in docs/source/features/sampling.md. - Update unit tests: enum skip properties, resolve_advanced_sampling_filters, NO_TOPK == FULL bit-for-bit via the unified op, native greedy argmax on mixed batches, and base-config acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
8265a20 to
7c56183
Compare
|
PR_Github #62482 [ run ] completed with state |
e2ff292 to
1107598
Compare
|
/bot run |
|
PR_Github #62483 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62514 [ run ] triggered by Bot. Commit: |
|
PR_Github #62515 [ run ] triggered by Bot. Commit: |
|
PR_Github #62514 [ run ] completed with state |
|
PR_Github #62515 [ run ] completed with state
|
1107598 to
467cfe0
Compare
|
/bot run |
|
PR_Github #62535 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #62535 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62551 [ run ] triggered by Bot. Commit: |
|
PR_Github #62548 [ run ] triggered by Bot. Commit: |
|
PR_Github/16561-467cfe0 #62551 was force-killed by a newer pipeline run. |
|
PR_Github #62548 [ run ] completed with state
|
467cfe0 to
c8dd668
Compare
…els in the one-model speculative sampler Add a deploy-time `advanced_sampling_mode` enum on `DecodingBaseConfig` (default FULL) that lets the one-model speculative sampler skip flashinfer's top-k mask and top-p renorm kernels when the deployment disables those filters. When a skipped filter is already disabled the output matches FULL bit-for-bit, so this is a lossless throughput optimization for fixed-config / temperature-only deployments (measured ~8-28% OTPS/gpu gain on Qwen3.6-35B-A3B-NVFP4, B200, MTP draft_len=3). - Enum end-to-end: `AdvancedSamplingMode.skips_top_k` / `.skips_top_p` are the single source of truth; pydantic parses the config string into the enum, `SpecMetadata` carries it, and `resolve_advanced_sampling_filters` resolves it once on the speculative side so the ops stay mode-agnostic. - Unified sampler: `sample_from_logits_op` (top_k mask -> softmax -> top_p sample / plain sample) replaces the `sampling_batch_spec_dec_one_model` wrapper; a `None` top_k / top_p skips that filter's kernel. - Greedy rows are handled natively via the `DISABLE_TEMP_VAL` sentinel temperature (softmax collapses to a one-hot argmax), so any mode supports mixed greedy + sampling batches with no special-case guard. - All modes work with `use_rejection_sampling` on or off; the two are independent. - Regenerated `tensorrt_llm/usage/llm_args_golden_manifest.json` (adds `speculative_config.advanced_sampling_mode`). - Unit tests in `tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py` and docs in `docs/source/features/sampling.md`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
c8dd668 to
4fc76dc
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62574 [ run ] triggered by Bot. Commit: |
|
PR_Github #62574 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62693 [ run ] triggered by Bot. Commit: |
|
PR_Github #62693 [ run ] completed with state |
Description
For one-model speculative decoding (MTP-Eagle one-model), the advanced sampler runs
flashinfer's top-k mask (
RadixTopKMaskLogits) and top-p renorm (AirTopPRenorm*)kernels on every draft layer and the target verification — even when the
deployment's requests disable those filters (
top_k=0,top_p=1). In that case thefilters are mathematical no-ops (top-k mask at
k=vocab, top-p renorm atp=1), butthe kernels still launch: measured at ~290 µs (top-k mask) + ~1.34 ms (top-p
renorm) per decode step at
draft_len=3(nsys, Qwen3.6-35B-A3B-NVFP4, B200). Forfixed-config / temperature-only deployments this is pure overhead.
Add a deploy-time
advanced_sampling_modeenum onDecodingBaseConfig(availableto any speculative config; default
FULL) that tells the one-model sampler whichfilters the deployment disables, so their flashinfer kernels are skipped:
FULL(default)NO_TOPKNO_TOPPNO_TOPK_NO_TOPPDesign:
AdvancedSamplingMode.skips_top_k/.skips_top_pproperties are the only place the mode→filter mapping lives. The enum is carried
end-to-end: pydantic parses the YAML/JSON string into the enum on the config,
SpecMetadatastores the enum, and there are no bare-string predicates downstream.resolve_advanced_sampling_filters(mode, top_k, top_p)returns effective tensors with a disabled filter set toNone, so thesampler/prob ops stay mode-agnostic.
sample_from_logits_op(logits, temperatures, top_k, top_p, seed, offset)(mirroringcompute_probs_from_logits_op) is the single one-model sampler:a
Nonetop_kskips the mask kernel; aNonetop_psamples viasampling_from_probsinstead oftop_p_sampling_from_probs. The rejection path'scompute_probs_from_logitsskips the top-k mask / top-p renorm identically.(
DISABLE_TEMP_VAL) set in_scan_one_model_sampling, sosoftmax(logits / DISABLE_TEMP_VAL)collapses to a one-hot at the argmax and sampling returns theargmax token. Every mode therefore supports mixed greedy + sampling batches with no
special-case error.
advanced_sampling_modeis a deploy-time setting and is notpart of the CUDA-graph key.
NO_TOPKwithtop_k=0), the probs — and therefore the sampled tokens and rejection acceptance —are bit-for-bit identical to
FULL; only the kernel launch is removed.use_rejection_samplingon or off; the two are independent.Benchmark
Qwen3.6-35B-A3B-NVFP4, TRT-LLM (1×TP1, MTP
draft_len=3), B200, rejection sampling ON,SPEED-Bench low-entropy (ISL≈2197, OSL=1024, temp=0.7), prefix-cache OFF. Values are
OTPS/gpu (output tok/s/gpu). Plot:
plots/pareto_rejection_conc.png.* Speedup =
NO_TOPK_NO_TOPP/FULL (temp only)— same sampling config(
top_k=0, top_p=1), isolating the effect of skipping the two filter kernels.Acceptance length is unchanged across all rows, confirming identical sampling behavior.
FULLwith realtop_k=20 / top_p=0.95is shown as a reference for the actualfiltered-sampling throughput.
Result: for a temperature-only deployment,
NO_TOPK_NO_TOPPgives a consistent~8–28% throughput gain (≈+12% at high concurrency) at zero accuracy cost, by
removing the two redundant per-draft-layer filter kernels.
Test coverage
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py(12 cases, all passing on B200):
{full, no_topk, no_topp, no_topk_no_topp}and theskips_top_k/skips_top_pproperties; the field lives onDecodingBaseConfig(default
FULL); all four modes construct with or withoutuse_rejection_sampling(no gating).
resolve_advanced_sampling_filterssets the disabled filter toNoneper mode (full→neither,no_topk→top_k,no_topp→top_p,no_topk_no_topp→both) and passes kept filters through unchanged.NO_TOPKviasample_from_logits_opyields theexact same tokens as
FULLwhen top_k is disabled, for the same seed/offset, acrosstop_p ∈ {1.0, 0.9}(the top_k mask is a no-op atk=vocab).returns its argmax token under
NO_TOPKandNO_TOPK_NO_TOPP, validating theguard-free design.
FULLmodes construct on any spec-decoding path(config-time acceptance; unsupported paths fall back to
FULLbehavior at runtime).Telemetry manifest / CODEOWNERS
Because this adds a new
DecodingBaseConfigfield,tensorrt_llm/usage/llm_args_golden_manifest.jsonis regenerated via
scripts/generate_llm_args_golden_manifest.py(addsspeculative_config.advanced_sampling_mode,kind: categorical). The manifest--checkCI gate passes. Touching the usage-telemetry path requires approval from@NVIDIA/trt-llm-usage-telemetry-devs,@NVIDIA/trt-llm-oss-compliance, and@NVIDIA/trt-llm-noncommitted-api-review-committee. The new field is additive withdefault
FULL(API-compatible).Note — alternative design
Instead of a deploy-time setting that changes which sampler the single advanced CUDA
graph captures, each mode could be a separately captured CUDA graph selected at
replay by a graph key — the same mechanism
is_all_greedy_samplealready uses to switchbetween the greedy (argmax) and advanced-sampling graphs. That would allow per-batch
mode switching, but it multiplies the number of captured decode graphs (one set per
mode) and lengthens warmup. We chose the deploy-time enum (0 extra graphs) because the
sampling config is fixed per deployment, which is the target use case.
Test Coverage
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py(11 cases, all passing on B200):
{full, no_topk, no_topp, no_topk_no_topp};NO_TOPKis accepted without rejection sampling;NO_TOPPandNO_TOPK_NO_TOPPraise aValidationErrorunlessuse_rejection_sampling=True(and are accepted with it).
compute_probs_from_logitspassestop_k=Noneand/ortop_p=Noneinto the flashinfer op (checked with a spy),so the right kernel is skipped:
full→neither,no_topk→top_k,no_topp→top_p,no_topk_no_topp→both.sampling_batch_spec_dec_one_modelinNO_TOPKyields the exact same tokens asFULLwhen top_k is disabled, for thesame seed/offset, across
top_p ∈ {1.0, 0.9, 0.5}(the top_k mask is a no-opat
k=vocab). Skipped on CPU-only CI.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.Summary by CodeRabbit
New Features
AdvancedSamplingModefor one-model MTP speculative decoding, exposed asDecodingBaseConfig.advanced_sampling_mode(defaultFULL) and plumbed viaSpecMetadata.FULL— apply both top-k and top-p filtersNO_TOPK— skip top-k filteringNO_TOPP— skip top-p renormalizationNO_TOPK_NO_TOPP— skip both filtersNone:resolve_advanced_sampling_filters(...)converts(advanced_sampling_mode, top_k, top_p)into effective optional tensors where disabled filters becomeNoneNonesample_from_logits_op(...):top_kis providedtop_pis providedtop_k/top_pfromSpecMetadata.advanced_sampling_modesample_from_logits_op(...)for non-greedy target token samplingTests
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.pycovering:skips_top_k/skips_top_pflagsDecodingBaseConfig/MTPDecodingConfigresolve_advanced_sampling_filters(...)identity vs disabled-to-NonebehaviorNO_TOPKmatchesFULLwhentop_kis disabled)FULLmodes across speculative/spec paths (including removal of the prior “one-model-only” gate)tensorrt_llm/usage/llm_args_golden_manifest.jsonto addspeculative_config.advanced_sampling_modeas a categorical argument.Docs / Config
docs/source/features/sampling.md, including semantics, deploy-time behavior (not in CUDA-graph key), independence from rejection sampling, and a usage example.Dev Engineer Review
Correctness & API consistency
resolve_advanced_sampling_filters(...)and consistently threaded into:compute_probs_from_logits)top_k/top_pasNonealigns with the stated goal: redundant FlashInfer kernels can be skipped while keeping mathematical no-op behavior when filters are disabled.CUDA-graph / keying
advanced_sampling_modeis excluded from the CUDA-graph key; functional selection is driven by resolved optional filter tensors (Nonevs provided tensors), matching the intended graph-stability requirement.Config / manifest
AdvancedSamplingModeadded totensorrt_llm/llmapi/llm_args.pyand defaulted toFULLonDecodingBaseConfig.tensorrt_llm/usage/llm_args_golden_manifest.jsonupdated to includespeculative_config.advanced_sampling_modewith allowed values:full,no_topk,no_topp,no_topk_no_topp.QA Engineer Review
Test changes
tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.pyAdvancedSamplingModevalues and skip flags)DecodingBaseConfig,MTPDecodingConfig, independent of rejection sampling)resolve_advanced_sampling_filtersmaps disabled inputs toNone)no_topk/no_topk_no_topp)FULLmodes allowed and stored)tests/integration/test_lists/): not indicated in the provided context.