feat(kda/sm90): FlashKDA prefill port + intracard-CP prefill

BENCHMARK_H200.md

@@ -55,5 +55,5 @@ Summary (28 configs): **avg=1.58x**, min=1.02x, max=2.51x.
 To reproduce:
 
 ```bash
-python benchmarks/bench_kda_fused_fwd.py --mode both
+python benchmarks/bench_kda_sm90_prefill.py --mode both
 ```

README.md

@@ -149,23 +149,25 @@ python benchmarks/generate_benchmark_hopper_md.py
 
 ```bash
 # Tests for modular KDA forward against FLA Triton implementation
-python -m pytest tests/test_kda_compare_fla.py -v
+python -m pytest tests/test_kda_sm100_chunk_vs_fla.py -v
 # Tests for modular KDA forward against naive KDA reference
-python -m pytest tests/test_kda.py -v
-# Tests for KDA fused forward
+python -m pytest tests/test_kda_sm100_chunk_vs_naive.py -v
+# Tests for the KDA fused forward (SM90 CUDA C++)
 python -m pytest tests/test_kda_fused_fwd.py -v
-# Tests for Lightning Attention fused forward
-python tests/test_lightning_attn.py
+# Tests for the SM90 CuTeDSL two-kernel prefill + intracard CP (vs FLA)
+python -m pytest tests/test_kda_sm90_prefill_vs_fla.py tests/test_kda_sm90_intracard_cp.py -v
+# Tests for Lightning Attention prefill
+python tests/test_lightning_sm100_prefill.py
 # Tests for Lightning Attention decode
-python -m pytest tests/test_la_decode.py -v
+python -m pytest tests/test_lightning_decode.py -v
 
-# test_kda.py and test_kda_compare_fla.py support a fast/slow split.
+# test_kda_sm100_chunk_vs_naive.py and test_kda_sm100_chunk_vs_fla.py support a fast/slow split.
 # Fast (default) — representative correctness paths for default CI and local iteration
-python -m pytest tests/test_kda.py tests/test_kda_compare_fla.py -v
+python -m pytest tests/test_kda_sm100_chunk_vs_naive.py tests/test_kda_sm100_chunk_vs_fla.py -v
 # Slow — broader stress coverage for nightly or manual runs
-python -m pytest -m kda_slow tests/test_kda.py tests/test_kda_compare_fla.py -v
+python -m pytest -m kda_slow tests/test_kda_sm100_chunk_vs_naive.py tests/test_kda_sm100_chunk_vs_fla.py -v
 # Full sweep (fast + slow) — run before submitting a PR
-python -m pytest -m kda_full tests/test_kda.py tests/test_kda_compare_fla.py -v
+python -m pytest -m kda_full tests/test_kda_sm100_chunk_vs_naive.py tests/test_kda_sm100_chunk_vs_fla.py -v
 ```
 
 <details>

REPO_LAYOUT.md

@@ -7,7 +7,6 @@ cuLA/
 ├── cula/                         # Python package (pip install -e .)
 │   ├── __init__.py
 │   ├── utils.py                  # arch asserts, get_pre_scan, cu_seqlens helpers, ...
-│   ├── cudac.py                  # re-export shim over the compiled C++ extension(s)
 │   │
 │   ├── kda/                      # KDA PUBLIC API + autograd + dispatch (NO kernels)
 │   │   ├── __init__.py           # lazy PUBLIC API: chunk_kda, kda_prefill_hopper,
@@ -16,7 +15,7 @@ cuLA/
 │   │   ├── chunk_fwd.py          # chunk_kda_fwd — fwd orchestration (lazy-imports kernels)
 │   │   ├── chunk_intra.py        # fwd intra (C++ ext) + bwd intra (Triton)
 │   │   ├── chunk_bwd.py          # chunk_kda_bwd — Triton + FLA + CuTeDSL + C++ mix
-│   │   └── hopper_fused_fwd.py   # cula_kda_prefill (=kda_prefill_hopper) — SM90 prefill via the C++ kernel (cula.cudac)
+│   │   └── flashkda.py   # cula_kda_prefill (=flashkda_prefill) — SM90 two-kernel K1+K2 prefill, fwd-only
 │   │
 │   ├── lightning/                # [non-KDA] Lightning Attention operator (LinearAttentionChunkwiseDecay, lightning_attn_fwd, linear_attention_decode)
 │   │   └── __init__.py
