[Lightning Attention] Add Hopper SM90 CuTe DSL prefill

[Aharrypotter]

📌 Description

This PR adds a Hopper SM90 Lightning Attention prefill backend implemented in
CuTe DSL.

• adds a 384-thread, three-warp-group TMA/WGMMA kernel under
  cula/ops/lightning/sm90/, with the FP32 recurrent state resident in
  registers across 64-token chunks;
• extends the public lightning_attn_fwd and
  lightning_attn_fwd_varlen APIs with architecture dispatch across the
  existing SM100/SM103 backend and the new SM90 backend;
• supports fixed-length and packed variable-length BF16 prefill, optional
  initial/final state, MHA and GVA head mappings, explicit state-pool indices,
  and persistent or non-persistent packed scheduling;
• uses a per-stream TensorMap workspace and TMA stores for packed ragged tails;
• keeps the default packed path free of dispatcher-visible device-value
  validation and host synchronization;
• adds an independent FP32 tokenwise reference, focused correctness and
  dispatch tests, a prefill-specific benchmark, and Hopper pipeline
  documentation.

The supported SM90 specialization uses BF16 Q/K/V/O, FP32 decay and recurrent
state, key/value dimension 128, and chunk size 64. Unsupported architectures
or specializations fail explicitly rather than silently falling back.

The packed-tail TMA optimization is a separate commit so its implementation
and evidence can be reviewed independently.

🔍 Related Issues

This work is intentionally separate from #76, which tracks Gated DeltaNet
rather than Lightning Attention.

🧪 Final-source validation

Validation environment:

GPU:                 NVIDIA H20-3e (SM90a, 78 SMs)
PyTorch:             2.9.1+cu130
Torch CUDA:          13.0
nvidia-cutlass-dsl:  4.5.1
Codegen tools:       CUDA 13.1

• Final focused product suite: 40/40 pass.
• Final same-campaign C++ comparison: 10/10 precision rows and 10/10
  timing-admission rows pass three-way checks against an independent FP32
  tokenwise reference.
• Packed persistent/non-persistent output and state are bitwise equal in
  the affected-path matrix; unselected state-pool slots remain unchanged.
• Public default-path host-sync audit: four packed calls, zero
  dispatcher-visible violations, including a cold TensorMap-workspace
  allocation.
• Two-stream stress: persistent and non-persistent modes, 20 launches per
  mode, bitwise concurrent-versus-serialized parity, and stable distinct
  per-stream TensorMap workspaces.
• Final affected-path Compute Sanitizer matrix: 9/9 pass
  (5 memcheck and 4 synccheck).
• Fresh source-bound codegen: 10/10 specializations, 9 unique cubins,
  24 WGMMA per specialization, and zero spill/stack/local memory.
• Ruff lint/format and git diff --check pass.

🔬 Accuracy vs original C++ implementation

We added a source-bound comparison against the original FlashInfer SM90
C++/CUTLASS implementation (7086b148, CUTLASS f3fde583).

The comparison freezes the common semantic subset: BF16 Q/K/V/output, FP32
state, D=V=128, C64, MHA (H=HV=64), scale=1,
lambda=exp(-slope) with exponent offset 1, unique sequential state slots,
fixed/packed execution, and final-state output.

The common subset is necessary because the original and cuLA public contracts
differ in GVA ownership, non-unit scale ownership, optional final state,
packed scheduling, and indirect/shared state-pool ownership. This is therefore
a numerical comparison on the common subset, not a whole-API equivalence
claim.

The final-source comparison is bound to pushed PR head 93da019.

CuTe vs C++ metric	Final-source result
Rows passing all three comparisons	10/10
Worst output max abs	1.2207e-4
Worst output relative RMS	6.6605e-5
Minimum output cosine similarity	0.9999999978
Minimum output bitwise agreement	99.9409%
Worst state max abs	2.4214e-8
Worst state relative RMS	2.0957e-7
Minimum state cosine similarity	0.9999999999999855
Minimum state bitwise agreement	13.9045%
Minimum elementwise tolerance pass rate	100%

The implementations are not fully bitwise identical because the recurrent
accumulations use different instruction schedules. The supported conclusion is
numerical alignment under the frozen tolerance contract, not complete bitwise
parity.

⚡ Performance vs original C++ implementation

The primary denominator is the original C++ Operation::run after one-time
workspace allocation and operation initialization. Every row uses five fresh
processes per implementation after triggering the pre-registered three-percent
escalation rule, rotated order, 20 warmups, 100 CUDA-event timed iterations,
unique CuTe caches, and audited receipts. A CuTe/C++ ratio below 1.0x means
CuTe is faster.

