Cast/normalize attn_bias in scaled_dot_product_efficient_attention - #4636
Open
apbose wants to merge 2 commits into
Open
Cast/normalize attn_bias in scaled_dot_product_efficient_attention#4636apbose wants to merge 2 commits into
apbose wants to merge 2 commits into
Conversation
attn_bias was assigned straight to attention_layer.mask with no dtype cast or rank normalization, unlike the sibling scaled_dot_product_attention() converter which already does this for attn_mask. When PyTorch's SDPA dispatcher routes through the efficient backend with an int32 padding mask (e.g. BERT bfloat16), TRT's AttentionInput layer rejected the mask (must be Float/Half/BFloat16/ Bool), producing a broken output tensor whose shape then crashes whatever converter reads it next (seen as permute's issue #4496 and today as matmul.py's "Could not get tensor shape"). Apply the same BOOL-check/cast/rank-normalize treatment attn_mask already gets, for both the causal+attn_bias (additive bias combined via elementwise add) and plain attn_bias branches.
github-actions
Bot
added
component: conversion
apbose
force-pushed
the
abose/fix-sdpa-efficient-attention-mask-cast
branch
from
August 28, 2026 23:46
537a499 to
c2ea7b5
Compare
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/conversion/test_attention.py 2026-08-28 23:46:43.316968+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/conversion/test_attention.py 2026-08-28 23:47:09.397118+00:00
@@ -126,13 +126,11 @@
# straight to TensorRT's IAttention layer with no cast, which TRT's
# AttentionInput layer rejects for anything but Float/Half/BFloat16/
# Bool matching the attention dtype. int32 (not e.g. float32) is used
# here since PyTorch's own eager kernel accepts an int32 bias but
# rejects a mismatched float dtype outright.
- attn_bias = torch.zeros(
- (query_shape[0], 1, 1, key_shape[2]), dtype=torch.int32
- )
+ attn_bias = torch.zeros((query_shape[0], 1, 1, key_shape[2]), dtype=torch.int32)
inputs.extend([query, key, value, attn_bias])
self.run_test(
SDPA(),
inputs,
rtol=1e-2,Covers the case a dtype-mismatched attn_bias (e.g. float32 bias with float16 query/key/value) is passed to _scaled_dot_product_efficient_attention, which previously reached TensorRT's IAttention layer uncast and could be rejected.
apbose
force-pushed
the
abose/fix-sdpa-efficient-attention-mask-cast
branch
from
August 28, 2026 23:53
c2ea7b5 to
4077390
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
attn_bias was assigned straight to attention_layer.mask with no dtype cast or rank normalization, unlike scaled_dot_product_attention() which already does this for attn_mask. TRT's AttentionInput layer rejects non-Float/Half/BFloat16/Bool masks (e.g. int32 padding masks), leaving a broken output tensor that crashes whatever converter reads it next.
Fixes the root cause of #4496.
Type of change
Checklist: