[TRTLLM-13230][feat] support min_p sampling for TorchSampler#16590
[TRTLLM-13230][feat] support min_p sampling for TorchSampler#16590lori-ren wants to merge 6 commits into
Conversation
Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run |
📝 WalkthroughWalkthroughAdds min-p sampling across parameter propagation, strategy selection, probability renormalization, FlashInfer execution, greedy-decoding inference, grouped sampling, and unit tests. ChangesMin-p sampling
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Request as Sampling request
participant Utils as UtilsSamplingParams
participant Resolver as resolve_sampling_strategy
participant Sampler as MinPSampleOnly
participant FlashInfer as min_p_sampling_from_probs_op
Request->>Utils: provide min_p
Utils->>Resolver: resolve min-p strategy
Resolver->>Sampler: select MinPSampleOnly
Sampler->>FlashInfer: sample from probabilities with min_p
FlashInfer-->>Sampler: return sampled token
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/sampling_params.py (1)
1-1: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMissing
min_prange validation lets NaN reach probability renormalization.SamplingParams._validate()validatestop_p(0<=top_p<=1) andtop_k(>=0) but never validatesmin_p, even though this PR wiresmin_pinto greedy-decoding resolution and probability renormalization; an out-of-rangemin_p > 1makesmin_p_renorm_probszero an entire probability row and divide by zero.
tensorrt_llm/sampling_params.py#L433-445: add amin_pbound check in_validate()(near the existingtop_p/top_kchecks), e.g.if self.min_p is not None and (self.min_p < 0 or self.min_p > 1): raise ValueError(...).tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py#L53-67: once validated upstream,min_p_renorm_probs's unguardedprobs / probs.sum(...)is safe; no code change needed here beyond the upstream fix.🤖 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/sampling_params.py` at line 1, Update SamplingParams._validate() to validate non-None min_p values using the same [0, 1] bounds as top_p, raising ValueError when min_p is below 0 or above 1. Place the check alongside the existing top_p and top_k validation; do not modify min_p_renorm_probs.tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py (1)
70-97: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
top_k=0sentinel from the new min-p strategy will trip this function's own assert.
top_k_top_p_sampling_batch's contract (per its own docstring) is that onlyNoneorvocab_sizedisable top-k; passing0hitsassert top_k > 1a few lines down and raisesAssertionError.resolve_sampling_strategy's new"min_p"branch insampling_utils.pydeliberately returns0as the "disabled top-k" sentinel (to dodge int32 overflow from the2**31greedy-probe vocab size), but the standalonesample()dispatcher forwards that0straight into this function without thesanitize_top_kstep that the FlashInfer-based_compute_probspath applies. See thesampling_utils.pycomment for the root-cause fix; flagging here since this is where the crash actually manifests.🤖 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/pyexecutor/sampler/ops/vanilla.py` around lines 70 - 97, Update top_k_top_p_sampling_batch to accept the top_k=0 sentinel used by resolve_sampling_strategy for min_p sampling by normalizing it to the existing disabled-top-k representation before the top_k validation and need_top_k calculation. Preserve the current behavior for None, vocab_size, and positive top_k values, and keep the existing assertions for invalid sampling parameters.
🤖 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/ops/vanilla.py`:
- Around line 53-67: Update SamplingParams._validate() to validate min_p with
the same source-level range contract as top_p, requiring values between 0 and 1
inclusive before they reach min_p_renorm_probs. Support the existing scalar or
per-request representation as appropriate, and preserve valid min_p values
unchanged.
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py`:
- Around line 166-171: Update top_k_top_p_sampling_batch to treat top_k values
of None and 0 as disabled, matching sanitize_top_k’s contract, while preserving
existing validation and sampling behavior for enabled top_k values. This ensures
the sample() dispatcher’s "min_p" path and draft-model rejection sampling handle
an unset top_k without assertion failures.
In `@tensorrt_llm/sampling_params.py`:
- Around line 433-445: Update SamplingParams._validate() to validate min_p like
the other probability thresholds, requiring it to remain within the inclusive
range 0 to 1. Preserve existing top_p/top_k validation and ensure invalid min_p
values are rejected before min_p_renorm_probs or greedy-decoding logic consumes
them.
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Line 1257: Update the SamplingParams reconstruction in the VaryParams branch
to preserve the original request’s min_p value by carrying param.min_p into the
reconstructed parameters. Keep the existing top-p, top-k, and temperature values
unchanged so randomized mixed min-p batching remains covered.
- Around line 395-404: The existing test only validates strategy tuple
resolution for min_p combinations, not runtime behavior. Extend the min_p
coverage in test_probs and test_samples with executable cases for min_p combined
with top_k, top_p, and both, using representative non-neutral values. Verify
filtering order and FlashInfer inputs while preserving the existing standalone
min_p cases.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py`:
- Around line 70-97: Update top_k_top_p_sampling_batch to accept the top_k=0
sentinel used by resolve_sampling_strategy for min_p sampling by normalizing it
to the existing disabled-top-k representation before the top_k validation and
need_top_k calculation. Preserve the current behavior for None, vocab_size, and
positive top_k values, and keep the existing assertions for invalid sampling
parameters.
In `@tensorrt_llm/sampling_params.py`:
- Line 1: Update SamplingParams._validate() to validate non-None min_p values
using the same [0, 1] bounds as top_p, raising ValueError when min_p is below 0
or above 1. Place the check alongside the existing top_p and top_k validation;
do not modify min_p_renorm_probs.
🪄 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: 2043a377-fdc6-4869-8f77-edcea06800fc
📒 Files selected for processing (6)
tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.pytensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.pytensorrt_llm/sampling_params.pytests/unittest/_torch/sampler/test_torch_sampler.py
| min_p: Optional[float] = None, | ||
| ): | ||
| return (not use_beam_search) and ( | ||
| (temperature is None and top_p is None and top_k is None) | ||
| or top_k == 1 | ||
| or top_p == 0.0 | ||
| or temperature == 0 | ||
| ) | ||
| if use_beam_search: | ||
| return False | ||
| # Explicit greedy triggers collapse sampling to a single token, | ||
| # regardless of min_p. | ||
| if top_k == 1 or top_p == 0.0 or temperature == 0: | ||
| return True | ||
| # Otherwise sampling is greedy only when no sampling knob is set. A | ||
| # positive min_p on its own still selects among multiple tokens. | ||
| params_unset = temperature is None and top_p is None and top_k is None | ||
| min_p_set = min_p is not None and min_p > 0.0 | ||
| return params_unset and not min_p_set |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Consider validating min_p's range alongside top_p/top_k.
The new greedy-decoding logic here is correct and well covered by tests. Separately, _validate() (a few lines above, unchanged by this diff) validates top_p (0 <= top_p <= 1) and top_k (>= 0) but has no equivalent bound for min_p. Since min_p is now load-bearing for both greedy-decoding resolution (here) and probability renormalization (min_p_renorm_probs in tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py), an out-of-range min_p > 1 will zero an entire probability row and divide by zero (NaN) downstream. See the more detailed comment on vanilla.py lines 53-67 for the concrete failure mode and a proposed _validate() fix.
Also applies to: 454-454
🤖 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/sampling_params.py` around lines 433 - 445, Update
SamplingParams._validate() to validate min_p like the other probability
thresholds, requiring it to remain within the inclusive range 0 to 1. Preserve
existing top_p/top_k validation and ensure invalid min_p values are rejected
before min_p_renorm_probs or greedy-decoding logic consumes them.
| TopP: SamplingParams(top_p=0.42, temperature=0.2), | ||
| TopK: SamplingParams(top_k=27, temperature=0.5), | ||
| TopKTopP: SamplingParams(top_k=27, top_p=0.6, temperature=0.5), | ||
| MinP: SamplingParams(min_p=0.05, temperature=1.0), |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Preserve min_p in randomized mixed batches.
Adding MinP here feeds it into the later VaryParams branch, whose reconstructed SamplingParams includes only top-p, top-k, and temperature. Those MinP requests therefore become temperature-only, leaving randomized mixed min-p batching untested. Carry param.min_p into the reconstructed parameters.
As per path instructions, coverage is insufficient for randomized min-p requests.
🤖 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/sampler/test_torch_sampler.py` at line 1257, Update the
SamplingParams reconstruction in the VaryParams branch to preserve the original
request’s min_p value by carrying param.min_p into the reconstructed parameters.
Keep the existing top-p, top-k, and temperature values unchanged so randomized
mixed min-p batching remains covered.
Source: Path instructions
|
PR_Github #60261 [ run ] triggered by Bot. Commit: |
…led sentinel in vanilla sampling Signed-off-by: Lori Ren <lorir@nvidia.com>
|
PR_Github #60261 [ run ] completed with state |
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
According to coderabbit, the one engine spec dec path is currently ignoring min p. That's fine for now, but we should document the limitation and fix it in a follow up.
We should fix that as soon as possible. We've had issues with spec dec ignoring sampling params in the past, users find it confusing
Signed-off-by: Lori Ren <lorir@nvidia.com>
… merge A main merge dropped the min_p override in params_imply_greedy_decoding: min_p stayed in the signature but the body only handled top_p_decay, so a min_p-only request silently resolved to greedy (resolve_sampling_strategy returned GREEDY before reaching the min_p branch). Restore the override so min_p and top_p_decay both keep sampling multi-token. Also fix a stale test_min_p_only assertion (the disabled-top_k sentinel is 0, not vocab_size), add a regression test for the draft-model rejection path (min_p>0 with the top_k=0 sentinel through sample()), and document that one-engine speculative decoding does not yet propagate min_p. Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run |
|
PR_Github #61513 [ run ] triggered by Bot. Commit: |
|
PR_Github #61513 [ run ] completed with state
|
…coding One-model speculative decoding ignored min_p (no min_p buffer; the sampling kernels took only temperature/top_k/top_p), so a min_p-only request decoded greedily. Thread min_p through mirroring top_p: add per-token min_ps / per-request request_min_p buffers and a skip_min_p flag, read min_p in _scan_one_model_sampling, and apply it as a min_p_renorm_probs stage in compute_probs_from_logits (flashinfer has no fused min-p kernel). The tokens-only sampler now shares the probs pipeline, so min_p affects both the sampled tokens and the rejection-acceptance distributions across the MTP / Eagle3 / dynamic-tree one-model paths. Signed-off-by: Lori Ren <lorir@nvidia.com>
Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run |
|
Hi @mikeiovine, I've extended min_p support for one-eng spec path. Pls review at your convenience. |
|
PR_Github #61539 [ run ] triggered by Bot. Commit: |
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
In sample_and_accept_draft_tokens (the target-probs path) you correctly short-circuit via skip_min_p, passing None so the if min_p is not None branch in compute_probs_from_logits gets traced away and no min-p op enters the graph:
interface.py ~L1732
min_ps = (None if spec_metadata.skip_min_p else spec_metadata.min_ps[gen_start:gen_end])
But the draft-sampling call sites pass the tensor unconditionally instead of honoring skip_min_p:
interface.py ~L1552, ~L1580, ~L1960, ~L1973, ~L2246, ~L2259
min_ps = spec_metadata.request_min_p[:batch_size] # never None
sampling_batch_spec_dec_one_model(..., min_p=min_ps, ...)
Because request_min_p / min_ps is a torch.zeros(...) buffer slice (never None), the if min_p is not None guard in compute_probs_from_logits is always True at torch.compile/CUDA-graph capture time. This bakes min_p_renorm_probs into the captured graph and runs it on every replay, including for the very common case where no request enables min-p.
min_p_renorm_probs is a numerical no-op when min_p == 0 (thresholds become 0, where keeps everything, sum→div by 1.0), but it is not free: it adds a full-vocab row-max reduction + elementwise where + row-sum + divide per draft step, per replay. On the draft path this fires once per draft token, so the wasted work scales with draft length.
And please test the spec decoding performance with https://gitlab-master.nvidia.com/ftp/spec-benchmarks
|
Please take a look of #16561 we need add a skip min-p filter. |
|
PR_Github #61539 [ run ] completed with state
|
4 participants
Summary by CodeRabbit
New Features
Tests
Description
The default PyTorch TorchSampler does not apply min_p — the param is accepted by SamplingParams but silently ignored (only the TRT backend / deprecated TRTLLMSampler handle it via the C++ executor SamplingConfig.min_p). Min-P is a common sampling parameter, so this is a feature-parity gap for the PyTorch backend.
Scope:
Apply Min-P in TorchSampler: keep only tokens whose probability ≥ min_p × p_max (where p_max is the highest token probability), then sample from the renormalized distribution.
Wire it into the existing strategy resolution / FlashInfer fast path (resolve_sampling_strategy, sampling_utils_flashinfer.py) — FlashInfer exposes min_p_sampling_from_probs / a min-p renorm op that can be composed with the current top-k/top-p kernels.
Respect existing semantics (SamplingParams definition / C++ runtime default min_p=0.0 = no-op) so behavior matches the legacy backend.
Support combination with temperature / top-k / top-p, and work for speculative-decoding expanded shapes.
Reference: vLLM and SGLang both support Min-P in their GPU samplers.
Test Coverage
test_torch_sampler.pyto add tests for min-p samplingPR 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.