Row	CuTe ms	C++ steady ms	CuTe / C++
f_b1_t512_z	0.039393	0.038846	1.014086x
f_b2_t512_z	0.074610	0.073885	1.009814x
f_b1_t4096_z	0.257587	0.251742	1.023216x
f_b2_t4096_z	0.507993	0.497066	1.021983x
f_b1_t16384_z	1.001478	0.978733	1.023239x
f_b1_t512_h0	0.040810	0.040749	1.001476x
f_b1_t4096_h0	0.260324	0.255249	1.019886x
f_b1_t16384_h0	1.006083	0.983927	1.022518x
p_uniform_4096_h0	0.298915	0.291469	1.025548x
p_skewed_8192_h0	0.467908	0.461152	1.014649x
Combined geomean	—	—	1.017616x

The pre-registered limits were a ten-row geomean <=1.10x, every row
<=1.25x, and the critical skewed packed row <=1.10x; all pass. There were
zero contention-invalid receipts, fallbacks, skips, or missing rows. The final
CuTe DSL source is therefore about 1.76% slower in geometric mean and within
2.56% of the prepared original C++ kernel on every row of this matrix.

⚡ Canonical CuTe vs FLA baseline

Before the packed-tail TMA optimization, the canonical H20 matrix contained
38 fixed and packed-variable-length workloads: 34 native-MHA rows used for
the primary geometric mean and four GVA rows reported separately. All rows use
BF16 Q/K/V, FP32 decay/state, D=128, and chunk size 64.

This frozen P7-B matrix is retained as the broad CuTe-vs-FLA baseline. The
final packed-tail TMA source is faster on the two same-boundary packed rows
screened above, so these packed FLA numbers are conservative rather than
final-source claims.

For every row and implementation:

1. run in three independent fresh processes with a unique CuTe cache;
2. record compilation and first-call setup separately;
3. execute 20 warmups outside timing;
4. capture 100 CUDA-event timed iterations;
5. use the median of the three per-process averages;
6. record GPU utilization, clocks, power, and temperature around each process.

Check	Reference limit	Baseline result
CuTe / FLA geomean over 34 native-MHA rows	<= 1.10x	0.487849x
Worst primary row	<= 1.25x	0.960543x
Complete row coverage	38 rows	38/38, all ratios < 1.0x

Workload group	Rows	CuTe / FLA geomean	Speedup	Min ratio	Max ratio
Fixed, no state (MHA)	10	0.738844x	1.35x	0.278828x	0.958029x
Fixed, initial/final state (MHA)	6	0.704192x	1.42x	0.289854x	0.960543x
Packed uniform (MHA)	6	0.338190x	2.96x	0.184498x	0.710279x
Packed random (MHA)	6	0.350382x	2.85x	0.199508x	0.711503x
Packed skewed (MHA)	6	0.339864x	2.94x	0.193118x	0.677091x
All fixed MHA	16	0.725654x	1.38x	0.278828x	0.960543x
All packed MHA	18	0.342770x	2.92x	0.184498x	0.711503x
All native MHA	34	0.487849x	2.05x	0.184498x	0.960543x
GVA, outside primary geomean	4	0.544746x	1.84x	0.257395x	0.942551x

Full fixed-length CuTe/FLA baseline (18 rows)

Row	CuTe ms	FLA ms	CuTe / FLA	Speedup
f_ns_b1_t512	0.042767	0.153382	0.278828x	3.59x
f_ns_b1_t1024	0.079612	0.154749	0.514456x	1.94x
f_ns_b1_t4096	0.292828	0.310630	0.942691x	1.06x
f_ns_b1_t8192	0.576859	0.613204	0.940729x	1.06x
f_ns_b1_t16384	1.144923	1.242369	0.921564x	1.09x
f_ns_b2_t512	0.082020	0.152919	0.536362x	1.86x
f_ns_b2_t1024	0.154230	0.164080	0.939967x	1.06x
f_ns_b2_t4096	0.583828	0.609406	0.958029x	1.04x
f_ns_b2_t8192	1.160640	1.249189	0.929115x	1.08x
f_ns_b2_t16384	2.310179	2.507152	0.921436x	1.09x
f_st_b1_t512	0.045991	0.158671	0.289854x	3.45x
f_st_b1_t4096	0.298916	0.321602	0.929458x	1.08x
f_st_b1_t16384	1.163152	1.268672	0.916827x	1.09x
f_st_b2_t512	0.089200	0.159130	0.560548x	1.78x
f_st_b2_t4096	0.596951	0.621473	0.960543x	1.04x
f_st_b2_t16384	2.333013	2.544457	0.916900x	1.09x
g_ns_b1_t4096_hv	0.292447	0.310271	0.942551x	1.06x
g_st_b1_t4096_h	0.297945	0.322211	0.924688x	1.08x

Full packed-variable-length CuTe/FLA baseline (20 rows)