@@ -25,15 +24,19 @@ cuLA/
 │       ├── __init__.py           # exports kda_decode, fused_sigmoid_..., linear_attention_decode
 │       ├── inv.py / ptx.py       # shared low-level helpers
 │       ├── sm100/                # SM100 shared helper only
+│       │   ├── __init__.py
 │       │   └── ptx.py            #   shared PTX helpers (used by KDA + lightning kernels)
 │       │
-│       ├── kda/                  # ★ KDA Python backends — by arch (sm100 today)
-│       │   ├── policy.py         # SM100 CP dispatch policy: use_intracard_cp:"auto"|bool
+│       ├── kda/                  # ★ ALL KDA backend kernels — by arch (sm90 / sm100)
+│       │   ├── policy.py         # CP dispatch policy: sm90/sm100 decisions, use_intracard_cp:"auto"|bool
 │       │   ├── sm100/            # SM100 (Blackwell) modular-chunk kernels
 │       │   │   ├── delta_h.py    #   recurrence (chunk_gated_delta_rule_fwd_h)
 │       │   │   ├── fwd_o.py      #   output (chunk_gla_fwd_o)
 │       │   │   ├── bwd_wy_dqkg.py#   backward wy/dqkg fused (used by chunk_bwd)
 │       │   │   └── cp/           #   SM100 intracard-CP: chunk_delta_h, pre_scan, merge
+│       │   ├── sm90/             # SM90 (Hopper) two-kernel FlashKDA prefill, fwd-only
+│       │   │   ├── fwd.py  k1.py  k2.py     #   flash_kda_fwd → launch_k1 (prepare) + launch_k2 (recurrence)
+│       │   │   └── cp/           #   SM90 intracard-CP: flashkda, pre_scan, merge, plan
 │       │   ├── decode/           #   single-token + MTP decode
 │       │   │   ├── cute.py       #     kda_decode / fused_sigmoid_gating_delta_rule_update (CuTe DSL)
 │       │   │   ├── mtp.py        #     kda_decode_mtp recurrent / recurrent_ws MTP verify (CuTe DSL)
@@ -49,10 +52,9 @@ cuLA/
 │       └── experimental/
 │           └── linear_attn_prototype.py     # [non-KDA] unwired normalized-linear-attn prototype
 │
-├── csrc/                         # CUDA C++ / CUTLASS
-│   ├── api/{kda_sm90.cu, kda_sm100.cu}        # PyBind11 (cula.cudac): SM90 prefill + SM100 chunk intra/recompute_w_u
-│   ├── kda/sm90/                 # SM90 (Hopper) KDA C++ kernels (CUTLASS 3.x, TMA/wgmma)
-│   ├── kda/sm100/                # Blackwell KDA C++ kernels (CUTLASS 3.x + UMMA)
+├── csrc/                         # CUDA C++ / CUTLASS — SM100 ONLY
+│   ├── api/{pybind.cu, kda_sm100.cu}        # PyBind11 (module cula.cudac): chunk intra + recompute_w_u
+│   ├── kda/sm100/                # Blackwell KDA C++ kernels (CUTLASS 3.x + UMMA); kda_fwd_sm100.cu is the only .cu
 │   └── kerutils/include/         # shared C++ headers (generic device helpers sm80/sm90/sm100, host)
 │
 ├── benchmarks/  tests/  docs/
@@ -66,8 +68,8 @@ cuLA/
 
 | Directory | Language | Description |
 |-----------|----------|-------------|
