Block-mask attention: express shared-prefix and strip-thinking multi-turn training that cu_seqlens cannot

Summary

Two training patterns we want need attention masks that cu_seqlens cannot express, and both need the same primitive: a block-sparse mask plumbed from the data path into the attention backend.

1. Shared-prefix training — many suffixes (rollouts, candidates, branches) share one prompt prefix. Pack the prefix once and let each suffix attend to prefix ∪ own_suffix, instead of materializing the prefix n times.
2. Strip-thinking multi-turn training — reasoning models strip history <think> blocks from every turn but the current one at inference. Training the whole conversation in one forward pass requires each turn's tokens to attend to prior turns' answers but not prior turns' thinking.

Today xorl can only express block-diagonal masking. position_ids → prepare_fa_kwargs_from_position_ids → cu_seq_lens_q/k (src/xorl/data/collators/packing_concat_collator.py, src/xorl/utils/seqlen_pos_transform_utils.py) gives each packed document a causal square on the diagonal and nothing off it. Both patterns above are off-diagonal by construction.

Status: the kernel question is settled. Everything below the "Measured" heading was verified on an H100 (SM90). The comments on this issue carry the raw numbers and methodology. The remaining work is plumbing and CP/SP, not kernels.

Why this is unblocked now

Shared-prefix attention was removed in #66 (e12b12b) for a dependency reason, not a design objection: its backend imported flash_attn_interface (FA3), which is not shipped for the CUDA 13 / FA4 profile pyproject pins. The PR framed the choice as "build FA3 for CUDA 13, or delete it."

There is a third resolution. FlexAttention + FlashAttention-4 gives FlexAttention an FA4 backend, and FA4 is what this tree already pins (flash-attn-4==4.0.0b19). Verified working here: kernel_options={"BACKEND": "FLASH"} generates CuTeDSL that instantiates FA4, with zero Triton in the generated code.

Current state in this tree

A partial foundation exists but is unreachable:

• backend/flex_attention.py has make_causal_block_mask() composing causal + document + padding into a BlockMask.
• It is registered in CAUSAL_MASK_FUNCTIONS["flex_attention"] but not in ATTENTION_FUNCTIONS, so flex_attention is never actually called.
• attn_implementation is Literal["eager", "sdpa", "native", "flash_attention_3", "flash_attention_4"] (src/xorl/arguments.py:458) — "flex_attention" is not a legal value.
• AttentionKwargs carries only cu_seq_lens_q/k and max_length_q/k — no field for a block mask or per-token segment metadata.

Mask specifications

Shared prefix

Sequence [P, S₁, …, Sₙ], segment_id = 0 for P and i for Sᵢ:

def shared_prefix(b, h, qi, ki):
    return (qi >= ki) & ((seg[ki] == 0) | (seg[ki] == seg[qi]))

position_ids restart at len(P) for every suffix. Saves (n−1) × len(P) tokens of KV compute and memory.

Strip-thinking multi-turn

Turns t = 1…N, each (userₜ, thinkₜ, answerₜ):

def strip_thinking(b, h, qi, ki):
    return (qi >= ki) & ~((role[ki] == THINK) & (turn[ki] < turn[qi]))

Own-turn thinking stays visible, so answerₜ still attends thinkₜ; only earlier turns' thinking is cut. The naive alternative is N forward passes per conversation — O(N²) redundant tokens and no gradient sharing.

The exactness question — still needs a decision

With one copy of each token, answerⱼ's hidden states are computed with thinkⱼ visible (correct — it is the query's own turn). But when a later turn attends answerⱼ as a key, those keys were produced under a context inference never has. This is why One-Pass to Reason adds token duplication rather than just masking.

• (a) Mask only. Accept contaminated history keys. One copy, one pass, cheapest.
• (b) Mask + token duplication. Duplicate each answer span: one copy attends its own thinking and carries the loss, a second attends the stripped context and serves as history. Exact w.r.t. inference, longer sequence, more intricate mask.

Recommendation: land (a) first — a strict superset of today's expressiveness, and shared with the shared-prefix case — then evaluate whether (b) changes downstream quality enough to justify it. This is the one open design question that measurement cannot settle.

Measured (H100 SM90, B=1 H=16 D=128 bf16, median of 30)

Correctness

All routes land at one bf16 ULP vs an FP32 dense reference, on out, dq, dk, dv, for all three masks (relative mean error 1.42e-03 – 1.70e-03). FA4's block-sparse backward is not measurably worse than Triton's or SDPA's.

Speed, fwd / fwd+bwd ms

