varlen backward hangs on SM90/H100 at head_dim=512 with torch 2.13/triton 3.7.1 (returns correctly on torch 2.11/triton 3.6)

[dengoswei]

Summary

On H100 (SM90), ffpa_attn_varlen_func(...).sum().backward() never returns for head_dim=512. The forward is fine and fast; only the backward hangs. It reproduces at every shape I tried, including one as small as total_q=128, num_heads=2, so it is not a size or tiling boundary.

The dense entry point ffpa_attn_func with an explicit TritonBackend(backward=True, enable_tma=True, persist_dkdv=True) works perfectly on the same machine, in the same process, at the same geometry — 29.2 ms fwd+bwd for 3x2048, 32q/4kv, D512, bf16, causal, which is ~8x faster than what we run today. So this is specifically the varlen path.

Environment

ffpa-attn	0.2.2 (latest on PyPI)
torch	2.13.0+cu130
triton	3.7.1
GPU	NVIDIA H100 80GB HBM3, sm90
python	3.12

Your own test suite passes here unmodified: pytest tests/test_ffpa_bwd.py -k sm90 → 8 passed in 4.87 s. So this is not a broken install or an environment/version skew.

Minimal repro

import torch
from ffpa_attn import ffpa_attn_varlen_func

nseq, S, H, D = 1, 128, 2, 512          # smallest shape that reproduces
dt = torch.bfloat16
cu = torch.arange(0, (nseq + 1) * S, S, device="cuda", dtype=torch.int32)
q = torch.randn(nseq * S, H, D, device="cuda", dtype=dt, requires_grad=True)
k = torch.randn(nseq * S, H, D, device="cuda", dtype=dt, requires_grad=True)
v = torch.randn(nseq * S, H, D, device="cuda", dtype=dt, requires_grad=True)

o = ffpa_attn_varlen_func(q, k, v, cu, cu, S, S, causal=True)
torch.cuda.synchronize()                # returns; 6.79 s cold (compile), 0.07 s warm
o.sum().backward()
torch.cuda.synchronize()                # never returns

Shape sweep

Each shape in its own child process, killed at 120 s. head_dim=512, bf16, causal, H100.

shape	forward	backward
nseq=1 S=128 HQ=2 HKV=2	ok 6.79 s (cold compile)	hang
nseq=1 S=128 HQ=32 HKV=4	ok 0.07 s	hang
nseq=1 S=512 HQ=32 HKV=4	ok 0.07 s	hang
nseq=1 S=2048 HQ=32 HKV=4	ok 0.08 s	hang
nseq=2 S=2048 HQ=32 HKV=4	ok 0.09 s	hang
nseq=3 S=2048 HQ=32 HKV=4	ok 0.08 s	hang

Where it stops

Python stack, captured by a watchdog thread after 120 s (identical at every shape):

