Skip to content

[qwen2-vl] fix FA2 inference #39121

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
zucchini-nlp merged 9 commits into from
Jul 1, 2025
Merged

[qwen2-vl] fix FA2 inference #39121

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
4301faa
fix FA2
zucchini-nlp Jun 30, 2025
2cb7f21
update is causal flag and remove mask for FA2
zucchini-nlp Jun 30, 2025
7823ffc
update for FA2 with varlen path
zucchini-nlp Jun 30, 2025
5640fe1
Merge remote-tracking branch 'upstream/main' into qwen2-fix
zucchini-nlp Jun 30, 2025
ebb46c3
how the tests were passing with different devices?
zucchini-nlp Jun 30, 2025
f2248b7
add comment and ref to the PR
zucchini-nlp Jul 1, 2025
c567fe6
move mask preparation to base pretrained model
zucchini-nlp Jul 1, 2025
3a3afe3
seq len is the first dim, not second
zucchini-nlp Jul 1, 2025
0e1d686
fix copies to fix GLM4V
zucchini-nlp Jul 1, 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
25 changes: 17 additions & 8 deletions src/transformers/modeling_flash_attention_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -508,6 +508,22 @@ def _flash_attention_forward(
query_states, key_states, value_states, target_dtype
)

# We will use `flash_attn_varlen_func` to prevent cross-example attention and also allow padding free approach
# under two cases:
# Case 1. If position_ids is provided and check all examples do not contain only 1 sequence, If tensor in increasing
# then we probably have one sequence, otherwise it is packed. Additionally check we are in pre-fill/training stage.
# Case 2. Some models pass directly pre-computed `cu_seqlens` so we don't need to infer it from position ids. It is safe to
# use `flash_attn_varlen_func` knowing we already have all necessary the kwargs. NOTE: it is user's responsibility
# to take care of flattenning `position_ids` if that's needed by the model. See #39121 for more information
is_fa2_with_position_ids = (
position_ids is not None
and query_states.shape[0] == 1
and (max_length_q is not None or (query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all()))
)
is_fa2_with_varlen_kwargs = all(
kwarg is not None for kwarg in (cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k)
)
Comment on lines +518 to +525

@vasqu vasqu Jun 30, 2025

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.

I'm a bit unsure about this since we would allow cu_seq and max_seq only but on most models we also have RoPE so it's breaking those models silently if we eff up not passing position_ids (due to RoPE positions being bound to position_ids as well). We should imo add at least a warning on only varlen kwargs to give some discretion here.

On another note, what do the integration tests use? Are they still working as expected 👀 seems a bit sus

@zucchini-nlp zucchini-nlp Jul 1, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

The integration tests always use position_ids which is the first is_fa2_with_position_ids, I just copied it from existing code and moved up here. The second case is added for Qwen only, afaik no other model passes pre-computed cu_lens form attention layers

In qwen we don't need any position ids, because they are 3D and won't help at all in inferring cu_lens

@vasqu vasqu Jul 1, 2025 •
edited
Loading

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.

I meant the general integration tests like e.g.

transformers/tests/models/qwen2_vl/test_modeling_qwen2_vl.py

Line 507 in d53518c

def test_small_model_integration_test_batch_flashatt2(self):

For why, I'm concerned about the second case on qwen is future model additions and general usability, not the validity of qwen. For developers, is_fa2_with_varlen_kwargs indicates that this suffices for varlen - before, we (unintentionally) checked for the existence of (correct flattened) position ids that RoPE models need when using varlen. Maybe #35941 helps for reference on what I mean.

Imo, it would help to add at least comments that for varlen most models need correct flattened position ids (from e.g. a collator), especially RoPE models which make up the majority of newer models.

@zucchini-nlp zucchini-nlp Jul 1, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Hmm, am I right that users might be passing only cu_lens without correct position_ids? I believe that would be users' responsibility to take care that RoPE is applied correctly, but I will add a comment in code explaining it, sure

