From 4563f82050f4302890c7e33e0e7afa61ae303c00 Mon Sep 17 00:00:00 2001 From: David Bertoin Date: Tue, 1 Sep 2026 11:10:44 +0000 Subject: [PATCH] Don't materialise the PRX attention mask `PRXAttnProcessor2_0` builds the joint [text | image] mask and then expands it to the full `[B, heads, L_img, L_all]` before handing it to attention. Every backend broadcasts a mask itself, so the expansion only costs bandwidth: * `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 * `sage` / `aiter` / `_native_npu`: reject `attn_mask` outright At PRX-1B's training shape (batch 32, 1024x1024, 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. Passing the unexpanded `[B, 1, 1, L_all]` is bitwise identical -- verified with `torch.equal` on both the block output and the full gradient vector, and against an fp32 unfused-MATH reference the relative error is unchanged to 6 significant figures. Measured on an H200, PRX-1B, batch 32 @ 1024px, bf16 autocast, 5 warmup / 20 timed steps (fake tensors: no dataloader, no text encoder, no loss terms), with `set_attention_backend("_native_cudnn")`: 8 GPU DDP, compiled 446.1 -> 427.5 ms/step peak 110.2 -> 75.2 GiB 1 GPU, compiled 383.8 -> 373.8 ms/step peak 65.6 -> 63.4 GiB 1 GPU, eager 702.3 -> 649.4 ms/step peak 132.7 -> 97.7 GiB On the default `native` backend the same change is 809.1 -> 762.1 ms/step eager. --- src/diffusers/models/transformers/transformer_prx.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/transformers/transformer_prx.py b/src/diffusers/models/transformers/transformer_prx.py index 2676db2e7158..33e1a44dad92 100644 --- a/src/diffusers/models/transformers/transformer_prx.py +++ b/src/diffusers/models/transformers/transformer_prx.py @@ -161,7 +161,10 @@ def __call__( ones_img = torch.ones((bs, l_img), dtype=torch.bool, device=device) attention_mask = attention_mask.to(device=device, dtype=torch.bool) joint_mask = torch.cat([attention_mask, ones_img], dim=-1) - attn_mask_tensor = joint_mask[:, None, None, :].expand(-1, attn.heads, l_img, -1) + # Every attention backend broadcasts the mask itself, so materialising + # [B, heads, L_img, L_all] only costs bandwidth: at batch 32, 1024 image tokens + # and 28 heads that is a 1.1 GiB mask read per block. Keep it broadcastable. + attn_mask_tensor = joint_mask[:, None, None, :] # Apply attention using dispatch_attention_fn for backend support # Reshape to match dispatch_attention_fn expectations: [B, L, H, D]