ffpa_attn/cute/__init__.py:919        _varlen_fwd_backward
  dq, dk, dv = torch.ops.ffpa_attn._varlen_bwd_cute(
ffpa_attn/cute/__init__.py:828        _varlen_bwd_custom
ffpa_attn/cute/_ffpa_bwd_sm90.py:447  _ffpa_attn_backward_sm90
  _call_with_tvm_ffi_current_stream(         # <- the dQ launch
ffpa_attn/cute/_utils.py:165          _call_with_tvm_ffi_current_stream
  return fn(*args)
nvidia_cutlass_dsl/python_packages/cutlass/cutlass_dsl/tvm_ffi_provider.py:629

Line 447 is the dQ launch, i.e. the dKdV launch at ~431 already returned. But since _call_with_tvm_ffi_current_stream does not synchronise (_utils.py:159-165), the launches are async and I cannot tell from this which of the two kernels is actually spinning — I have not yet run with CUDA_LAUNCH_BLOCKING=1 to disambiguate. Happy to, if useful.

GPU stays at 100% utilisation throughout.

What is ruled out

• Not compilation. The first forward compiles in 6.79 s and every later one takes 0.07 s, so compilation is cached and costs seconds. The backward got 120 s at every shape, including that tiny one.
• Not the compile cache. Reproduced with FLASH_ATTENTION_CUTE_DSL_CACHE_DIR on fast local disk (an earlier attempt on a network mount had its own unrelated problem; this run did not).
• Not shape. See the sweep — the smallest shape hangs too.
• Not a version skew. Your sm90 backward tests pass here, 8/8.
• Not the call site, as far as I can tell: ffpa_attn_varlen_func exposes no backward_backend, and functional.py:278-284 requires forward_backend == backward_backend. Since the varlen entry dispatches straight to _ffpa_attn_varlen_cute, its backward is pinned to CuTeDSL with no way to select the Triton backward that the dense tests exercise. If there is a supported way to steer this, I would be glad to be corrected.

Possibly relevant: test coverage of this path

I went looking for an existing test to run before filing, and could not find one for this combination. Across all 9 files in tests/, exactly one test exercises varlen together with a gradient: tests/test_ffpa_cute_sm80.py::test_sm80_cutedsl_varlen_autograd_matches_sdpa. It differs from the failing case on four axes:

	that test	this report
path	SM80	SM90
head_dim	320	512
causal	no	yes
heads	2q/2kv (MHA)	2q/2kv and 32q/4kv (GQA) — both hang

tests/test_ffpa_bwd.py has 29 backward tests and 0 mentions of varlen; all 10 of its head_dim=512 backward tests pass an explicit backward_backend=TritonBackend(...).

So the SM90 varlen backward looks uncovered rather than regressed — which would explain why this survived a release. Mentioning it only because it might help you decide where to look; apologies if I have missed a test elsewhere.

Context

This is relevant to #276, where it is stated that "Hopper works great currently for packed/padded training because its backward is cu_seqlens-aware and compiles once". That does not match what I measure on H100 with head_dim=512 — the forward is indeed fast and compiles once, but the backward does not return. It is possible the difference is head_dim (320 vs 512) or causal; I have not yet isolated which, and would run whichever variant is most useful to you.

We were evaluating FFPA for packed-sequence RL training with Gemma-4's 512-dim global-attention layers, where it benchmarks ~8x faster than our current backend. Very happy to run further diagnostics on H100 — this is a machine we have access to.

[dengoswei]

Follow-up: I cloned the repo and read the kernels, and I think I can point at the exact asymmetry. Everything in the first two sections below is read straight off main (19f04dbdba41c663e695c352e816bf5a05244afc, v0.2.2-12-g19f04db); the mechanism in the third is my inference from it, not something I have instrumented yet.

src/ffpa_attn/cute/ is untouched since v0.2.2

None of the 12 commits since the tag touch that directory — they are all cuda/triton backend work — so this reproduces on main, not only on the release.

In the D512 dQ kernel, one of the three warp roles does not honour process_tile

_dq_d512_sm90.py has three warp-role loops. process_tile appears in only two of them:

role	def	while work_tile.is_valid_tile	process_tile
load (producer / TMA)	740	784	absent
mma_wg1	1007	1088	1092, gated at 1141
mma_wg2	1334	1419	1424, gated at 1478

_dkdv_d512_sm90.py has the same three-role structure and gates all three (776/780, 1067/1071, 1345/1349).

And the value is compile-time constant in the dense case:

process_tile = const_expr(not self.is_varlen_k) or n_block_min < n_block_max

Dense ⇒ const_expr(True) ⇒ always True ⇒ all three roles run every tile body, no asymmetry.
Varlen ⇒ const_expr(False) ⇒ process_tile = n_block_min < n_block_max, which can be False.

Inferred mechanism (not yet instrumented)

When process_tile is False under varlen, mma_wg1 and mma_wg2 skip the tile body — and with it the pipeline acquire/release and barrier arrivals inside it — while load has no such guard and keeps pushing Q/K/dO/V for that tile. Producer and consumers desynchronise and someone waits on an mbarrier that never arrives. That matches the symptom: 100% GPU utilisation, no progress, no error.

Where the empty tiles come from, if this is right: the varlen-only grid padding at _ffpa_bwd_sm90.py:283-285,

total_q_rounded_padded = (
  total_q + cu_seqlens_q.shape[0] * m_block_size - 1
) // m_block_size * m_block_size

For the minimal repro (total_q=128, cu_seqlens.shape[0]=2, m_block_size=64) that gives 192 — three m_blocks covering a 128-row segment, so the third is entirely past the end and n_block_min >= n_block_max. That would explain why even the smallest shape hangs rather than only ragged or large ones.

This also fits the stack in the original report sitting at the dQ launch (_ffpa_bwd_sm90.py:447) rather than the dKdV launch — though as I noted there, the launches are async so the stack alone did not establish that. The source asymmetry is a second, independent line pointing the same way.

What I have not done

I have not confirmed at runtime that process_tile is actually False for these shapes, nor which kernel is the one spinning. Two cheap checks would settle both — CUDA_LAUNCH_BLOCKING=1, and printing n_block_min/n_block_max per work tile. I am re-acquiring an H100 and will run them, but I did not want to sit on the source finding meanwhile in case it is immediately obvious to you.

If the diagnosis holds, the fix looks like it would be adding the same process_tile computation and guard to the load role in _dq_d512_sm90.py that _dkdv_d512_sm90.py already has in all three roles — but you would know far better than I whether the producer can simply skip the tile or needs to drain something first.

[dengoswei]

Correction to my original report: I was wrong that this is not a version issue. It is. Sorry for the noise — here is the data.

I re-ran the same minimal case on a second H100 with a different software stack, and the varlen D512 backward returns and produces correct gradients there.

	first report	this run
GPU	H100 80GB HBM3, sm90	H100 80GB HBM3, sm90
torch	2.13.0+cu130	2.11.0
triton	3.7.1	3.6.0
ffpa-attn	0.2.2	0.2.2
varlen D512 backward	never returns, every shape	returns: 12.99 s cold, 7.840 ms steady

Gradient correctness on the run that works, 3x2048 packed, 32 q heads / 4 kv heads, head_dim=512, bf16, causal, checked against autograd through fp32 SDPA on the same values:

dq: max|g|=4.9062e+00  rel_err_vs_sdpa=2.970e-03
dk: max|g|=7.5000e+00  rel_err_vs_sdpa=3.758e-03
dv: max|g|=1.4625e+01  rel_err_vs_sdpa=2.592e-03

Non-zero, non-None, bf16-level agreement — so it is genuinely computing the backward, not returning early.

What this changes

• Retract: "Not a version skew. Your sm90 backward tests pass here, 8/8." The suite passing was real, but it does not cover varlen (the point of my first report), so it was never evidence about this path. Version is exactly the variable.
• Keep: everything about the hanging stack itself, and the shape sweep — those reproduce on the torch 2.13.0+cu130 / triton 3.7.1 stack.
• The source asymmetry I described in my previous comment (the load warp role in _dq_d512_sm90.py not honouring process_tile, while both mma_wg roles and all three _dkdv_d512_sm90.py roles do) is unchanged — it is the same source on both machines. If that reading is right, it would be a latent race that a different compiler/runtime happens to schedule around, which is consistent with a version-sensitive hang rather than a deterministic bug. But I want to be clear that this is now my inference stacked on my inference, and I have not instrumented it.

What I could not isolate

The two stacks differ in both torch and triton, so I cannot say which one moves it. If it would help, I am happy to run a specific combination on H100 — pinning one and varying the other is a few minutes of work on my side, and I would rather do that than have you chase it.

CUDA_LAUNCH_BLOCKING=1 on the hanging stack (to name which of the dKdV / dQ kernels is the one spinning) is still on my list; the machine that reproduced the hang was reclaimed before I got to it, and the replacement runs the working stack.

[DefTruth]

@dengoswei Hi, thank you for bringing this issue to our attention. Regrettably, we do not have consistent access to H100/B200 hardware for development and testing purposes. As such, we welcome contributors to submit pull requests for bug fixes. ffpa‑attn is currently independently maintained. We sincerely apologize for not having sufficient resources to address community‑reported issues in a timely manner.