In slow integration tests we don't pass position_ids, not that of I know. For most LLMs the fa2 path integration tests fallback to inferring cu_lens from the mask, and in Qwen the position_ids are constructed on-the-fly during forward call. The model has a requirement for adding rope deltas on top of 3D positions and I don't think users would be doing all that manually

@vasqu vasqu Jul 1, 2025 •
edited
Loading

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.

"Hmm, am I right that users might be passing only cu_lens without correct position_ids?" - Yes, not only users but possibly us as well because it's something that's harder to figure out when done wrong imo :D


# Contains at least one padding token in the sequence
if attention_mask is not None:
batch_size = query_states.shape[0]
Expand All @@ -531,14 +547,7 @@ def _flash_attention_forward(
)
attn_output = _pad_input(attn_output_unpad, indices_q, batch_size, query_length)

# If position_ids is provided and check all examples do not contain only 1 sequence, If tensor in increasing
# then we probably have one sequence, otherwise it is packed. Additionally check we are in pre-fill/training stage.
# Use `flash_attn_varlen_func` to prevent cross-example attention and also allow padding free approach
elif (
position_ids is not None
and query_states.shape[0] == 1
and (max_length_q is not None or (query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all()))
):
elif is_fa2_with_varlen_kwargs or is_fa2_with_position_ids:
batch_size = query_states.size(0)

