Add attention_backend="flex" for packed sequences

[timothyngo]

Summary

Adds model.attention_backend = "flex", routing packed batches through FlexAttention's block-sparse mask instead of the dense (B, 1, S, S) mask the packed path builds today.

Stacked on #205 — this PR targets fix/reject-packing-under-pp, so review that one first. GitHub will show only this PR's own diff.

Why

attention.py:214-227 materializes a dense boolean mask in every layer and hands it to F.scaled_dot_product_attention. An explicit attn_mask is not a FlashAttention-2 shape, so SDPA silently falls back to the mem-efficient/math kernel. Packing therefore gave up the fast kernel in order to save padding, and rebuilt an O(B·S²) mask per layer to do it — at seq_len=8192, B=4 that is a 256 MB bool tensor constructed 32 times per forward. The comment already at attention.py:201 names FlexAttention as the intended fix; this is it.

What

• kempnerforge/model/masking.py (new) — build_doc_causal_block_mask (block-diagonal causal; H=None so the mask broadcasts over heads, which is what keeps it correct under TP's sharded head counts) and flex_attention_fn, which compiles flex_attention once per process on CUDA and stays eager on CPU.
• attention.py — a block_mask branch placed after the capture_attention_weights check and before the dense-mask path. Ordering is load-bearing: MoMaBlock builds this same class and passes both doc_ids and key_padding_mask, so a flex branch placed first would silently drop the padding mask. GQA uses enable_gqa rather than repeat_interleave, so repeated K/V heads are never materialized.
• transformer.py — builds the BlockMask once per forward and threads it to every block, replacing the per-layer dense mask. Raises if doc_ids is shorter than the sequence reaching attention (the VLM image-prefix case).
• config/model.py / config/job.py — the attention_backend field, FLEX_BLOCK_SIZE, and validation.

Behavior change

Default "sdpa" is bit-identical to before — verified, not assumed: test_no_doc_ids_is_bit_identical_to_sdpa and test_doc_ids_none_is_exact_noop. The flag only does anything when data.pack_sequences is on; unpacked batches take the is_causal SDPA fast path under either backend.

The seq_len >= 128 restriction

This one is worth a close look. Under torch.compile, flex silently leaks across documents when seq_len is below FlexAttention's 128 block size, while eager stays correct. Measured drift in document-1 logits when document-0 tokens change (must be exactly 0):

S	32	64	96	120	127	128	129	130	200	300	1000	2048
drift	3.3e-2	2.5e-2	2.5e-2	2.9e-2	2.8e-2	0	0	0	0	0	0	0

Ruled out along the way: a torch bug (a minimal pure-torch repro is clean in every variant, ~7e-7), the nested torch.compile (raw and inner-compiled flex fail identically), the dynamo graph break (building the mask outside the model fails the same), out-of-bounds indexing in the padded tail (S=160 has a padded tail and is exact; clamping indices changed nothing), and layer count (1 layer fails like 2).

JobConfig.validate rejects the configuration rather than papering over it. The restriction costs nothing real: every shipped config uses seq_len 512/1024/4096, the default is 2048, and below one block there is no block sparsity for flex to exploit, so it would be pure overhead even if it were correct.

Numerics

Checked against an fp64 CPU reference rather than trusting a tolerance: flex is closer to exact than the dense-mask SDPA path on every parameter (relative error 3.1e-7..5.4e-7 vs SDPA's 3.2e-7..9.1e-7). The gradient test therefore bounds disagreement relative to each parameter's own scale — gradients reach ~1.8e3 here, so a fixed atol is meaningless — with both an elementwise and a whole-tensor-norm bound, so a single bad element cannot hide in the norm.

Testing

• uv run ruff check / ruff format --check pass
• uv run pyright kempnerforge/ — only the pre-existing video_io.py import av error (optional video group, not installed); zero in touched files
• uv run pytest tests/unit/ — 1843 passed, 13 skipped (1808 on main)
• pytest tests/integration/test_compile.py on 1×H200 — 14 passed, including:
  • test_cross_document_isolation[compiled] — direct regression test for the bug above
  • test_sequence_lengths_around_the_block_size[128/200/300] — ragged lengths with a partial trailing block
  • test_no_recompilation_across_document_layouts — the per-step mask_mod closure must not retrigger compilation, since that would blow cache_size_limit and fall back to eager flex, which is slower and hungrier than what it replaces, with no error to notice
• pre-commit (ruff, ruff-format) on all changed files

Not run: tests/distributed/ — this PR adds no distributed code path. TP interaction is argued from H=None (head-broadcast mask) and n_rep being computed from global head counts, but a 2-GPU TP+flex parity test is in the follow-up rather than asserted here.

Known limitations

• FlexAttention has no CPU backward in torch 2.11 and raises as soon as an input requires grad, so "flex" is CUDA-only in practice; CPU coverage is forward-only under torch.no_grad().
• capture_attention_weights is unsupported under flex (the kernel never forms an attention-weight matrix) and raises with a clear message.
• key_padding_mask (VLM video) and MoT self-attention stay on the dense SDPA path, by scope.
• RoPE positions still run continuously across a packed window — unchanged from today, so loss curves stay comparable. Per-document position reset is separate work.

Follow-ups

• Throughput benchmarks and the AC × compile matrix, plus a TP+flex distributed case.
• Flipping the default to "flex" for packed runs, once benchmarked.

Refs #204

[timothyngo]

Pushed 79a269a, correcting two things this PR got wrong. Both were found by measuring on H200, and both are worth calling out because the original text stated them with more confidence than I had earned.

head_dim <= 256 was false, and has been removed

That bound was never verified — it came from a design review plus a comment in a sibling repo about the lower limit. Measured:

   8  RAISED   NYI: embedding dimension
  16  OK       32  OK       64  OK       96  OK      128  OK
 192  RAISED   No valid triton configs / OutOfMemoryError
 256  OK      320  OK      384  OK      512  OK

head_dim 320, 384 and 512 all compute correctly, forward and backward, to ~2e-6 against a dense reference. The cap was rejecting working configurations — including the larger heads this repo will want.

It is deliberately not replaced with a wider range. The usable set is not an interval: 192 fails while 128 and 256 on either side are fine, and it is not a power-of-two rule either (96, 320, 384 all work). A config-time check cannot predict which sizes Triton can tile on a given GPU. That failure is also a loud compile-time error rather than silent corruption, so 192 is documented as known-bad rather than guarded — a guard's job is catching silent wrongness, not duplicating an error that already surfaces.

The floor stays at 16: head_dim 8 fails to lower.

The seq_len >= 128 bound is right; my explanation of it was wrong

This PR claimed "below one block the compiled kernel silently returns incorrect results". That is false. Bare compiled flex_attention matches a dense reference at seq_len 32/64/96 across head_dim 16/32/64/128 — all 20 combinations correct, ~1e-6. The kernel handles sub-block sequences fine.

The leak appears only once the call sits inside a compiled Transformer graph, and it reproduces across head dimensions:

dim	head_dim	S=32	S=128	S=256
128	32	leak 3.3e-02	0.0	0.0
256	64	leak 8.6e-02	0.0	0.0
512	128	leak 3.7e-01	0.0	0.0

So it is a sequence-length effect, not a head-dimension one, and the bound guards the right variable. But the cause is not established — and there is an unresolved contradiction behind it: a minimal torch.compiled wrapper calling flex_attention at S=32 is correct, while the full model at S=32 is not. Something about the larger fused graph matters and I did not isolate what.

The error message, the FLEX_BLOCK_SIZE comment, the masking.py docstring, the CHANGELOG and validation-rules.md now all say what was observed, state that the bound is empirical, and say the root cause is unknown. That distinction is load-bearing for whoever revisits this: "the kernel is broken" sends someone upstream to PyTorch, when the kernel is provably fine and the unexplored territory is our compiled graph.

Tests: 1846 unit, 14 integration on H200.