Support SM80 (Ampere) with dual-arch build

[Leslie360]

FlashKDA SM80 (Ampere) Support

Summary

Add SM80 (Ampere/A800) support to FlashKDA via a dual-architecture build. The existing SM90+ TMA path is preserved unchanged; a new cooperative-copy path is added under csrc/sm80/.

• SM80 fwd e2e: ~3.40 ms (T=8192, H=96, D=128, A800)
• fla Triton chunk_kda: ~6.13 ms
• Speedup: ~1.8×

Design

• csrc/sm80/: new SM80 cooperative-copy implementation (no TMA, no cluster barriers)
• csrc/smxx/: existing SM90+ TMA implementation (unchanged)
• setup.py: builds flash_kda_C_sm80 and flash_kda_C_sm90 as separate extensions based on FLASH_KDA_CUDA_ARCHS
• flash_kda/__init__.py: runtime dispatch by torch.cuda.get_device_capability()
• csrc/flash_kda.cpp: dispatch_fwd selects flash_kda::sm80::launch_fwd or flash_kda::sm90::launch_fwd; each extension is compiled with FLASH_KDA_SM80_ONLY / FLASH_KDA_SM90_ONLY

SM80-Specific Optimizations

• 2-stage cp.async pipeline in K2 to hide gmem latency
• Vectorized cooperative copies (coop_copy_2d_vec8 / coop_copy_1d_vec4, 16B per thread-op)
• SharedStorageK2 specialized on StateFP32: the 64 KB fp32 conversion scratch is only reserved when the state dtype is actually fp32; the bf16 path uses 71.2 KB smem (→ 2 CTA/SM on A800) instead of ~96 KB
• Pipeline stages unified to the values the kernel actually uses (kK2InputStages = 2, kK2OutputStages = 1)
• K1 __launch_bounds__ adjusted from (256, 8) to (256, 4); ptxas reports 0 spills

Benchmark

See BENCHMARK_A800.md for full data.

Implementation	Mean	vs flash_kda
flash_kda (SM80, bf16 state)	3.40 ms	—
chunk_kda (fla Triton)	6.13 ms	1.80× slower
chunk_gated_delta_rule (fla Triton)	3.63 ms	1.07× slower

Kernel breakdown (PyTorch Profiler): K2 recurrence 2.35 ms (58.8%), K1 prepare 1.63 ms (40.8%).

chunk_gated_delta_rule is included as a reference only; it implements Gated DeltaNet (scalar per-head gate), not KDA.

Correctness

• tests/test_fwd.py: 4 passed (test_fwd, test_fwd_varlen, test_fwd_vs_fla, test_fwd_varlen_vs_fla)
• Bit-exact against torch_ref for tested shapes.
• vs fla chunk_kda: err_ratio ≈ 3–5e-3 (bf16 noise level).

Build Notes

• FLASH_KDA_CUDA_ARCHS=auto / 80 / 80,90a verified on A800 (CUDA 12.2).
• FLASH_KDA_CUDA_ARCHS=all requires CUDA 12.9+ for compute_100a; not verified in the current environment.
• SM90 path is untouched; no functional change expected on Hopper/Blackwell.

Files Changed

• setup.py
• README.md
• csrc/flash_kda.cpp
• csrc/fwd.h
• flash_kda/__init__.py
• csrc/sm80/ (new)
• BENCHMARK_A800.md (new)