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

[qywu]

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. Remove shared-prefix attention #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. Remove shared-prefix attention #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 Remove shared-prefix attention #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

• Remove shared-prefix attention #66 / e12b12b — the shared-prefix removal and its FA3 rationale; Reorganize src/xorl/ops: quarantine vendored code, separate kernels / layers / objectives, one home for the exact-contract family #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

[qywu]

Verified on H100: FA4 block masks work, fwd+bwd are correct, and the determinism limit is real but fails closed

Measured on 1× H100 80GB SXM (SM90), B=1 H=16 D=128 bf16, median of 30 iterations. This settles the "verify the FLASH backend is actually present and functional" item that was listed as the first thing to check.

Environment caveat: this box runs the sglang profile — torch 2.11.0+cu130, FA4 4.0.0b15, nvidia-cutlass-dsl 4.5.0. The default profile pins torch 2.12.1+cu132, FA4 4.0.0b19, cutlass-dsl 4.5.2. Numbers should be re-taken on the pinned pair before being used as a target.

1. FA4 exposes block masks directly, including varlen

flash_attn/cute/interface.py — both entry points take them:

flash_attn_func(..., score_mod=None, score_mod_bwd=None, mask_mod=None,
                aux_tensors=None, block_sparse_tensors=None,
                block_sparse_tensors_bwd=None, deterministic=False)

flash_attn_varlen_func(..., cu_seqlens_q=..., cu_seqlens_k=...,
                       mask_mod=None, block_sparse_tensors=None, ...)

BlockSparseTensorsTorch carries (mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx, cu_total_m_blocks, cu_block_idx_offsets, block_size, dq_write_order, dq_write_order_full, spt), and compute_block_sparsity(tile_m, tile_n, ..., mask_mod, aux_tensors, cu_seqlens_q, cu_seqlens_k, ...) builds it from a mask_mod. Varlen block sparsity is supported and needs cu_total_m_blocks — relevant because our packed path is varlen.

Note flash_attn_varlen_func accepts block_sparse_tensors but not block_sparse_tensors_bwd, unlike the dense entry point.

Important for design: FA4's mask_mod is a CuTeDSL callable — (b_idx, h_idx, q_idx, kv_idx, seqlen_info, aux_tensors) over TensorSSA, not a Python predicate. Authoring against FA4 directly means writing CuTeDSL. FlexAttention's value is that Inductor generates that CuTeDSL from ordinary Python, so we should go through FlexAttention, not call FA4 directly.

2. The FLASH backend really does lower to FA4

Verified with a fresh TORCHINDUCTOR_CACHE_DIR per backend and a scan of generated code:

backend	files	cute.jit	flash_attn.cute	@triton.jit	tl.dot
default (Triton)	7	0	0	7	6
kernel_options={"BACKEND": "FLASH"}	6	6	6	0	0

Cleanly disjoint. Both rows below are the same flex_attention call with the same Python mask_mod; only the flag differs.

3. Forward and backward are correct on both target patterns

Against an FP32 dense reference (the methodology from the PyTorch post), at S=4096, H=8 — relative mean error:

mask	blockmask sparsity	tensor	flex Triton	flex FA4	SDPA dense mask
causal	48.4%	out	1.43e-03	1.44e-03	1.44e-03
shared_prefix	67.2%	out	1.42e-03	1.43e-03	1.44e-03
strip_thinking	67.2%	out	1.43e-03	1.44e-03	1.45e-03

dq / dk / dv land at 1.64e-03 – 1.70e-03 for every implementation; max abs error is one bf16 ULP (0.015625) across the board. FA4's block-sparse backward is not measurably worse than Triton's or SDPA's.

The two mask_mods used, for the record:

# shared prefix: [P, S1..Sn]; suffix i sees P and itself
def shared_prefix(b, h, qi, ki):
    return (qi >= ki) & ((seg[ki] == 0) | (seg[ki] == seg[qi]))

