[https://nvbugs/6322045][fix] In triton_context, when max_q_len == max_kv_len (so cache_lens=0) and… - #15479
[https://nvbugs/6322045][fix] In triton_context, when max_q_len == max_kv_len (so cache_lens=0) and…#15479tensorrt-cicd wants to merge 3 commits into
Conversation
|
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 (1)
💤 Files with no reviewable changes (1)
WalkthroughThe SDPA path now uses fused causal attention for large equal-length workloads when estimated scratch usage exceeds 2 GiB. Other workloads retain the combined padding and causal mask. Two related L40S test waiver entries are removed. ChangesSDPA Causal Masking and Test Waiver
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py`:
- Around line 1181-1182: The sdpa_scratch_bytes calculation on line 1181
incorrectly uses q.element_size() which reflects the query tensor's dtype (e.g.,
2 bytes for fp16/bf16), but the actual SDPA scratch memory is allocated for fp32
score tensors which use 4 bytes per element. To fix this, replace the
q.element_size() call with a hardcoded 4 bytes (or equivalent like the size of a
float32 element) in the sdpa_scratch_bytes calculation to accurately estimate
the memory requirement and prevent premature selection of the explicit-mask
backend that could cause OOM on memory-constrained GPUs.
🪄 Autofix (Beta)
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: 07606216-ae91-4e99-93cf-6accb9a55b37
📒 Files selected for processing (2)
tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
| sdpa_scratch_bytes = num_seq * n_heads * max_q_len * max_kv_len * q.element_size() | ||
| use_is_causal = max_q_len == max_kv_len and sdpa_scratch_bytes > 2 * 1024**3 |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# Find PyTorch version pinning
find . -type f \( -name "requirements*.txt" -o -name "pyproject.toml" -o -name "setup.py" \) | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 952
🏁 Script executed:
# Read the specific file segment
cat -n tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | sed -n '1174,1200p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1788
🏁 Script executed:
# Also check surrounding comments for context on SDPA behavior
cat -n tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | sed -n '1160,1185p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1464
🏁 Script executed:
# Find main project configuration
cat pyproject.toml | grep -A 10 "torch\|pytorch" | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1219
🏁 Script executed:
# Also check if there's documentation about SDPA backend behavior
rg -i "sdpa|scaled_dot_product" --type py -B 2 -A 2 | grep -i "dtype\|fp32\|float32" | head -30Repository: NVIDIA/TensorRT-LLM
Length of output: 3723
🏁 Script executed:
# Check PyTorch documentation link or comments that might clarify the backend
rg "math backend\|flash\|fused kernel" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/ -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check the main pyproject.toml for torch version
cat pyproject.toml | grep -i "torch" | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 1219
🏁 Script executed:
# Look for PyTorch version in the base requirements
find . -name "requirements.txt" -exec grep -l "torch" {} \; | head -5Repository: NVIDIA/TensorRT-LLM
Length of output: 259
🏁 Script executed:
# Search for any tests or documentation about SDPA math backend behavior
find . -path "./tests/*" -name "*.py" -exec rg "scaled_dot_product_attention" {} + | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 3178
🏁 Script executed:
# Check if there are related PRs or comments about fp32 score tensors
rg "fp32.*score|score.*fp32" --type py tensorrt_llm/_torch/Repository: NVIDIA/TensorRT-LLM
Length of output: 497
🏁 Script executed:
# Check main requirements.txt for torch version
cat requirements.txt | grep -i torchRepository: NVIDIA/TensorRT-LLM
Length of output: 624
🏁 Script executed:
# Look for test that might reveal SDPA backend behavior with fp16/bf16
cat tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention.py | grep -A 20 "def test\|torch.float16\|torch.bfloat16" | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 3109
🏁 Script executed:
# Check if there's documentation about SDPA math backend in the repo
rg "torch.nn.functional.scaled_dot_product_attention" tensorrt_llm/_torch/ -B 5 -A 5 | grep -E "dtype|float32|fp32|math.*backend" | head -20Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Look at the actual test to see if fp16/bf16 queries are tested
cat tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention.py | sed -n '1,150p'Repository: NVIDIA/TensorRT-LLM
Length of output: 5217
🏁 Script executed:
# Search for PyTorch SDPA documentation or implementation details in the codebase
rg "math.*backend|cuda_flash\|efficient_attention" tensorrt_llm/_torch/ -B 2 -A 5 | head -40Repository: NVIDIA/TensorRT-LLM
Length of output: 3371
🏁 Script executed:
# Look for any comments about SDPA dtype behavior in tests or implementations
rg -i "sdpa.*dtype|score.*dtype|attention.*math" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/ -B 3 -A 3Repository: NVIDIA/TensorRT-LLM
Length of output: 6500
🏁 Script executed:
# Check if there are any references to PyTorch SDPA backend selection
cat tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | sed -n '1100,1200p' | grep -i "backend\|dtype\|flash"Repository: NVIDIA/TensorRT-LLM
Length of output: 611
🏁 Script executed:
# Verify the proposed constants follow naming guidelines
# Check if similar constants are defined in the file
rg "THRESHOLD|BYTES|SIZE" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1681
🏁 Script executed:
# Verify the fix formatting matches the codebase style
# Check how constants are defined elsewhere in the file
cat tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | grep -E "^[A-Z_]+\s*=" | head -10Repository: NVIDIA/TensorRT-LLM
Length of output: 82
🏁 Script executed:
# Verify there are no existing fp32/fp64 element size constants already defined
rg "element_size\|ELEMENT_SIZE" tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
# Check the context of the _MIN_TL_DOT_K constant to see naming style
cat tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | grep -B 5 -A 5 "_MIN_TL_DOT_K"Repository: NVIDIA/TensorRT-LLM
Length of output: 730
🏁 Script executed:
# Check the exact lines and verify the fix makes sense with surrounding code
cat tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py | sed -n '1175,1200p'Repository: NVIDIA/TensorRT-LLM
Length of output: 1526
Use fp32 bytes for the SDPA scratch threshold.
Line 1181 uses q.element_size(), but the preceding comment describes an fp32 score tensor. For fp16/bf16 queries, this halves the estimate, so the explicit-mask math backend can still be selected when the real fp32 score exceeds the 2 GiB guard—defeating the OOM prevention on memory-tight GPUs like L40S.
Proposed fix
- sdpa_scratch_bytes = num_seq * n_heads * max_q_len * max_kv_len * q.element_size()
- use_is_causal = max_q_len == max_kv_len and sdpa_scratch_bytes > 2 * 1024**3
+ FP32_ELEMENT_SIZE_BYTES = 4
+ SDPA_SCRATCH_THRESHOLD_BYTES = 2 * 1024**3
+ sdpa_scratch_bytes = (
+ num_seq * n_heads * max_q_len * max_kv_len * FP32_ELEMENT_SIZE_BYTES
+ )
+ use_is_causal = (
+ max_q_len == max_kv_len and sdpa_scratch_bytes > SDPA_SCRATCH_THRESHOLD_BYTES
+ )🤖 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/auto_deploy/custom_ops/attention/triton_attention.py`
around lines 1181 - 1182, The sdpa_scratch_bytes calculation on line 1181
incorrectly uses q.element_size() which reflects the query tensor's dtype (e.g.,
2 bytes for fp16/bf16), but the actual SDPA scratch memory is allocated for fp32
score tensors which use 4 bytes per element. To fix this, replace the
q.element_size() call with a hardcoded 4 bytes (or equivalent like the size of a
float32 element) in the sdpa_scratch_bytes calculation to accurately estimate
the memory requirement and prevent premature selection of the explicit-mask
backend that could cause OOM on memory-constrained GPUs.
6e60bc4 to
6e13e27
Compare
| cache_lens.view(num_seq, 1, 1, 1) + q_positions.view(1, 1, max_q_len, 1) | ||
| ) | ||
| attn_mask = attn_mask & causal_mask | ||
| # An explicit attn_mask forces SDPA onto the math backend, which |
There was a problem hiding this comment.
The target test is for trtllm attention too, but it is fixing only triton attention. Seems that this is not the proper fix.
6dd0238 to
3094b9b
Compare
c16425c to
3f754d0
Compare
3f754d0 to
7fcfa38
Compare
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
LGTM — AutoDeploy (beta) triton_context: the explicit-mask→is_causal substitution only fires when max_q_len==max_kv_len AND SDPA scratch >2 GiB (the giant fresh-prefill warmup where cache_lens==0 makes the mask exactly lower-triangular). Smaller/normal prefill and any cache-reuse path keep the explicit-mask numerics.
|
Please resolve @taylor-yb-lee's waiver concern first: this PR only fixes the triton path, but it also removes the waiver for |
|
Waiting on code owner review from NVIDIA/trt-llm-torch-autodeploy-devs. |
…attn The triton context-attention SDPA fast path passes an explicit attn_mask to torch.nn.functional.scaled_dot_product_attention, which forces the math backend and materializes [num_seq, n_heads, max_q_len, max_kv_len] fp32 score and softmax tensors. During piecewise cudagraph warmup at 8192 tokens for Llama-3.1-8B (1*32*8192*8192*4 = 8 GiB per tensor), this OOMs on memory-tight GPUs such as L40S (44 GiB). When the explicit mask reduces to a plain lower-triangular mask (every sequence has the same q_len under the existing all_same_q_len gate, and max_q_len == max_kv_len implies all cache_lens are 0), and the projected scratch would exceed 2 GiB, drop the mask and dispatch via is_causal=True so SDPA can use the fused flash kernel instead. The original explicit-mask path is preserved for shorter sequences and any case with cache reuse, so numerics are unchanged outside the OOM-prone regime. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
7fcfa38 to
369cc0b
Compare
|
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 --only-qa-verify test accuracy/test_llm_api_autodeploy.py::TestLlama3_1_8B::test_auto_dtype[triton-False-1] |
|
PR_Github #65470 [ run ] triggered by Bot. Commit: |
StanleySun639
left a comment
StanleySun639
left a comment
There was a problem hiding this comment.
LGTM if the target test case can pass.
|
PR_Github #65470 [ run ] completed with state |
|
Thanks for the reviews and apologies for the belated response. Auto Deploy will be deprecated in the TRTLLM, so there is no point to make change here. Let me close this PR and request close the original ticket too. |
9 participants
Summary
Test plan
Links
Summary
triton_contextattention path during Llama-3.1-8B prefill warmup at 8192 tokens on L40S GPUs.test_auto_dtypewaivers for thetriton-False-1andtrtllm-False-1configurations.Dev Engineer Review
test_auto_dtypeentries.trtllm-False-1waiver because that backend does not exercisetriton_context. The full L40S pipeline should validate both configurations.QA Engineer Review
tests/integration/test_lists/waives.txt.TestLlama3_1_8B::test_auto_dtypeentries fortriton-False-1andtrtllm-False-1.