-| `cula/kda/` | Python | KDA **public API only** — autograd + dispatch, no kernels. Two prefill entries: modular chunk `chunk_kda` (SM100) and `kda_prefill_hopper` (SM90, driving the C++ kernel). |
-| `cula/ops/kda/` | Python (CuTe DSL) | **KDA Python backends**, by arch: `sm100/` (+cp), `decode/`, `experimental/`, plus `policy.py` (CP dispatch). |
+| `cula/kda/` | Python | KDA **public API only** — autograd + dispatch, no kernels. Two prefill entries: modular chunk `chunk_kda` (SM100) and two-kernel K1+K2 `flashkda_prefill` (SM90; the CUDA C++ fused path stays on `kda_prefill_hopper`). See [`cula/kda/README.md`](cula/kda/README.md). |
+| `cula/ops/kda/` | Python (CuTe DSL) | **All KDA backends**, by arch: `sm100/` (+cp), `sm90/` (+cp), `decode/`, `experimental/`, plus `policy.py` (CP dispatch). Both prefill backends are chunked forward; arch is the discriminator (1 impl each). |
 | `cula/ops/lightning/` · `cula/ops/experimental/` | Python (CuTe DSL) | `[non-KDA]` Lightning/linear attention kernels. |
 | `cula/ops/{inv,ptx}.py`, `cula/ops/sm100/ptx.py` | Python | Shared low-level helpers (kept in place; not KDA-specific). |
 | `csrc/kda/{sm90,sm100}/` · `csrc/api/` | CUDA C++ | Hopper SM90 prefill + Blackwell SM100 (chunk intra + recompute_w_u), exposed as `cula.cudac`. |

USAGE.md

@@ -12,6 +12,7 @@ cuLA provides two KDA kernel implementations targeting different GPU architectur
 |---|---|---|
 | Modular Forward | Blackwell (SM100) | `from cula.kda import chunk_kda` |
 | Fused Forward | Hopper (SM90) | `from cula.kda import kda_prefill_hopper` |
