[Feature] Integrate Lightning SM90 Implementation via CuTe DSL

[zheyang0825]

[Feature] Integrate Lightning SM90 Implementation via CuTe DSL

Background

FlashInfer has a production-ready Lightning Attention SM90 (Hopper) implementation. For API completeness, we could port it into cuLA using CuTe DSL, along with corresponding benchmarks and unit tests.

Upstream PR: flashinfer-ai/flashinfer#2276 from guangyunh-nv

Note: The kernel design and algorithm credits belong to the original authors (guangyunh-nv). This issue tracks a porting effort, not an original implementation.

Upstream Implementation Overview

FlashInfer PR #2276 implements Lightning Attention Prefill on Hopper architecture:

• SM90 Optimized: TMA warp-specialized architecture with asynchronous copy and warp group scheduling
• Optional gating and final-state output
• High-level Python API with chunked prefill support

Task Checklist

• Port kernels via CuTe DSL: Rewrite the upstream C++ Hopper lightning kernels using CuTe DSL, following cuLA's build system and coding conventions
• Adapt Python interface: Align with cuLA's existing API style and expose the Lightning SM90 Python bindings
• Add benchmarks: Add performance benchmarks under cuLA's framework, covering the same settings as the KDA benchmarks:
  • Fixed-length: B={1,2}, T={512, 1024, 4096, 8192, 16384}, H=64, D=128, dtype=bf16
  • Varlen: num_seqs={10, 20}, total_len={4096, 8192, 16384}, distributions: uniform / random / skewed
• Performance validation: Verify that the CuTe DSL version achieves comparable performance to the upstream C++ version across all benchmark settings above
• Add unit tests: Port the reference implementation and test cases from upstream tests/gdn/ to ensure correctness
• Update documentation: Document the newly added Lightning SM90 support

References

• Upstream PR: flashinfer-ai/flashinfer#2276
• Related Issue: flashinfer-ai/flashinfer#1690

[Aharrypotter]

I'd like to take this. I've been studying #2276's C++ collective mainloop closely.
I'll open a draft PR this week. Please assign it to me.

[zheyang0825]

I'd like to take this. I've been studying #2276's C++ collective mainloop closely. I'll open a draft PR this week. Please assign it to me.

Assigned to you. Thanks for volunteering!

If you run into any questions or need clarification while working on this, feel free to ask here. Looking forward to your PR.