[None][fix] Fix max_context_length value for attention workspace sizing#15156
Conversation
|
/bot run |
📝 WalkthroughWalkthroughThis PR removes an off-by-one subtraction from context length calculations in the TensorRT-LLM attention backend. The ChangesContext length off-by-one removal
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 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)
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)
tensorrt_llm/_torch/attention_backend/trtllm_gen.py (1)
741-742: ⚡ Quick winUse
metadata.max_context_lengthto avoid contract drift.Line 741 duplicates the metadata contract already defined in
TrtllmAttentionMetadata.max_context_length. Reusing the property here keeps both paths locked to one source of truth.♻️ Proposed refactor
- max_context_length = min(metadata.max_seq_len, metadata.max_num_tokens) + max_context_length = metadata.max_context_length🤖 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/attention_backend/trtllm_gen.py` around lines 741 - 742, Replace the ad-hoc computation of context length with the metadata property to avoid contract drift: set max_context_length from metadata.max_context_length (instead of min(metadata.max_seq_len, metadata.max_num_tokens)) and use metadata.max_context_length as the fallback for attention_window_size when forward_args.attention_window_size is falsy (i.e., change attention_window_size to use forward_args.attention_window_size or metadata.max_context_length); update references in this block (max_context_length and attention_window_size) to rely on TrtllmAttentionMetadata.max_context_length as the single source of truth.
🤖 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/attention_backend/trtllm_gen.py`:
- Around line 741-742: Replace the ad-hoc computation of context length with the
metadata property to avoid contract drift: set max_context_length from
metadata.max_context_length (instead of min(metadata.max_seq_len,
metadata.max_num_tokens)) and use metadata.max_context_length as the fallback
for attention_window_size when forward_args.attention_window_size is falsy
(i.e., change attention_window_size to use forward_args.attention_window_size or
metadata.max_context_length); update references in this block
(max_context_length and attention_window_size) to rely on
TrtllmAttentionMetadata.max_context_length as the single source of truth.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a69ef961-72fb-495c-886f-08efe3ca69d4
📒 Files selected for processing (2)
tensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/attention_backend/trtllm_gen.py
|
PR_Github #53044 [ run ] triggered by Bot. Commit: |
|
PR_Github #53044 [ run ] completed with state
|
ee70236 to
a95d292
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53244 [ run ] triggered by Bot. Commit: |
a95d292 to
f5fefc3
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53283 [ run ] triggered by Bot. Commit: |
|
PR_Github #53244 [ run ] completed with state |
|
PR_Github #53283 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53312 [ run ] triggered by Bot. Commit: |
|
PR_Github #53312 [ run ] completed with state
|
f5fefc3 to
86c4352
Compare
TRTLLM attention used max_context_length only to estimate context workspace size. For visualGen, every token is a context token so max_context_length should be equal to max_seq_len. Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com>
86c4352 to
6a68216
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53521 [ run ] triggered by Bot. Commit: |
|
PR_Github #53521 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53591 [ run ] triggered by Bot. Commit: |
|
PR_Github #53591 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53759 [ run ] triggered by Bot. Commit: |
|
PR_Github #53759 [ run ] completed with state |
3 participants
TRTLLM attention used max_context_length only to estimate context workspace size. For VisualGen, every token is a context token so max_context_length should be equal to max_seq_len.
This PR is to address a small hidden OOB write of attention workspace.
Summary by CodeRabbit
Description
Test 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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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.