+| Two-Kernel Prefill | Hopper (SM90) | `from cula.kda import flashkda_prefill` |
 
 Both are drop-in replacements for [FLA](https://github.com/fla-org/flash-linear-attention)'s `chunk_kda` — just change the import.
 
@@ -115,9 +116,77 @@ print(f'Final state shape: {final_state.shape}')  # [2, 32, 128, 128]
 
 ---
 
-## Intra-Card Context Parallel (chunk_delta_h)
+### Two-Kernel Prefill (SM90 — Hopper)
 
-cuLA includes an intra-card context parallel (CP) path for `chunk_gated_delta_rule_fwd_h`. Long sequences are split into sub-sequences, processed independently in parallel, then merged via a prefix-scan step — unlocking sequence-dimension parallelism on a single GPU.
+The SM90 prefill is a **two-kernel pipeline** — K1 (Prepare) → K2 (Recurrence) — *not* a single fused kernel. Intra-chunk attention, inter-chunk state propagation, and output are computed across the two kernels. **Forward-only; backward is not yet implemented.**
+
+#### Example
+
+```python
+import torch
+from cula.kda import flashkda_prefill
+
+B, T, H, K, V = 2, 2048, 32, 128, 128
+device = 'cuda'
+
+q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16)
+k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16)
+v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16)
+g = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) * 0.1
+beta = torch.randn(B, T, H, device=device, dtype=torch.bfloat16).sigmoid()
+A_log = torch.randn(H, device=device, dtype=torch.float32) * 0.01
+dt_bias = torch.zeros(H * K, device=device, dtype=torch.float32)
+init_state = torch.zeros(B, H, K, V, device=device, dtype=torch.float32)
+
+o, final_state = flashkda_prefill(
+    q=q, k=k, v=v, g=g, beta=beta,
+    A_log=A_log, dt_bias=dt_bias,
+    initial_state=init_state,
+    output_final_state=True,
+    use_qk_l2norm_in_kernel=True,
+    use_gate_in_kernel=True,
+    safe_gate=True,
+    lower_bound=-5.0,
+)
+
+print(f'Output shape: {o.shape}')             # [2, 2048, 32, 128]
+print(f'Final state shape: {final_state.shape}')  # [2, 32, 128, 128]
+```
+
+**Notes**
+
+- Mainly **suitable for large-batch inference**. When both batch size and head count are small, throughput on long sequences is recovered by enabling **intra-card CP** (`use_intracard_cp`), which parallelizes the sequence-length dimension on a single GPU — see [Intra-Card Context Parallel](#intra-card-context-parallel).
+- **Matrix inversion uses fp16 precision**, which is faster and occupies less shared memory but introduces minor numerical differences compared to tf32 inversion.
+- **Intra-subchunk attention uses g-first as anchor**, which causes some numerical differences compared with the FLA Triton implementation (FLA uses g-half as anchor in the diagonal).
+
+---
+
+## Intra-Card Context Parallel
+
+Long sequences can be split into sub-sequences, processed in parallel on one GPU, and merged via a prefix scan — unlocking sequence-dimension parallelism. **Default off; inference-only.** Two surfaces:
+
+### SM90 — via `flashkda_prefill(use_intracard_cp=...)`
+
+Pass `use_intracard_cp` (alias `use_cp`) to the Hopper prefill:
+
+- **`"auto"`** — the auto-router enables CP only when it predicts a speedup, otherwise runs serial (no error, no regression).
+- **`True`** — force CP; raises if the shape cannot be meaningfully split.
+- **`False`** (default) — CP off.
+
+Works with **any sequence length** (non-CHUNK-aligned is handled internally) and **dense or varlen** input. The auto decision uses the device SM count plus two tunables in `cula/ops/kda/sm90/cp/plan.py` (`CULA_KDA_CP_RERUN_RATIO`, `CULA_KDA_CP_AUTO_MIN_SEG_TILES`).
+
+```python
+o, final_state = flashkda_prefill(
+    q=q, k=k, v=v, g=g, beta=beta, A_log=A_log, dt_bias=dt_bias,
+    cu_seqlens=cu_seqlens,              # varlen packed (int32)
+    output_final_state=True, use_gate_in_kernel=True, safe_gate=True, lower_bound=-5.0,
+    use_intracard_cp="auto",           # "auto" | True | False  (default False)
+)
+```
+
+### SM100 — low-level `chunk_gated_delta_rule_fwd_h`
+
+cuLA also exposes intra-card CP directly on the SM100 recurrence `chunk_gated_delta_rule_fwd_h`, gated by an environment variable.
 
 **Requirements**
 

benchmarks/bench_kda_sm90_cp.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+# Copyright 2025-2026 Ant Group Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+bench_kda_sm90_cp.py — Benchmark: SM90 intracard CP speedup (CP-on vs CP-off)
+
+Measures the speedup of the SM90 intracard context-parallel path against
+the serial K1+K2 baseline across varlen configurations.
+
+Usage:
+  python bench_kda_sm90_cp.py [--ncu] [--sanitizer]
+
+With --ncu, warmup=1 and iters=1 for ncu profiling:
+  ncu --set full -o report python bench_kda_sm90_cp.py --ncu
+"""
+
+import argparse
+import pathlib
+import sys
+
+sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
+
+import torch
+
+from benchmarks.utils import SEED, exclusive_cumsum, prepare_safe_gate_inputs, set_seed
+from cula.kda import flashkda_prefill as cula_kda_prefill
+from cula.utils import assert_hopper, get_device_sm_count
+
+D = 128
+H_VALUES = [4, 8]
+WARMUP = 10
+N_ITERS = 100
+NCU_MODE = False
+SANITIZER_MODE = False
+
+# (tag, seq_lens) — each entry is tested at every H in H_VALUES.
+# SM90 CHUNK=16, so sequences need to be long enough (>= ~8K tiles) for CP to pay off.
+CONFIGS = [
+    ("T=4K", [4096]),
+    ("T=8K", [8192]),
+    ("T=16K", [16384]),
+    ("T=32K", [32768]),
+    ("T=64K", [65536]),
+    ("2x16K", [16384, 16384]),
+    ("32K+4K", [32768, 4096]),
+    ("32K+1K", [32768, 1024]),
+    ("64K+1K", [65536, 1024]),
+    ("64K+2x1K", [65536, 1024, 1024]),
+    ("64K+5x1K", [65536] + [1024] * 5),
+]
+
+
+def time_kernel(fn, warmup=None, n_iters=None):
+    if warmup is None:
+        warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP
+    if n_iters is None:
+        n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start_evt = torch.cuda.Event(enable_timing=True)
+    end_evt = torch.cuda.Event(enable_timing=True)
+    start_evt.record()
+    for _ in range(n_iters):
+        fn()
+    end_evt.record()
+    torch.cuda.synchronize()
+    return start_evt.elapsed_time(end_evt) / n_iters
+
+
+def run_kernel(q, k, v, g, beta, scale, A_log, dt_bias, cu_seqlens, lower_bound, *, use_cp):
+    cula_kda_prefill(
+        q,
+        k,
+        v,
+        g,
+        beta,
+        scale=scale,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        cu_seqlens=cu_seqlens,
+        output_final_state=False,
+        safe_gate=True,
+        lower_bound=lower_bound,
+        use_intracard_cp="auto" if use_cp else False,
+    )
+
+
+def bench_cp(h_values, configs):
+    print("\n" + "=" * 100)
+    print(" SM90 Intracard CP Benchmark: CP-on vs CP-off")
+    print("=" * 100)
+
+    device = torch.device("cuda")
+    assert_hopper(device)
+    num_sms = get_device_sm_count(device)
+    print(f" [Device] {torch.cuda.get_device_name(0)}  SMs={num_sms}")
+    results = []
+
+    for H in h_values:
+        for tag, seq_lens in configs:
+            set_seed(SEED)
+            torch.cuda.empty_cache()
+
+            total_T = sum(seq_lens)
+            cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device)
+            inputs = prepare_safe_gate_inputs(1, total_T, H, D, device, cu_seqlens=cu_seqlens, seed=SEED)
+            q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"]
+            A_log, dt_bias = inputs["A_log"], inputs["dt_bias"]
+            scale, lower_bound = inputs["scale"], inputs["lower_bound"]
+
+            common = dict(
+                q=q,
+                k=k,
+                v=v,
+                g=g,
+                beta=beta,
+                scale=scale,
+                A_log=A_log,
+                dt_bias=dt_bias,
+                cu_seqlens=cu_seqlens,
+                lower_bound=lower_bound,
+            )
+
+            ms_off = time_kernel(lambda: run_kernel(**common, use_cp=False))
+            ms_on = time_kernel(lambda: run_kernel(**common, use_cp=True))
+
+            speedup = ms_off / ms_on if ms_on > 0 else float("inf")
+            r = dict(tag=tag, H=H, total_T=total_T, ms_off=ms_off, ms_on=ms_on, speedup=speedup)
+            results.append(r)
+
+            del q, k, v, g, beta, A_log, dt_bias, inputs
+            torch.cuda.empty_cache()
+
+    return results
+
+
+def print_report(results, h_values):
+    sep = "=" * 95
+    print(f"\n\n{sep}")
+    print("               BENCHMARK REPORT: SM90 Intracard CP")
+    print("               CP-on vs CP-off (intracard_prefill vs flash_kda_fwd)")
+    print(f"               D={D}  dtype=bf16  safe_gate=True")
+    wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP
+    ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS
+    mode_tag = "  [NCU mode]" if NCU_MODE else ("  [Sanitizer mode]" if SANITIZER_MODE else "")
+    print(f"               Warmup={wu}  Iters={ni}{mode_tag}")
+    print(sep)
+
+    for H_val in h_values:
+        h_results = [r for r in results if r["H"] == H_val]
+        if not h_results:
+            continue
+        print(f"\n  [H={H_val}]")
+        print(f"  {'─' * 80}")
+        print(f"  {'config':<20s} {'T':>7s}  │  {'CP_off(ms)':>10s}  {'CP_on(ms)':>10s}  {'Speedup':>8s}")
+        print(f"  {'─' * 80}")
+        for r in h_results:
+            print(f"  {r['tag']:<20s} {r['total_T']:>7d}  │  {r['ms_off']:>10.4f}  {r['ms_on']:>10.4f}  {r['speedup']:>7.2f}x")
+        print(f"  {'─' * 80}")
+
+    speedups = [r["speedup"] for r in results]
+    if speedups:
+        geo = 1.0
+        for s in speedups:
+            geo *= s
+        geo = geo ** (1 / len(speedups))
+        print(f"\n  All configs: geo-mean={geo:.2f}x  best={max(speedups):.2f}x  worst={min(speedups):.2f}x")
+
+    print(f"\n{sep}\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="bench_kda_sm90_cp: SM90 intracard CP speedup")
+    parser.add_argument("--ncu", action="store_true", help="NCU profiling mode: warmup=1, iters=1")
+    parser.add_argument("--sanitizer", action="store_true", help="Sanitizer mode: warmup=1, iters=1")
+    args = parser.parse_args()
+
+    global NCU_MODE, SANITIZER_MODE
+    if args.ncu:
+        NCU_MODE = True
+        print("[NCU mode] warmup=1, iters=1")
+    if args.sanitizer:
+        SANITIZER_MODE = True
+        print("[Sanitizer mode] warmup=1, iters=1")
+
+    results = bench_cp(H_VALUES, CONFIGS)
+    print_report(results, H_VALUES)
+    return results
+
+
+if __name__ == "__main__":
+    main()