Sdpa for owlvit #42136

Aravind-11 · 2025-11-10T23:38:23Z

What does this PR do?

Implements SDPA for OWL VIT.

Fixes #28103

Before submitting

Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
[] Did you write any new necessary tests?

Who can review?

@vasqu @younesbelkada

Aravind-11 · 2025-11-11T00:35:19Z

What does this PR do?

Implements SDPA for OWL VIT. Revamp of #28818

Fixes #28103

Before submitting

Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.

Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.

[] Did you write any new necessary tests?

Who can review?

@vasqu @younesbelkada

I ran the RUN_SLOW=1 python -m pytest tests/models/owlvit/test_modeling_owlvit.py for the original owlvit implementation and it seemed to fail the same tests as my current implementation. I'm not sure how to infer from that.

vasqu

Sorry but I've got to be strict about this. We no longer implement separate classes for all the attention flavors but one unified one. I think ViT is a good example in this case, e.g. see https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit/modeling_vit.py

Before changing this to these standards I won't take a proper look for now.

Aravind-11 · 2025-11-11T18:51:20Z

Sorry but I've got to be strict about this. We no longer implement separate classes for all the attention flavors but one unified one. I think ViT is a good example in this case, e.g. see https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit/modeling_vit.py

Before changing this to these standards I won't take a proper look for now.

Got it. Thanks a lot!

Aravind-11 · 2025-11-12T06:42:48Z

Sorry but I've got to be strict about this. We no longer implement separate classes for all the attention flavors but one unified one. I think ViT is a good example in this case, e.g. see https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit/modeling_vit.py

Before changing this to these standards I won't take a proper look for now.

I made similar changes as in the vit and removed the seperate sdpa class. Let me know what you think!

vasqu

Added some comments but in general it would be best to have a green CI before requesting a review. Atm, things are likely not working as expected

vasqu · 2025-11-12T12:43:42Z

src/transformers/models/owlvit/modeling_owlvit.py

-        causal_attention_mask = _create_4d_causal_attention_mask(
-            input_shape, hidden_states.dtype, device=hidden_states.device
+        # OWL-ViT uses a bidirectional (non-causal) encoder.
+        attention_mask = create_bidirectional_mask(
+            config=self.config,
+            input_embeds=hidden_states,
+            attention_mask=attention_mask,
        )
-        # expand attention_mask
-        if attention_mask is not None:
-            # [num_samples, seq_len] -> [num_samples, 1, tgt_seq_len, src_seq_len]
-            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)


This seems to suffer from the same issue as in #41750

It does not use a bidirectional mask, but a causal mask:

The first mask is a based causal mask

The second is a padding mask

These are added on top creating a causal mask with padding included

This also may need to adjust the is_causal argument dynamically as in the PR I linked - although I'm not sure if it's just causal in general

Thanks! I made some changes to the code after referring to CLIP - removing the output_attention, return dict and casual_attention_mask. Also copied the eager attention part, attention reshaping from CLIP. Added the flash and flex attn too.

I think that the current CI is failing because the OWL VIT config file is conflicting with the current encoder implementation. Could you guide me here? Thanks a lot!

Thanks! I made some changes to the code after referring to CLIP - removing the output_attention, return dict and casual_attention_mask. Also copied the eager attention part, attention reshaping from CLIP. Added the flash and flex attn too.

I think that the current CI is failing because the OWL VIT config file is conflicting with the current encoder implementation. Could you guide me here? Thanks a lot!

Hi, I investigated the failing OwlViTForObjectDetectionTest::test_eager_matches_sdpa_inference_09_fp32_pad_left.

The failure is due to the test invoking OwlViTForObjectDetection.forward() without providing pixel_values.

OwlViTForObjectDetection requires pixel_values (image tensors) for its vision backbone. When the test omits them, the model raises a ValueError: 'pixel_values' is None.

vasqu · 2025-11-12T12:44:51Z

src/transformers/models/owlvit/modeling_owlvit.py

    base_model_prefix = "owlvit"
    input_modalities = ["image", "text"]
    supports_gradient_checkpointing = True
+    _supports_sdpa = True


It should also support flash attn and flex attention then

vasqu · 2025-11-12T12:46:29Z

src/transformers/models/owlvit/modeling_owlvit.py



-# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->OwlViT
+# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->OwlViT


This doesn't work, output_attentions would be ignored --> see the check_model_inputs decorator in other models like CLIP

vasqu · 2025-11-12T12:48:20Z

src/transformers/models/owlvit/modeling_owlvit.py

+            # Eager attention implementation
+            # Scale query


Please follow what is done in other models, we should be able to copy the eager function from bert (see clip as example)

vasqu · 2025-11-12T12:48:45Z

src/transformers/models/owlvit/modeling_owlvit.py

+        # Prepare attention mask - OWL-ViT uses bidirectional (non-causal) attention
+        attention_mask = create_bidirectional_mask(
+            config=self.config,
+            input_embeds=hidden_states,
+            attention_mask=attention_mask,
+        )


The mask is prepared outside, not within the attention module

vasqu · 2025-11-12T12:49:16Z

src/transformers/models/owlvit/modeling_owlvit.py

+        # Get query, key, value projections
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)

-        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
-        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
-        key_states = key_states.view(*proj_shape)
-        value_states = value_states.view(*proj_shape)
+        # Reshape to (bsz, num_heads, seq_len, head_dim)
+        query_states = query_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)


Please follow llama / clip, the -1 dim is also important for example

github-actions · 2025-11-16T05:06:10Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: owlv2, owlvit

vasqu reviewed Nov 11, 2025

View reviewed changes

Aravind-11 force-pushed the sdpa_for_OWL_ViT branch from d519ced to 77f5221 Compare November 12, 2025 05:15

vasqu reviewed Nov 12, 2025

View reviewed changes

nileshkokane01 and others added 24 commits November 15, 2025 18:22

Added sdpa attention

f9498f3

Added Changes to OWL-Vit as suggested

ddd9eeb

Fixed nits

c20d934

removed unwanted files

a75f357

Fixed nits

2c07d0f

Fixed past_key_values_length to length 0

a1d25d1

Fixed nits

62f47ab

Fixed dim issue

9c6dfbf

fixed nits

2ccdc99

Fixed dim issue

ec085b3

add sdpa for owlvit

a1c22e4

fixes

8cff000

fixe

b12a5fc

fixes

9ca23ae

fixes

9923d24

fixes

7a7977d

remove return_dict from config file

a08d71d

remove return_dict from ovlv2 config

de20883

commit same changes to owlv2

5d5cf7d

make fix-copies

b281a10

fixes

9b9020a

fixes

4dbbca2

fixes

90d32e9

fixes

57e88a3

Aravind-11 added 4 commits November 15, 2025 18:22

fixes

83e244e

fixes

8e32ea9

revert changes

aea9460

fixes

42971a3

Aravind-11 force-pushed the sdpa_for_OWL_ViT branch from 41a76bf to 42971a3 Compare November 16, 2025 01:43

Aravind-11 added 4 commits November 15, 2025 18:52

fixes

752540e

fixes

e64aa6f

fixes

3597dd5

make style

7ad3918



		# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->OwlViT
		# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->OwlViT

Sdpa for owlvit #42136

Are you sure you want to change the base?

Sdpa for owlvit #42136

Conversation

Aravind-11 commented Nov 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

Aravind-11 commented Nov 11, 2025

What does this PR do?

Before submitting

Who can review?

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Aravind-11 commented Nov 11, 2025

Uh oh!

Aravind-11 commented Nov 12, 2025

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Nov 16, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Aravind-11 commented Nov 10, 2025 •

edited

Loading