Skip to content

Fix SigLIP casual mask bug #25360

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nenad1002 merged 10 commits into from
Jul 15, 2025
Merged

Fix SigLIP casual mask bug #25360

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
f87251a
Fix SigLIP casual mask bug
nenad1002 Jul 10, 2025
487c5ff
Check causal_mask_nodes
nenad1002 Jul 11, 2025
919ee48
Move causal mask nodes outside local scope
nenad1002 Jul 12, 2025
b991613
Move causal mask nodes outside local scope
nenad1002 Jul 12, 2025
e8884cf
Add a new test
nenad1002 Jul 15, 2025
49af722
Update onnxruntime/test/python/transformers/test_phi_vision.py
nenad1002 Jul 15, 2025
fc8f467
Update onnxruntime/test/python/transformers/test_phi_vision.py
nenad1002 Jul 15, 2025
1cec7b6
Update onnxruntime/test/python/transformers/test_phi_vision.py
nenad1002 Jul 15, 2025
7a633be
Update onnxruntime/test/python/transformers/test_phi_vision.py
nenad1002 Jul 15, 2025
fb36760
Trigger CLA check
nenad1002 Jul 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 38 additions & 32 deletions onnxruntime/python/tools/transformers/fusion_attention_clip.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,42 +269,48 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
attention_last_node = reshape_qkv

add_qk = ""
causal_mask_nodes_1 = None
causal_mask_nodes_2 = None
if add_mask is not None:
# 4D Add after Q x K'
add_qk_nodes = self.model.match_parent_path(
add_mask,
[
"Where",
"Sub",
"Cast",
"Expand",
"Unsqueeze",
"Unsqueeze",
"Reshape",
"Reshape",
"Cast",
],
[1, 2, 1, 0, 0, 0, 0, 0, 0],
)
if add_qk_nodes is not None:
if add_mask.input[1] == "attention_mask":
add_qk = add_mask.input[1]
else:
# Here we do not match the whole subgraph since it is very complex. Instead, we just check whether a key path
# of computing causal mask.
causal_mask_nodes_1 = self.model.match_parent_path(
add_mask,
["Concat", "Expand", "Unsqueeze", "Unsqueeze", "Where", "Less"],
[causal_mask_input_index, 0, 0, 0, 0, 0],
)
# If the model is exported with batch_size == 1, there is no Concat node
causal_mask_nodes_2 = self.model.match_parent_path(
# 4D Add after Q x K'
add_qk_nodes = self.model.match_parent_path(
add_mask,
["Expand", "Unsqueeze", "Unsqueeze", "Where", "Less"],
[causal_mask_input_index, 0, 0, 0, 0],
[
"Where",
"Sub",
"Cast",
"Expand",
"Unsqueeze",
"Unsqueeze",
"Reshape",
"Reshape",
"Cast",
],
[1, 2, 1, 0, 0, 0, 0, 0, 0],
)
if causal_mask_nodes_1 is None and causal_mask_nodes_2 is None:
logger.debug("fuse_attention: failed to match causal mask subgraph")
return
if add_qk_nodes is not None:
add_qk = add_mask.input[1]
else:
# Here we do not match the whole subgraph since it is very complex. Instead, we just check whether a key path
# of computing causal mask.
causal_mask_nodes_1 = self.model.match_parent_path(
add_mask,
["Concat", "Expand", "Unsqueeze", "Unsqueeze", "Where", "Less"],
[causal_mask_input_index, 0, 0, 0, 0, 0],
)
# If the model is exported with batch_size == 1, there is no Concat node
causal_mask_nodes_2 = self.model.match_parent_path(
add_mask,
["Expand", "Unsqueeze", "Unsqueeze", "Where", "Less"],
[causal_mask_input_index, 0, 0, 0, 0],
)

if causal_mask_nodes_1 is None and causal_mask_nodes_2 is None:
logger.debug("fuse_attention: failed to match causal mask subgraph")
return

new_node = self.create_attention_node(
mask_index=None,
Expand All @@ -320,7 +326,7 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
output=attention_last_node.output[0],
add_qk_str=add_qk,
scale=None,
causal=(add_mask is not None),
causal=(causal_mask_nodes_1 is not None) or (causal_mask_nodes_2 is not None),
)
if new_node is None:
logger.debug("fuse_attention: failed to create fused node")
Expand Down
Binary file added onnxruntime/test/python/transformers/test_data/models/phi-4-v-instruct-vision-attention.onnx
View file
Open in desktop
Binary file not shown.
70 changes: 59 additions & 11 deletions onnxruntime/test/python/transformers/test_phi_vision.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ def __init__(self):
self.attn = PhiVCLIPAttention()
self.ln = torch.nn.LayerNorm(20, eps=1e-05)

def forward(self, x):
def forward(self, x, attention_mask=None):
# SkipLayerNorm ------+
# | |
# Attention |
Expand All @@ -163,8 +163,7 @@ def forward(self, x):
x = self.ln(x)
residual = x

# Attention + MatMul
x = self.attn(x)
x = self.attn(x, attention_mask=attention_mask)

# SkipLayerNorm
x = residual + x
Expand Down Expand Up @@ -194,14 +193,31 @@ def verify_fusion(self, optimized_model, expected_model_filename):
)

def export(self, model, inputs):
torch.onnx.export(
model,
args=inputs,
f=os.path.join(os.path.dirname(__file__), "export.onnx"),
export_params=True,
opset_version=14,
do_constant_folding=True,
)
path = os.path.join(os.path.dirname(__file__), "export.onnx")

if len(inputs) == 2:
torch.onnx.export(
model,
args=inputs,
f=path,
export_params=True,
opset_version=14,
do_constant_folding=True,
input_names=["input", "attention_mask"],
dynamic_axes={
"input": {0: "batch", 1: "seq"},
"attention_mask": {0: "batch", 2: "seq", 3: "seq"},
},
)
else:
torch.onnx.export(
model,
args=inputs,
f=path,
export_params=True,
opset_version=14,
do_constant_folding=True,
)

def tearDown(self):
path = os.path.join(os.path.dirname(__file__), "export.onnx")
Expand Down Expand Up @@ -249,6 +265,38 @@ def test_phi_vision_attention(self):
)
self.verify_fusion(optimized_model, "phi-3.5-v-instruct-vision-attention.onnx")

def test_phi_vision_attention_with_mask(self):
model = PhiVCLIPAttentionAndLayerNorm()

batch, seq_len, dim = 1, 2, 20
mask = torch.zeros(batch, 1, seq_len, seq_len)
mask[:, 1:] = float("-inf")

inputs = (torch.randn(batch, seq_len, dim), mask)
self.export(model, inputs)
original_model = onnx.load(os.path.join(os.path.dirname(__file__), "export.onnx"))
options = FusionOptions("clip")
optimized_model = optimize_model(
original_model,
model_type="clip",
num_heads=2,
hidden_size=20,
optimization_options=options,
opt_level=0,
use_gpu=True,
)
self.verify_fusion(optimized_model, "phi-4-v-instruct-vision-attention.onnx")

graph = optimized_model.model.graph
attention_node = next((n for n in graph.node if n.name == "Attention_0"), None)
self.assertIsNotNone(attention_node, "Could not find the Attention fused node")
attr_names = [attr.name for attr in attention_node.attribute]
self.assertNotIn(
"unidirectional",
attr_names,
f"The attention node should not have a 'unidirectional' attribute: {attr_names}",
)


if __name__ == "__main__":
unittest.main()
Loading