[lewisword]

Summary

The NVFP4 (CUDABackend(enable_fp4=True)) sm_120 persist-D path aborts with a CUTLASS TMA assertion for long attention sequences. The delta_s gmem layout in the fp4 launchers is built with int32 extents; the batch-mode block stride Nkv_pad * (Nq_pad/kBr) * Nh overflows INT32_MAX once the sequence exceeds ~70k tokens (at Nh=56, D=128), wraps negative, and trips the TMA descriptor assertion.

Fixes #386

Error

cute/atom/copy_traits_sm90_tma.hpp:983: ... Assertion
`(gmem_prob_stride[4]) < (uint64_t(1) << 40)' failed.
Aborted (core dumped)

Root cause

csrc/cuffpa/launch/cute_fp4_persist_d.cuh (same pattern in the split-D and split-D m4n2 launchers):

auto layout_DS =
    tile_to_shape(SmemLayoutAtomDS{}, make_shape(Nq_pad, Nkv_pad, Nh, Nb),
                  Step<_2, _1, _3, _4>{});

tile_to_shape -> blocked_product -> make_ordered_layout accumulates the block strides as products of these int32 values. The batch-mode stride is Nkv_pad * (Nq_pad/kBr) * Nh, which scales as N^2 and crosses INT32_MAX at roughly 70k attention tokens (B=1, H=56, D=128). The wrapped negative stride is reinterpreted as uint64 inside make_tma_copy_desc and aborts on the assertion.

Note the true stride (2.9e9 elements for an 81k-token sequence, ~11.6 GB in bytes) is far below the 2^40-byte TMA hardware limit - this is purely an int32 accumulation overflow, not a genuine TMA limit.

Minimal reproduction

import os
os.environ["FFPA_CUDA_ALLOW_SMALL_D"] = "1"
import torch
from ffpa_attn import CUDABackend, ffpa_attn_func

backend = CUDABackend(forward=True, backward=False, tensor_layout="NHD",
                      enable_fp4=True, fp4_smooth_v=False, fp4_hadamard=False)
# B=1, S=77000, H=56, D=128 (long-sequence video DiT self-attention)
q = torch.randn(1, 77000, 56, 128, dtype=torch.bfloat16, device="cuda")
k = torch.randn_like(q); v = torch.randn_like(q)
with torch.no_grad():
    ffpa_attn_func(q, k, v, forward_backend=backend)   # aborts

S=42000 succeeds; S=77000 aborts deterministically.

Fix

Promote the four dynamic extents to int64_t in make_shape in all three fp4 launchers (cute_fp4_persist_d.cuh, cute_fp4_split_d.cuh, cute_fp4_split_d_m4n2.cuh). The TMA atom's static type layout is unchanged, so no kernel-side changes are required.

Verification (RTX PRO 5000, sm_120f, torch 2.13.0+cu129, CUDA 12.9)

• The minimal repro above aborts before the change and passes after it.
• Numerics unchanged vs SDPA (cos ~0.982, the inherent NVFP4 quantization error), identical between short (42k) and long (77k) sequences.
• End-to-end: a 362-frame (15s) 768x1344 video-DiT generation (MiniMax-H3, ~117k tokens) that previously aborted at the first denoise step now completes.

[Copilot]

🟢 Approval recommended

The narrowly scoped casts address overflow before stride accumulation consistently across all three launchers, with no unresolved findings.

No open issues 🎉

What changed in this PR

Fixes #386 by preventing 32-bit stride overflow in the CUDA backend’s NVFP4 delta_s layouts, which caused TMA descriptor assertions for long attention sequences.

Changes:

• Cast all four dynamic layout extents to int64_t before stride accumulation in all three FP4 launchers.
• Document the overflow cause while preserving layout geometry and kernel interfaces.

File	Description
csrc/​cuffpa/​launch/​cute_fp4_split_d.cuh	Widens split-D layout extents to prevent stride overflow.
csrc/​cuffpa/​launch/​cute_fp4_split_d_m4n2.cuh	Applies the same protection to the M4N2 layout.
csrc/​cuffpa/​launch/​cute_fp4_persist_d.cuh	Widens persist-D layout extents and explains the assertion’s cause.

_{Review effort: Balanced}

[DefTruth]

LGTM~ Thanks for this fix!