[BUG] CuTe DSL >= 4.6: FlashAttention-4 SM100 backward with block-sparse masks spills registers, up to 1.8x slower than 4.5.3 (same kernel source) #3707

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug

With the same kernel source, CuTe DSL 4.6.0.dev0 and later generate PTX for the FlashAttention-4 SM100 backward kernel (flash_attn.cute.flash_bwd_sm100.FlashAttentionBackwardSm100) that spills registers. The spills appear when the kernel is specialized for block-sparse attention with a mask_mod. That is the path torch.nn.attention.flex_attention(..., kernel_options={"BACKEND": "FLASH"}) takes. With CuTe DSL 4.5.3, the same source compiles to PTX with no spills, or almost none. The kernel is at the 128-register cap (512 threads per CTA), so the extra pressure becomes local-memory traffic. On B200 the backward kernel is 22% slower with a sliding-window mask and 82% slower with a document mask. The dense, causal, GQA and native window_size specializations of the same kernel are not affected.

This looks like the backward counterpart of #3701 (forward kernel, 4.7/4.8 vs 4.5.1), and it narrows down where the regression landed. 4.5.3 is fine. 4.6.0.dev0 already regresses, and it is the first release whose PTX header says Based on NVVM 23.0.0 (4.5.3: NVVM 20.0.0).

Backward kernel GPU time, ms (B200, bf16, head dim 128, 16 heads; median of 50; main backward kernel only, from the torch profiler; a repeat run agrees within 1% for the block-sparse rows and within 5% for the others)

The 4.5.3 and 4.6.0.dev0 columns run the same FA4 source (b20, see "Same source" below).
¹ The 4.8.0 column runs FA4 4.0.0b32, the latest release at the time of testing (b33 has a byte-identical backward kernel). FA4 changed between b20 and b32, so this column shows that the regression persists; it is not a clean A/B.

The same PTX files assembled by one ptxas (CUDA 13.0, V13.0.88, ptxas -arch=sm_100a -O3 -v)

All of these are at 128 registers (.reqntid 512), with warp-specialized budgets (setmaxnreg.inc 136/152, setmaxnreg.dec 88/24). The PTX files have no .local declarations, so ptxas introduces the spills from register pressure in the NVVM 23 PTX. The SASS that 4.8.0 ships through its own ptxas has 50 STL / 42 LDL for the document mask and 35 / 26 for the sliding window (causal: 2 / 2).

Likely pressure region. CUTE_DSL_COMPILER_OPT="remarks{ptx}" CUTE_DSL_LINEINFO=1 on 4.8.0 marks flash_attn/cute/block_sparse_utils.py:1564 in get_block_sparse_iteration_info_bwd. That line loads the block-sparse count and index tensors: curr_q_cnt = q_cnt[batch_idx, head_idx, n_block] and q_idx[batch_idx, head_idx, n_block, None]. #3701 found the same kind of small metadata-tensor load for the forward kernel (cu_seqlens in the varlen scheduler). Here the IMAD count rises together with the spills (document mask: 653 → 1049). We suspect the address arithmetic of these indexed loads, but have not verified it.

Steps/Code to reproduce bug

The attached nvvm23-sm100-bwd-repro.zip has the scripts, the nine PTX files and our raw results. Its README has more detail.

1. No GPU needed. Assemble the attached PTX pair:

ptxas -arch=sm_100a -O3 -v ptx/bwd.document_mask.fa4-b20.cutlass-dsl-4.5.3.nvvm20.ptx -o /dev/null
    16 bytes stack frame, 28 bytes spill stores, 16 bytes spill loads
ptxas -arch=sm_100a -O3 -v ptx/bwd.document_mask.fa4-b20.cutlass-dsl-4.6.0.dev0.nvvm23.ptx -o /dev/null
    152 bytes stack frame, 204 bytes spill stores, 204 bytes spill loads

2. On a B200: time the backward with the script below (FA4 through flex_attention, about 50 lines), once per environment:

python3.12 -m venv nvvm20 && . nvvm20/bin/activate
pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cu130
pip install torch==2.14.0 numpy flash-attn-4==4.0.0b19 nvidia-cutlass-dsl==4.5.3 quack-kernels==0.5.0 apache-tvm-ffi==0.1.11
python minimal_repro.py

python3.12 -m venv nvvm23 && . nvvm23/bin/activate
pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cu130
pip install torch==2.14.0 numpy flash-attn-4==4.0.0b20 nvidia-cutlass-dsl==4.6.0.dev0 quack-kernels==0.5.3 apache-tvm-ffi==0.1.12
python minimal_repro.py

import importlib.metadata as md
import statistics

import torch
import torch._functorch.config
from torch.nn.attention.flex_attention import create_block_mask, flex_attention

B, H, S, D = 1, 16, 32768, 128

# Causal attention within documents of random length (2^9 .. 2^12 tokens), a standard flex attention mask.
generator = torch.Generator().manual_seed(0)
lengths = []
while sum(lengths) < S:
    lengths.append(int(torch.randint(512, 4096, (1,), generator=generator)))
