[TRTLLM-13215][perf] Skip redundant one-model sampling-param refills

[zhaoyangwang-nvidia]

Description

populate_sampling_params_for_one_model rebuilt and re-uploaded every sampling-parameter buffer on every step, even though a request's sampling params are fixed for its lifetime. Per step that is a host expansion of batch_size * (draft_len + 1) list appends plus six H2D copies, all ahead of the forward on the critical path. This PR records what the buffers hold and skips the refill when the step reproduces it.

• The two buffer groups are tracked with separate signatures because they depend on different inputs. request_* (length batch) is a function of the ordered sampling values; the expanded temperatures / top_ks / top_ps (length batch * (draft_len + 1)) additionally depend on each request's token count, which sets their layout. A context request becoming a generation request grows its span from one row to draft_len + 1 and shifts every later request, invalidating the expanded buffers while leaving the per-request ones valid; the same applies when runtime_draft_len changes. A single merged signature would refill request_* needlessly on every such transition.
• Slot ids are deliberately not part of either signature. Every consumer reads these buffers by batch position, so the order of the normalized values already encodes the batch ordering, and a reshuffle changes the signature on its own. Slot ids index batch_slot_ids, which is copied separately; including them would only force refills when a slot changes hands between requests that happen to sample identically.
• top_k_max derives from the per-request values, so it now refreshes exactly when those buffers do rather than on every call.
• The signatures live in a list that create_cuda_graph_metadata's shallow copy keeps sharing. The graph views and the eager view write the same tensors, so the record of what those tensors hold has to be shared too — plain fields would give each view its own stale answer and let one skip a fill another view invalidated, silently sampling with another batch's params.

An earlier revision of this PR instead moved the expansion to the device, keeping the params in slot-indexed tables and gathering them with index_select. That cut host time substantially (batch 128: 186 us to 54 us, and flat in the row count rather than linear) but end-to-end throughput did not improve — it regressed slightly, up to -2.24% per-user at batch 16. The host work is submitted asynchronously (pinned + non_blocking=True), so it is already hidden whenever the GPU is the bottleneck, while the three added index_select kernels are real GPU work on the critical path. Skipping the work entirely avoids both.

Test Coverage

Validated on H200 NVL (SM 9.0).

• Signature behavior per buffer group, reported as (need_update_sampler_param, need_update_expanded_sampler_param): first call (True, True); steady-state repeat (False, False); context to generation, which changes the token count, (False, True); a runtime_draft_len change (False, True); a sampling param change (True, True); slot order reshuffled (True, True); batch growth (True, True); explicit invalidate (True, True). A graph copy shares the signature list, and after the graph view fills, the eager view on the same batch reports (False, False).
• Equivalence against a forced full refill, which is the check that matters here: driving a request sequence through populate_sampling_params_for_one_model and comparing every buffer plus top_k_max against a metadata object that refills unconditionally. Identical across the initial fill, a steady-state skip, a generation-to-context transition, a return to steady state, and a sampling-param change. The skip path produces bit-identical buffers to always refilling.
• Edge cases: an all-greedy step interleaved between two identical non-greedy batches still hits the cache correctly; buffer reallocation forces a refill; the top_k sentinel (INT32_MAX) is distinguished from 0.
• End-to-end test_llama_eagle3_rejection_sampling_modes passes (4 cases), covering use_cuda_graph x use_dynamic_tree with temperature=1.0 / top_p=0.9 / top_k=50.
• MT-Bench full-batch output throughput, draft_len=5, 80 prompts, one run per point: batch 1, 420.82 to 421.44 tok/s (+0.15%); batch 2, 841.52 to 841.13 (-0.05%); batch 4, 1614.47 to 1620.39 (+0.37%); batch 8, 2907.33 to 2909.33 (+0.07%); batch 16, 5640.86 to 5592.40 (-0.86%). All within noise — the deltas swing both ways with no trend, and a single run per point gives no variance estimate. Three runs of the unmodified baseline at batch 16 span 3.222 to 3.374 acceptance length and 5442 to 5641 tok/s, a 3.6% spread from sampling alone, well above anything this change could contribute. Acceptance length tracks the baseline at every batch size, which is the check that matters: the sampling distribution is untouched.
• Changed-file pre-commit hooks pass.

No measurable throughput change is the expected shape. The skipped work is host-side and submitted asynchronously, so it is already hidden whenever the GPU is the bottleneck; this removes redundant work rather than shortening the critical path.

JIRA: https://jirasw.nvidia.com/browse/TRTLLM-13215

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[zhaoyangwang-nvidia]

Hi @lori-ren @allisonlim-nv could you help to review this PR, thanks~

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 379e7b3d-df79-4c95-a06d-d7789e595de0

📥 Commits

Reviewing files that changed from the base of the PR and between f20561a and af6f657.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/speculative/interface.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tensorrt_llm/_torch/speculative/interface.py

---

Walkthrough

SpecMetadata now shares sampling signatures across metadata copies, detects changed sampling buffer groups, updates them independently, and invalidates signatures after reallocation or through a public method.

Changes

Speculative sampling metadata

Layer / File(s)	Summary
Sampling cache state and lifecycle tensorrt_llm/_torch/speculative/interface.py	SpecMetadata stores shared signatures for per-request and expanded per-token sampling buffers. CUDA-graph copies retain the shared list. Buffer reallocation and the new public method clear both signatures.
Sampling buffer update decisions tensorrt_llm/_torch/speculative/interface.py	Sampling signatures include normalized values, request order, and token-count layout. Population refills only changed buffer groups and refreshes top_k_max with the per-request group.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: lori-ren, allisonlim-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the ticket, performance focus, and primary change: skipping redundant one-model sampling-parameter refills.
Description check	✅ Passed	The description explains the problem and solution, documents test coverage, and includes the required checklist with relevant items reviewed.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 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 1105-1144: Remove slot_ids from
_sampling_params_buffers_need_update signatures: use values alone for the
per-request signature and (values, num_tokens) for the expanded-buffer
signature. Keep the existing comparison, recording, and return behavior
unchanged so slot changes no longer trigger sampling-buffer updates.

🪄 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: 9244481a-efe2-46c2-a3f8-d930b62e514d

📥 Commits

Reviewing files that changed from the base of the PR and between 07b3e82 and f20561a.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/speculative/interface.py

[zhaoyangwang-nvidia]

/bot run --disable-fail-fast

[zhaoyangwang-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65554 [ run ] triggered by Bot. Commit: af6f657 Link to invocation

[tensorrt-cicd]

PR_Github #65554 [ run ] completed with state FAILURE. Commit: af6f657
/LLM/main/L0_MergeRequest_PR pipeline #53291 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[zhaoyangwang-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65741 [ run ] triggered by Bot. Commit: af6f657 Link to invocation

[tensorrt-cicd]

PR_Github #65741 [ run ] completed with state SUCCESS. Commit: af6f657
/LLM/main/L0_MergeRequest_PR pipeline #53457 completed with status: 'SUCCESS'

CI Report

Link to invocation