Fix Pixtral Flash Attention image boundaries after axial RoPE refactor by oliverholworthy · Pull Request #49373 · huggingface/transformers

oliverholworthy · 2026-10-06T16:37:16Z

What does this PR do?

Pixtral passes axial RoPE coordinates shaped [total_patches, 2] to the Flash Attention packed-sequence detector. It flattens both coordinates, producing false image boundaries and a terminal cumulative length twice the number of Q/K/V tokens. A checkpoint-free Pixtral model raises a CUDA illegal memory access on Transformers 5.17.0, 5.18.0, and main at 571c5cc; smaller configurations can return incorrect output without crashing. The same reproduction passes on 5.16.1.

This change derives cumulative lengths and maximum lengths from the per-image patch grids, and passes them explicitly in the flash attention case. Keeping axial coordinates for rotary embeddings and preserve the existing block mask for other attention implementations.

This follows the Qwen2-VL/Qwen3-VL vision-attention approach:

transformers/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py

Lines 287 to 298 in 571c5cc

    
               key_states, 
        
               value_states, 
        
               attention_mask=None, 
        
               scaling=self.scaling, 
        
               dropout=0.0 if not self.training else self.attention_dropout, 
        
               cu_seq_lens_q=cu_seqlens, 
        
               cu_seq_lens_k=cu_seqlens, 
        
               max_length_q=max_seqlen, 
        
               max_length_k=max_seqlen, 
        
               is_causal=False, 
        
               **kwargs, 
        
           )

Related work / coordination

🚨 Vision (2d/3d) rotary embeddings #48105 introduced the axial-position change.
Open Fix _is_packed_sequence for 1D and 3D position_ids #47342 normalizes 1D and 3D position IDs. Its current diff leaves this 2D [total_patches, 2] representation unchanged, so it does not fix this reproduction. This proposal instead supplies explicit image boundaries within Pixtral.
Open Update Pixtral to use SDPA #34542 adds SDPA support to Pixtral; it addresses a different problem.
Searched open issues and PRs on 2026-10-06. No matching fix found

Who can review?

@zucchini-nlp

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>

zucchini-nlp · 2026-10-07T11:03:18Z

-            # We only rely on position_ids when using flash attention
+            # Axial RoPE positions have two coordinates per patch, so they cannot
+            # describe the packed image boundaries expected by Flash Attention.
+            sequence_lengths = [p.shape[-2] * p.shape[-1] for p in patch_embeds_list]
+            cu_seqlens = torch.tensor([0, *sequence_lengths], device=patch_embeds.device, dtype=torch.int32)
+            cu_seqlens = cu_seqlens.cumsum(dim=0, dtype=torch.int32)
+            kwargs.update(
+                cu_seq_lens_q=cu_seqlens,
+                cu_seq_lens_k=cu_seqlens,
+                max_length_q=max(sequence_lengths),
+                max_length_k=max(sequence_lengths),
+            )
            attention_mask = None


hmm quite the same as vision utils but oke, not really worth materializing image_grid only to infer FA kwargs

transformers/src/transformers/vision_utils.py

Lines 68 to 78 in e598fba

def get_vision_attention_seqlens(

grid_thw: torch.Tensor,

config: PreTrainedConfig,

merge_temporal: bool = False,

kwargs: dict | None = None,

) -> tuple[torch.Tensor, int | None]:

"""Get cumulative and maximum sequence lengths for packed vision attention. ``merge_temporal`` is

forwarded to [`get_vision_cu_seqlens`] (``True`` for clip-level attention, e.g. kimi_k25)."""

cu_seqlens = get_vision_cu_seqlens(grid_thw, merge_temporal=merge_temporal, kwargs=kwargs)

max_seqlen = get_max_seqlen(cu_seqlens, config, kwargs=kwargs)

return cu_seqlens, max_seqlen

zucchini-nlp · 2026-10-07T11:04:50Z