# strip-thinking: history thinking is invisible, own-turn thinking is not
def strip_thinking(b, h, qi, ki):
    return (qi >= ki) & ~((role[ki] == 1) & (turn[ki] < turn[qi]))

4. Speed: FA4 wins long, loses short; SDPA-with-a-mask is not viable

Times in ms, fwd / fwd+bwd:

mask	seqlen	flex Triton	flex FA4	SDPA dense mask
causal (48% sparse)	2048	0.127 / 0.473	0.188 / 0.733	0.151 / 0.512
	4096	0.281 / 0.865	0.289 / 0.900	0.478 / 1.652
	8192	0.840 / 2.718	0.610 / 2.033	1.821 / 6.170
	16384	2.781 / 9.590	1.797 / 6.826	7.874 / 25.718
shared_prefix (68%)	2048	0.098 / 0.390	0.189 / 0.731	0.152 / 0.512
	4096	0.203 / 0.672	0.259 / 0.880	0.479 / 1.653
	8192	0.542 / 1.865	0.470 / 1.542	1.843 / 6.182
	16384	1.783 / 6.342	1.233 / 4.704	7.880 / 25.690
strip_thinking (68%)	2048	0.124 / 0.456	0.207 / 0.810	0.150 / 0.512
	4096	0.209 / 0.679	0.273 / 0.851	0.478 / 1.650
	8192	0.549 / 1.898	0.496 / 1.551	1.840 / 6.147
	16384	1.822 / 6.416	1.300 / 4.732	7.883 / 25.754

Four conclusions:

• Crossover at ~4–8k. FA4 is ~2× slower than Triton at 2048 (launch/dispatch overhead with no work to hide it) and 1.35–1.45× faster at 16384. The backend has to be a tunable, not a hardcoded choice. Below the blog's 1.41–1.89× on H200, consistent with SM90 + older FA4 + no LPT scheduling.
• SDPA with an explicit mask does not exploit sparsity at all. Its 16384 row is 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. It is 5.5× slower than flex-FA4 on shared_prefix fwd+bwd, and allocates an S×S mask besides. This is the number that justifies the feature.
• Block sparsity pays off as expected. Same kernel and shape, only density differs: causal at 48% sparse costs 6.826 fwd+bwd vs shared_prefix at 68% sparse at 4.704 — a 1.45× saving tracking the sparsity ratio. The masks are cheap, not merely correct.
• Generic block-mask causal ≈ hand-specialized causal. At 16384 fwd: sdpa_is_causal 1.784, FA4 native causal=True 1.691, flex-FA4 block-mask causal 1.797. Expressing causal as a generic block mask costs ~6%, so the block-sparse machinery imposes no structural tax.

5. Backward determinism: the limitation is real, but it fails closed

Same input, 4 repeated backward passes, checking bit-exactness (S=4096, H=16):

mask	impl	dq	dk	dv
causal	flex Triton	bit-exact	bit-exact	bit-exact
causal	flex FA4	NONDETERMINISTIC (max diff 2.44e-04)	bit-exact	bit-exact
shared_prefix	flex Triton	bit-exact	bit-exact	bit-exact
shared_prefix	flex FA4	bit-exact (this run)	bit-exact	bit-exact
strip_thinking	flex Triton	bit-exact	bit-exact	bit-exact
strip_thinking	flex FA4	NONDETERMINISTIC (max diff 1.49e-08)	bit-exact	bit-exact

Only dq is affected — consistent with dQ being accumulated across KV blocks in a schedule-dependent order. Note shared_prefix came out bit-exact in this run: the nondeterminism is intermittent, so a naive determinism test can pass by luck. Any test we write must repeat enough to be meaningful, or assert the guard instead.

The good news is the layering. FA4 the kernel does support deterministic block-sparse backward — interface.py asserts:

"deterministic block-sparse backward requires dq_write_order in block_sparse_tensors"
"deterministic block-sparse backward requires dq_write_order_full when full blocks are present"
"deterministic block-sparse backward requires block_sparse_tensors.spt ..."

