feat: refactor four attention converters with TRT add_attention_v2()#4416
feat: refactor four attention converters with TRT add_attention_v2()#4416zewenli98 wants to merge 8 commits into
add_attention_v2()#4416Conversation
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py 2026-07-17 18:09:43.975250+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py 2026-07-17 18:10:04.835218+00:00
@@ -2,11 +2,10 @@
import torch
from torch.fx import GraphModule, Node
from .pass_utils import clean_up_graph_after_modifications
-
_INT64_MAX = 2**63 - 1
_SYM_MIN = getattr(torch, "sym_min", None)
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py 2026-07-17 18:09:44.412631+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/eliminate_sym_min_int64_max.py 2026-07-17 18:10:05.940612+00:00
@@ -2,11 +2,10 @@
import torch
from torch.fx import GraphModule, Node
from .pass_utils import clean_up_graph_after_modifications
-
_INT64_MAX = 2**63 - 1
_SYM_MIN = getattr(torch, "sym_min", None)
| out = out[0] | ||
|
|
||
| torch.testing.assert_close(ref, out, rtol=1e-1, atol=2e-1) | ||
| torch.testing.assert_close(ref, out, rtol=0.3, atol=0.3) |
There was a problem hiding this comment.
This is quite high? Will it be stable in practice?
There was a problem hiding this comment.
The relaxation is mainly for TRT-RTX. As I talked with @SandSnip3r, TRT-RTX 1.6 is going to fix the accuracy issue but it's not released yet. So I skipped these tests for now and changed to rtol=1e-1, atol=3e-1 for the single slightly mismatched element on standard TRT:
Mismatched elements: 1 / 9723904 (0.0%)
Greatest absolute difference: 0.2071533203125 at index (0, 0, 80751) (up to 0.2 allowed)
Greatest relative difference: 2.931640625 at index (0, 0, 80751) (up to 0.1 allowed)
0ed13a5 to
0c27738
Compare
| return mask | ||
|
|
||
|
|
||
| def _restore_static_query_heads( |
There was a problem hiding this comment.
Should this be a lowering pass?
There was a problem hiding this comment.
I'm not sure if this is general enough to make it a lowering pass now because I only observed in the query of attention in some model implementations like HF Qwen. Maybe we can keep it in the attention converters now. If we observe more cases later, we can make it a lowering pass then.
Description
This PR refactors four attention converters with TRT
add_attention_v2(). Additionally, since newer torch rejects passing bothis_causal=Trueand explicitattn_maskto SDPA, which aligns with TRT, some scenarios are gated in validators and the unreachable branches are pruned.Fixes # (issue)
Type of change
Checklist: