[FA4] [head_dim=256] dV is incorrect when using non-contiguous q, k, v from a single qkv projection

Hey team, I am trying to use the recent release https://github.com/Dao-AILab/flash-attention/releases/tag/fa4-v4.0.0.beta11 that adds experimental support for head_dim=256 and hitting issues with the dV tensor being corrupted when q, k, v are non contiguous. This does not happen with head_dim=64 or 128. I am also using `torch.use_deterministic_algorithms(True)` to force nans instead of garbage data. ``` device NVIDIA B300 SXM6 AC head_dim=64 random contiguous: ok (128/128 tokens finite) three separate Linears: ok (128/128 tokens finite) fused QKV Linear, sliced views: ok (128/128 tokens finite) head_dim=128 random contiguous: ok (128/128 tokens finite) three separate Linears: ok (128/128 tokens finite) fused QKV Linear, sliced views: ok (128/128 tokens finite) head_dim=256 random contiguous: ok (128/128 tokens finite) three separate Linears: ok (128/128 tokens finite) fused QKV Linear, sliced views: NaN 32/128 tokens finite, finite_idx[:8]=[0, 4, 8, 12, 16, 20, 24, 28], nan_idx[:8]=[1, 2, 3, 5, 6, 7, 9, 10] ``` To reproduce: ```python import os os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") import torch import torch.nn as nn from flash_attn.cute.interface import flash_attn_varlen_func def _pkg_version(name): import importlib.metadata as md try: return md.version(name) except md.PackageNotFoundError: return "?" print(f"torch {torch.__version__} (cuda {torch.version.cuda})") print(f"flash_attn_4 {_pkg_version('flash_attn_4')}") print(f"cutlass_dsl {_pkg_version('nvidia_cutlass_dsl')}") print(f"device {torch.cuda.get_device_name(0)}") print() # to force nans torch.use_deterministic_algorithms(True, warn_only=False) device = torch.device("cuda:0") dtype = torch.bfloat16 n_tokens = 128 n_q_heads, n_kv_heads = 4, 2 seg_lens = [68, 60] # cu_seqlens=[0, 68, 128] cu_seqlens = torch.tensor( [0, *torch.cumsum(torch.tensor(seg_lens), 0).tolist()], dtype=torch.int32, device=device, ) max_seqlen = max(seg_lens) def fa4_backward_dv(q, k, v, head_dim): captured = {} v.register_hook(lambda g: captured.__setitem__("dv", g.detach().cpu())) out = flash_attn_varlen_func( q, k, v, cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens, max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen, softmax_scale=head_dim ** -0.5, causal=True, window_size=(-1, -1), num_splits=0, deterministic=False, ) if isinstance(out, (list, tuple)): out = out[0] out.sum().backward() return captured["dv"] def random_contig(d): q = torch.randn(n_tokens, n_q_heads, d, dtype=dtype, device=device, requires_grad=True) k = torch.randn(n_tokens, n_kv_heads, d, dtype=dtype, device=device, requires_grad=True) v = torch.randn(n_tokens, n_kv_heads, d, dtype=dtype, device=device, requires_grad=True) return fa4_backward_dv(q, k, v, d) def three_linears(d): hidden = n_q_heads * d x = torch.randn(n_tokens, hidden, dtype=dtype, device=device, requires_grad=True) proj_q = nn.Linear(hidden, n_q_heads * d, bias=False).to(device=device, dtype=dtype) proj_k = nn.Linear(hidden, n_kv_heads * d, bias=False).to(device=device, dtype=dtype) proj_v = nn.Linear(hidden, n_kv_heads * d, bias=False).to(device=device, dtype=dtype) q = proj_q(x).reshape(n_tokens, n_q_heads, d) k = proj_k(x).reshape(n_tokens, n_kv_heads, d) v = proj_v(x).reshape(n_tokens, n_kv_heads, d) return fa4_backward_dv(q, k, v, d) def fused_qkv(d): hidden = n_q_heads * d x = torch.randn(n_tokens, hidden, dtype=dtype, device=device, requires_grad=True) proj_qkv = nn.Linear( hidden, (n_q_heads + 2 * n_kv_heads) * d, bias=False, ).to(device=device, dtype=dtype) qkv = proj_qkv(x) q_end = n_q_heads * d k_end = q_end + n_kv_heads * d q = qkv[:, :q_end].reshape(n_tokens, n_q_heads, d) k = qkv[:, q_end:k_end].reshape(n_tokens, n_kv_heads, d) v = qkv[:, k_end:].reshape(n_tokens, n_kv_heads, d) return fa4_backward_dv(q, k, v, d) def summarize(label, dv): finite = torch.isfinite(dv).flatten(1).all(dim=-1) bad = (~finite).nonzero().flatten().tolist() if not bad: print(f" {label}: ok ({dv.shape[0]}/{dv.shape[0]} tokens finite)") return good = finite.nonzero().flatten().tolist() print( f" {label}: NaN {len(good)}/{dv.shape[0]} tokens finite, " f"finite_idx[:8]={good[:8]}, nan_idx[:8]={bad[:8]}" ) for d in (64, 128, 256): print(f"head_dim={d}") for fn, label in ( (random_contig, "random contiguous"), (three_linears, "three separate Linears"), (fused_qkv, "fused QKV Linear, sliced views"), ): torch.manual_seed(0) summarize(label, fn(d)) ``` The interesting pattern is that finite gradients can be seen only at positions that are %4==0: 0, 4, 8, 12, 16, 20, 24, 28 etc.