[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 indexed by stable logical CTA slots,
  plus 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, stable TensorMap-slot repair, and decay-LUT
warp-ordering repair are separate commits so their implementations 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

The final evidence is bound to pushed PR head
247941ee0addb6de4c279fda31f2389fa5863c19.

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

• Focused local product suite: 43/43 pass.
• Complete H20 FLA semantic matrix: 38/38 rows 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.
• Commit-1/final-source output and state are bitwise equal across fixed,
  stateful, GVA, packed non-persistent, and packed persistent paths.
• Ragged, stage-wrap, two-stream, and 80-work-unit/78-CTA reuse checks
  produce the same output/state hashes as commit 1.
• Public default-path host-sync audit: four warmed packed calls and zero
  dispatcher-visible violations.
• Final Compute Sanitizer matrix: 17/17 pass — 5 memcheck, 4 synccheck,
  and 8 racecheck attempts, with zero racecheck hazards.
• Fresh source-bound codegen: 10/10 specializations, 9 unique cubins,
  exactly two new warp barriers per specialization, 24 WGMMA per
  specialization, no %smid / SR_VIRTUALSMID, and zero spill/stack/local
  memory.
• Final tracked-source manifest: 195/195 files match locally and on H20.
• Ruff lint/format and git diff --check pass.

🔬 Accuracy vs original C++ implementation

The source-bound comparison uses 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.

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, 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.039275	0.038971	1.007817x
f_b2_t512_z	0.074635	0.073614	1.013862x
f_b1_t4096_z	0.257677	0.251751	1.023541x
f_b2_t4096_z	0.508240	0.497178	1.022250x
f_b1_t16384_z	1.002423	0.978645	1.024296x
f_b1_t512_h0	0.040428	0.040775	0.991493x
f_b1_t4096_h0	0.260528	0.254996	1.021695x
f_b1_t16384_h0	1.007592	0.984179	1.023789x
p_uniform_4096_h0	0.298660	0.291589	1.024249x
p_skewed_8192_h0	0.466860	0.461181	1.012313x
Combined geomean	—	—	1.016481x

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.65% slower in geometric mean and within
2.43% of the prepared original C++ kernel on every row of this matrix.

⚡ Final-source CuTe vs FLA

This table reseals the comparison on exact pushed PR head
247941ee0addb6de4c279fda31f2389fa5863c19.

The reference is FLA chunk_simple_gla at
cca521e9efa50d2859f27702923ad1ff46305dee (FLA 0.5.1, Triton 3.5.1).
Both implementations ran on the same NVIDIA H20-3e GPU in BF16 with D=128 and
chunk size 64. Ratios are CuTe latency / FLA latency, so lower than 1 means
CuTe is faster.

Acceptance and protocol

• Correctness: 38/38 rows passed the frozen tolerance contract. Packed
  output/state parity and untouched state-pool slots passed on 20/20 rows.
• Worst output error: max abs 7.62939453e-06, relative RMS
  3.14672361e-04.
• Worst final-state/state-pool error: max abs 7.86771998e-05, relative RMS
  7.14503476e-05.
• Timing: 20 warmups plus 100 CUDA-event iterations in each fresh process;
  three process replicas per implementation, with fixed rows alternating
  FLA/CuTe, packed rows using a three-way round-robin, and compile/first-call
  latency excluded.
• Audit accounting: 288/288 timing receipts accepted, no fallback, skip,
  contention-invalid receipt, or missing row. No ratio entered the frozen
  ±3% band, so five-process escalation was not triggered.
• Primary scope: 16 fixed native-MHA rows plus 18 packed native-MHA rows using
  the final default non-persistent scheduler. The explicit persistent scheduler
  and four GVA rows are reported as secondary results.

Scope	Rows	Geomean CuTe / FLA	Reciprocal speedup	Min ratio	Max ratio
Fixed native MHA	16	0.638699x	1.57x	0.254419x	0.839356x
Packed native MHA, default non-persistent	18	0.200602x	4.98x	0.085684x	0.504235x
Primary native MHA	34	0.345958x	2.89x	0.085684x	0.839356x
Packed native MHA, explicit persistent	18	0.207320x	4.82x	0.087398x	0.538515x
GVA, outside primary geomean	4	0.404571x	2.47x	0.154392x	0.827973x
All logical rows, using each row's primary path	38	0.351705x	2.84x	0.085684x	0.839356x

The final-source CuTe path is faster on all 34/34 primary native-MHA rows
and all 38/38 logical rows. The worst primary row is
f_ns_b2_t4096 at 0.839356x, or about 16.1% lower latency (1.19x
reciprocal speedup).

Full fixed-length final-source results (18 rows)

Row	Scope	CuTe ms	FLA ms	CuTe / FLA
f_ns_b1_t512	primary	0.038185	0.150085	0.254419x
f_ns_b1_t1024	primary	0.070295	0.149481	0.470259x
f_ns_b1_t4096	primary	0.256095	0.308653	0.829716x
f_ns_b1_t8192	primary	0.503510	0.605031	0.832205x
f_ns_b1_t16384	primary	0.997657	1.221316	0.816871x
f_ns_b2_t512	primary	0.072723	0.151736	0.479273x
f_ns_b2_t1024	primary	0.134844	0.163311	0.825689x
f_ns_b2_t4096	primary	0.504956	0.601600	0.839356x
f_ns_b2_t8192	primary	0.998426	1.223236	0.816217x
f_ns_b2_t16384	primary	1.987451	2.499582	0.795113x
f_st_b1_t512	primary	0.040460	0.154245	0.262311x
f_st_b1_t4096	primary	0.260280	0.319954	0.813494x
f_st_b1_t16384	primary	1.005222	1.258422	0.798796x
f_st_b2_t512	primary	0.079401	0.169517	0.468393x
f_st_b2_t4096	primary	0.513464	0.616820	0.832438x
f_st_b2_t16384	primary	2.002498	2.524512	0.793222x
g_ns_b1_t4096_hv	GVA secondary	0.255810	0.308959	0.827973x
g_st_b1_t4096_h	GVA secondary	0.260104	0.319813	0.813302x

Full packed-variable-length final-source results (20 rows)

Row	Scope	Non-persistent ms	Persistent ms	FLA ms	Non-persistent / FLA	Persistent / FLA
p_uniform_n10_t4096	primary	0.299094	0.304160	1.905937	0.156928x	0.159585x
p_random_n10_t4096	primary	0.273203	0.308375	1.944183	0.140523x	0.158614x
p_skewed_n10_t4096	primary	0.292760	0.296557	1.960785	0.149307x	0.151244x
p_uniform_n10_t8192	primary	0.509239	0.519610	1.914239	0.266027x	0.271445x
p_random_n10_t8192	primary	0.514724	0.549892	1.929898	0.266710x	0.284933x
p_skewed_n10_t8192	primary	0.518548	0.525644	1.933433	0.268201x	0.271871x
p_uniform_n10_t16384	primary	0.959634	0.977369	1.914227	0.501317x	0.510581x
p_random_n10_t16384	primary	0.944285	1.042516	1.935911	0.487773x	0.538515x
p_skewed_n10_t16384	primary	0.968080	0.982405	1.919900	0.504235x	0.511696x
p_uniform_n20_t4096	primary	0.363404	0.365287	3.696972	0.098298x	0.098807x
p_random_n20_t4096	primary	0.324556	0.339438	3.787835	0.085684x	0.089613x
p_skewed_n20_t4096	primary	0.327394	0.327566	3.747992	0.087352x	0.087398x
p_uniform_n20_t8192	primary	0.561136	0.568076	3.737222	0.150148x	0.152005x
p_random_n20_t8192	primary	0.526246	0.553435	3.794610	0.138682x	0.145848x
p_skewed_n20_t8192	primary	0.545861	0.547915	3.869205	0.141078x	0.141609x
p_uniform_n20_t16384	primary	0.955606	0.974590	3.759780	0.254165x	0.259215x
p_random_n20_t16384	primary	0.968760	1.034116	3.844807	0.251966x	0.268964x
p_skewed_n20_t16384	primary	0.980969	0.989274	3.790462	0.258799x	0.260990x
g_puni_n10_t4096_hv	GVA secondary	0.296281	0.301916	1.919022	0.154392x	0.157328x
g_pskew_n20_t16384_h	GVA secondary	0.980723	0.986702	3.805924	0.257683x	0.259254x

⚡ Direct synchronization-cost check

The final decay-LUT synchronization commit was also compared directly against
the accepted stable-slot commit on eight fixed and packed paths, using five
fresh processes per source and alternating order.

Check	Result
Output/state bitwise paths	8/8
Timing receipts	80/80
Final source / commit 1 geomean	0.999511x
Worst path	1.000424x

The explicit warp synchronization is therefore inside the frozen <=1.03x
geomean and <=1.05x per-path guardrails.

This is kernel-only H20 BF16 prefill evidence. It does not claim decode,
backward, other head dimensions, other GPUs, or end-to-end serving performance.

🚀 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, stable logical TensorMap-slot ownership, and the persistent CTA
  rendezvous between work units.
• Packed non-persistent scheduling indexes TensorMap slots by
  sequence_idx * value_heads + value_head_idx; persistent scheduling uses
  stable logical cta_idx. Device code no longer uses %smid.
• 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]

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.

