Skip to content

feat: refactor four attention converters with TRT add_attention_v2()#4416

Open
zewenli98 wants to merge 8 commits into
mainfrom
evanli/refactor_attention_impl
Open

feat: refactor four attention converters with TRT add_attention_v2()#4416
zewenli98 wants to merge 8 commits into
mainfrom
evanli/refactor_attention_impl

Conversation

@zewenli98

Copy link
Copy Markdown
Collaborator

Description

This PR refactors four attention converters with TRT add_attention_v2(). Additionally, since newer torch rejects passing both is_causal=True and explicit attn_mask to SDPA, which aligns with TRT, some scenarios are gated in validators and the unreachable branches are pruned.

Fixes # (issue)

Type of change

  • New feature (non-breaking change which adds functionality)

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@zewenli98
zewenli98 requested a review from narendasan July 17, 2026 18:09
@zewenli98 zewenli98 self-assigned this Jul 17, 2026
@meta-cla meta-cla Bot added the cla signed label Jul 17, 2026
@github-actions github-actions Bot added component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: converters Issues re: Specific op converters component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Jul 17, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

@github-actions github-actions Bot added the component: lowering Issues re: The lowering / preprocessing passes label Jul 17, 2026
@github-actions github-actions Bot added the component: tests Issues re: Tests label Jul 20, 2026
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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is quite high? Will it be stable in practice?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

@zewenli98 zewenli98 added the Force All Tests[L0+L1+L2] For run all the L0, L1, L2 tests label Jul 22, 2026
@github-actions
github-actions Bot requested a review from cehongwang July 22, 2026 05:36
@zewenli98
zewenli98 force-pushed the evanli/refactor_attention_impl branch from 0ed13a5 to 0c27738 Compare July 22, 2026 05:39
@zewenli98
zewenli98 requested a review from narendasan July 22, 2026 22:39
return mask


def _restore_static_query_heads(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this be a lowering pass?

@zewenli98 zewenli98 Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests Force All Tests[L0+L1+L2] For run all the L0, L1, L2 tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants