Skip to content

Commit 1be4125

Browse files
zewenli98zewenli98
authored
feat: enhance validators of attention converters and add single-token tests (#4203)
1 parent a5be9d2 commit 1be4125

2 files changed

Lines changed: 107 additions & 65 deletions

File tree

‎py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py‎

Lines changed: 40 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -3961,34 +3961,57 @@ def aten_ops_linear(
39613961
)
39623962

39633963

3964-
def scaled_dot_product_attention_validator(
3965-
node: Node, settings: Optional[CompilationSettings] = None
3966-
) -> bool:
3967-
if node.kwargs.get("enable_gqa", False):
3968-
_LOGGER.debug(
3969-
"enable_gqa is not yet supported by the converter. Please try setting decompose_attention=True in the compilation settings."
3970-
)
3971-
return False
3972-
3964+
def _attention_qkv_shapes_supported(node: Node) -> bool:
39733965
query_shape, key_shape, value_shape = None, None, None
39743966
if "val" in node.args[0].meta:
39753967
query_shape = node.args[0].meta["val"].size()
39763968
if "val" in node.args[1].meta:
39773969
key_shape = node.args[1].meta["val"].size()
39783970
if "val" in node.args[2].meta:
39793971
value_shape = node.args[2].meta["val"].size()
3980-
if (
3981-
query_shape != key_shape
3982-
or query_shape != value_shape
3983-
or key_shape != value_shape
3984-
):
3972+
3973+
# If shape metadata is unavailable, defer to runtime/converter checks.
3974+
if query_shape is None or key_shape is None or value_shape is None:
3975+
return True
3976+
3977+
if len(query_shape) != len(key_shape) or len(query_shape) != len(value_shape):
39853978
_LOGGER.debug(
3986-
"query, key, and value have different shapes. Please try setting decompose_attention=True in the compilation settings."
3979+
"query, key, and value must have the same rank. Please try setting decompose_attention=True in the compilation settings."
39873980
)
39883981
return False
3982+
3983+
# TensorRT IAttention layer supports different sequence lengths for query and key/value
3984+
# ([B, Nq, Sq, H] vs [B, Nkv, Skv, H]), but K and V must still agree on all dims.
3985+
seq_dim = len(query_shape) - 2
3986+
for dim, (query_dim, key_dim, value_dim) in enumerate(
3987+
zip(query_shape, key_shape, value_shape)
3988+
):
3989+
if dim == seq_dim:
3990+
if key_dim != value_dim:
3991+
_LOGGER.debug(
3992+
"key and value must have the same sequence length. Please try setting decompose_attention=True in the compilation settings."
3993+
)
3994+
return False
3995+
else:
3996+
if query_dim != key_dim or query_dim != value_dim or key_dim != value_dim:
3997+
_LOGGER.debug(
3998+
"query, key, and value differ on a non-sequence dimension. Please try setting decompose_attention=True in the compilation settings."
3999+
)
4000+
return False
39894001
return True
39904002

39914003

4004+
def scaled_dot_product_attention_validator(
4005+
node: Node, settings: Optional[CompilationSettings] = None
4006+
) -> bool:
4007+
if node.kwargs.get("enable_gqa", False):
4008+
_LOGGER.debug(
4009+
"enable_gqa is not yet supported by the converter. Please try setting decompose_attention=True in the compilation settings."
4010+
)
4011+
return False
4012+
return _attention_qkv_shapes_supported(node)
4013+
4014+
39924015
@dynamo_tensorrt_converter(
39934016
torch.ops.aten.scaled_dot_product_attention.default,
39944017
supports_dynamic_shapes=True,
@@ -4024,24 +4047,7 @@ def scaled_dot_product_flash_attention_validator(
40244047
if args_bounds_check(node.args, 5, False):
40254048
_LOGGER.debug("return_debug_mask is not yet supported.")
40264049
return False
4027-
4028-
query_shape, key_shape, value_shape = None, None, None
4029-
if "val" in node.args[0].meta:
4030-
query_shape = node.args[0].meta["val"].size()
4031-
if "val" in node.args[1].meta:
4032-
key_shape = node.args[1].meta["val"].size()
4033-
if "val" in node.args[2].meta:
4034-
value_shape = node.args[2].meta["val"].size()
4035-
if (
4036-
query_shape != key_shape
4037-
or query_shape != value_shape
4038-
or key_shape != value_shape
4039-
):
4040-
_LOGGER.debug(
4041-
"query, key, and value have different shapes. Please try setting decompose_attention=True in the compilation settings."
4042-
)
4043-
return False
4044-
return True
4050+
return _attention_qkv_shapes_supported(node)
40454051

40464052

40474053
@dynamo_tensorrt_converter(
@@ -4078,24 +4084,7 @@ def scaled_dot_product_efficient_attention_validator(
40784084
if args_bounds_check(node.args, 4, False):
40794085
_LOGGER.debug("compute_log_sumexp is not yet supported.")
40804086
return False
4081-
4082-
query_shape, key_shape, value_shape = None, None, None
4083-
if "val" in node.args[0].meta:
4084-
query_shape = node.args[0].meta["val"].size()
4085-
if "val" in node.args[1].meta:
4086-
key_shape = node.args[1].meta["val"].size()
4087-
if "val" in node.args[2].meta:
4088-
value_shape = node.args[2].meta["val"].size()
4089-
if (
4090-
query_shape != key_shape
4091-
or query_shape != value_shape
4092-
or key_shape != value_shape
4093-
):
4094-
_LOGGER.debug(
4095-
"query, key, and value have different shapes. Please try setting decompose_attention=True in the compilation settings."
4096-
)
4097-
return False
4098-
return True
4087+
return _attention_qkv_shapes_supported(node)
40994088

41004089

41014090
@dynamo_tensorrt_converter(

‎tests/py/dynamo/conversion/test_attention_aten.py‎

Lines changed: 67 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ class TestScaledDotProductAttention(DispatchTestCase):
2020
torch.float16,
2121
0.0,
2222
False,
23-
), # flash attention
23+
),
2424
(
2525
(4, 8, 32, 16),
2626
(4, 8, 32, 16),
@@ -31,7 +31,7 @@ class TestScaledDotProductAttention(DispatchTestCase):
3131
torch.float16,
3232
0.0,
3333
False,
34-
), # flash attention
34+
),
3535
(
3636
(4, 8, 32, 16),
3737
(4, 8, 32, 16),
@@ -42,7 +42,7 @@ class TestScaledDotProductAttention(DispatchTestCase):
4242
torch.float32,
4343
0.0,
4444
False,
45-
), # flash attention
45+
),
4646
(
4747
(4, 8, 32, 16),
4848
(4, 8, 32, 16),
@@ -53,7 +53,7 @@ class TestScaledDotProductAttention(DispatchTestCase):
5353
torch.float32,
5454
0.0,
5555
False,
56-
), # flash attention
56+
),
5757
(
5858
(4, 8, 32, 16),
5959
(4, 8, 32, 16),
@@ -64,7 +64,7 @@ class TestScaledDotProductAttention(DispatchTestCase):
6464
torch.float16,
6565
0.0,
6666
False,
67-
), # efficient attention
67+
),
6868
(
6969
(2, 8, 128, 64),
7070
(2, 8, 128, 64),
@@ -75,7 +75,7 @@ class TestScaledDotProductAttention(DispatchTestCase):
7575
torch.float16,
7676
0.0,
7777
False,
78-
), # efficient attention
78+
),
7979
(
8080
(2, 8, 128, 64),
8181
(2, 8, 128, 64),
@@ -86,7 +86,18 @@ class TestScaledDotProductAttention(DispatchTestCase):
8686
torch.float32,
8787
0.0,
8888
False,
89-
), # efficient attention
89+
),
90+
(
91+
(4, 8, 1, 64),
92+
(4, 8, 4, 64),
93+
(4, 8, 4, 64),
94+
(1, 1, 4),
95+
False,
96+
None,
97+
torch.float16,
98+
0.0,
99+
False,
100+
), # decoder-style single-token attention
90101
]
91102
)
92103
def test_sdpa_bool_mask(
@@ -145,7 +156,7 @@ def forward(self, query, key, value, attn_mask=None):
145156
torch.float16,
146157
0.0,
147158
False,
148-
), # flash attention
159+
),
149160
(
150161
(4, 8, 32, 16),
151162
(4, 8, 32, 16),
@@ -156,7 +167,7 @@ def forward(self, query, key, value, attn_mask=None):
156167
torch.float16,
157168
0.0,
158169
False,
159-
), # flash attention
170+
),
160171
(
161172
(4, 8, 32, 16),
162173
(4, 8, 32, 16),
@@ -167,7 +178,7 @@ def forward(self, query, key, value, attn_mask=None):
167178
torch.float32,
168179
0.0,
169180
False,
170-
), # flash attention
181+
),
171182
(
172183
(4, 8, 32, 16),
173184
(4, 8, 32, 16),
@@ -178,7 +189,7 @@ def forward(self, query, key, value, attn_mask=None):
178189
torch.float32,
179190
0.0,
180191
False,
181-
), # flash attention
192+
),
182193
(
183194
(4, 8, 32, 16),
184195
(4, 8, 32, 16),
@@ -189,7 +200,7 @@ def forward(self, query, key, value, attn_mask=None):
189200
torch.float16,
190201
0.0,
191202
False,
192-
), # efficient attention
203+
),
193204
(
194205
(2, 8, 128, 64),
195206
(2, 8, 128, 64),
@@ -200,7 +211,7 @@ def forward(self, query, key, value, attn_mask=None):
200211
torch.float16,
201212
0.0,
202213
False,
203-
), # efficient attention
214+
),
204215
(
205216
(2, 8, 128, 64),
206217
(2, 8, 128, 64),
@@ -211,7 +222,18 @@ def forward(self, query, key, value, attn_mask=None):
211222
torch.float32,
212223
0.0,
213224
False,
214-
), # efficient attention
225+
),
226+
(
227+
(4, 8, 1, 64),
228+
(4, 8, 4, 64),
229+
(4, 8, 4, 64),
230+
(1, 1, 4),
231+
False,
232+
None,
233+
torch.float16,
234+
0.0,
235+
False,
236+
), # decoder-style single-token attention
215237
]
216238
)
217239
def test_sdpa_fp_mask(
@@ -304,6 +326,17 @@ def forward(self, query, key, value, attn_mask=None):
304326
0.0,
305327
False,
306328
),
329+
(
330+
[(2, 8, 1, 64), (4, 8, 1, 64), (8, 8, 1, 64)],
331+
[(2, 8, 4, 64), (4, 8, 4, 64), (8, 8, 4, 64)],
332+
[(2, 8, 4, 64), (4, 8, 4, 64), (8, 8, 4, 64)],
333+
[(1, 1, 4), (1, 1, 4), (1, 1, 4)],
334+
False,
335+
None,
336+
torch.float16,
337+
0.0,
338+
False,
339+
), # decoder-style single-token attention
307340
]
308341
)
309342
def test_dynamic_sdpa_fp_mask(
@@ -416,6 +449,16 @@ class TestScaledDotProductEfficientAttention(DispatchTestCase):
416449
torch.float32,
417450
0.0,
418451
),
452+
(
453+
(4, 8, 1, 64),
454+
(4, 8, 4, 64),
455+
(4, 8, 4, 64),
456+
(1, 1, 4),
457+
False,
458+
None,
459+
torch.float16,
460+
0.0,
461+
), # decoder-style single-token attention
419462
]
420463
)
421464
def test_efficient_sdpa(
@@ -518,6 +561,16 @@ def forward(self, query, key, value, attn_bias=None):
518561
torch.float32,
519562
0.0,
520563
),
564+
(
565+
(4, 8, 1, 64),
566+
(4, 8, 4, 64),
567+
(4, 8, 4, 64),
568+
(1, 1, 4),
569+
False,
570+
None,
571+
torch.float16,
572+
0.0,
573+
), # decoder-style single-token attention
521574
]
522575
)
523576
def test_efficient_sdpa_random_attn_bias(

0 commit comments

Comments
 (0)