Support SM80 (Ampere) with dual-arch build

BENCHMARK_A800.md

@@ -0,0 +1,79 @@
+# FlashKDA Benchmark on NVIDIA A800 (SM80)
+
+> Test date: 2026-07-30
+> GPU: NVIDIA A800-SXM4-80GB (SM80, Ampere)
+> CUDA: 12.2
+> PyTorch: 2.4+
+> flash_kda: dual-arch branch `dual-arch-sm80`
+
+## Summary
+
+FlashKDA SM80 forward kernel achieves **~1.8× speedup** over the fla Triton
+`chunk_kda` implementation on A800 for the KDA prefill shape (T=8192, H=96,
+D=128).
+
+| Implementation | Mean | Min | Max | vs flash_kda |
+|---|---|---|---|---|
+| **flash_kda (SM80, bf16 state)** | **3.40 ms** | 3.10 ms | 4.46 ms | — |
+| chunk_kda (fla Triton) | 6.13 ms | 6.09 ms | 6.52 ms | **1.80× slower** |
+| chunk_gated_delta_rule (fla Triton) | 3.63 ms | 3.59 ms | 4.49 ms | 1.07× slower |
+
+## Kernel Breakdown (PyTorch Profiler)
+
+| Kernel | Time / call | Share |
+|---|---|---|
+| `_flash_kda_fwd_recurrence_sm80` (K2) | 2.35 ms | 58.8% |
+| `_flash_kda_fwd_prepare_sm80` (K1) | 1.63 ms | 40.8% |
+
+> ncu is unavailable in the test environment (driver resource restricted);
+> PyTorch Profiler is used for the kernel-level breakdown.
+
+## Shared Memory / Occupancy
+
+`SharedStorageK2` is specialized on `StateFP32` so the fp32 conversion scratch
+buffer (64 KB) is only reserved when actually needed:
+
+| Path | Shared Memory / CTA | Theoretical CTAs / SM (A800 164 KB) |
+|---|---|---|
+| bf16 state | **71.2 KB** | **2** |
+| fp32 state | 96.0 KB | 1 |
+
+Pipeline stages are unified to the values the kernel actually uses:
+`kK2InputStages = 2`, `kK2OutputStages = 1` (previously declared as 3/2 in the
+launch code while the kernel only ever used 2/1).
+
+## Test Configuration
+
+- Shape: `B=1, T=8192, H=96, D=128`
+- `initial_state` / `final_state`: bf16
+- Benchmark script: `benchmarks/bench_fwd.py --mode fixed --H 96 --D 128 --warmup 5 --iters 20 --repeats 3`
+- `chunk_gated_delta_rule` is included as a reference point only; it implements
+  Gated DeltaNet (scalar per-head gate), not KDA.
+
+## Notes
+
+- SM80 path uses cooperative copies and a 2-stage `cp.async` pipeline instead of
+  TMA; numerical output is bit-exact against the torch reference for the tested
+  shapes.
+- `no state` and `fp32 state` configurations show the same min latency
+  (~3.0–3.2 ms); occasional max outliers are first-iteration noise.
+- `FLASH_KDA_CUDA_ARCHS=all` build requires CUDA 12.9+ for `compute_100a`;
+  dual-arch (`80,90a`) build verified on CUDA 12.2.
+
+## Correctness
+
+- `tests/test_fwd.py`: 4 passed (`test_fwd`, `test_fwd_varlen`,
+  `test_fwd_vs_fla`, `test_fwd_varlen_vs_fla`)
+- Output matches `torch_ref` bit-exactly.
+- vs fla `chunk_kda`: err_ratio ≈ 3–5e-3 (bf16 noise level).
+
+## Re-run on 2026-08-17
+
+- **Kernel breakdown reproduced** (PyTorch Profiler): K2 recurrence 2.35 ms (58.8%),
+  K1 prepare 1.63 ms (40.8%) — identical to the table above.
+- End-to-end ~4.0 ms (mean) under heavy GPU co-tenancy (~70 GB used by other
+  tenants); clean-env figure remains ~3.40 ms.
+- `compute-sanitizer --tool memcheck`: **0 errors** on the full A800 shape,
+  both `fixed` and `varlen`.
+- ncu/nsys remain unavailable in this environment (driver resource held,
+  `perf_event_paranoid=4`, no Nsight Systems). See `PROFILING_A800_REFRESH.md`.

README.md

@@ -7,8 +7,8 @@ FlashKDA: Flash Kimi Delta Attention — high-performance KDA kernels built on C
 - **2026-04-22** — Deep-Dive Blog: the design decisions behind FlashKDA v1, read it [here](docs/20260420-flashkda-v1-deep-dive.md).
 
 ## Requirements
-- SM90 and above
-- CUDA 12.9 and above
+- SM80 and above (SM80 Ampere path uses cooperative copies; SM90+ Hopper/Blackwell path uses TMA)
+- CUDA 12.2 and above
 - PyTorch 2.4 and above
 
 ## Installation