doc_id = torch.repeat_interleave(torch.arange(len(lengths)), torch.tensor(lengths))[:S].cuda()


def document_causal(b, h, q_idx, kv_idx):
    return (doc_id[q_idx] == doc_id[kv_idx]) & (q_idx >= kv_idx)


def sliding_window(b, h, q_idx, kv_idx):
    return (q_idx >= kv_idx) & (q_idx - kv_idx < 1024)


def backward_ms(mask_mod):
    """Median time of one attention backward (FA4 kernels via flex attention's FLASH backend)."""
    block_mask = create_block_mask(mask_mod, None, None, S, S, BLOCK_SIZE=(256, 128))
    attention = torch.compile(lambda q, k, v: flex_attention(q, k, v, block_mask=block_mask, kernel_options={"BACKEND": "FLASH"}))
    torch.manual_seed(0)
    q, k, v = (torch.randn(B, H, S, D, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3))
    out = attention(q, k, v)
    dout = torch.randn_like(out)
    times = []
    for i in range(60):
        start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        start.record()
        torch.autograd.grad(out, (q, k, v), dout, retain_graph=True)
        end.record()
        torch.cuda.synchronize()
        if i >= 10:
            times.append(start.elapsed_time(end))
    return statistics.median(times)


def main():
    """Prints the backward time of two block-sparse attention masks."""
    # replaying one compiled backward (retain_graph=True) needs donated buffers off
    torch._functorch.config.donated_buffer = False
    print(" ".join(f"{p}=={md.version(p)}" for p in ("torch", "flash-attn-4", "nvidia-cutlass-dsl")), torch.cuda.get_device_name())
    for mask_mod in (document_causal, sliding_window):
        print(f"{mask_mod.__name__:16s} backward {backward_ms(mask_mod):.3f} ms")


if __name__ == "__main__":
    main()

Output on B200 (the 4.8.0 line uses FA4 b32 with quack-kernels 0.6.5 and apache-tvm-ffi 0.1.14.post1):

torch==2.14.0+cu130 flash-attn-4==4.0.0b19 nvidia-cutlass-dsl==4.5.3 NVIDIA B200
document_causal  backward 1.549 ms
sliding_window   backward 1.132 ms
torch==2.14.0+cu130 flash-attn-4==4.0.0b20 nvidia-cutlass-dsl==4.6.0.dev0 NVIDIA B200
document_causal  backward 2.588 ms
sliding_window   backward 1.323 ms
torch==2.14.0+cu130 flash-attn-4==4.0.0b32 nvidia-cutlass-dsl==4.8.0 NVIDIA B200
document_causal  backward 2.480 ms
sliding_window   backward 1.316 ms

3. Same source. FA4 b20 is b19 ported to the CuTe DSL 4.6 API. The port's changes to the backward kernel are the .ptr accessors on the TMEM allocator's struct fields and assume_tensor_aligned on LSE/dPsum. b19's kernels do not compile on 4.6, and b20's do not compile on 4.5.3, because b20 calls the 4.6 signature of nvvm.fmax. Restoring b19's fmax helper, which picks the signature by CUDA version (fa4-b20-fmax-for-cutlass-dsl-4.5.patch in the zip, one function), makes the same b20 source run on 4.5.3. It then gives the 4.5.3 numbers above: 1.279 ms for the document mask, matching b19's 1.278 ms. The *.cutlass-dsl-4.5.3.nvvm20.ptx files come from that setup.

Expected behavior

The block-sparse specializations of the backward kernel should compile with CuTe DSL >= 4.6 to code that spills no more, and runs no slower, than with 4.5.3, as the dense and causal specializations already do.

Environment details (please complete the following information):

• Environment location: Linux x86_64 container on a B200 node; CUDA 13.0 toolkit for the ptxas cross-check
• GPU: NVIDIA B200 (sm_100a), driver 595.91.07
• Python 3.12.13, torch 2.14.0+cu130
• nvidia-cutlass-dsl 4.5.3 / 4.6.0.dev0 / 4.8.0 (default install, no [cu13] extra)
• flash-attn-4 4.0.0b19 and b20 / b20 / b32; quack-kernels 0.5.0 / 0.5.3 / 0.6.5; apache-tvm-ffi 0.1.11 / 0.1.12 / 0.1.14.post1

Additional context

• FA_DISABLE_2CTA=1 (forces FA4's 1-CTA backward) gives the same block-sparse times on both toolchains, so this is not a kernel-dispatch difference.
• We also saw the slowdown with 4.6.2, 4.6.3 and 4.7.1, on other block-sparse shapes.
• On 4.8.0 the forward kernel's block-sparse specialization also reports spills (20 B stores / 480 B loads). The remark points at the analogous metadata load in the forward (block_sparse_utils.py:68, mask_block_cnt[batch_idx, head_idx, m_block]). We measured no forward slowdown from it.
• We can share CUTE_DSL_KEEP=ir dumps if they help.
• Questions: is this a known NVVM 23 / 4.6 toolchain regression? Is there a compile option in 4.6+ that restores the 4.5 behavior for kernels like this one until a fix ships?

nvvm23-sm100-bwd-repro.zip