Don't materialise the PRX attention mask

[DavidBert]

What does this PR do?

PRXAttnProcessor2_0 builds the joint [text | image] attention mask and then expands it to a dense [B, heads, L_img, L_all] before handing it to dispatch_attention_fn:

attn_mask_tensor = joint_mask[:, None, None, :].expand(-1, attn.heads, l_img, -1)

This PR drops the .expand(...) and passes the broadcastable [B, 1, 1, L_all] instead.

Two reasons:

1. It's the documented shape. .ai/references/models.md ("Attention masks") specifies "Padding mask → bool (B, L) or (B, 1, 1, L)". The current dense (B, heads, L_q, L_k) is neither, so this brings PRX in line with the convention rather than introducing a new one.

2. The expansion is pure bandwidth. Every backend broadcasts the mask itself:

• native / _native_* — torch SDPA broadcasts attn_mask natively
• flex — _native_flex_attention does attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv) on a 4-D mask before building the block mask
• xformers — same, attn_mask.expand(...) for a 4-D mask
• flash / _flash_3 / sage / aiter / _native_npu — reject attn_mask outright
• the *_varlen paths reduce it via _normalize_attn_mask, whose 4-D branch is attn_mask.any(dim=(1, 2)) — so [B, 1, 1, L] normalizes fine and keeps the varlen path (a dense additive-float mask would not, per models.md)

At PRX-1B's training shape (batch 32, 1024×1024, patch 32 → 1024 image + 256 text tokens, 28 heads) the expanded mask is 1120 MiB per block, read once per block per forward, for 16 blocks.

Why [B, 1, 1, L] and not the 2-D [B, L_all]

joint_mask is already (B, L_all), so passing it directly looks tempting and _native_attention would handle it (it reshapes a 2-D mask to [B, 1, 1, L_k]). But _native_flex_attention treats dim 1 of a 2-D mask as the query axis — attn_mask.view(attn_mask.size(0), 1, attn_mask.size(1), 1) — which is wrong here, since L_all != seq_len_q for PRX. The 4-D form is unambiguous on every backend, so that's what this uses.

Is it equivalent?

Bitwise, yes. Verified with torch.equal on both the block output and the concatenated gradient vector, across batch sizes 1 and 3, with and without a mask, using a realistically padded mask (half the text tokens masked). Against an fp32 unfused-MATH reference the relative error is unchanged to 6 significant figures on the forward and to the 8th digit on gradients.

Benchmarks

H200, PRX-1B (hidden_size=1792, num_heads=28, depth=16, patch_size=32, bottleneck_size=256), batch 32 @ 1024×1024, 256 text tokens, fp32 master weights + bf16 autocast, 5 warmup / 20 timed steps, torch.cuda.synchronize() around the timed region. Fake tensors — no dataloader, no text encoder, no loss terms — so this isolates the model step.

With set_attention_backend("_native_cudnn"):

	before	after
8 GPU DDP, compiled	446.1 ms/step · peak 110.2 GiB	427.5 ms/step · peak 75.2 GiB
1 GPU, compiled	383.8 ms/step · peak 65.6 GiB	373.8 ms/step · peak 63.4 GiB
1 GPU, eager	702.3 ms/step · peak 132.7 GiB	649.4 ms/step · peak 97.7 GiB

On the default native backend the same change is 809.1 → 762.1 ms/step eager.

The memory is the larger effect: −35 GiB at eager batch 32, −35 GiB on 8-GPU compiled.

Self-review notes

Ran the self-review skill on the diff against .ai/references/review-rules.md, code_style.md and models.md.

Blocking issues: none.

Non-blocking / additional info:

1. No # Copied from linkage. PRXAttnProcessor2_0 carries no # Copied from header and nothing else in src/ or tests/ references it, so no make fix-copies propagation is needed and there is no ripple to other models.
2. Ephemeral context check. The added comment states the reason and the magnitude rather than referring to this PR or its review, so it stands alone for a future reader.
3. Out of scope, but worth more than this PR: models.md also says "Only pass when the batch actually contains padding… See pipeline_qwenimage.py encode_prompt for the pattern: if mask.all(): mask = None." PRXPixelPipeline always passes a mask. Dropping it when it carries no information is worth roughly another 20% of the eager step in my measurements (a no-mask forward also regains the flash path). That is a pipeline behaviour change, so I have deliberately left it out of this PR — happy to open a separate one if you'd like it.
4. Docs. No user-facing behaviour or API change, so no docs/ update applies. The reasoning lives as a comment at the site.

Verdict: READY.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? Yes — Claude Code.
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment? Yes — see above.
• Did you read the contributor guideline?
• Did you read our philosophy doc?
• Was this discussed/approved via a GitHub issue or the forum? No — it is a self-contained perf fix with no behaviour change.
• Did you make sure to update the documentation with your changes? No user-facing behaviour or API change; the reasoning is a code comment at the site.
• Did you write any new necessary tests? No new test — the change is bitwise identical, so the existing PRX tests cover it. Happy to add an explicit mask-shape assertion if you'd prefer one.
• Are you the author (or part of the team) of the model/pipeline? Yes — I authored the PRX model and the PRXPixelPipeline integration (Prx #12525, Add PRXPixelPipeline: pixel-space PRX text-to-image pipeline #13928).

Who can review?

@sayakpaul @yiyixuxu — git log on transformer_prx.py and the models entry in the template. Also related: #12787 (PRX model compilation), since this change matters most under torch.compile.

[HuggingFaceDocBuilderDev]

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.

[yiyixuxu]

thanks