[None][perf] Add batched-pybind fast-path in TorchSampler.update_requests for gpt-oss-120b

[chenfeiz0326]

Summary

• Root cause: On disagg-gen_only-* workloads for gpt-oss-120b-fp4 on GB200,
  the gen-worker per-iter bubble (~8.8 ms) is dominated by the host-side sampler
  tail in TorchSampler.update_requests — a Python loop that fires 256 per-request
  pybind chatter calls (req.add_new_token etc.) after state.sampler_event.synchronize().
• Fix: Add a batched-pybind fast-path at sampler.py:3654. When the whole
  batch requires no beam search, no logprobs, no draft tokens, no stop-word
  lists, and no speculative-tree manager, collapse the per-request loop into a
  single batched add_new_tokens_to_requests(...) call plus one linear finish-reason
  scan. Cached eligibility flag (_batch_fastpath_eligible) is set once in
  __init__ to avoid a property call per iter. Falls through to the unchanged
  per-request path on any invariant break; runs entirely post-sampler_event.synchronize()
  so CUDA-graph capture/replay is unaffected.
• Perf metric: output_token_throughput (higher is better)
• Baseline (reproduce): 8934.14 tok/s
• Fix median (3-rep same-node verify): 9135.99 tok/s — [9002.96, 9135.99, 9318.70], CV=1.73%
• Gain vs baseline: +2.26%; gain vs verify-side ToT: +3.83%

Test plan

• Reproduce baseline on Lyris GB200 (3 cross-node reps; all within 0.5% of median).
• Apply fix, rebuild TRT-LLM wheel, run 3-rep same-node perf.
• Confirm gain ≥ 2% and same-node CV ≤ 5%.
• Programmatic anti-pattern scan clean (no debug artifacts, no test-infra edits,
  no pytest.skip/time.sleep, no CUDA-graph disable, no except: fallback).
• Correctness argument: finish-reason filter matches finish_if_reason's
  {END_ID, LENGTH, STOP_WORDS} set; bookkeeping (py_num_accepted_draft_tokens=0,
  py_rewind_len=0, py_decoding_iter += 1) mirrors the existing MTP=0 branch.
• add_new_tokens_to_requests batched pybind is already used by TRTLLMSampler
  (see sampler.py:5125); no new C++ binding introduced.

Automation

Automated fix generated by repair-bot auto_perf_stabilize workflow.
Run key: perfstab-perf-test-perf-sanity-py-test-e2-c87588
Test case: perf/test_perf_sanity.py::test_e2e[disagg-gen_only-gb200_gpt-oss-120b-fp4_8k1k_con512_ctx1_tp1_gen1_dep2_eplb0_mtp0_ccb-NIXL]

Summary by CodeRabbit

• Performance
  • Improved generation efficiency for eligible decoding batches by processing sampled tokens in a single batched operation.
  • Reduced per-request processing overhead during decoding, which may improve throughput and response speed.
  • Preserved handling for supported completion conditions, including end-of-sequence, length limits, and stop-word termination.

[coderabbitai]

📝 Walkthrough

Walkthrough

TorchSampler caches single-beam eligibility and adds an early greedy batch-update path that processes eligible requests together, applies selected finish reasons, resets draft-tracking state, and advances the decoding iteration.

Changes

Greedy sampler batch fast path

Layer / File(s)	Summary
Cached eligibility and batched greedy updates tensorrt_llm/_torch/pyexecutor/sampler/sampler.py	TorchSampler caches max_beam_width == 1 eligibility and, when batch invariants hold, batch-adds sampled tokens, applies END_ID, LENGTH, and STOP_WORDS finish reasons, resets draft-token state, increments py_decoding_iter, and skips per-request processing.

Estimated code review effort: 3 (Moderate) | ~15–30 minutes

Suggested reviewers: zhaoyangwang-nvidia

🚥 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 is concise and accurately describes the main performance fast-path change.
Description check	✅ Passed	The description covers the problem, fix, performance results, and test plan, with only minor template sections omitted.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[tongyuantongyu]

LGTM. I updated the description a bit: Non beam search does not mean greedy only, sampling such as top-k is also eligible.

[chenfeiz0326]

Empirical A/B across gpt-oss GB200 cases (fix wheel vs origin/main ToT wheel)