@@ -25,7 +25,7 @@ By default, the build detects the current CUDA device and compiles for that arch
 FLASH_KDA_CUDA_ARCHS=all pip install -v --no-build-isolation .
 ```
 
-Supported values are `auto` (default), `all`, or a comma-separated arch list such as `90a,100a`.
+Supported values are `auto` (default), `all`, or a comma-separated arch list such as `80,90a,100a`.
 
 ## Using FlashKDA as an FLA backend
 

csrc/flash_kda.cpp

@@ -25,6 +25,53 @@ int64_t get_workspace_size(
     return H * total_tiles * per_tile_bytes + tile_prefix_bytes;
 }
 
+// Runtime dispatch between the SM80 cooperative-copy implementation and the
+// SM90+ TMA implementation.  Each extension is compiled with exactly one of
+// FLASH_KDA_SM80_ONLY / FLASH_KDA_SM90_ONLY so it only references its own path.
+template <bool HasStateIn, bool HasStateOut, bool StateFP32, bool IsVarlen>
+static void dispatch_fwd(
+    cutlass::bfloat16_t const* q_ptr,
+    cutlass::bfloat16_t const* k_ptr,
+    cutlass::bfloat16_t const* v_ptr,
+    cutlass::bfloat16_t const* g_ptr,
+    cutlass::bfloat16_t const* beta_ptr,
+    void const* initial_state_raw,
+    float scale_f,
+    void* final_state_raw,
+    cutlass::bfloat16_t* out_ptr,
+    void* workspace_ptr,
+    int total_tiles,
+    int T_total,
+    int H,
+    int N_val,
+    int64_t const* cu_seqlens_dev,
+    float const* A_log_ptr,
+    float const* dt_bias_ptr,
+    float gate_scale,
+    cudaStream_t stream,
+    int arch_major
+) {
+#if defined(FLASH_KDA_SM80_ONLY)
+    TORCH_CHECK(arch_major == 8, "flash_kda_C_sm80 requires an SM80 (Ampere) device");
+    flash_kda::sm80::launch_fwd<128, HasStateIn, HasStateOut, StateFP32, IsVarlen>(
+        q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr,
+        initial_state_raw, scale_f, final_state_raw, out_ptr,
+        workspace_ptr, total_tiles,
+        T_total, H, N_val, cu_seqlens_dev,
+        A_log_ptr, dt_bias_ptr, gate_scale, stream);
+#elif defined(FLASH_KDA_SM90_ONLY)
+    TORCH_CHECK(arch_major >= 9, "flash_kda_C_sm90 requires an SM90+ device");
+    flash_kda::sm90::launch_fwd<128, HasStateIn, HasStateOut, StateFP32, IsVarlen>(
+        q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr,
+        initial_state_raw, scale_f, final_state_raw, out_ptr,
+        workspace_ptr, total_tiles,
+        T_total, H, N_val, cu_seqlens_dev,
+        A_log_ptr, dt_bias_ptr, gate_scale, stream);
+#else
+    #error "Must define FLASH_KDA_SM80_ONLY or FLASH_KDA_SM90_ONLY"
+#endif
+}
+
 void fwd(
     torch::Tensor q,
     torch::Tensor k,
@@ -135,6 +182,17 @@ void fwd(
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
 
+    // Detect current device compute capability for runtime dispatch.
+    int current_device = c10::cuda::current_device();
+    cudaDeviceProp prop;
+    cudaGetDeviceProperties(&prop, current_device);
+    int const arch_major = prop.major;
+    TORCH_CHECK(
+        arch_major == 8 || arch_major >= 9,
+        "FlashKDA requires SM80 (Ampere) or newer; got compute capability ",
+        arch_major, ".", prop.minor
+    );
+
     constexpr int CHUNK = 16;
 
     // Get state pointers (nullptr if not present)
@@ -182,12 +240,12 @@ void fwd(
 
     // Dispatch based on state configuration and varlen
     #define LAUNCH(HI, HO, FP32, VL) \
-        launch_fwd<128, HI, HO, FP32, VL>( \
+        dispatch_fwd<HI, HO, FP32, VL>( \
             q_ptr, k_ptr, v_ptr, g_ptr, beta_t_ptr, \
             initial_state_raw, scale_f, final_state_raw, out_ptr, \
             workspace_ptr, total_tiles, \
             int(T_total), int(H), int(N_val), cu_seqlens_dev, \
-            A_log_ptr, dt_bias_ptr, gate_scale, stream)
+            A_log_ptr, dt_bias_ptr, gate_scale, stream, arch_major)
 
     #define DISPATCH_STATE(VL) \
         if (!has_state_in && !has_state_out) { \

csrc/fwd.h

@@ -3,6 +3,9 @@
 
 #include <cutlass/bfloat16.h>
 
+namespace flash_kda {
+namespace sm80 {
+
 template <int D, bool HasStateIn = true, bool HasStateOut = true, bool StateFP32 = false, bool IsVarlen = true>
 void launch_fwd(
     cutlass::bfloat16_t const* q_ptr,
@@ -25,3 +28,33 @@ void launch_fwd(
     float gate_scale,
     cudaStream_t stream
 );
+
+} // namespace sm80
+
+namespace sm90 {
+
+template <int D, bool HasStateIn = true, bool HasStateOut = true, bool StateFP32 = false, bool IsVarlen = true>
+void launch_fwd(
+    cutlass::bfloat16_t const* q_ptr,
+    cutlass::bfloat16_t const* k_ptr,
+    cutlass::bfloat16_t const* v_ptr,
+    cutlass::bfloat16_t const* g_bf16_ptr,
+    cutlass::bfloat16_t const* beta_ptr,
+    void const* initial_state_ptr,
+    float scale,
+    void* final_state_ptr,
+    cutlass::bfloat16_t* out_ptr,
+    void* workspace_ptr,
+    int total_tiles,
+    int T_total,
+    int H,
+    int N,
+    int64_t const* cu_seqlens_ptr,
+    float const* A_log_ptr,
+    float const* dt_bias_ptr,
+    float gate_scale,
+    cudaStream_t stream
+);
+
+} // namespace sm90
+} // namespace flash_kda