+    @require_flash_attn
+    @require_torch_gpu
+    def test_flash_attention_2_packed_images(self):
+        config = PixtralVisionConfig(
+            hidden_size=256,
+            intermediate_size=512,


actually we have a test tests/models/pixtral/test_modeling_pixtral.py::PixtralVisionModelModelTest::test_flash_attn_2_inference_equivalence

and i see it failing on main so lets just make sure it passes now, this new test is not needed

yeah, the workflow is all green so lets just delete the test and i'll merge

Thanks for the review and running the tests. I've removed the new test.

zucchini-nlp · 2026-10-07T11:05:07Z

run-slow: pixtral

github-actions · 2026-10-07T11:06:46Z

AMD CI

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs on AMD:

models: ["models/pixtral"]

github-actions · 2026-10-07T11:06:47Z

Nvidia CI

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs on Nvidia:

models: ["models/pixtral"]
quantizations: []

github-actions · 2026-10-07T11:12:22Z

CI Results (AMD)

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	0f9114b5	workflow commit (merge commit)
PR	621dcd13	branch commit (from PR)
main	cbde22f4	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

HuggingFaceDocBuilderDev · 2026-10-07T11:14:44Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

github-actions · 2026-10-07T11:15:06Z

CI Results (Nvidia)

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	0f9114b5	workflow commit (merge commit)
PR	621dcd13	branch commit (from PR)
main	cbde22f4	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>

github-actions · 2026-10-07T11:31:44Z

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

run-slow: pixtral

github-actions · 2026-10-07T11:36:02Z

CI recap

Dashboard: View test results in Grafana
Latest run: 37611695664:1
Result: success | Jobs: 1 | Tests: 46 | Failures: 0 | Duration: 1m 52s

oliverholworthy added 2 commits October 6, 2026 15:39

Fix Pixtral Flash Attention packed image boundaries

5454956

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>

Cover unequal image lengths in Pixtral Flash Attention regression

621dcd1

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>

zucchini-nlp self-requested a review October 6, 2026 16:38

zucchini-nlp approved these changes Oct 7, 2026

View reviewed changes

Use existing Pixtral Flash Attention equivalence coverage

3752a6f

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>

zucchini-nlp enabled auto-merge October 7, 2026 11:38

zucchini-nlp added this pull request to the merge queue Oct 7, 2026

Merged via the queue into huggingface:main with commit ff704ff Oct 7, 2026
39 checks passed

akoumpa mentioned this pull request Oct 8, 2026

fix(retrieval): backport Pixtral FlashAttention image boundaries NVIDIA-NeMo/Automodel#4204

Merged

3 tasks done

stevhliu mentioned this pull request Oct 23, 2026

[docs] Translation workflow #49889

Closed

Navigation Menu

Repository navigation

Fix Pixtral Flash Attention image boundaries after axial RoPE refactor - #49373

Fix Pixtral Flash Attention image boundaries after axial RoPE refactor#49373
zucchini-nlp merged 3 commits into
huggingface:mainfrom
oliverholworthy:oholworthy/fix-pixtral-flash-attention-image-boundaries

oliverholworthy commented Oct 6, 2026 •

edited by github-actions Bot

Loading

Uh oh!

zucchini-nlp Oct 7, 2026

Uh oh!

zucchini-nlp Oct 7, 2026

Uh oh!

zucchini-nlp Oct 7, 2026

Uh oh!

oliverholworthy Oct 7, 2026

Uh oh!

zucchini-nlp commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	key_states,
	value_states,
	attention_mask=None,
	scaling=self.scaling,
	dropout=0.0 if not self.training else self.attention_dropout,
	cu_seq_lens_q=cu_seqlens,
	cu_seq_lens_k=cu_seqlens,
	max_length_q=max_seqlen,
	max_length_k=max_seqlen,
	is_causal=False,
	**kwargs,
	)

	def get_vision_attention_seqlens(
	grid_thw: torch.Tensor,
	config: PreTrainedConfig,
	merge_temporal: bool = False,
	kwargs: dict \| None = None,
	) -> tuple[torch.Tensor, int \| None]:
	"""Get cumulative and maximum sequence lengths for packed vision attention. ``merge_temporal`` is
	forwarded to [`get_vision_cu_seqlens`] (``True`` for clip-level attention, e.g. kimi_k25)."""
	cu_seqlens = get_vision_cu_seqlens(grid_thw, merge_temporal=merge_temporal, kwargs=kwargs)
	max_seqlen = get_max_seqlen(cu_seqlens, config, kwargs=kwargs)
	return cu_seqlens, max_seqlen

Conversation

oliverholworthy commented Oct 6, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Related work / coordination

Who can review?

Uh oh!

zucchini-nlp Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

zucchini-nlp Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

zucchini-nlp Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

oliverholworthy Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

zucchini-nlp commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

AMD CI

Uh oh!

github-actions Bot commented Oct 7, 2026

Nvidia CI

Uh oh!

github-actions Bot commented Oct 7, 2026

CI Results (AMD)

Commit Info

Uh oh!

HuggingFaceDocBuilderDev commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

CI Results (Nvidia)

Commit Info

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

CI recap

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

oliverholworthy commented Oct 6, 2026 •

edited by github-actions Bot

Loading