Lyris GB200. Same two wheels reused across every rep — trtllm-682c38d7c (ToT at PR filing time) vs same commit + this PR's diff applied and built. Comparison uses mean-of-N, Welch t-statistic. |t|>2.78 at N=3 → p<0.05.

Case	Mode	Conc	N	ToT mean	Fix mean	Gain%	CV_tot	CV_fix	t-stat	Verdict
gpt-oss ctx_only con=4	ctx_only	4	0/0	—	—	—	—	—	—	incomplete
gpt-oss ctx_only con=512	ctx_only	512	0/0	—	—	—	—	—	—	incomplete
gpt-oss e2e con=1024	e2e	1024	3/3	17004.44	17107.86	+0.61%	0.99%	1.97%	+0.48 ×	in_noise
gpt-oss gen_only con=1024	gen_only	1024	6/6	13187.43	13055.73	-1.00%	1.36%	1.12%	-1.39 ×	in_noise
gpt-oss gen_only con=128	gen_only	128	3/3	7509.22	7585.61	+1.02%	0.41%	0.65%	+2.28 ×	marginal_gain_noisy
gpt-oss gen_only con=4	gen_only	4	3/3	901.68	927.56	+2.87%	3.40%	3.00%	+1.08 ×	success_noisy
PENALTY deepseek-v32 MTP1 con=2048	gen_only	2048	3/3	27260.28	27341.83	+0.30%	1.22%	0.99%	+0.33 ×	in_noise

(Skipped: gen_only con=512 dep2 — already measured at +2.26% in the original PR body.)

Key takeaways

• No statistically significant regression on any gpt-oss case, including gen_only con=1024 (which came in -1% but t=-1.39 at N=6 fails p<0.05).
• gen_only con=128 shows the cleanest signal: +1.02% with tight CV (0.41% / 0.65%) and t=+2.28 (borderline significant at N=3).
• e2e con=1024 and gen_only con=4 trend positive but within noise. con=4 at CV 3% is expected — batch=4 sampler tail is only tens of μs, so the win is small in absolute terms and drowned by iter-time variance.
• ctx_only cases are unmeasurable via output_token_throughput — the harness runs prefill only, no gen loop → no bench log → no throughput metric. The fast-path is theoretically ≤ 0.3% here (sampler fires once per request for the first-token emit), well below the noise floor.
• Penalty probe (deepseek-v32 MTP1 con=2048): +0.30%, t=+0.33. The eligibility check runs and rejects at check 3 (get_spec_tree_manager() != None) on every gen iter — no detectable cost. See threaded reply for the detailed breakdown.

Method

For every (case, side) pair I:

1. Provisioned a per-rep dir on lyris (no rebuild — same two wheels).
2. Ran a parametrized slurm-disaggregated-ci-pytest/trigger.sh with --wheel-dir <tot|fix> --run-count 1.
3. Grepped Output token throughput (tok/s) from disagg-*/trtllm-benchmark.0.0.log.
4. Computed mean, stdev, CV, and Welch t on the resulting per-rep lists.

Raw per-rep numbers and the launcher/parser scripts are in out/gpt-oss-120b-seeds/comparison/ in my repair-bot checkout.

[chenfeiz0326]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60602 [ run ] triggered by Bot. Commit: d08c23b Link to invocation

[tensorrt-cicd]

PR_Github #60602 [ run ] completed with state SUCCESS. Commit: d08c23b
/LLM/main/L0_MergeRequest_PR pipeline #48910 completed with status: 'FAILURE'

CI Report

⚠️ 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

CI Agent Failure Analysis

Link to invocation

[chenfeiz0326]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60681 [ run ] triggered by Bot. Commit: d08c23b Link to invocation

[tensorrt-cicd]

PR_Github #60681 [ run ] completed with state FAILURE. Commit: d08c23b
/LLM/main/L0_MergeRequest_PR pipeline #48976 completed with status: 'FAILURE'

CI Report

⚠️ 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

CI Agent Failure Analysis

Link to invocation

[chenfeiz0326]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60856 [ run ] triggered by Bot. Commit: d08c23b Link to invocation

[tensorrt-cicd]

PR_Github #60856 [ run ] completed with state SUCCESS. Commit: d08c23b
/LLM/main/L0_MergeRequest_PR pipeline #49128 completed with status: 'SUCCESS'

CI Report

Link to invocation