[LiRunGuo]

Motivation

flash_attn_func / flash_attn_varlen_func (and flex attention's FLASH backend, which calls them) take no tile sizes. On SM90 the forward tile comes from _tile_size_fwd_sm90, which only adjusts tile_m for the sparse Q block when head_dim <= 96. It never looks at the sparse KV block. SM90 has no KV subtiling (kv_subtile_factor must be 1), and a sparse Q block smaller than tile_m cannot be subtiled either. So any BlockSparseTensorsTorch.block_size that differs from the heuristic tile raises a ValueError in the forward, even when the kernel runs that tile fine through _flash_attn_fwd(..., tile_mn=...):

• hdim 192: flex's default BLOCK_SIZE=128 fails (tile_n=112).
• hdim 256: every KV block except 80 fails (tile_n=80). The only accepted size, (128, 80), is then rejected by the backward (tile_n=64), so SM90 hdim-256 block sparsity through autograd is not usable at all.
• hdim <= 128: (64, 128), (128, 64), (64, 64), (128, 112), ... all fail.

It also accounts for all 120 SM90 failures of test_mask_mod.py::test_parameterized_masks[True-128-112-*] / [True-64-128-*] (use_autograd=True with block sparsity). Those have failed on Hopper since the autograd variant was added in #2485. FA4 CI runs only on B200, which skips non-128x128 tiles, so nothing caught them.

Change

_fit_sm90_fwd_tile_to_block_sparsity runs after the SM90 heuristic when the caller supplies block_size:

• If the sparse Q block isn't a multiple of tile_m but is a multiple of 64, use tile_m=64. Larger sparse Q blocks still use q_subtile_factor.
• If the sparse KV block is smaller than tile_n and a multiple of 16, use tile_n = sparse KV block.

The tile only shrinks, which never adds smem or register pressure. The rule fires only when the heuristic tile would already have been rejected by normalize_block_sparse_config, so configurations that work today are unchanged: same tiles, same kernels, same performance. Also fixes the SM90 KV-subtile error message, which printed tile_n where it meant the sparse KV block size.

Results (H200, bf16, B=2 H=4 S=1024, block-aligned causal BlockMask vs fp32 flex)

(head_dim x block size) combinations accepted by flash_attn_func, out of 40 (hdim 64/96/128/192/256 x 8 block sizes):

	main	this PR
forward	8 / 40	32 / 40
forward + backward	6 / 40	11 / 40

Forward max abs error vs the fp32 flex reference is <= 0.0039 for every newly accepted case. The remaining backward rejections come from the SM90 backward's own tile_n (128 for hdim <= 128, 96 for 192, 64 for 256). That is a separate heuristic and is left for a follow-up.

Tests

• New test_sm90_block_sparse_fwd_tile_fits_block_size (7 cases, forward and backward where supported, via flash_attn_func): 7/7 fail on main, 7/7 pass here.
• tests/cute/test_mask_mod.py (H200): main 1256 passed / 120 failed -> this PR 1383 passed / 0 failed (+7 new).
• tests/cute/test_block_sparsity.py (H200): 4883 passed / 0 failed.
• tests/cute/test_mask_mod_varlen.py + tests/cute/test_flash_attn_fast.py (H200): 784 passed, 2 failed. The two failures, test_varlen_static_masks[block_diagonal-mha-*], were CUDA OOMs while four suites shared the node. Rerun alone, they pass on both main and this PR (17/17).
• Two-pass run (FLASH_ATTENTION_FAKE_TENSOR=1 compile, then GPU), separate cache dir for the patched tree. ruff check: no new findings.

This overlaps the function that #2733 moves into typed configs. The change is a self-contained post-pass and is easy to rebase onto that stack, whichever lands first.

[LiRunGuo]

@drisspg could you take a look when you have a moment? You own most of the block-sparsity code and the typed-config stack (#2732/#2733) that moves _tile_size_fwd_sm90, so you're the natural reviewer here.

Suggested path, whichever you prefer:

1. Land this first. It's a self-contained post-pass after the SM90 heuristic and only changes configurations that currently raise, so it should rebase onto [CuTe] Route forward through typed configs #2733 trivially. I'm happy to do that rebase myself once [CuTe] Route forward through typed configs #2733 lands.
2. Or fold it into [CuTe] Route forward through typed configs #2733. If you'd rather keep all tile selection in the typed configs, I can close this and send the same logic plus the test as a small PR on top of your stack.

Either way, the follow-up I'd like to send next is the matching change to the SM90 backward tile selection, so that more of the 40 (head_dim x block size) combinations work end to end. Right now 11 do. I'll keep that separate, since the backward tiles are more tightly tuned.

I don't think FA4 CI runs on Hopper (it's B200 only), so I can post H200 results for any revision of this.

[drisspg]

looks good thanks!