mask	seqlen	direct FA4	flex → FA4	flex → Triton	SDPA dense mask
shared_prefix (68% sparse)	4096	0.181 / 0.566	0.259 / 0.880	0.203 / 0.672	0.479 / 1.653
	16384	1.146 / 4.371	1.233 / 4.704	1.783 / 6.342	7.880 / 25.690
strip_thinking (68%)	16384	1.162 / 4.559	1.300 / 4.732	1.822 / 6.416	7.883 / 25.754

• SDPA with an explicit mask is not viable. Its 16384 times are 25.718 / 25.690 / 25.754 across three masks of very different density — identical, because it materializes the full S×S mask and computes every entry. 5.9× slower than direct FA4. This is the number that justifies the feature.
• Block sparsity pays as expected: causal at 48% sparse costs 6.826 fwd+bwd vs shared_prefix at 68% at 4.704 — a 1.45× saving tracking the sparsity ratio.
• FA4-vs-Triton crosses over at ~4–8k. FA4 is ~2× slower at 2048 (dispatch overhead) and 1.35–1.45× faster at 16384. The backend must be a tunable, not hardcoded.
• Generic block-mask causal ≈ hand-specialized causal (1.797 vs 1.691 fwd at 16k) — the block-sparse machinery imposes no structural tax.

mask_mod vs score_mod — why this must be a mask

The same shared-prefix pattern expressed as a score_mod returning -inf is numerically identical (max diff 0.001953, bf16 rounding) and dramatically slower, because score_mod does not skip tiles:

seqlen	as block_mask	as score_mod(-inf)	penalty
4096	0.854 ms	3.196 ms	3.7×
8192	1.507 ms	11.433 ms	7.6×

Determinism

deterministic mode	backend	block_mask	score_mod only	plain
off	Triton	exact	exact	exact
off	FA4	(intermittent)	not exact	not exact
strict	Triton	exact	exact	exact
strict	FA4	refused	exact	exact
warn_only	FA4	warns, runs	not exact, no warning	—

• Triton is deterministic by default, in every configuration, without flags.
• FA4 is non-reproducible by default even without a block mask.
• FA4 + block mask under strict mode raises NotImplementedError naming Triton as the alternative — correct fail-closed behavior per docs/k3/ATTENTION_CONTRACT.md.
• warn_only=True silently disables FA4 determinism everywhere and warns only for the block-mask case. A trap, since it is a common setting for unrelated ops.

Batch / packing invariance

Forward: fully invariant, both backends, 24/24 cells bit-exact, including a layout axis where block sparsity swung 67.2% → 90.6% and one ragged split aligned to no tile edge. This is the packing scenario we run.

Backward: Triton exactly 0 everywhere (zero noise floor). FA4's cross-batch diffs exactly equal its run-to-run noise floor (6.10e-05 / 7.63e-06) — no detectable neighbour dependence, but not reproducible, so contract-grade invariance is unachievable on FA4 for these masks.

Mechanism: dQ reproducibility tracks contributors per dQ tile. Block-diagonal document masks are bit-exact over 12 reps; causal and our long-range masks are not (1.22e-04). Our patterns are unavoidably long-range.

Verified FA4 gaps (checked against b19 wheels and b27, the latest)

gap	b19 (our pin)	b27 (latest)
transposed backward metadata helper	TODO	TODO, file byte-identical
varlen backward accepts block sparsity	no	no
deterministic block-sparse backward, SM90	hard assert	hard assert

1. No backward metadata helper. compute_block_sparsity returns only the forward M-major orientation. Omit the transposed view and ctx.block_sparse_tensors_bwd is None, so the backward silently runs dense — no error, correct gradients, 3.2–3.3× slower at 16384. Buildable by hand (~15 lines); needs an assertion so the silent path is impossible.
2. Varlen backward cannot take block sparsity at all. FlashAttnVarlenFunc.forward accepts block_sparse_tensors but never stores it on ctx; backward passes only mask_mod. There is no parameter to supply one.
3. Deterministic block-sparse backward is absent on Hopper. dq_write_order appears 0 times in flash_bwd_sm90.py; the block-sparse branch opens assert not self.deterministic. The assert guards genuinely missing code — patching it out would produce non-reproducible results while claiming determinism. flash_bwd_sm100.py does implement it (_dq_semaphore_lock_value consumes dq_write_order), so Blackwell likely works; untested here.