if cu_seq_lens_q is None or cu_seq_lens_k is None:
Expand Down
67 changes: 37 additions & 30 deletions src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -634,14 +634,17 @@ def forward(
value_states = value_states.transpose(0, 1).unsqueeze(0)
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

attention_mask = torch.full(
[1, 1, seq_length, key_states.shape[-2]],
torch.finfo(query_states.dtype).min,
device=query_states.device,
dtype=query_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
# Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
attention_mask = None
if self.config._attn_implementation != "flash_attention_2":
attention_mask = torch.full(
[1, 1, seq_length, key_states.shape[-2]],
torch.finfo(query_states.dtype).min,
device=query_states.device,
dtype=query_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0

@vasqu vasqu Jun 30, 2025

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 probably also fixes #39067? Might be nice to change the flash attention integration path in the future to prioritize pos+cu_seq (even with an attention mask).

@zucchini-nlp zucchini-nlp Jun 30, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

No, it doesn't, that issue was there before refactor and I found it's related to FA2 using ragged inputs which the SDPA/eager path doesn't support in the same way

I will comment under #39067 when I have a clear fix. Let's keep this PR open until then, better to fix it once and forever


attention_interface: Callable = eager_attention_forward
if self.config._attn_implementation != "eager":
Expand All @@ -652,13 +655,13 @@ def forward(
query_states,
key_states,
value_states,
attention_mask,
dropout=0.0 if not self.training else self.dropout,
attention_mask=attention_mask,
dropout=0.0,
scaling=self.scaling,
cu_seqlens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
cu_seq_lens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seq_lens_k=cu_seqlens,
max_length_q=max_seqlen,
max_length_k=max_seqlen,
is_causal=False,
**kwargs,
)
Expand Down Expand Up @@ -928,6 +931,7 @@ def __init__(self, config: Qwen2_5OmniVisionEncoderConfig = None) -> None:
self.scaling = self.head_dim**-0.5
self.num_key_value_groups = 1 # needed for eager attention
self.config = config
self.is_causal = False

def forward(
self,
Expand All @@ -943,18 +947,21 @@ def forward(
query_states = apply_rotary_pos_emb_vision(query_states.unsqueeze(0), rotary_pos_emb).squeeze(0)
key_states = apply_rotary_pos_emb_vision(key_states.unsqueeze(0), rotary_pos_emb).squeeze(0)

attention_mask = torch.full(
[1, 1, seq_length, seq_length],
torch.finfo(query_states.dtype).min,
device=query_states.device,
dtype=query_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
# Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
attention_mask = None
if self.config._attn_implementation != "flash_attention_2":
attention_mask = torch.full(
[1, 1, seq_length, seq_length],
torch.finfo(value_states.dtype).min,
device=value_states.device,
dtype=value_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0

query_states = query_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
key_states = key_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
value_states = value_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
query_states = query_states.transpose(0, 1).unsqueeze(0)
key_states = key_states.transpose(0, 1).unsqueeze(0)
value_states = value_states.transpose(0, 1).unsqueeze(0)
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

attention_interface: Callable = eager_attention_forward
Expand All @@ -966,13 +973,13 @@ def forward(
query_states,
key_states,
value_states,
attention_mask,
attention_mask=attention_mask,
dropout=0.0,
scaling=self.scaling,
cu_seqlens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
cu_seq_lens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seq_lens_k=cu_seqlens,
max_length_q=max_seqlen,
max_length_k=max_seqlen,
is_causal=False,
**kwargs,
)
Expand Down
67 changes: 37 additions & 30 deletions src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1638,14 +1638,17 @@ def forward(
value_states = value_states.transpose(0, 1).unsqueeze(0)
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

attention_mask = torch.full(
[1, 1, seq_length, key_states.shape[-2]],
torch.finfo(query_states.dtype).min,
device=query_states.device,
dtype=query_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
# Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
attention_mask = None
if self.config._attn_implementation != "flash_attention_2":
attention_mask = torch.full(
[1, 1, seq_length, key_states.shape[-2]],
torch.finfo(query_states.dtype).min,
device=query_states.device,
dtype=query_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0

attention_interface: Callable = eager_attention_forward
if self.config._attn_implementation != "eager":
Expand All @@ -1656,13 +1659,13 @@ def forward(
query_states,
key_states,
value_states,
attention_mask,
dropout=0.0 if not self.training else self.dropout,
attention_mask=attention_mask,
dropout=0.0,
scaling=self.scaling,
cu_seqlens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
cu_seq_lens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seq_lens_k=cu_seqlens,
max_length_q=max_seqlen,
max_length_k=max_seqlen,
is_causal=False,
**kwargs,
)
Expand Down Expand Up @@ -1906,6 +1909,7 @@ def __init__(self, config: Qwen2_5OmniVisionEncoderConfig = None) -> None:
self.scaling = self.head_dim**-0.5
self.num_key_value_groups = 1 # needed for eager attention
self.config = config
self.is_causal = False

def forward(
self,
Expand All @@ -1921,18 +1925,21 @@ def forward(
query_states = apply_rotary_pos_emb_vision(query_states.unsqueeze(0), rotary_pos_emb).squeeze(0)
key_states = apply_rotary_pos_emb_vision(key_states.unsqueeze(0), rotary_pos_emb).squeeze(0)

attention_mask = torch.full(
[1, 1, seq_length, seq_length],
torch.finfo(query_states.dtype).min,
device=query_states.device,
dtype=query_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
# Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
attention_mask = None
if self.config._attn_implementation != "flash_attention_2":
attention_mask = torch.full(
[1, 1, seq_length, seq_length],
torch.finfo(value_states.dtype).min,
device=value_states.device,
dtype=value_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0

query_states = query_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
key_states = key_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
value_states = value_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
query_states = query_states.transpose(0, 1).unsqueeze(0)
key_states = key_states.transpose(0, 1).unsqueeze(0)
value_states = value_states.transpose(0, 1).unsqueeze(0)
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

attention_interface: Callable = eager_attention_forward
Expand All @@ -1944,13 +1951,13 @@ def forward(
query_states,
key_states,
value_states,
attention_mask,
attention_mask=attention_mask,
dropout=0.0,
scaling=self.scaling,
cu_seqlens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
cu_seq_lens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seq_lens_k=cu_seqlens,
max_length_q=max_seqlen,
max_length_k=max_seqlen,
is_causal=False,
**kwargs,
)
Expand Down
36 changes: 20 additions & 16 deletions src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,7 @@ def __init__(self, config: Qwen2_5_VLVisionConfig) -> None:
self.proj = nn.Linear(self.dim, self.dim)
self.scaling = self.head_dim**-0.5
self.config = config
self.is_causal = False

def forward(
self,
Expand Down Expand Up @@ -233,18 +234,21 @@ def forward(
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)

attention_mask = torch.full(
[1, 1, seq_length, seq_length],
torch.finfo(value_states.dtype).min,
device=value_states.device,
dtype=value_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
# Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
attention_mask = None
if self.config._attn_implementation != "flash_attention_2":
attention_mask = torch.full(
[1, 1, seq_length, seq_length],
torch.finfo(value_states.dtype).min,
device=value_states.device,
dtype=value_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0

query_states = query_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
key_states = key_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
value_states = value_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
query_states = query_states.transpose(0, 1).unsqueeze(0)
key_states = key_states.transpose(0, 1).unsqueeze(0)
value_states = value_states.transpose(0, 1).unsqueeze(0)
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

attention_interface: Callable = eager_attention_forward
Expand All @@ -256,13 +260,13 @@ def forward(
query_states,
key_states,
value_states,
attention_mask,
attention_mask=attention_mask,
dropout=0.0,
scaling=self.scaling,
cu_seqlens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
cu_seq_lens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seq_lens_k=cu_seqlens,
max_length_q=max_seqlen,
max_length_k=max_seqlen,
is_causal=False,
**kwargs,
)
Expand Down
36 changes: 20 additions & 16 deletions src/transformers/models/qwen2_vl/modeling_qwen2_vl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -324,6 +324,7 @@ def __init__(self, config: Qwen2VLVisionConfig) -> None:
self.proj = nn.Linear(self.dim, self.dim)
self.scaling = self.head_dim**-0.5
self.config = config
self.is_causal = False

def forward(
self,
Expand Down Expand Up @@ -351,18 +352,21 @@ def forward(
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)

attention_mask = torch.full(
[1, 1, seq_length, seq_length],
torch.finfo(value_states.dtype).min,
device=value_states.device,
dtype=value_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
# Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
attention_mask = None
if self.config._attn_implementation != "flash_attention_2":
attention_mask = torch.full(
[1, 1, seq_length, seq_length],
torch.finfo(value_states.dtype).min,
device=value_states.device,
dtype=value_states.dtype,
)
for i in range(1, len(cu_seqlens)):
attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0

query_states = query_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
key_states = key_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
value_states = value_states.transpose(0, 1).unsqueeze(0) # unsqueeze batch_dim
query_states = query_states.transpose(0, 1).unsqueeze(0)
key_states = key_states.transpose(0, 1).unsqueeze(0)
value_states = value_states.transpose(0, 1).unsqueeze(0)
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

attention_interface: Callable = eager_attention_forward
Expand All @@ -374,13 +378,13 @@ def forward(
query_states,
key_states,
value_states,
attention_mask,
attention_mask=attention_mask,
dropout=0.0,
scaling=self.scaling,
cu_seqlens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
cu_seq_lens_q=cu_seqlens, # pass cu seq lens for FA2
cu_seq_lens_k=cu_seqlens,
max_length_q=max_seqlen,
max_length_k=max_seqlen,
is_causal=False,
**kwargs,
)
Expand Down
2 changes: 1 addition & 1 deletion tests/models/qwen2_5_vl/test_modeling_qwen2_5_vl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ def prepare_config_and_inputs_for_common(self):
input_ids[:, self.num_image_tokens - 1] = self.vision_start_token_id
inputs_dict = {
"pixel_values": pixel_values,
"image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size),
"image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size, device=torch_device),
"input_ids": input_ids,
"attention_mask": attention_mask,
}
Expand Down
2 changes: 1 addition & 1 deletion tests/models/qwen2_vl/test_modeling_qwen2_vl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ def prepare_config_and_inputs_for_common(self):

inputs_dict = {
"pixel_values": pixel_values,
"image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size),
"image_grid_thw": torch.tensor([[1, 1, 1]] * self.batch_size, device=torch_device),
"input_ids": input_ids,
"attention_mask": attention_mask,
}
Expand Down