[https://nvbugs/5615248][fix] Beam history copies only on terminal steps#13799
Conversation
PR NVIDIA#13799 (commits 'Defer per-step beam-history D2H to terminal step' and 'Cast seq_slots/seq_lens to CUDA exactly once per step') introduced synchronous .cpu() calls inside _prepare_beam_history's _builder closure on terminal steps, which blocks the host thread and breaks the overlap scheduler's CPU-GPU pipelining. This commit restores overlap by: * Routing the deferred terminal-step D2H copies through a dedicated side stream so they no longer block the main stream. * Recording the side-stream completion via a SamplerEvent secondary event so the main stream is not gated on it. * Batching all per-request side-stream copies inside a single 'with torch.cuda.stream(side)' context to amortize Python overhead. * Adding a host-side predictor (_predict_beam_history_terminal) that skips D2H copies entirely on definitely-non-terminal steps based on py_max_new_tokens / max_seq_len / py_stop_words / lagged first_finish_reasons. The predictor preserves correctness via a fallback synchronous .cpu() path on predictor misses. Robust pooled per-request results (5 runs x 16 requests, Llama-3.1 8B, beam_width=4, ISL 100/OSL 20): * Median E2E: -0.32 to -0.47 ms (-0.4 to -0.55 percent) across two independent benchmark sessions. * Median ITL: -0.0015 to -0.0017 ms (~ -0.4 percent). * Mann-Whitney U two-sided p < 1e-11 in both runs. * TTFT: no significant effect (correctly so; predictor only affects the decode loop's per-step terminal check). Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
PR NVIDIA#13799 (commits 'Defer per-step beam-history D2H to terminal step' and 'Cast seq_slots/seq_lens to CUDA exactly once per step') introduced synchronous .cpu() calls inside _prepare_beam_history's _builder closure on terminal steps, which blocks the host thread and breaks the overlap scheduler's CPU-GPU pipelining. This commit restores overlap by: * Routing the deferred terminal-step D2H copies through a dedicated side stream so they no longer block the main stream. * Recording the side-stream completion via a SamplerEvent secondary event so the main stream is not gated on it. * Batching all per-request side-stream copies inside a single 'with torch.cuda.stream(side)' context to amortize Python overhead. * Adding a host-side predictor (_predict_beam_history_terminal) that skips D2H copies entirely on definitely-non-terminal steps based on py_max_new_tokens / max_seq_len / py_stop_words / lagged first_finish_reasons. The predictor preserves correctness via a fallback synchronous .cpu() path on predictor misses. Robust pooled per-request results (5 runs x 16 requests, Llama-3.1 8B, beam_width=4, ISL 100/OSL 20): * Median E2E: -0.32 to -0.47 ms (-0.4 to -0.55 percent) across two independent benchmark sessions. * Median ITL: -0.0015 to -0.0017 ms (~ -0.4 percent). * Mann-Whitney U two-sided p < 1e-11 in both runs. * TTFT: no significant effect (correctly so; predictor only affects the decode loop's per-step terminal check). Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
09f8feb to
8999b4b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #47463 [ run ] triggered by Bot. Commit: |
|
PR_Github #47463 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #47538 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR optimizes beam-search history finalization by introducing side-stream-based asynchronous device-to-host transfers that avoid synchronization on non-terminal steps. The core insight is predicting terminal steps using a lagged snapshot of finish reasons, then selectively issuing non-blocking D2H copies only when needed. CUDA tensor casting is consolidated to reduce redundant transfers. ChangesSide-Stream Async Beam-History D2H Optimization
🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/sampler.py (2)
3514-3518: ⚡ Quick winReset the lagged finish-reason cache on slot reuse.
_predict_beam_history_terminal()keys this state byseq_slot, but the snapshot written here is never cleared when a slot is reassigned. A recycled slot can inherit a finished beam from the previous request and force eager beam-history D2H until the new request reaches its firstupdate_requests(). Clearing the entry during new-request setup, or when the request completes, keeps the predictor aligned with the current occupant of the slot.🤖 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.py` around lines 3514 - 3518, The snapshot cache _prev_first_finish_reasons_host is never cleared when a seq_slot is reused, so update the new-request setup (the code path that assigns req.py_seq_slot and/or the request completion path) to reset the entry for that slot (indexed by req.py_seq_slot) before the new request begins or when a request finishes; this keeps _predict_beam_history_terminal keyed by seq_slot aligned with the current occupant and prevents inheriting stale finish reasons—clear the entry (e.g., set to None/zeros/empty array) in the code that assigns req.py_seq_slot or in the request teardown/complete routine and ensure update_requests() semantics remain correct.
3428-3429: ⚡ Quick winKeep the predictor enabled for single-token stop words.
This branch disables the optimization for every request with stop words, even though
_prepare_beam_history()only forces unconditional finalization for multi-token stop words via_check_stop_words_length(). Using that narrower check here preserves correctness via the.cpu()miss fallback and avoids per-step eager copies for the single-token case.🤖 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.py` around lines 3428 - 3429, The current branch unconditionally disables the predictor whenever request.py_stop_words_list is present, which is too broad; change the condition so the predictor is only disabled for multi-token stop words (use the same logic as _prepare_beam_history() via _check_stop_words_length(request.py_stop_words_list) or equivalent) so single-token stop words still use the predictor and rely on the .cpu() miss fallback instead of forcing eager copies.
🤖 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.py`:
- Around line 3514-3518: The snapshot cache _prev_first_finish_reasons_host is
never cleared when a seq_slot is reused, so update the new-request setup (the
code path that assigns req.py_seq_slot and/or the request completion path) to
reset the entry for that slot (indexed by req.py_seq_slot) before the new
request begins or when a request finishes; this keeps
_predict_beam_history_terminal keyed by seq_slot aligned with the current
occupant and prevents inheriting stale finish reasons—clear the entry (e.g., set
to None/zeros/empty array) in the code that assigns req.py_seq_slot or in the
request teardown/complete routine and ensure update_requests() semantics remain
correct.
- Around line 3428-3429: The current branch unconditionally disables the
predictor whenever request.py_stop_words_list is present, which is too broad;
change the condition so the predictor is only disabled for multi-token stop
words (use the same logic as _prepare_beam_history() via
_check_stop_words_length(request.py_stop_words_list) or equivalent) so
single-token stop words still use the predictor and rely on the .cpu() miss
fallback instead of forcing eager copies.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a2bdc6b2-42a6-4935-9a98-7c6b54518700
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/sampler.py
|
PR_Github #47538 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #47542 [ run ] triggered by Bot. Commit: |
|
PR_Github #47542 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #47619 [ run ] triggered by Bot. Commit: |
|
PR_Github #47619 [ run ] completed with state |
ixlmar
left a comment
There was a problem hiding this comment.
The changes make sense, but we need to figure out how to keep the overall sampler code maintainable and how to not regress any existing features.
Since the side-stream D2H copies are supposed to overlap with compute, do we even need the speculative D2H-optout with fallback to .cpu() and resulting potential syncs that need to be covered in unit tests and may alter performance characteristics unexpectedly (for big enough models the risk of breaking overlap scheduling by .cpu() may outweigh any perf benefit)?
Generally, this is a form of speculative execution with the usual performance implications. Although I agree that the .cpu() branch would appear to be hit only rarely, this might be unexpected to users running benchmarks etc. and thus perhaps it would be better to gate this optimization by some configuration option (having "spec" in its name)?
PR NVIDIA#13799 (commits 'Defer per-step beam-history D2H to terminal step' and 'Cast seq_slots/seq_lens to CUDA exactly once per step') introduced synchronous .cpu() calls inside _prepare_beam_history's _builder closure on terminal steps, which blocks the host thread and breaks the overlap scheduler's CPU-GPU pipelining. This commit restores overlap by: * Routing the deferred terminal-step D2H copies through a dedicated side stream so they no longer block the main stream. * Recording the side-stream completion via a SamplerEvent secondary event so the main stream is not gated on it. * Batching all per-request side-stream copies inside a single 'with torch.cuda.stream(side)' context to amortize Python overhead. * Adding a host-side predictor (_predict_beam_history_terminal) that skips D2H copies entirely on definitely-non-terminal steps based on py_max_new_tokens / max_seq_len / py_stop_words / lagged first_finish_reasons. The predictor preserves correctness via a fallback synchronous .cpu() path on predictor misses. Robust pooled per-request results (5 runs x 16 requests, Llama-3.1 8B, beam_width=4, ISL 100/OSL 20): * Median E2E: -0.32 to -0.47 ms (-0.4 to -0.55 percent) across two independent benchmark sessions. * Median ITL: -0.0015 to -0.0017 ms (~ -0.4 percent). * Mann-Whitney U two-sided p < 1e-11 in both runs. * TTFT: no significant effect (correctly so; predictor only affects the decode loop's per-step terminal check). Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>
b85c8cf to
0b44394
Compare
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
74da52a to
1216a56
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48634 [ run ] triggered by Bot. Commit: |
ixlmar
left a comment
There was a problem hiding this comment.
Had a superficial look at the tests as well (looking good)
|
/bot run --disable-fail-fast |
|
PR_Github #48664 [ run ] triggered by Bot. Commit: |
|
PR_Github #48634 [ run ] completed with state |
|
PR_Github #48664 [ run ] completed with state
|
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
ebdd632 to
cf4d7cd
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48705 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #48706 [ run ] triggered by Bot. Commit: |
|
PR_Github #48705 [ run ] completed with state |
|
PR_Github #48706 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #48784 [ run ] triggered by Bot. Commit: |
|
PR_Github #48784 [ run ] completed with state |
…eps (NVIDIA#13799) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Description
This addresses https://nvbugspro.nvidia.com/bug/5615248. Two host-side optimizations for the beam-search decode path:
_prepare_beam_history- only copy beam history on terminal steps.SamplerEventand awaited host-side inupdate_requests..cpu()so correctness is preserved for any unforeseen finish pattern.seq_slots/seq_lensH2D - cast once per step at the top of_process_requestsand reused. Eliminates 2 redundant H2D launches/step on the typical concurrency=1 / single-strategy beam-search path. Multi-strategy keeps the existing per-group fall-back.Workload: TinyLlama-1.1B-Chat-v1.0 on L40S, ISL=100 / OSL=20, beam_width=10,
max_batch_size=1, piecewise CUDA graphs, 5 runs × 16 requests pooled per side (n = 80). Sign convention: negative Δ = improvement.Reproduction
A companion validation branch carries the bench harness used to produce the table above (kept off this MR to keep the diff minimal):
user/brb/nvbug-5615248-beam-host-overheadnvbugs_5615248/trtllm_bench/REPRO_V6_V7.mdTest Coverage
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)
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.Summary by CodeRabbit