[TRTLLM-13231][feat] Support top_p_decay in the PyTorch TorchSampler - #16184
Conversation
185d692 to
49e2e80
Compare
49e2e80 to
f38d44a
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 full TorchSampler and FlashInfer support for per-sequence top-p decay, including parameter validation, slot state, grouped metadata, compiled GPU gather/update operations, request lifecycle handling, and updated sampler test instrumentation. ChangesTop-P decay sampling
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant TorchSampler
participant FlashInferGroupedStrategySampler
participant Fusions
TorchSampler->>FlashInferGroupedStrategySampler: pass TopPDecayMetadata
FlashInferGroupedStrategySampler->>Fusions: gather runtime top-p by slot
Fusions-->>FlashInferGroupedStrategySampler: return decayed top-p values
FlashInferGroupedStrategySampler-->>TorchSampler: return sampled tokens
TorchSampler->>Fusions: update sampled decay slots
Fusions-->>TorchSampler: mutate runtime top-p state
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: 1
🧹 Nitpick comments (2)
tests/unittest/_torch/sampler/test_top_p_decay.py (2)
332-358: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRejection coverage only exercises speculative decoding, not beam search or non-FlashInfer sampling.
The PR objectives state decay combined with speculative decoding, beam search, or non-FlashInfer sampling should be rejected/unsupported, but
TestRejectiononly covers the speculative-decoding path (test_reject_speculative) plus the accept/no-op cases. Beam-search and non-FlashInfer-backend rejection are untested here.Suggest adding
test_reject_beam_search(decay +use_beam_search=True) andtest_reject_non_flashinfer_backend(decay with a non-FlashInfer_grouped_sampler_cls) totests/unittest/_torch/sampler/test_top_p_decay.py, or confirm these paths are already covered elsewhere and add a comment/reference here.🤖 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_top_p_decay.py` around lines 332 - 358, Extend TestRejection with coverage for both unsupported decay combinations: add test_reject_beam_search using decay and use_beam_search=True, and test_reject_non_flashinfer_backend using decay with a non-FlashInfer _grouped_sampler_cls. Assert each request is rejected through sampler.validate_request, reusing the existing rejection pattern and helpers.Source: Path instructions
190-327: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftNo end-to-end test exercises
_build_top_p_runtime_override/_apply_top_p_runtime_overridethroughsample_async.
TestRuntimeUpdatevalidates_update_top_p_decay_after_sampleandtop_p_decay_gatherin isolation (good unit coverage), but nothing in this file ortest_torch_sampler.pydrives a decay-active request through the fullsample_asyncpipeline to confirmTopPDecayOverrideis actually constructed and consumed end-to-end (the wiring shown insampling_utils.py/sampler.pycontext snippets).TestBatchedSampling._build_test_casesintest_torch_sampler.pyalso doesn't include a decay-activeBASE_CASESentry.Since this is the integration point tying together admission, per-slot state, and grouped sampling, an end-to-end regression here would catch wiring breaks that the isolated unit tests can't. Suggest adding an integration test (e.g. in
tests/unittest/_torch/sampler/test_top_p_decay.py) that runssample_asyncacross multiple steps with a decay-active request and asserts the runtime top-p actually shrinks between steps, or flag this as a follow-up if it's intentionally out of scope for this PR.🤖 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_top_p_decay.py` around lines 190 - 327, Add an end-to-end test in TestRuntimeUpdate that drives a decay-active request through sample_async for multiple steps, verifying _build_top_p_runtime_override and _apply_top_p_runtime_override are wired correctly and the request’s runtime top-p decreases between samples. Include the required request/admission setup and assert the observed sampling top-p values, while preserving the existing isolated unit tests.Source: Path instructions
🤖 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/sampler.py`:
- Around line 4180-4249: The _build_top_p_runtime_override method must not raise
ValueError when a decay-active row has req_num_steps greater than one, since
this aborts the entire batch. Replace the exception path in the per-request
layout construction with request-scoped handling: skip the top-p decay override
for that row or otherwise fail only the affected request, while preserving
override behavior for valid single-token decay rows.
---
Nitpick comments:
In `@tests/unittest/_torch/sampler/test_top_p_decay.py`:
- Around line 332-358: Extend TestRejection with coverage for both unsupported
decay combinations: add test_reject_beam_search using decay and
use_beam_search=True, and test_reject_non_flashinfer_backend using decay with a
non-FlashInfer _grouped_sampler_cls. Assert each request is rejected through
sampler.validate_request, reusing the existing rejection pattern and helpers.
- Around line 190-327: Add an end-to-end test in TestRuntimeUpdate that drives a
decay-active request through sample_async for multiple steps, verifying
_build_top_p_runtime_override and _apply_top_p_runtime_override are wired
correctly and the request’s runtime top-p decreases between samples. Include the
required request/admission setup and assert the observed sampling top-p values,
while preserving the existing isolated unit tests.
🪄 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: 4627d58b-0abc-4344-aee9-b093cc208387
📒 Files selected for processing (10)
cpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.cucpp/tensorrt_llm/kernels/samplerKernels/toppDecayKernels.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/samplerOp.cpptensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.pytensorrt_llm/sampling_params.pytests/unittest/_torch/sampler/test_top_p_decay.pytests/unittest/_torch/sampler/test_torch_sampler.py
2abe199 to
3abbb4c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59692 [ run ] triggered by Bot. Commit: |
|
PR_Github #59692 [ run ] completed with state
|
3abbb4c to
6a1de75
Compare
Add Top-P Decay (top_p_decay / top_p_min / top_p_reset_ids) support to the
PyTorch backend TorchSampler, with semantics matching the C++ computeToppDecay
kernel: after each sampled token,
runtime_top_p = initial_top_p if token == reset_id (reset_id >= 0)
= max(runtime_top_p * top_p_decay, top_p_min) otherwise
- Per-slot runtime state in resident CUDA buffers, gated on-device by an
is_decay_slot mask plus a host-side slot set for O(1) early-outs; the hot
path needs no host/device sync.
- Two fused custom ops (thop/samplerOp.cpp, kernels/samplerKernels/):
top_p_decay_gather (pre-sample per-row top-p, replaces index_select x2 +
where) and top_p_decay_update (post-sample in-place decay with in-kernel
token gather from the strided new_tokens view).
- Active decay forces a top-p-capable strategy even when top_p is 1.0/unset;
an explicit greedy control (top_k==1 / top_p==0 / temperature==0) wins over
decay (single-source predicate SamplingParams.params_imply_explicit_greedy).
- Out-of-range decay params follow the C++ warn-and-default policy.
- Unsupported combinations (beam search, speculative draft tokens through
TorchSampler) are rejected per-request at admission via validate_request,
with a req_num_steps == 1 guard at sample time; finished requests retire
their slot so the early-outs re-arm.
- Unit tests (test_torch_sampler.py::TestTopPDecay): strategy routing,
runtime-update parity with the C++ reference (cases ported from
topPSamplingLayerTest.cpp), and per-request rejection of unsupported
combinations.
Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…pdate Sampled token ids are non-negative, so the negative reset_ids sentinel (-1, reset disabled) never matches without an explicit reset_id >= 0 test; compare directly, matching the legacy computeToppDecay kernel. Update the recurrence described in the op/kernel docs and the test reference accordingly. Addresses review feedback from @ixlmar. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…ent cleanups - Extract all top-p-decay bookkeeping from setup_sampler_step into _setup_top_p_decay_for_new_requests; rename decay_params -> sampling_params. - Encapsulate TopPDecayStore allocation as TopPDecayStore.create(). - _retire_top_p_decay_slot: move the GENERATION_COMPLETE check to the call sites so it is not re-checked redundantly. - Narrow _build_top_p_runtime_override's strategy_key type to FlashInferGroupedStrategySampler.STRATEGY_KEY_TYPE. - Inline the float_tensor helper; generalize the SamplingConfig unwrap comment to cover all sampling fields; document that decay param range validation already lives in the executor::SamplingConfig constructor. - Tests: hoist the mock-request helper in TestTopPDecay. Addresses review feedback from @ixlmar. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…ct invalid decay params Address review feedback from @ixlmar: - Pass the top-p-decay override through the existing StrategyMetadata mechanism (TopPDecayMetadata, mirroring BeamSearchMetadata) instead of a decay_override kwarg on sample_grouped_strategies; the override is applied inside the TopP*/TopKTopP* strategy impls via TopPDecayMixin (which declares _top_p, removing the getattr/setattr workaround). - Vectorize the per-STEP slot layout with torch.repeat_interleave; guard the single-token invariant in an if __debug__ block; use .item() over int(). - BREAKING: SamplingParams rejects out-of-range decay params (top_p_decay / top_p_min outside (0, 1], negative top_p_reset_ids) instead of warn-and-default, mirroring the executor::SamplingConfig hard checks. - Single-source the greedy/decay resolution: params_imply_top_p_decay_active and an extended params_imply_greedy_decoding(top_p_decay=...) on SamplingParams, reused by _greedy_decoding and resolve_sampling_strategy (decay predicate now evaluated lazily). - Consolidate the feature documentation: TopPDecayStore's docstring is the single source for semantics and the slot lifecycle; other sites keep their local contract and point back to it. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
_unbatch_sampling_results gained a required seq_slots_cuda kwarg (the resident device copy used by the post-sample top-p-decay update), but the direct caller in test_unbatch_sampling_results was not updated. Precompute the device copy outside the no-sync region. Full test_torch_sampler.py now passes (205/205) in a matching container. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
48e2f7f to
8e2dea1
Compare
…rch.compile'd ops Per review discussion (@ixlmar): keep the new sampler kernel-free. The two decay ops (pre-sample per-row gather, post-sample in-place update) are now native torch functions fused with torch.compile in sampler/ops/top_p_decay.py; the custom CUDA kernels, the thop op wrappers, and their build wiring are removed. Compile configuration and rationale (benchmarked on H200, bs=32, per-step sampler cost beyond a fake model forward): - mode=max-autotune-no-cudagraphs: cudagraphs is unsafe here (in-place per-slot state, output consumed outside the compiled region) and the cudagraphs-enabled modes were either slower or produced overwritten outputs. - torch._dynamo.mark_dynamic on the batch-varying dims avoids recompilation as batch composition changes. - With a >=2ms forward, this matches the removed fused kernels (~141us vs ~142us per step, both ~35us ahead of eager); with a ~1ms forward it costs ~40us/step over eager (~90us over the fused kernels), the fixed compiled-entry cost being exposed. The kernel-free maintainability trade-off was judged worth it for the non-critical-path decay feature. - Compilation is lazy: first decay-active request pays ~1s; non-decay workloads never trigger it. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
….Fusions
They are pure-torch compiled ops; Fusions in ops/vanilla.py is the existing
home for torch.compile'd fused helpers (same impl + mark_dynamic wrapper
pattern as gather_log_softmax), so the standalone ops/top_p_decay.py module
is folded into it and call sites use Fusions.top_p_decay_{gather,update}.
Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py (1)
399-419: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winKeep the static top-p tensor immutable across sampling calls.
Line 414 persists gathered runtime values in
self._top_p. On a later group with reordered or newly non-decay rows, Line 418 uses those stale values asstatic_top_p, so plain top-p requests can sample with another request’s decayed threshold. Preserve_static_top_p, reset_top_pfrom it on every call, and gather from_static_top_p.🤖 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/sampling_utils.py` around lines 399 - 419, The _maybe_apply_top_p_decay method currently mutates self._top_p, allowing decayed values to leak into later sampling calls. Preserve the immutable static tensor in _static_top_p, reset self._top_p from it on every invocation, and pass _static_top_p as static_top_p to top_p_decay_gather while retaining the existing metadata validation and non-decay behavior.tensorrt_llm/_torch/pyexecutor/sampler/sampler.py (1)
2589-2595: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftPrevent late speculative-decoding mismatches from aborting the batch.
get_draft_token_length()is explicitly best-effort at admission because drafter tokens may be attached later. Such a request can reach the assertion at Lines 4326-4338 withreq_num_steps > 1; the assertion then raises during grouped metadata construction and aborts the whole sampler batch. Reject the request when draft tokens become known, or skip/fail only that request instead of asserting in the batch path. This is the same late top-p-decay mismatch previously reported.Also applies to: 4326-4338
🤖 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/sampler.py` around lines 2589 - 2595, Prevent late speculative-decoding requests from aborting grouped metadata construction: update the admission/metadata flow around get_draft_token_length and _build_top_p_decay_metadata so requests whose draft tokens become known with req_num_steps > 1 are rejected or skipped individually before the batch-level assertion. Remove or replace the assertion path in _build_top_p_decay_metadata with per-request handling while preserving valid non-speculative requests.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py (1)
72-94: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the public tensor contracts with
Args:sections.These exported functions have shape, dtype, mutation, and indexing requirements, but their docstrings do not document each argument in Google style. Add explicit
Args:entries, includingsampled_slots/slotsindexing and the in-place mutation ofruntime_top_p.Also applies to: 108-124
🤖 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/top_p_decay.py` around lines 72 - 94, Add Google-style Args sections to the exported functions top_p_decay_update and the related function at the referenced later section, documenting each tensor’s shape, dtype, indexing semantics, and role. Explicitly describe sampled_slots/slots as indices into per-slot tensors and state that runtime_top_p is mutated in place while the functions return None.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.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 2589-2595: Prevent late speculative-decoding requests from
aborting grouped metadata construction: update the admission/metadata flow
around get_draft_token_length and _build_top_p_decay_metadata so requests whose
draft tokens become known with req_num_steps > 1 are rejected or skipped
individually before the batch-level assertion. Remove or replace the assertion
path in _build_top_p_decay_metadata with per-request handling while preserving
valid non-speculative requests.
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py`:
- Around line 399-419: The _maybe_apply_top_p_decay method currently mutates
self._top_p, allowing decayed values to leak into later sampling calls. Preserve
the immutable static tensor in _static_top_p, reset self._top_p from it on every
invocation, and pass _static_top_p as static_top_p to top_p_decay_gather while
retaining the existing metadata validation and non-decay behavior.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.py`:
- Around line 72-94: Add Google-style Args sections to the exported functions
top_p_decay_update and the related function at the referenced later section,
documenting each tensor’s shape, dtype, indexing semantics, and role. Explicitly
describe sampled_slots/slots as indices into per-slot tensors and state that
runtime_top_p is mutated in place while the functions return None.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 269e9fdb-3ca0-4fe7-ab04-308a160806e0
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/sampler/ops/top_p_decay.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
…space change - TopPDecayStore lifecycle step 1 still said 'fused kernels'; the CUDA kernels were replaced by torch.compile'd ops. - Restore a blank line in setup_sampler_step accidentally dropped from the upstream code, keeping the diff against main minimal. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py (1)
389-389: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAdd
fullgraph=Trueto the two@torch.compile(...)decorators
mode="max-autotune-no-cudagraphs"can be combined withfullgraph=True, and these in-place index updates/gathers are supported bytorch.compile. Keepingfullgraph=Truehere would make any unexpected graph break fail loudly instead of silently falling back to eager on this per-step path.🤖 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` at line 389, Update both `@torch.compile` decorators in vanilla.py at lines 389-389 and 447-447 to include fullgraph=True alongside mode="max-autotune-no-cudagraphs". Apply the same change to both decorator sites and leave the wrapped functions unchanged.
🤖 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.
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py`:
- Line 389: Update both `@torch.compile` decorators in vanilla.py at lines 389-389
and 447-447 to include fullgraph=True alongside
mode="max-autotune-no-cudagraphs". Apply the same change to both decorator sites
and leave the wrapped functions unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cad5f9e6-79e1-4ae5-8024-699589097c8a
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
|
/bot run --disable-fail-fast |
|
PR_Github #60927 [ run ] triggered by Bot. Commit: |
|
PR_Github #60927 [ run ] completed with state
|
…dexing mypy rejected slicing seq_slots_cuda with Tensor.item() results (Slice index must be an integer). Cast first_req/last_req to int, matching the existing int()/cast(int, ...) pattern used elsewhere in this file. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #60972 [ run ] triggered by Bot. Commit: |
|
PR_Github #60972 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61001 [ run ] triggered by Bot. Commit: |
|
PR_Github #61001 [ run ] completed with state |
…VIDIA#16184) Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
6 participants
Implement Top-P Decay on the FlashInfer sampling path, porting the semantics of the legacy C++ computeToppDecay kernel. The three SamplingParams fields top_p_decay / top_p_min / top_p_reset_ids were previously dropped by the PyTorch backend; they are now plumbed through and honored.
Each decode step, after the token is sampled, the per-request runtime top-p is updated as runtime_p = max(runtime_p * top_p_decay, top_p_min), or reset to its initial value when the sampled token equals top_p_reset_ids (gated on reset_id >= 0). The runtime top-p is carried as a per-row tensor into the FlashInfer batched sampling ops, never encoded into the cached strategy tuple.
Scope is limited to the FlashInfer path and single-token decode steps; decay combined with speculative decoding, beam search, or the non-FlashInfer backend is rejected early. Out-of-range top_p_decay / top_p_min warn-and-default, matching the C++ runtime.
Adds unit tests covering activation/routing, warn-and-default validation, numeric parity with the C++ recurrence, slot-reuse lifecycle, and mixed decay / non-decay batches.
Dev Engineer Review
TorchSampler, including per-slot runtime state (top_p_decay,top_p_min,top_p_reset_ids, initial/runtime values, and reset handling) gated by a device-side “decay-active slot” mask, with slot lifecycle updates on request admission/retirement.runtime_top_pfor FlashInfer sampling and (b) updateruntime_top_pin-place after each sampled token usingtop_p_decaywithtop_p_minclamping andreset_idsrecurrence (parity w/ legacy C++computeToppDecaysemantics).top_p_decay < 1.0activates top-p-capable strategies at runtime even when initial/statictop_pimplies greedy; added greedy-resolution predicates inSamplingParamsto ensure explicit greedy intent can override decay activity.tensorrt_llm/sampling_params.pyto be within supported ranges; out-of-range values raiseValueError. Requests combining Top-p decay with unsupported flows (e.g., speculative draft tokens / multi-token decode steps) are rejected per-request during admission.QA Engineer Review
tests/unittest/_torch/sampler/test_torch_sampler.pyby addingTestTopPDecaywith these new test methods:test_strategy_routingtest_runtime_update_paritytest_out_of_range_decay_params_rejectedtest_reject_speculative_draft_tokenstests/integration/test_lists/(test-db/ and qa/) was not provided; Verdict: needs follow-up.Description
Adds Top-P Decay (
top_p_decay/top_p_min/top_p_reset_idsinSamplingParams)to the PyTorch backend's
TorchSampler. These parameters were previously only honoredby the legacy TRT sampler; on the torch path they were silently dropped.
Semantics (parity with the legacy C++
computeToppDecaykernel)After each sampled token of a decay-active request:
top_p_decayis set and< 1.0;top_p_min/top_p_reset_idsalone do not activate it. Defaults mirror the C++ runtime (
min=1e-6,reset_id=-1,a negative sentinel that never matches since token ids are non-negative).
top_pis 1.0/unset(initial top-p defaults to 1.0), so the decayed value can take effect on later steps.
An explicit greedy control (
top_k==1,top_p==0.0,temperature==0) wins overdecay. The greedy/decay resolution is single-sourced in
SamplingParams.params_imply_greedy_decodingand shared by_greedy_decodingandresolve_sampling_strategy.top_p_decay/top_p_minoutside(0, 1],negative
top_p_reset_ids) are now rejected bySamplingParamswith a clear error,replacing the former warn-and-default behavior and mirroring the hard checks in the
executor::SamplingConfigconstructor.Implementation (kernel-free)
No custom CUDA code. Per-slot runtime state lives in resident CUDA buffers
(
TorchSampler.TopPDecayStore, the single documentation point for the feature'ssemantics and slot lifecycle), gated by a device-side
is_decay_slotmask plus ahost-side slot set used only for O(1) early-outs — the hot path needs no host/device
sync or host-side filtering.
are native torch functions fused with
torch.compileinops.vanilla.Fusions,following the existing
Fusionspattern (gather_log_softmax).mode="max-autotune-no-cudagraphs"— cudagraphs is deliberately disabled: the updatemutates persistent per-slot state in place and the gather's output is consumed outside
the compiled region, which is unsafe with cudagraph static buffers (verified: outputs
get overwritten by subsequent replays).
torch._dynamo.mark_dynamicon thebatch-varying dims avoids recompilation as batch composition changes. Compilation is
lazy (~1 s on the first decay-active request; non-decay workloads never trigger it).
StrategyMetadatamechanism(
TopPDecayMetadata, mirroringBeamSearchMetadata) and consumed inside theTopP*/TopKTopP* strategy impls via
TopPDecayMixin— the generic sampling entrypoint is unchanged.
TorchSampler) are rejected per-request at admission viaSampler.validate_request,with a debug-only invariant check at sample time.
Performance
Measured with an isolated sampler benchmark (H200, bs=32, steady-state decode, per-step
sampler cost beyond a fake model forward of the given length; p50 over 2048 steps):
With a forward of ≥ 2 ms the compiled ops match the hand-written fused kernels (both
~20 µs decay overhead, ~35 µs ahead of eager): the compiled-function entry cost hides
behind the forward and the Inductor-fused kernels keep the GPU tail as short as the
custom kernels did. With a
1 ms forward the fixed per-call entry cost is exposed+40 µs vs. eager). We took the kernel-free maintainability trade-off, as the decay(
feature is not on the critical benchmarking path; an earlier revision of this PR
contains the fused-kernel implementation should that regime ever matter.
A broader sweep (all four
torch.compilemodes × dynamic/static/mark_dynamic,12 configurations) and profiling notes are in the PR discussion.
Test Coverage
tests/unittest/_torch/sampler/test_torch_sampler.py::TestTopPDecay:decay == 1.0inactive).from
topPSamplingLayerTest.cppTopPDecay, plus thetop_p_min > initialrise case),including reset behavior.
SamplingParamsconstruction.validate_request.Additionally verified: a 300-step pipelined (no per-step sync) trajectory check of the
compiled ops against the analytic recurrence, and the full
test_torch_sampler.pysuite (205 passed) in a matching container.
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.