with compute_dq_write_order() assigning semaphore lock values for ordered dQ accumulation. PyTorch's Flex lowering simply does not build that metadata yet, and rather than degrade silently it raises (torch/_inductor/kernel/flex/flex_attention.py):

raise NotImplementedError(
    "Deterministic backward for flex_attention with block_mask using the FLASH backend "
    "is not yet implemented. The TRITON backend supports deterministic backward."
)

Confirmed empirically under torch.use_deterministic_algorithms(True):

impl	result
flex Triton	ran fine
flex FA4	InductorError: Deterministic backward … not yet implemented. The TRITON backend supports deterministic backward.

So this reconciles with XORL_FLASH_ATTN_DETERMINISTIC and docs/k3 better than feared: deterministic runs route to the Triton backend, throughput runs route to FA4, and the combination we cannot honor raises instead of silently producing irreproducible gradients. That is exactly the fail-closed behavior ATTENTION_CONTRACT.md requires. Three options for later, in preference order: wait for PyTorch to plumb dq_write_order; call FA4 directly with the metadata built ourselves; or accept Triton-only for exact lanes.

Practical blocker found along the way

flash_attn.cute is shadowed in this environment. flash-attn-4 installs only flash_attn/cute/ (no flash_attn/__init__.py) into ~/.local/.../site-packages, while flash_attn 2.7.4 owns the flash_attn package in /usr/local/.../dist-packages and wins the import — so import flash_attn.cute fails with ModuleNotFoundError even though the files are present. Every measurement above needed:

import flash_attn
flash_attn.__path__.append("/home/qywu/.local/lib/python3.12/site-packages/flash_attn")

Worth checking whether the pinned default profile has the same collision, since backend/flash_attention.py gates FA4_AVAILABLE on from flash_attn.cute import ... and would silently report FA4 as unavailable.

Net effect on this issue

The dependency blocker from #66 is gone: block-sparse fwd+bwd is reachable today on the FA4 the tree already pins, it is numerically correct on both target patterns, and it is 5.5× faster than the only alternative that can express these masks. The remaining work is the plumbing and CP/SP items already listed above — not the kernel.

[qywu]

Follow-up: driving FA4 directly, without FlexAttention or torch.compile

Same H100, B=1 H=16 D=128 bf16, median of 30. This bypasses Flex entirely — hand-written CuTeDSL mask_mod plus compute_block_sparsity, calling flash_attn_func directly. It changes the recommendation from the comment above, so worth recording.

Writing a CuTeDSL mask_mod is easier than expected — for uniform geometry

No aux_tensors needed if segment boundaries are expressible as index arithmetic:

