[0z5a]

Closes #3773.

Adds variable sequence lengths to the Hopper CuTe DSL FMHA example
(examples/python/CuTeDSL/cute/hopper/kernel/attention/fmha.py), following the
interface of the Blackwell example.

Interface

• __call__ takes cum_seqlen_q, cum_seqlen_k (int32, [0, s_0, s_0 + s_1, ...])
  and max_seqlen_q. Pass None for a fixed-length side. Q/K/V/O/LSE of a
  variable-length side are packed along the sequence dimension with batch size 1.
• run() and the CLI accept per-sequence lengths as in the Blackwell example,
  e.g. --q_shape "3,(1,129,700),8,128" --k_shape "3,(65,300,700),2,128".
  Q and K can be variable-length independently. GQA/MQA works as before.

Implementation

• The grid covers the longest query sequence (as in the Blackwell example).
  All warp groups skip the same tiles: those that start past the end of their
  query sequence, or whose key sequence is empty. Rows of a query sequence
  with no keys are therefore not written, as in the Blackwell example.
• Per tile, Q/K/V and LSE are offset to their sequence with domain_offset,
  and the mask helpers get that sequence's lengths.
• Tails: a tail tile inside the packed buffer reads rows of the next sequence;
  TMA zero fill only applies past the end of the whole buffer. Extra K columns
  are masked: variable-length keys use the residual mask when no window is
  set, and the causal/window masks already cover the last K tile. Extra Q rows
  are computed but not stored.
• Store: the O tile can extend into the next sequence, so with
  variable-length queries the epilogue copies the staged shared-memory tile to
  global memory with 128-bit stores and a row predicate instead of the TMA
  store (the Blackwell example also drops its TMA store for variable-length
  queries).
• Fixed-length inputs compute the same thing as before; the varlen-only parts
  are behind const_expr.
• FP8: can_implement now requires key lengths divisible by 16. FP8 V is
  s-major, and its head stride is s_k bytes, which TMA needs to be a multiple
  of 16 bytes, so fixed-length shapes that fail this check could not run
  before either. With variable lengths, key offsets not divisible by 16 made
  the kernel fault on GH200.
• run(): the default softmax scale uses the head dimension (q_shape[3]),
  not q_shape[1], which is a tuple for variable-length input. [Hopper CuTeDSL] Add tests for the dense GEMM, FMHA and CTA norm examples #3770 makes the
  same one-line change. This is a behavior change for a caller that passes
  scale_softmax=0.0 with a fixed-length, non-square S_q != head_dim: that
  graph now uses 1/sqrt(head_dim). default_softmax_scale names the value and
  a host test pins it for non-square shapes.
• The variable-length row store asserts the epilogue-tile and warp-group
  divisibility, and (for a static layout) D contiguity, that its 128-bit copy
  assumes, so an unsupported (out_dtype, head_dim) fails at trace time instead
  of dropping rows.

Tests

test/examples/CuTeDSL/hopper/test_fmha_varlen.py (separate from the
fixed-length tests in #3770). Output and LSE are checked per sequence against
a PyTorch reference (fp16/bf16 atol 1e-2, fp8 atol 0.13, rtol 1e-5):

• lengths (1, 0, 63, 64, 65, 128, 129, 700): empty and length-1 sequences,
  both sides of the 64- and 128-row tiles, unaligned sequence starts;
• no mask / causal / window (128, 0) x GQA 8:8, 8:2, 8:1 x persistent and
  non-persistent;
• different q/k lengths per sequence (no mask, bottom-right causal);
  variable-length Q only and K only;
• d=64, d=256, bf16, fp8;
• a sequence with queries but no keys (completes; not reference-checked).

On GH200: 32/32 pass, and the #3770 fixed-length suite (53 cases) passes on
this branch.

Performance (GH200)

fp16, h=16, d=128, mma_tiler_mn=(64, 128), 16384 total tokens, TFLOPS, mean
of two runs of 50 iterations. FLOPs count the attended (q, k) pairs of each
sequence. Neither branch includes #3399.

fixed 8 x 2048	main	this PR
no mask, persistent	516.3	516.7
causal, persistent	364.6	365.3
no mask	509.3	509.6
causal	394.6	394.8

varlen (this PR)	no mask, persistent	causal, persistent	no mask	causal
8 x 2048 packed	520.5	396.1	498.7	393.3
4096, 512, 3072, 1024, 2560, 1536, 2048, 1536	454.8	311.7	508.4	419.7

• Fixed-length performance is unchanged.
• With equal lengths, varlen is within about 2% of fixed length. In the
  persistent causal case, fixed length on main runs below non-persistent
  (365 vs 395), and the varlen path does not show that gap. I haven't looked
  into why.
• With skewed lengths, non-persistent throughput holds. Persistent drops
  (455 / 312) because its static round-robin assignment over the max-length
  grid balances the work poorly. A length-aware persistent schedule could be
  a follow-up.

[brandon-yujie-sun]

@arvin-chou

[arvin-chou]

The varlen path looks correct to me. A few things I'd want to settle before merging:

1. fmha.py:2460 — the default softmax scale changed from 1/sqrt(S_q) to 1/sqrt(head_dim). That's the right value and varlen needs it (S_q is a tuple now), but it silently changes results for a caller that passes scale_softmax=0.0 with S_q != head_dim. Worth a line in the PR description and a test that pins the scale for a non-square shape.
2. fmha.py:1480 — store_o_rows derives thr_n = epi_tile[1] // copy_elems and a thread layout from it without guarding that copy_elems divides epi_tile[1],
   that thr_n divides 128, or that o is d-contiguous. Fine for the tested configs, but an accepted-but-untested (out_dtype, head_dim) could hit a zero-division or silently drop rows — a can_implement check or an assert would make the safe set explicit.

The rest I saw are efficiency/cleanup. follow-ups, not blockers.

[0z5a]

Thanks — addressed both points in c6cfa35:

1. Default softmax scale: the default now goes through default_softmax_scale(q_shape) = 1/sqrt(head_dim), and a host test pins it for non-square shapes (S_q packed tuple of 129 or fixed 700 vs head_dim=64; 129 vs 128). The PR description now calls out the behavior change for a caller that passes scale_softmax=0.0 with a fixed-length non-square S_q != head_dim.
2. store_o_rows guards: it now asserts epi_tile[1] % copy_elems == 0, num_threads_per_warp_group % thr_n == 0, and, when the layout is static, gO.stride[0] == 1 (D-contiguous). An unsupported (out_dtype, head_dim) now fails at trace time instead of dividing by zero or dropping rows.

Validation on GH200 (SM90, CuTe DSL 4.8.1): test_fmha_varlen.py 33 passed (57.7 s), including the new host test and all varlen cases through the guarded store. This is a Hopper example, so no SM100 run was applicable.