[#12702][feat] Autodeploy deprecate the legacy triton attention - #14194
Conversation
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
f3346ea to
3573846
Compare
📝 WalkthroughWalkthroughThis PR consolidates the Triton attention backend infrastructure by unifying separate paged and unpaged implementations into a single ChangesTriton Attention Backend Migration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention.py (1)
16-1744: QA list updates are unnecessary for this change scope.This file is unit-test migration only; no integration test-list update is required in this PR slice.
As per coding guidelines: "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional."
🤖 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 `@tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention.py` around lines 16 - 1744, Summary: The PR only modifies unit tests under the unittest/ scope so QA list updates are unnecessary; add an explicit note to that effect. Fix: add a single-line statement (e.g., "QA list updates unnecessary for unit-test-only changes") either to the PR description/commit message or as a short comment in the test module docstring near the top of this test file (the module containing tests like TestTritonMHAIntegration and functions prepare_triton_metadata/triton_mha_with_cache), so reviewers see the intent; keep the note minimal and do not change test logic.
🤖 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 1514-1535: The code currently silently ignores unsupported
attention semantics (e.g., non-default dropout_p, sinks, logit_cap) causing
incorrect routing to triton_mha_with_cache; inside the same function (the
get_constants/extract_op_args handling for source_attn_node) replace the
ad_logger.debug for dropout_p with a fail-fast RuntimeError and add similar
checks that raise if sinks or logit_cap (obtained via
extract_op_args(source_attn_node, "sinks", "logit_cap") or equivalent) are
present/non-default; keep the existing layout check, retain the scale type-check
behavior, and only return [scale, sliding_window] when all unsupported arguments
are default—otherwise raise with a clear message referencing source_attn_node so
the backend is not selected for incompatible semantics.
- Around line 1211-1225: The current assert only checks
custom_attn_mask.shape[1] == 1 but does not validate other dimensions used by
_paged_context_masked_kernel; replace the single assert with explicit boundary
checks: verify custom_attn_mask.shape[0] == num_seq (batch dim),
custom_attn_mask.shape[2] >= max_q_len (q-length dim must cover the largest
per-sequence query length computed from qo_indptr / max_q_len), and
custom_attn_mask.shape[3] >= total_kv_tokens (kv-length dim must cover all KV
positions indexed by kv_positions/runtime metadata); if any check fails, raise a
clear ValueError referencing custom_attn_mask, num_seq, max_q_len, and
total_kv_tokens so the kernel cannot be launched with undersized mask. Ensure
these checks occur before the dtype/contiguous conversion and before launching
_paged_context_masked_kernel.
In
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py`:
- Around line 491-492: The regex strings passed to pytest.raises(match=...) are
regular expressions but are written as normal Python strings, so escape
sequences like .* are interpreted and trigger RUF043; update both message
arguments (the one starting with "Cached attention backend 'flashinfer'..." and
the one containing ".*gemma4_multimodal_mask.*Supported backends: torch,
triton") to be raw string literals by prefixing them with r (e.g., r"Cached
attention..." and r".*gemma4_multimodal_mask.*Supported backends: torch,
triton") so the regex metacharacters are treated correctly.
---
Nitpick comments:
In
`@tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention.py`:
- Around line 16-1744: Summary: The PR only modifies unit tests under the
unittest/ scope so QA list updates are unnecessary; add an explicit note to that
effect. Fix: add a single-line statement (e.g., "QA list updates unnecessary for
unit-test-only changes") either to the PR description/commit message or as a
short comment in the test module docstring near the top of this test file (the
module containing tests like TestTritonMHAIntegration and functions
prepare_triton_metadata/triton_mha_with_cache), so reviewers see the intent;
keep the note minimal and do not change test logic.
🪄 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: 9570cd79-7902-44b2-ae42-8cf8f3f0bbec
📒 Files selected for processing (21)
examples/auto_deploy/cookbooks/gemma_4_trtllm_cookbook.ipynbexamples/auto_deploy/model_registry/configs/gemma4_dense.yamlexamples/auto_deploy/model_registry/configs/gemma4_moe.yamlexamples/auto_deploy/model_registry/configs/gemma4_moe_base.yamltensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.pytensorrt_llm/_torch/auto_deploy/custom_ops/README.mdtensorrt_llm/_torch/auto_deploy/custom_ops/attention/__init__.pytensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention_with_kv_cache.pytensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.pytensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma4.pytensorrt_llm/_torch/auto_deploy/transform/library/kvcache.pytests/integration/defs/accuracy/test_llm_api_autodeploy.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_h100.ymltests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_attention_op.pytests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention.pytests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention_with_kv_cache.pytests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py
💤 Files with no reviewable changes (5)
- tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention_with_kv_cache.py
- tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_attention_op.py
- tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py
- tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py
- tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention_with_kv_cache.py
|
/bot run --disable-fail-fast |
|
PR_Github #48632 [ run ] triggered by Bot. Commit: |
|
PR_Github #48632 [ run ] completed with state
|
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
…iton_attention_0515 Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> # Conflicts: # tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_attention_op.py # tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_attention_with_kv_cache.py
|
/bot run --disable-fail-fast |
|
PR_Github #48945 [ run ] triggered by Bot. Commit: |
|
PR_Github #48945 [ run ] completed with state
|
…iton_attention_0515 Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #49974 [ run ] triggered by Bot. Commit: |
|
PR_Github #49974 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #50374 [ run ] triggered by Bot. Commit: |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #50376 [ run ] triggered by Bot. Commit: |
|
PR_Github #50376 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #50432 [ run ] triggered by Bot. Commit: |
|
PR_Github #50432 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
…iton_attention_0515 Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #50837 [ run ] triggered by Bot. Commit: |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #50887 [ run ] triggered by Bot. Commit: |
|
PR_Github #50837 [ run ] completed with state |
|
PR_Github #50887 [ run ] completed with state
|
|
/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1,DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #51081 [ run ] triggered by Bot. Commit: |
|
PR_Github #51081 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51382 [ run ] triggered by Bot. Commit: |
|
PR_Github #51382 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51432 [ run ] triggered by Bot. Commit: |
|
PR_Github #51432 [ run ] completed with state |
4 participants
#12702
As the triton paged attention is added to the codebase, deprecate the legacy triton attentions. Also remove the related test. Rename the attention backend: triton_paged to triton.
Summary by CodeRabbit
Configuration Changes
tritonattention backend instead oftriton_paged.Improvements
Documentation
Tests
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.