[https://nvbugs/6571220][fix] Correct unfused attention context workspace sizing - #17026
[https://nvbugs/6571220][fix] Correct unfused attention context workspace sizing#17026pranav-nvidia wants to merge 5 commits into
Conversation
…KV length Runner::getWorkspaceSize passed cross_kv_length=0, so the unfused-path buffers were sized at zero while enqueueContext carved them at the real length. The carved views overrun the allocation, surfacing as CUBLAS_STATUS_EXECUTION_FAILED at the QK^T GEMM. FP32 enc-dec hits it, since FP32 takes the unfused path. Pass the context sequences' max past-KV length through the size query. Every changed buffer is !mEnableContextFMHA-gated, so fused FP16/BF16 is unchanged. Also drops the model_engine guard that disabled fp32 enc-dec CUDA graphs pending this fix, and flips the Whisper case that pinned it. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
…ce-v2 Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
/bot run --disable-fail-fast |
WalkthroughThe change corrects unfused attention workspace sizing, propagates maximum cross-KV length for cross-attention, removes the FP32 encoder-decoder CUDA graph guard, and updates T5 and Whisper integration coverage. ChangesAttention workspace and FP32 graph execution
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: 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.
🧹 Nitpick comments (1)
cpp/tensorrt_llm/thop/attentionOp.cpp (1)
348-349: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse lower camelCase for the new workspace-length identifier.
cpp/tensorrt_llm/thop/attentionOp.cpp#L348-L349: Rename the virtual parameter tomaxCrossKvLength.cpp/tensorrt_llm/thop/attentionOp.cpp#L414-L418: Rename the override parameter and forwarded argument.cpp/tensorrt_llm/thop/attentionOp.cpp#L1369-L1375: Rename the local variable and call-site argument.As per coding guidelines, use lowercase camelCase for C++ parameters and locals.
🤖 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 348 - 349, Rename the workspace-length identifier from max_cross_kv_length to maxCrossKvLength throughout cpp/tensorrt_llm/thop/attentionOp.cpp: update the virtual parameter at lines 348-349, the override parameter and forwarded argument at lines 414-418, and the local variable and call-site argument at lines 1369-1375. Preserve the existing behavior and all argument forwarding.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 348-349: Rename the workspace-length identifier from
max_cross_kv_length to maxCrossKvLength throughout
cpp/tensorrt_llm/thop/attentionOp.cpp: update the virtual parameter at lines
348-349, the override parameter and forwarded argument at lines 414-418, and the
local variable and call-site argument at lines 1369-1375. Preserve the existing
behavior and all argument forwarding.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dfbd59cf-1ca7-45ee-bd67-64cb3181348c
📒 Files selected for processing (6)
cpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/thop/attentionOp.cpptensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/defs/llmapi/test_llm_api_pytorch_t5.pytests/integration/defs/llmapi/test_llm_api_pytorch_whisper.pytests/integration/test_lists/test-db/l0_h100.yml
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/pyexecutor/model_engine.py
|
PR_Github #64109 [ run ] triggered by Bot. Commit: |
|
PR_Github #64109 [ run ] completed with state
|
|
Verified the root cause independently — the size query and the
while One thing I checked so nobody else has to: Not approving yet, and only for want of evidence: the run at |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The sizing changes hold up — I walked every term in getWorkspaceSizeForContext against the carve in enqueueContext (attentionOp.cpp:1504-1600) and after this PR the remaining differences all over-allocate (max_num_tokens ≥ context tokens, mMaxContextLength ≥ max_context_q_len, max_num_requests ≥ num_contexts). Note the quadratic terms qk_buf/qk_buf_float were already sized by batch_size * input_seq_length * kv_seq_length, so the three buffers moved off max_num_tokens here don't change the allocation's order of magnitude.
Two notes on the write-up:
- The title says "cross-KV length", but three of the four
common/attentionOp.cppchanges (attention_mask,qkv_buf_2,padding_offset) are a separate packed-vs-dense fix that applies to self-attention with padding removal too. The description covers it under "align the remaining packed-token versus dense-layout buffer sizes", but the title undersells the blast radius. [None]on a fix for a memory under-allocation reaching cuBLAS. If there's an NVBug behind this, please tag it; if it was found by inspection while removing the #16141 workaround, saying so in the description is enough.
Also worth confirming in CI that the newly enabled FP32 enc-dec graph path is exercised with an encoder length longer than anything seen during warmup capture, since the workspace size now varies with the batch's encoder length.
| // The unfused-MHA buffers below must upper-bound the enqueueContext carve, which sizes them by | ||
| // batch_size * input_seq_length (not num_tokens): with padding removal the actual token count can be | ||
| // smaller than batch_size * max(context q length), so sizing by max_num_tokens underestimates. | ||
| size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_length * kv_seq_length; |
There was a problem hiding this comment.
This is the right fix, but it patches one of two independently written copies of the same arithmetic — enqueueContext recomputes all ~25 of these sizes at line 1504 onward, and both then fill an AttentionContextWorkspaceSizes and call buildContextLayout. The bug being fixed here is that drift, and nothing prevents the next one: no assert compares the planned size against the carved layout.
Since both sides already funnel through the same struct, a follow-up that extracts one computeContextWorkspaceSizes(batch, input_seq_length, cross_kv_length, num_tokens, total_kv_len, ...) called from both — the size query passing engine maxima, the enqueue passing runtime dims — would make this class of bug structurally impossible. Failing that, a TLLM_CHECK(carvedTotal <= plannedTotal) in enqueueContext would at least turn a silent OOB into a clear error instead of CUBLAS_STATUS_EXECUTION_FAILED.
There was a problem hiding this comment.
Agreed on the shared computeContextWorkspaceSizes. The TLLM_CHECK needs a new EnqueueParams field too. It currently carries void* workspace (attentionOp.h:119) with no allocated byte count to compare against. I will take both as a follow-up PR rather than widening this one.
| int32_t max_cross_kv_length = 0; | ||
| if (op->isCrossAttention() && num_contexts > 0) | ||
| { | ||
| max_cross_kv_length = host_past_key_value_lengths.slice(0, 0, num_contexts).max().item<int32_t>(); |
There was a problem hiding this comment.
This is character-for-character the same expression as line 901, which sets enqueue_params.cross_kv_length for the context run (same tensor, same slice(0, 0, num_contexts)). Keeping two copies in sync by convention is what produced the bug this PR fixes; consider hoisting it to a single const above the getWorkspaceSize call and passing it into run() (or at least a cross-reference comment on both).
Minor gating difference worth a note in the comment: line 901 additionally requires cross_kv.has_value(), so an isCrossAttention() op called without a cross_kv tensor gets a nonzero size here and a zero at enqueue. That direction is safe (over-allocates), but a reader checking the two for equality will trip on it.
There was a problem hiding this comment.
Kept the two computations but documented them in both directions, and called out the cross_kv.has_value() asymmetry as intentional safe over-allocation. I did try hoisting it and passing it into run(). It works and is strictly better, but it adds a parameter to that signature and reflows both call sites, which felt like scope creep here. Happy to fold it into the refactor follow-up.
cascade812
left a comment
cascade812
left a comment
There was a problem hiding this comment.
LGTM, thanks for fixing this!
…kspace sizing Rename the new workspace-length parameter to lower camelCase per the C++ coding guidelines, and cross-reference the two sites that derive the context-stage cross-KV length so the size query and the enqueue carve stay in agreement. Assert the T5 mixed-encoder-length cases actually capture decoder CUDA graphs, which needs the single-process worker to reach the engine. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
pranav-nvidia
changed the title
|
/bot run --disable-fail-fast |
|
Removed the "ci: full pre-merge approved" label because @pranav-nvidia could not be verified as an active member of NVIDIA/trt-llm-ci-approvers. Ask a member of that team to apply it. |
|
PR_Github #64663 [ run ] triggered by Bot. Commit: |
|
PR_Github #64663 [ run ] completed with state
|
|
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. |
|
/bot run --disable-fail-fast |
|
PR_Github #64718 [ run ] triggered by Bot. Commit: |
|
PR_Github #64718 [ run ] completed with state |
…ce-v2 Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
/bot run |
|
PR_Github #65072 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tensorrt_llm/_torch/pyexecutor/model_engine.py (3)
3824-3833: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRelease the pinned encoder staging pool during cleanup.
These lines allocate large pinned tensors and retain them in
_encoder_decoder_host_buffer_pool. The existingcleanup()path does not clear this pool. If the engine is cleaned up while its object remains alive, the pinned memory remains allocated and can accumulate across reloads. Release the buffers after their pending CUDA events complete.Proposed cleanup change
+ for buffers in getattr(self, "_encoder_decoder_host_buffer_pool", []): + event = buffers["event"] + if event is not None: + event.synchronize() + self._encoder_decoder_host_buffer_pool = NoneAlso applies to: 3835-3871
🤖 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/model_engine.py` around lines 3824 - 3833, Update cleanup() to release _encoder_decoder_host_buffer_pool after pending CUDA events complete. Ensure each pooled buffer’s tensors and associated event references are cleared or discarded only once their events have completed, then empty the pool so pinned encoder staging memory is not retained across reloads.
3685-3686: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winReplace runtime
assertchecks with explicit exceptions.Raise
TypeErrorwhencross_attn_metadatais notTrtllmAttentionMetadata. RaiseValueErrorwhennum_tokensexceedsencoder_max_num_tokens. Python omits both checks under optimization.🤖 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/model_engine.py` around lines 3685 - 3686, In the code around prepare_encoder_decoder_from_precomputed_lengths, replace the runtime assert on cross_attn_metadata with an explicit TypeError when it is not a TrtllmAttentionMetadata instance, and add an explicit ValueError when num_tokens exceeds encoder_max_num_tokens. Preserve the existing valid-input execution path without relying on assert-based validation.Source: Coding guidelines
1495-1500: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winCatch only expected prewarm failures.
The handler also suppresses tensor-allocation failures and
ValueErrorcontract errors from the fused CuTe DSL path. Catch only expected runtime failures from_project_and_quantize_q; let setup and contract errors propagate.🤖 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/model_engine.py` around lines 1495 - 1500, Update the prewarm exception handling around _project_and_quantize_q to catch only the expected runtime/prewarm failure type, while allowing tensor-allocation failures and ValueError contract/setup errors to propagate. Preserve the existing warning message for the handled prewarm failure.Sources: Coding guidelines, Linters/SAST tools
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
2683-2683: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winComplete the new type annotations.
Line 2683 declares
max_seq_lenasintwhile assigningNone. Ruff reports RUF013. The new_prepare_encoder_decoder_inputs_fastfunction also has no return annotation. Use Python 3.10 union syntax and annotate the returned tuple.Proposed typing fix
- max_seq_len: int = None, + max_seq_len: int | None = None, ... - resource_manager: Optional[ResourceManager]): + resource_manager: Optional[ResourceManager] + ) -> tuple[dict[str, Any], torch.Tensor | None]:As per coding guidelines, annotate every function and prefer Python 3.10
|unions.Also applies to: 3875-3879
🤖 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/model_engine.py` at line 2683, Update _prepare_encoder_decoder_inputs_fast so max_seq_len uses a Python 3.10 nullable annotation (int | None) and add the function’s return annotation for its returned tuple. Also complete the missing annotations in the additional function range noted by the review, ensuring every function parameter and return type follows the project’s Python 3.10 union style.Sources: Coding guidelines, Linters/SAST tools
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 3824-3833: Update cleanup() to release
_encoder_decoder_host_buffer_pool after pending CUDA events complete. Ensure
each pooled buffer’s tensors and associated event references are cleared or
discarded only once their events have completed, then empty the pool so pinned
encoder staging memory is not retained across reloads.
- Around line 3685-3686: In the code around
prepare_encoder_decoder_from_precomputed_lengths, replace the runtime assert on
cross_attn_metadata with an explicit TypeError when it is not a
TrtllmAttentionMetadata instance, and add an explicit ValueError when num_tokens
exceeds encoder_max_num_tokens. Preserve the existing valid-input execution path
without relying on assert-based validation.
- Around line 1495-1500: Update the prewarm exception handling around
_project_and_quantize_q to catch only the expected runtime/prewarm failure type,
while allowing tensor-allocation failures and ValueError contract/setup errors
to propagate. Preserve the existing warning message for the handled prewarm
failure.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Line 2683: Update _prepare_encoder_decoder_inputs_fast so max_seq_len uses a
Python 3.10 nullable annotation (int | None) and add the function’s return
annotation for its returned tuple. Also complete the missing annotations in the
additional function range noted by the review, ensuring every function parameter
and return type follows the project’s Python 3.10 union style.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ab0e0c94-d86c-4a52-88fc-ad36fdbc84c0
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/defs/llmapi/test_llm_api_pytorch_t5.pytests/integration/test_lists/test-db/l0_h100.yml
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/integration/test_lists/test-db/l0_h100.yml
- tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py
|
PR_Github #65072 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65108 [ run ] triggered by Bot. Commit: |
|
PR_Github #65108 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65153 [ run ] triggered by Bot. Commit: |
|
PR_Github #65153 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65209 [ run ] triggered by Bot. Commit: |
|
PR_Github #65209 [ run ] completed with state
|
7 participants
Description
AttentionOp::getWorkspaceSizeForContext(the size query) andAttentionOp::enqueueContext(the carve) disagreed on four unfused-MHA context buffers, so the query could return fewer bytes than the carve consumes and the carved views ran past the end of the allocation. The shortfall grows with batch size, so it presented as a batch-size threshold rather than a universal failure.attentionMaskmax_num_tokens * kv_seq_lengthbatch_size * input_seq_length * kv_seq_lengthqkvBufmax_num_tokens * local_hidden_units_qobatch_size * input_seq_length * local_hidden_units_qopaddingOffsetsizeof(int) * max_num_tokenssizeof(int) * batch_size * input_seq_lengthencoderPaddingOffsetsizeof(int) * max_num_tokenssizeof(int) * batch_size * cross_kv_lengthWith padding removal the real token count is smaller than
batch_size * max(context q length), so sizing bymax_num_tokensunderestimates. Separately,cpp/tensorrt_llm/thop/attentionOp.cpppassed a literal0forcross_kv_length; for cross-attentionkv_seq_lengthiscross_kv_length, soattentionMask,kBuf,vBuf,qkBufandqkFloatBufwere queried at zero bytes while the carve sized them from the real encoder KV length.Every affected buffer is guarded by
mEnableContextFMHA ? 0 : <size>, so fused FP16/BF16 paths are unaffected. The models that reach it are those falling back to unfused MHA: T5/mT5 at any dtype (relative position embedding) and anything at FP32. BART and Whisper at FP16/BF16 take the fused path.Two measured exposures:
CUBLAS_STATUS_EXECUTION_FAILEDduring warmup capture, then an illegal memory access. This is what the temporary guard inmodel_engine.pyworked around; the guard is removed here.CUBLAS_STATUS_INTERNAL_ERROR. FP32 is not required to hit this.Latent since ~
v0.11.0, because nothing exercised the unfused encoder-decoder cross-attention path above batch 2.The fix passes the context sequences' maximum past-KV length into the size query as
cross_kv_lengthand aligns the remaining packed-token buffers with the dense layout the carve uses, so the query upper-bounds the carve.Note on scope: three of the four
common/attentionOp.cppchanges are a packed-vs-dense fix that also applies to self-attention with padding removal, not only to cross-attention.Test Coverage
test_llm_api_pytorch_t5.py— new FP32 mixed-encoder-length, batch-2, KV-manager-v1 CUDA graph case, added tol0_h100.yml. The mixed-batch cases now assert decoder CUDA graphs actually captured, so an engine that silently declines capture fails instead of passing on output checks alone.test_llm_api_pytorch_whisper.py— the FP32 feature-combination case now expects decoder graph capture rather than the engine declining graphs.Verified on SM120 with a single-variable A/B (same binary otherwise),
t5-smallmixed-encoder-length suite: the FP32 case fails without the fix and passes with it; the BF16 greedy, BF16 beam-2 and KV-v2 cases pass in both arms.PR Checklist
Dev Engineer Review
enqueueContextbuffer carving.ci: full pre-merge approvedlabel, and the later L0 pipeline still failed.QA Engineer Review
test_llm_api_pytorch_t5.py.test_llm_api_pytorch_whisper.py.tests/integration/test_lists/test-db/l0_h100.yml.