[https://nvbugs/6565412][fix] Size trtllm-gen and thop decode buffers for beam search - #17014
Conversation
…search The trtllm-gen decode kernel keeps one multi-CTA KV counter per attention head per decoder sequence, but the buffer was sized from the request count alone. Beam search expands each request into beam_width sequences, so the kernel rejects the buffer at warmup with TVM_FFI_CHECK(multi_ctas_kv_counter_size >= counter_bytes) The max(..., multi_processor_count) floor hides the missing factor at small batch, so this only reproduces past a threshold: with 6 decoder heads and 148 SMs, beam 2 passes up to max_batch_size 8 and fails from 16 up. Every in-tree beam-search test runs below that threshold, so CI does not cover it. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
…ounter Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (6)
🚧 Files skipped from review as they are similar to previous changes (6)
WalkthroughThe attention API now accepts a separate maximum sequence count. Runner caching, semaphore capacity, generation workspace sizing, and FlashInfer KV counter allocation use this value, with fallback to the maximum request count. Tests cover beam-expanded sizing. ChangesSequence-capacity attention sizing
Estimated code review effort: 3 (Moderate) | ~20 minutes 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
🤖 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 `@tests/unittest/_torch/attention/test_fmha_page_index.py`:
- Around line 36-52: Add a test covering
FlashInferTrtllmGenFmha.prepare_workspace that uses metadata with distinct
max_num_requests and beam_width values, then verifies the counter sizing
reflects their product by asserting the helper receives the expanded batch or
the allocated counter buffer has the corresponding size. Keep the existing
direct helper tests unchanged.
🪄 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: 8897a81e-6bf4-4dfe-aabc-04bb1b97f09f
📒 Files selected for processing (2)
tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytests/unittest/_torch/attention/test_fmha_page_index.py
|
PR_Github #64110 [ run ] triggered by Bot. Commit: |
pranav-nvidia
changed the title
|
PR_Github #64110 [ run ] completed with state
|
…es, including thop Take the bound from metadata.max_num_sequences instead of recomputing max_num_requests * beam_width: beam_width is rewritten per step and is 1 during context, so it cannot size a buffer that must outlive CUDA graph capture. The thop attention path had the same gap and now takes the same bound for its semaphore and generation workspace sizing. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
pranav-nvidia
changed the title
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/tensorrt_llm/thop/attentionOp.cpp (1)
338-343: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winInitialize
RunnerBasestate.
RunnerBase::max_num_sequencesis a new externally visible scalar with no constructor or in-class initializer. Initialize it together with the other runner capacity fields so a default-constructedRunnerBasecannot expose indeterminate state.As per coding guidelines, initialize externally visible class members in the constructor.
🤖 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 `@cpp/tensorrt_llm/thop/attentionOp.cpp` around lines 338 - 343, Initialize RunnerBase::max_num_sequences in the constructor alongside beam_width, max_num_requests, and attention_window_size, ensuring default-constructed RunnerBase instances expose a defined value.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.
Nitpick comments:
In `@cpp/tensorrt_llm/thop/attentionOp.cpp`:
- Around line 338-343: Initialize RunnerBase::max_num_sequences in the
constructor alongside beam_width, max_num_requests, and attention_window_size,
ensuring default-constructed RunnerBase instances expose a defined value.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9c0f0b25-c806-4177-badb-2fc3dc7c298c
📒 Files selected for processing (6)
cpp/tensorrt_llm/nanobind/thop/bindings.cppcpp/tensorrt_llm/thop/attentionOp.cppcpp/tensorrt_llm/thop/attentionOp.htensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytests/unittest/_torch/attention/test_fmha_page_index.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
- tests/unittest/_torch/attention/test_fmha_page_index.py
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #64643 [ run ] triggered by Bot. Commit: |
|
PR_Github #64643 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64714 [ run ] triggered by Bot. Commit: |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
PR_Github #64714 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64725 [ run ] triggered by Bot. Commit: |
|
PR_Github #64725 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64736 [ run ] triggered by Bot. Commit: |
|
PR_Github #64736 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65070 [ run ] triggered by Bot. Commit: |
|
PR_Github #65070 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65107 [ run ] triggered by Bot. Commit: |
|
PR_Github #65107 [ run ] completed with state |
4 participants
Description
The trtllm-gen decode kernel keeps one multi-CTA KV counter per attention head per decoder sequence, and the generation workspace is sized per decoder sequence too. Both the flashinfer sizing helper and the C++ thop attention path sized these from
max_num_requestsalone, omitting the beam-width expansion, so beam search rejects the counter buffer at warmup with:The
max(..., multi_processor_count)floor hides the missing factor at small batch, so the failure appears as a threshold rather than a universal beam-search break. With 6 decoder heads on a 148-SM part at beam 2:max_batch_sizemax(6B, 148)6·B·2Not model-specific: any model reaching
num_heads · batch · beam_width > SM countunder beam search hits it. Present since the function was introduced in487330e8a03(2026-06-09); ships in v1.3.0rc19–rc22.Fix
Both paths size from
metadata.max_num_sequences, which the model engine sets tobatch_size * max_beam_width. That is a construction-time capacity, unlikemetadata.beam_width, which the model engine rewrites per step and holds at 1 during context — so it stays valid for a buffer that cannot grow under CUDA graph capture.thop.attentiongains an optionalmax_num_sequences, applied toreserveSemaphoreArrayandgetWorkspaceSizeForGeneration— whose first parameter is already namedmax_num_seqand expands tobatch_beam. Context sizing stays onmax_num_requests: beams share one prompt and fan out only at the first decode step. Callers that leavemax_num_sequencesunset run at beam width 1, so both paths fall back tomax_num_requests.Sizing up is safe — the kernel check is
>=, and for cross-attention (K/V not beam-packed) the only cost is4 · num_heads · batch · (beam − 1)bytes.Test Coverage
tests/unittest/_torch/attention/test_fmha_page_index.py— the old sizing under-allocates for a beam-expanded batch above the SM floor, the multi-processor floor is preserved at batch 1, andprepare_workspacepasses the beam-expanded bound to the sizing helper.max_batch_size16 and 32 go fail → pass, 8 stays passing. This covered the flashinfer counter sizing; the thop change landed after that run and is not covered by it, so it rests on CI.Why CI missed the original bug: every in-tree beam test sits in the masked region — the disagg beam test uses
max_beam_width=2at low concurrency,test_cli_flowbeam tests target the legacy CLI, and the Whisperbeam_searchcase usesmax_batch_size=2.PR Checklist
[JIRA/NVBUG/None][type] summaryDev Engineer Review
metadata.max_num_sequences.metadata.max_num_requests.attentionAPIs to accept optionalmax_num_sequences.metadata.max_num_sequences.QA Engineer Review
prepare_workspacepasses the beam-expanded maximum sequence count.tests/integration/test_lists/based on the provided changes.