Add optional key-only masks for NaFlexVit self-attention

[chenghuichen]

Summary

NaFlexVit already uses a [B, 1, 1, N] key-only mask for MAP attention pooling. This change adds an opt-in use_key_only_attn_mask setting to use the same broadcastable mask form in transformer self-attention, instead of materializing the symmetric [B, 1, N, N] mask.

For valid query tokens, both forms mask the same keys. The default remains unchanged, explicitly supplied attention masks are unaffected, and the generated mask is created after patch dropout from the updated patch_valid.

The option can be enabled with create_model(..., use_key_only_attn_mask=True).

Memory impact

For B=16, N=1024, and bf16, the generated mask is 32 KiB instead of 32 MiB.

In a small NVIDIA A10 benchmark with PyTorch 2.11 fused SDPA, the per-iteration peak allocated memory was:

Mode	Full-mask run	Compact-mask run	Reduction
Inference	152.000 MiB	120.031 MiB	31.969 MiB
Forward + backward	785.021 MiB	753.052 MiB	31.969 MiB

Benchmark outputs were bitwise identical. Median iteration time was 8.404 ms versus 7.980 ms for inference, and 29.312 ms versus 28.021 ms for forward + backward. The reproduction script:
benchmark_naflex_compact_key_mask.py

Tests

Added unit tests for mask shape, patch-dropout alignment, and exact output parity with average, token, and MAP pooling.

[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.