[BUG] CuTe DSL 4.7/4.8 regression: FlashAttention-4 SM100 forward kernel spills registers (0 -> 176 B) vs 4.5.1, ~3% slower on B300 #3701

Describe the bug

With identical kernel source, CuTe DSL 4.7.x/4.8.0 generates PTX for the FlashAttention-4 SM100 forward kernel (flash_attn.cute.flash_fwd_sm100.FlashAttentionForwardSm100) that spills registers, where CuTe DSL 4.5.1 generated PTX that does not. The kernel is at the 128-register cap in both cases (512 threads per CTA), so the extra pressure becomes local-memory spills. The varlen causal fwd+bwd step is about 2–3% slower on B300.

The regression is in the generated PTX, not in ptxas: the 4.8.0 PTX still spills when assembled with the older ptxas 13.2, while the 4.5.1 PTX assembled with that same ptxas does not.

Resource usage, same FA4 source (fa4-v4.0.0.beta16), same quack-kernels 0.4.1, sm_103a

MMA (UTCHMMA/UTCQMMA), MUFU, and packed FP32 (FFMA2/FMUL2/FADD2) instruction counts in the forward kernel are identical between the two versions, so the difference is spill traffic rather than different math.

PTX vs ptxas cross-check (forward kernel PTX from CUTE_DSL_KEEP=ptx)

Neither PTX contains .local declarations, so the spills are introduced by ptxas from register pressure in the 4.8.0 PTX.

CUTE_DSL_COMPILER_OPT="remarks{ptx}" with CUTE_DSL_LINEINFO=1 on 4.8.0 reports the spills for both forward and backward, and marks flash_attn/cute/tile_scheduler.py:923 (the varlen scheduler's _get_num_m_blocks, loading cu_seqlens) as the likely pressure region.

Timing (B300, varlen causal, bf16, head_dim 128, fwd+bwd, median of 20 steps after 3 warmups, runs interleaved on one GPU)

CuTe DSL 4.7.0 and 4.7.1 (measured with quack-kernels 0.6.5) show the same slowdown as 4.8.0. The current FA4 release (fa4-v4.0.0.beta32) with quack-kernels 0.6.5 on CuTe DSL 4.8.0 is within noise of beta16 on 4.8.0, so the difference tracks the DSL version, not the FA4 or quack version. I could not test 4.5.3 or 4.6.x with this FA4 source (import errors unrelated to the kernel), so the regression landed after 4.5.1 and by 4.7.0.

Steps/Code to reproduce bug

1. Environment: NVIDIA B300 SXM6 (sm_103a), driver 590.48.01, PyTorch 2.11 (CUDA 13.0), Python 3.12.
2. FA4 source: Dao-AILab/flash-attention at tag fa4-v4.0.0.beta16, installed with pip install --no-deps --no-build-isolation flash_attn/cute.
3. Install quack-kernels==0.4.1 (--no-deps) and one of:
   • nvidia-cutlass-dsl==4.5.1 with nvidia-cutlass-dsl-libs-cu13==4.5.1 (apache-tvm-ffi 0.1.6)
   • nvidia-cutlass-dsl==4.8.0 with nvidia-cutlass-dsl-libs-cu13==4.8.0 (apache-tvm-ffi 0.1.11)
4. Run the script below with CUTE_DSL_KEEP=ptx,cubin CUTE_DSL_DUMP_DIR=<dir>, then cuobjdump --dump-resource-usage <dir>/*flash_fwd_sm100*.cubin.

quack-kernels 0.4.1 annotates with cute.core.ThrMma/ThrCopy, which 4.8.0 exposes only on cute; the script aliases them so the same quack source imports on both DSL versions.

import importlib.metadata as md
import json
import sys

import torch
import cutlass.cute as cute

for _name in ("ThrMma", "ThrCopy", "TiledMma", "TiledCopy"):
    if not hasattr(cute.core, _name) and hasattr(cute, _name):
        setattr(cute.core, _name, getattr(cute, _name))

from flash_attn.cute.interface import flash_attn_varlen_func


def bench(total, heads, head_dim, docs):
    torch.manual_seed(0)
    q, k, v = (
        torch.randn(total, heads, head_dim, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3)
    )
    cu = torch.linspace(0, total, docs + 1, device="cuda").round().to(torch.int32)
    max_seqlen = int((cu[1:] - cu[:-1]).max())

    def step():
        out = flash_attn_varlen_func(
            q, k, v, cu_seqlens_q=cu, cu_seqlens_k=cu, max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen, causal=True
        )
        out = out[0] if isinstance(out, tuple) else out
        out.float().square().mean().backward()

    for _ in range(3):
        step()
    torch.cuda.synchronize()
    times = []
    for _ in range(20):
        start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        start.record()
        step()
        end.record()
        torch.cuda.synchronize()
        times.append(start.elapsed_time(end))
    return round(sorted(times)[10], 3)


shapes = [(8192, 32, 128, 4), (32768, 32, 128, 8), (32768, 16, 128, 1)]
print(json.dumps({
    "cutlass_dsl": md.version("nvidia-cutlass-dsl"),
    "gpu": torch.cuda.get_device_name(),
    "median_ms": {f"{t}x{h}h_d{d}_{n}docs": bench(t, h, d, n) for t, h, d, n in shapes},
}))

Expected behavior

CuTe DSL 4.8.0 generates PTX for this unchanged kernel source with register pressure comparable to 4.5.1, so the forward kernel does not spill and runs at 4.5.1 speed. If this is a deliberate lowering change, guidance on a compile option or kernel-side change that avoids the spill would help.

Environment details

• GPU: NVIDIA B300 SXM6 (sm_103a), driver 590.48.01
• CuTe DSL: 4.5.1 (good), 4.7.0 / 4.7.1 / 4.8.0 (spills), CUDA 13 libs
• ptxas used in the cross-check: 13.1, 13.2
• PyTorch 2.11, CUDA 13.0, Python 3.12, Linux x86_64
• FlashAttention-4: fa4-v4.0.0.beta16 (control) and fa4-v4.0.0.beta32
• quack-kernels: 0.4.1 (control) and 0.6.5

Additional context

Default compile options are identical in both versions (opt-level=3, empty ptxas options, no line info, assertions off), and none of the new 4.8 environment variables were set.