[icavan]

One portability note regarding the TensorMap workspace indexing:

[icavan]

LGTM

[Aharrypotter]

Follow-up after the TensorMap portability review: the PR is now resealed on exact pushed head 247941e.

• bb69d07 replaces %smid workspace ownership with stable logical CTA slots.
• 247941e adds explicit warp ordering to the three-segment decay-LUT scan.

Final-source validation:

• FLA: 38/38 correctness rows and 288/288 timing receipts; CuTe/FLA primary geomean 0.345958x, worst primary row 0.839356x.
• Prepared original C++: 10/10 correctness rows and 100/100 timing receipts; CuTe/C++ geomean 1.016481x, worst row 1.024296x.
• Direct commit-1/final-source check: 8/8 output/state bitwise paths and 80/80 timing receipts; geomean 0.999511x, worst path 1.000424x.
• Compute Sanitizer: 17/17 jobs pass (5 memcheck, 4 synccheck, 8 racecheck), with zero racecheck hazards.
• Codegen/resources: 10/10 specializations pass, %smid is absent, and stack/local/spill remain zero.

The PR body has been updated with the complete source-bound C++ and FLA tables and the final validation boundary.

[zheyang0825]

One packaging compatibility issue:

[Aharrypotter]

Additional context for the dependency constraint nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0:

4.5.0 is intentionally excluded because it contains known upstream CuTe DSL compiler regressions that were fixed in CUTLASS 4.5.1:

• NVIDIA/cutlass#3210 reports that indexing a vector loaded from register memory can introduce unexpected .local memory traffic and a severe performance regression.
• NVIDIA/cutlass#3201 reports incorrect scf.for loop-carried variable analysis, which can prevent modified variables from being propagated correctly.
• The CUTLASS 4.5.1 changelog lists both issues as fixed in 4.5.1.

These are relevant risks for this kernel because the Lightning SM90 implementation relies heavily on register-resident accumulators and DSL-lowered recurrent loops. To be precise, we did not claim a Lightning-specific reproduction on 4.5.0; we conservatively excluded the known-bad compiler release and qualified 4.5.1, 4.5.2, and 4.5.3 instead.

[zheyang0825]

LGTM