Broadcast SDPA attn_mask instead of materializing a per-head copy

[jannisringwald]

What

One-line change in Attention.forward (src/depth_anything_3/model/dinov2/layers/attention.py, fused path):

# before
attn_mask=(attn_mask)[:, None].repeat(1, self.num_heads, 1, 1) if attn_mask is not None else None
# after
attn_mask=attn_mask[:, None] if attn_mask is not None else None

F.scaled_dot_product_attention broadcasts the attention mask across the head dimension (docs), so the [B, 1, N, N] view is mathematically identical to the repeated [B, num_heads, N, N] tensor — the kernel sees the same mask values either way.

Why

For any caller that threads attn_mask through the backbone kwargs (DinoV2.forward(**kwargs) → get_intermediate_layers → the global-attention blocks), the .repeat() materializes a physical per-head copy before the kernel even runs. At DA3's multi-view global-attention token counts that is a very large single allocation, growing quadratically with view count (378×504 inputs, vitb: 973 tokens/view, 12 heads, fp32 mask):

views	N = views × 973	mask copy from .repeat()	broadcast view
8	7,784	2.7 GiB	0 bytes
14	13,622	8.3 GiB	0 bytes
32	31,136	43.3 GiB	0 bytes

A 14-view masked call would allocate 8.3 GiB just for the mask copy — more than half the memory of a 16 GB T4. We noticed the line while auditing the SDPA memory behavior of 14-view inference on 16 GB T4s.

Reproducible measurement of the line itself:

import torch
mask = torch.zeros(1, 13622, 13622, device="cuda")          # 0.7 GiB
torch.cuda.reset_peak_memory_stats()
m = mask[:, None].repeat(1, 12, 1, 1)                        # current code
print(torch.cuda.max_memory_allocated() / 2**30)             # ~8.3 GiB
del m; torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
m = mask[:, None]                                            # this PR
print(torch.cuda.max_memory_allocated() / 2**30)             # 0.0

Safety

• The in-repo inference pipelines call the backbone with attn_mask=None, so this is a no-op for them — verified A/B (patched vs unpatched) on a real 8-view T4 inference: identical peak memory (7.77 GB), outputs matching to run-to-run kernel noise (median/p95 relative depth error 0.000000).
• For mask-passing callers, broadcasting delivers the same mask values to the kernel, so results are unchanged.
• The non-fused (else) branch is untouched — it never materialized a per-head mask.