Row	Persistent ms	Non-persistent ms	FLA ms	P / FLA	Speedup
p_uniform_n10_t4096	0.503174	0.513231	1.928827	0.260871x	3.83x
p_uniform_n10_t8192	0.826287	0.832048	1.904565	0.433846x	2.30x
p_uniform_n10_t16384	1.351911	1.349822	1.903353	0.710279x	1.41x
p_uniform_n20_t4096	0.698982	0.726441	3.788571	0.184498x	5.42x
p_uniform_n20_t8192	0.941363	0.963633	3.815130	0.246745x	4.05x
p_uniform_n20_t16384	1.546499	1.562673	3.782814	0.408822x	2.45x
p_random_n10_t4096	0.527506	0.491917	1.966145	0.268295x	3.73x
p_random_n10_t8192	0.850628	0.822708	1.918113	0.443471x	2.25x
p_random_n10_t16384	1.401463	1.301126	1.969723	0.711503x	1.41x
p_random_n20_t4096	0.768505	0.768764	3.851995	0.199508x	5.01x
p_random_n20_t8192	1.031841	1.018108	3.856423	0.267564x	3.74x
p_random_n20_t16384	1.583002	1.563621	3.866134	0.409454x	2.44x
p_skewed_n10_t4096	0.537285	0.551372	1.976719	0.271807x	3.68x
p_skewed_n10_t8192	0.806166	0.814602	1.967009	0.409844x	2.44x
p_skewed_n10_t16384	1.339975	1.342001	1.979016	0.677091x	1.48x
p_skewed_n20_t4096	0.744314	0.770045	3.854193	0.193118x	5.18x
p_skewed_n20_t8192	1.017835	1.036325	3.831471	0.265651x	3.76x
p_skewed_n20_t16384	1.528866	1.556334	3.838823	0.398264x	2.51x
g_puni_n10_t4096_hv	0.498931	0.510662	1.938390	0.257395x	3.89x
g_pskew_n20_t16384_h	1.519689	1.570591	3.871482	0.392534x	2.55x

The benchmark entry point is:

python benchmarks/bench_lightning_attn_prefill.py \
  --modes no_state h0_ht varlen \
  --num-heads 64 \
  --iterations 100

🚀 Pull Request Checklist

• Repository-wide pre-commit hook set.
• Changed-file Ruff lint and format.
• git diff --check.
• Final source manifest and source-bound H20 evidence.
• Tests added or updated as needed.

Reviewer notes

• Please focus kernel review on Q/K/V TMA ownership, QK publication barriers,
  the register-resident state lifetime, packed tail TensorMap replacement/TMA
  stores, and the persistent CTA rendezvous between work units.
• The public recurrent state is FP32 in [V, K] orientation. The independent
  reference checks recurrence direction, same-token causality, GVA mapping,
  continuation, state-pool isolation, and packed persistent parity.
• Dynamic shared memory is 189,440 bytes, limiting the kernel to one CTA/SM.
  Final codegen reports 24 WGMMA/HGMMA operations and zero spill/stack/local
  memory for every validated specialization.
• This PR is forward-prefill only. Backward, decode changes, head dimensions
  other than 128, and chunk sizes other than 64 are out of scope.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Aharrypotter]

Thanks for asking. I reran the comparison on the exact pushed PR head 93da019 against the prepared original FlashInfer SM90 C++/CUTLASS implementation (7086b148, CUTLASS f3fde583) on an H20.

The comparison uses the shared semantic subset: BF16 Q/K/V/O, FP32 recurrent state, D=128, chunk size 64, 64-head MHA, scale=1, cpp_lambda=exp(-cula_slope), and final-state output. It is a direct kernel comparison, not a whole-API equivalence claim.

Accuracy:

• 10/10 rows pass CuTe/reference, C++/reference, and direct CuTe/C++ comparisons.
• The minimum per-row bitwise agreement is 99.9409% for BF16 output and 13.9045% for FP32 final state. We do not claim full bitwise identity because the recurrent accumulations use different instruction schedules.
• CuTe/C++ output max-abs / relative-RMS are 1.2207e-4 / 6.6605e-5. Final-state max-abs / relative-RMS are 2.4214e-8 / 2.0957e-7. Every output and state element passes the frozen tolerance.

Performance:

• The primary ratio is final CuTe DSL CUDA-event latency divided by one already-prepared C++ Operation::run; compilation, allocation, workspace setup, operation initialization, comparison, and host synchronization are outside both timing boundaries.
• Every row uses five fresh processes per implementation, 20 warmups, and 100 CUDA-event iterations. All ten rows triggered the pre-registered three-percent escalation rule. There were zero contention-invalid receipts, fallbacks, skips, or missing rows.
• CuTe/C++ ten-row geomean: 1.017616x.
• Worst row: 1.025548x.
• Critical skewed packed row: 1.014649x.

So on this frozen common-subset matrix, the final CuTe DSL source is about 1.76% slower in geometric mean and no more than 2.56% slower than the prepared original C++ kernel. The full run contains 10 source-bound correctness receipts and 100 timing receipts, and an independent auditor reproduced the aggregate exactly.