And we cannot upgrade past b19. b27 requires nvidia-cutlass-dsl>=4.6.2 and quack-kernels>=0.5.3; pyproject.toml pins 4.5.2 because "nvidia-cutlass-dsl 4.6.0 removed cutlass.cute.core.ThrCopy, which the Quack compatibility layer still needs at import time" and quack-kernels==0.5.0. b19 is a ceiling imposed by the vendored Quack tree (see #78), not a stale pin.

Decisions this settles

1. Use FlexAttention, not native FA4, as the default. It computes both metadata orientations, generates CuTeDSL from Python (so ragged geometry via captured tensors is tractable), and provides the deterministic-Triton fallback. Cost is 1.06–1.55×, and zero at 16384 causal.
2. Exact lanes → Triton. Throughput lanes → FA4. The only configuration that is both reproducible and invariant for our masks is Flex/Triton.
3. Direct FA4 is a documented escape hatch, not a non-option: fastest everywhere, removes the short-sequence regression, and on Blackwell may reach deterministic block-sparse backward that Flex refuses (PyTorch's guard is arch-blind). Costs hand-written CuTeDSL masks and hand-maintained transposed metadata.
4. Move document boundaries into the mask and out of cu_seqlens. Native FA4's varlen backward cannot be made sparse, and Flex does not use varlen — both viable routes point the same way. This is a larger data-path change than originally scoped and should be agreed before plumbing starts.

Proposed work

1. Plumb segment metadata. Emit per-token segment_ids / turn_ids / role_ids alongside position_ids in the packing collators; add an AttentionKwargs field for the tags or a prebuilt BlockMask.
2. Make flex_attention a real backend. Register in ATTENTION_FUNCTIONS, add to the attn_implementation literal, expose backend selection (FA4 vs Triton) as a tunable given the crossover, and route exact lanes to Triton.
3. Composable mask builders. Generalize make_causal_block_mask into causal / document / shared_prefix / strip_history_thinking mods, with create_block_mask compiled and cached across steps rather than rebuilt per microbatch.
4. Tests that actually execute. #66 noted all four shared-prefix tests carried importorskip guards and reported skipped, "which is easy to misread as covered." Each mask_mod should be checked against an explicit dense-mask eager reference at small shapes — CPU-runnable. Determinism/invariance tests need repeats, since the FA4 noise is intermittent.
5. Docs. A docs/k3 section: this is a new numerical program for attention, and the Triton/FA4 lane split belongs in the contract.

Open constraints (not settled by measurement)

• Block granularity. FA4's minimum sparse block is 128×128 on Hopper, 256×128 on Blackwell (q_stage=2). Arbitrary segment boundaries fall into partial blocks; interacts with pad_to_multiple_of.
• CP/SP — the hardest part. A BlockMask must shard consistently with src/xorl/distributed/sequence_parallel/. Ulysses shards heads and is probably tractable; ring attention is not obviously so, and sequence_shard_collator.py zigzag-reorders packed sequences (zigzag_reorder_packed_sequence) for causal load balance — a permutation whose premise is a causal triangle. #66 removed two ring-attention NotImplementedError guards that existed only to reject this; reinstating an explicit "unsupported with ring/hybrid CP" error is the right start.
• Compile interaction. Scalars captured in a mask_mod are baked into the compiled kernel. With enable_compile and varying segment layouts, segment data must be passed as tensors, not captured scalars, or every distinct value recompiles.
• Packaging. flash_attn.cute is shadowed where flash_attn 2.x owns the flash_attn package: flash-attn-4 ships only flash_attn/cute/ with no __init__.py, so import flash_attn.cute raises ModuleNotFoundError even when present. Since backend/flash_attention.py gates FA4_AVAILABLE on that import, an affected environment silently reports FA4 as unavailable. Needs checking on the pinned profile.

Non-goals

• Reviving the removed repack data path from #66 — this expresses the pattern in the mask instead.
• Inference/serving-side prefix caching.

Measurement caveats

All numbers came from the sglang profile — torch 2.11, FA4 4.0.0b15, nvidia-cutlass-dsl 4.5.0 — not the pinned default (torch 2.12.1, FA4 b19, 4.5.2). API facts were read directly from b19 and b27 wheels and are solid; timings should be re-taken on the pinned pair before being treated as targets. Single H100, bf16, dense-layout flex; the paged-KV path (XORL_FLASH_ATTN_PAGED_KVCACHE) is untested.

References

• #66 / e12b12b — the shared-prefix removal and its FA3 rationale; #78 — the vendored ops/quack tree blocking the FA4 upgrade
• FlexAttention + FlashAttention-4
• One-Pass to Reason (arXiv 2504.18246)
• Colfax: FlexAttention in FA CuTe DSL · attention-gym

Start with src/xorl/data/collators/packing_concat_collator.py, src/xorl/utils/seqlen_pos_transform_utils.py, backend/flex_attention.py, src/xorl/arguments.py, and AttentionKwargs. Run the existing shared-prefix tests and inspect their import guards, then compare each mask against a small dense reference. Done means metadata reaches an executable backend, mask tests run, and the Triton/FA4 and determinism behavior is documented.