Recover FlashAttention-2 on video training (the frame-mask key-padding disables the fast path)

[amazloumi]

What

Video training always takes the explicit-attn_mask SDPA branch in Attention (and MoTAttention), even for fully-decoded clips with no padding, so it loses the FlashAttention-2 fast path on every layer.

Why: VideoCollator always emits frame_mask, so _prefix_key_padding_mask always returns a non None mask (all-True when unpadded). Attention then builds a (B, 1, S, S) bool mask + a ~mask.any(-1) reduction and calls F.scaled_dot_product_attention(..., attn_mask=...). An explicit mask is not an FA2-supported shape, so SDPA falls back to the mem-efficient/math kernel. The result is identical (an all-True mask AND causal is causal), but FA2 is gone and the mask is materialized — a throughput/memory cost that grows with sequence length (more frames / longer video).

This was a deliberate choice in the frame-mask change (#132): always-masking is torch.compile/DP-friendly (one graph, no host sync). This issue tracks recovering the fast path. Two approaches:

(A) Skip the mask when a batch has no padding (cheap, partial). VideoCollator already knows which clips it padded, so it can emit a CPU-side has_padding: bool (no GPU sync); train.py passes frame_mask=None when not has_padding, restoring the is_causal → FA2 path for fully-decoded batches. Caveat: reintroduces a data-dependent branch — with torch.compile enabled (off by default in KF) it holds two graphs / recompiles when has_padding flips, which is exactly the compile-readiness the always-mask approach preserved. Recovers FA2 only for fully-unpadded batches.

(B) FlexAttention — the clean fix (depends on #97). FlexAttention applies an arbitrary mask inside a fused, compile-native kernel, so there's no FA2 fallback for padded or unpadded clips. Fixes the cost for all cases and removes the data-dependent branch; the right answer as video sequence length grows (where the FA2/memory gap matters most). Larger effort, gated on #97.

Scope

• kempnerforge/data/video_dataset.py (VideoCollator): emit a cheap CPU-side has_padding flag — option (A).
• scripts/train.py: pass frame_mask=None when not has_padding — option (A).
• kempnerforge/model/attention.py, kempnerforge/model/mot.py: move the masked SDPA path to FlexAttention — option (B), with Add FlexAttention as an alternative attention backend #97.
• No change to masking semantics — both options produce identical logits/loss; this is purely the attention kernel/perf path.

Backward compatibility

• Pure performance change — masking results are unchanged.
• Option (A) adds a data-dependent branch: benign for correctness (DP ranks may pick different kernels — fine), but interacts with torch.compile (two graphs / recompiles). Measure the actual gain at representative clip lengths first — it's narrow (unpadded batches only) and may not justify nicking compile-readiness.
• Option (B) is the durable fix and removes the branch entirely.