@cute.jit
def mm_strip_thinking(b_idx, h_idx, q_idx, kv_idx, seqlen_info, aux_tensors):
    PER, U, T = cutlass.Int32(S // 4), cutlass.Int32(S // 32), cutlass.Int32(S // 8)
    turn_q, turn_k = q_idx // PER, kv_idx // PER
    off_k = kv_idx - turn_k * PER
    kv_not_think = (off_k < U) | (off_k >= U + T)
    return (q_idx >= kv_idx) & (kv_not_think | (turn_k >= turn_q))

@cute.jit
def mm_shared_prefix(b_idx, h_idx, q_idx, kv_idx, seqlen_info, aux_tensors):
    P, L = cutlass.Int32(S // 4), cutlass.Int32((S - S // 4) // 3)
    return (q_idx >= kv_idx) & ((kv_idx < P) | (((q_idx - P) // L) == ((kv_idx - P) // L)))

Both produce block sparsity identical to Flex's create_block_mask (67.19% at 4096) and identical error vs the FP32 reference (out 1.44e-03, dq 1.69e-03), so the two routes are computing the same thing.

Gotcha worth knowing: ~ on a boolean vector is rejected — DSLRuntimeError: Unsupported DSL type: vector<1xi1>. Restructure to positive logic via De Morgan, as above. Real ragged turn/prefix lengths would still need aux_tensors loads, which is the genuinely harder path and the one the blog flags as a pipeline-stall risk ("pointer-chasing patterns, e.g. document masking with per-token metadata").

The real gap: FA4 ships no backward block-sparsity helper

compute_block_sparsity returns only the M-major forward orientation. Its own docstring says:

TODO: - transposed tensors for bwd pass

Pass block_sparse_tensors alone and ctx.block_sparse_tensors_bwd stays None, so the backward silently runs dense. No error, no warning — just a much slower backward. The cost is severe and grows with length:

seqlen	mask	fwd ms	fwd+bwd, dense bwd	fwd+bwd, transposed bwd	penalty
4096	causal	0.220	0.828	0.658	1.26×
4096	shared_prefix	0.181	1.143	0.566	2.02×
4096	strip_thinking	0.180	1.146	0.608	1.89×
16384	causal	1.725	10.554	6.826	1.55×
16384	shared_prefix	1.146	14.485	4.371	3.31×
16384	strip_thinking	1.162	14.588	4.559	3.20×

I built the transposed metadata by hand — _ordered_to_dense_simple → transpose → repack counts/indices, ~15 lines. Note the failure mode is silent: the forward is correctly sparse and the gradients are numerically correct, so nothing signals that you left 3× on the table. Anything we build on the direct path needs an assertion that block_sparse_tensors_bwd is present.

This is also precisely what Flex buys you. PyTorch's backward template passes Q_NUM_BLKS, Q_IDX, FULL_Q_NUM_BLKS, FULL_Q_IDX — the transposed view create_block_mask computes for free. Flex is not hiding the kernel; it is hiding the metadata bookkeeping in both orientations.

Full comparison, all four routes

fwd / fwd+bwd ms. Direct FA4 uses the hand-transposed backward metadata:

mask	seqlen	direct FA4	flex → FA4	flex → Triton	SDPA dense mask
causal	4096	0.220 / 0.658	0.289 / 0.900	0.281 / 0.865	0.478 / 1.652
	16384	1.725 / 6.826	1.797 / 6.826	2.781 / 9.590	7.874 / 25.718
shared_prefix	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	4096	0.180 / 0.608	0.273 / 0.851	0.209 / 0.679	0.478 / 1.650
	16384	1.162 / 4.559	1.300 / 4.732	1.822 / 6.416	7.883 / 25.754

• Direct FA4 is fastest everywhere, by 1.06–1.55× over flex→FA4. The gap is torch.compile/Inductor dispatch overhead, not kernel quality — at 16384 causal the two are identical (6.826 / 6.826), confirming the same kernel underneath and that the overhead amortizes at long sequence.
• Against the only alternative that can express these masks at all (SDPA with an explicit mask), direct FA4 is 5.9× faster on shared_prefix fwd+bwd at 16384.
• Direct FA4 also removes the short-sequence regression: at 4096 it beats flex→Triton (0.566 vs 0.672 on shared_prefix), where flex→FA4 lost to Triton.

Revised recommendation

Flex stays the default, for three reasons that outweigh 1.1–1.5×: it computes both metadata orientations, it generates the CuTeDSL from ordinary Python (so ragged geometry via captured tensors is tractable), and it gives us the deterministic-Triton fallback documented above.

Direct FA4 becomes a documented escape hatch rather than a non-option, which is a change from my earlier comment. It is worth reaching for if torch.compile interaction with enable_compile proves painful, or if short-sequence dispatch overhead matters — at the cost of hand-written CuTeDSL masks, hand-maintained transposed metadata, and no determinism story.

Caveats: 16384 correctness rows are missing because my FP32 dense reference OOMs at that size (my harness's limit, not FA4's) — 4096 correctness is verified for all three masks on both routes. Environment is still the sglang profile (torch 2.11, FA4 4.0.0b15, cutlass-dsl 4.5.0), not the pinned default.

[kiddyboots216]

Right, so I think the issue is that we don't have the proper bwd without Flex and writing that bwd pass by hand is nontrivial.

[qywu]

Correction: FA4 does not support deterministic block-sparse backward on Hopper

My first comment claimed "FA4 the kernel does support deterministic block-sparse backward — PyTorch's Flex lowering simply does not build that metadata yet." That is wrong. I inferred capability from the ValueError messages in interface.py demanding dq_write_order. Those are scaffolding for a path the SM90 kernel does not implement.

Tested directly: built the forward metadata, the transposed backward metadata, and dq_write_order / dq_write_order_full via FA4's own compute_dq_write_order, set spt both True and False, and passed deterministic=True. Every interface-level check passed. Then the kernel asserts:

AssertionError: Deterministic not implemented for block-sparse backward

flash_bwd_sm90.py has the deterministic dQ semaphore path (barrier.wait_eq / arrive_inc around the dQaccum store), but it is in the dense branch. The block-sparse branch is the else:

else:
    assert not self.deterministic, (
        "Deterministic not implemented for block-sparse backward"
    )
    dQaccum_store_block_sparse_bwd_sm90(...)

Verified present in b15, b19 (our pin), and b27 (latest). So compute_dq_write_order is currently unreachable machinery on Hopper, not a usable escape hatch.

flash_bwd_sm100.py has no equivalent assert, so Blackwell may support it — untested, this host is H100 only. SM100 does carry its own restriction: "Variable sequence length is not supported yet for blocksparse in bwd."

Measured non-determinism for scale: dq max run-to-run diff 4.88e-04 over 5 runs at S=4096 shared_prefix; dk / dv bit-exact. Roughly one bf16 ULP — the low bit of the gradient, not a large error, but not reproducible.

Effect on the recommendation: unchanged in direction, firmer in reason

On Hopper, Flex-with-Triton is the only way to get a deterministic block-sparse backward (verified bit-exact on all three masks). Native FA4 has no path to it. The earlier framing — "wait for PyTorch to plumb dq_write_order, or build it ourselves" — is not available; the kernel would have to implement it first.

So: exact lanes → Flex/Triton, throughput lanes → FA4 (either route). PyTorch's NotImplementedError when torch.use_deterministic_algorithms(True) meets FLASH + block_mask is therefore correct and load-bearing, not a temporary gap — it fails closed exactly as docs/k3/ATTENTION_CONTRACT.md requires.

Also: the latest FA4 does not change any of this, and we cannot upgrade to it anyway

flash-attn-4 ships weekly betas; b27 is current (2026-08-19). We pin b19 (2026-06-24).

	b19	b27
compute_block_sparsity.py	551 lines	551 lines — byte-identical
TODO: transposed tensors for bwd pass	present	still present
varlen accepts block_sparse_tensors_bwd	no	still no
assert not self.deterministic on block-sparse bwd (SM90)	present	still present

Two months of releases, no movement on any of the three gaps.

And the upgrade is blocked regardless: b27 requires nvidia-cutlass-dsl>=4.6.2 and quack-kernels>=0.5.3, while pyproject.toml pins nvidia-cutlass-dsl==4.5.2 for a stated reason —

# 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. So b19 is a ceiling imposed by the vendored Quack layer, not a stale pin. Anything wanting newer FA4 has to resolve ThrCopy first. Worth noting this is the same vendored ops/quack tree flagged in #78.

One more finding that matters more than the determinism question

b27's FlashAttnVarlenFunc.forward accepts block_sparse_tensors but never stores it on ctx, and its backward passes only mask_mod and deterministic — no sparsity metadata, and no block_sparse_tensors_bwd parameter to supply one.

So native FA4's varlen backward applies the mask element-wise but visits every block. Gradients are correct; sparsity skipping is entirely absent, and there is no API to fix it. Measured cost of a dense backward on the dense entry point, where metadata can be supplied: 3.2–3.3× at seqlen 16384.

This is the strongest argument for Flex, because our packed path is varlen today (position_ids → cu_seq_lens_q/k). Going native FA4 means either eating that penalty or moving to a dense layout with document boundaries encoded in the mask — which is what Flex does anyway. Whichever route we pick, the mask should absorb document boundaries and we should stop carrying them in cu_seqlens; that is now a design conclusion rather than an open question.

[qywu]

Batch / packing invariance measured — the property ATTENTION_CONTRACT.md actually requires

Run-to-run determinism (earlier comments) is weaker than what the contract asks for. The contract's concern is that "one request's result depend[s] on unrelated requests in the batch." That is a separate property and it is now measured. H100, S=4096, H=8, D=128, bf16.

Three axes, each holding the request under test byte-identical:

• A — batch size: same request in slot 0, same S, same mask, B ∈ {1,2,4,8} with random neighbours.
• B — neighbour data: document boundaries fixed, only the content of the other documents varies.
• C — neighbour layout: doc 0 pinned at [0, L), the rest of the row subdivided 4 ways (1 / 3 / 6 / ragged docs) — so a genuinely different block-sparsity pattern each time.

Forward: fully invariant, both backends

24/24 cells bit-exact, zero difference. Axis C is the strong one: block sparsity swung 67.2% → 85.9% → 90.6% → 80.9% across layouts (including a ragged 377/824 split aligning to no tile edge), so the kernel visited a materially different tile set each time, and doc 0's output did not move. That is the packing scenario we actually run — a sequence lands beside different companions every epoch.

Axis A passing on FA4 also speaks directly to the split-KV-from-batch-shape worry: on this path (dense-layout flex forward, SM90) that dependence does not appear.

Backward: Triton clean; FA4 has no neighbour dependence but is not reproducible

Because FA4's dQ is non-reproducible, a raw cross-batch diff is uninterpretable — so each run measures a 6-repeat run-to-run noise floor first and only flags neighbour dependence if cross-batch diffs exceed it.

backend	mask	noise floor	max cross-batch / neighbour diff	verdict
Triton	shared_prefix	0	0	invariant
Triton	strip_thinking	0	0	invariant
FA4	shared_prefix	6.10e-05	6.10e-05	≤ floor
FA4	strip_thinking	7.63e-06	7.63e-06	≤ floor

FA4's cross-batch diffs exactly equal the noise floor rather than merely falling under it — the signature of one shared cause: dQ accumulation order varies between any two invocations, and batch composition adds nothing. So there is no detectable neighbour dependence.

But the honest phrasing is "no detectable neighbour dependence," not "batch-invariant." Invariance in the contract's sense — this request always yields these bytes — is unachievable on FA4 for our masks, because it fails at the reproducibility step before batch composition is even in question.

Mechanism: dQ reproducibility tracks contributors per dQ tile

A 12-repeat control isolates why:

mask	block sparsity	dq bit-exact over 12 reps	max diff
document (block-diagonal)	85.9%	yes	0.0
causal	48.4%	no	1.22e-04

Block-diagonal masks give each dQ tile a small fixed contributor set, so the accumulation order is stable. Long-range masks give late tiles many contributors and the order varies. Our two target patterns are unavoidably long-range — every suffix reaches into the shared prefix; every turn reaches back across earlier answers — so they behave like causal, at ~1e-4 to 1e-6.

Note strip_thinking shows intermittency: three of five cells came back exactly 0 and one at 7.63e-06. A thin test can pass by luck, so any regression test here needs repeats.

Scope — do not over-read this

One H100, bf16, S=4096, one shape family, dense-layout flex only. It does not cover the paged-KV entry point (XORL_FLASH_ATTN_PAGED_KVCACHE), which is the path the contract's split-KV warning is actually aimed at, nor cross-device or cross-shape reproducibility.

Conclusion

Reinforces the split, now on the stronger property: exact lanes → Flex/Triton (reproducible and invariant, by default, no flags); throughput lanes → FA4 (no neighbour dependence, ~1e-4 non-reproducible on these mask shapes).