[https://nvbugs/6400533][fix] Cache per-request greedy flag to fast-exit spec-dec sampling scan

[trtllm-agent]

Summary

• Root cause: The one-model speculative-decoding sampler ran a per-step Python-level normalization of temperature/top_k/top_p across the entire batch inside _scan_one_model_sampling, even when every request was greedy. On gen-only workloads (e.g., DeepSeek-R1 FP4 with MTP3, con=1024), this per-step per-request work became a measurable host overhead and regressed generation throughput because the downstream populate_sampling_params_for_one_model fast-exits on all-greedy anyway, making the normalization wasted work.
• Fix: Cache an immutable py_is_greedy_sample bool on LlmRequest at admission time (sampling params don't change afterward), then add a hot-path fast-exit in SpecMetadata that does an O(batch_size) all(...) over cached bools and short-circuits with the greedy flags set, skipping the normalized-tuple construction and per-token H2D copies. The capture-time _force_non_greedy_for_capture path and the rejection-sampling slot-id list are preserved so CUDA-graph warmup and rejection-sampling configs remain correct.
• Perf metric: mean_gen_worker_per_iter_device_step_time (lower-is-better)
• Bad perf: 32.56
• Good perf: 31.33
• Perf after fix: 31.85
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6400533

Summary by CodeRabbit

• Performance Improvements
  • Improved generation efficiency for requests using greedy sampling.
  • Added faster handling for batches where all requests use greedy sampling, reducing unnecessary sampling-parameter processing.
  • Preserved compatibility with rejection sampling and capture configurations.

[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: 81d418b3-a79a-43ca-87b0-8bd95ad79a49

📥 Commits

Reviewing files that changed from the base of the PR and between e2ce7a6 and 5ecee80.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/llm_request.py
• tensorrt_llm/_torch/speculative/interface.py

---

📝 Walkthrough

Walkthrough

Caches whether each request uses greedy sampling and adds an all-greedy batch fast path in speculative sampling metadata scanning, including handling for rejection sampling and forced non-greedy capture.

Changes

Greedy sampling optimization

Layer / File(s)	Summary
Cache and apply greedy sampling state tensorrt_llm/_torch/pyexecutor/llm_request.py, tensorrt_llm/_torch/speculative/interface.py	LlmRequest derives py_is_greedy_sample from sampling parameters. _scan_one_model_sampling skips per-request normalization for all-greedy batches, preserves rejection-sampling slot IDs, and bypasses the fast path when non-greedy capture is forced.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, matches the performance fix, and follows the required ticket/type format.
Description check	✅ Passed	The description covers the issue, solution, test plan, and bug link, so it is mostly complete despite using different headings.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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