refactor(kda): replace SM100 native forward with CuTeDSL

.coderabbit.yaml

@@ -11,6 +11,10 @@ reviews:
     drafts: false
     auto_incremental_review: true
 
+  pre_merge_checks:
+    docstrings:
+      mode: "off"
+
   path_filters:
     - "!csrc/cutlass/**"
     - "!third_party/**"

README.md

@@ -20,13 +20,13 @@ Linear attention mechanisms reformulate standard attention to use linear-time st
 
 cuLA supports both **Hopper (SM90)** and **Blackwell (SM10X)** GPUs.
 
-> **Requirements (Hopper & Blackwell):** Python 3.12+, CUDA Toolkit 12.9+ (SM10X support), NVCC 12.9+, PyTorch 2.9.1+
+> **Requirements (Hopper & Blackwell):** Python 3.12+, CUDA Toolkit 12.9+, PyTorch 2.9.1+. Building the optional SM90 native extension also requires NVCC.
 
 > **Note:** The PyTorch CUDA version must match your system CUDA Toolkit version. Check with `nvcc --version` and `python -c "import torch; print(torch.version.cuda)"`.
 
 ### Pre-built Wheels
 
-Pre-built fat-binary wheels (SM90 + SM100 + SM103) are available on [GitHub Releases](https://github.com/inclusionAI/cuLA/releases). Linux wheels target `manylinux_2_28` and require glibc 2.28 or newer:
+Pre-built wheels with the SM90 native extension and CuTeDSL JIT kernels for SM100/SM103 are available on [GitHub Releases](https://github.com/inclusionAI/cuLA/releases). Linux wheels target `manylinux_2_28` and require glibc 2.28 or newer:
 
     pip install "cuda-linear-attention==<VERSION>+<CUDA_TAG>" -f https://github.com/inclusionAI/cuLA/releases/expanded_assets/<TAG>
 
@@ -57,7 +57,7 @@ pip install -e third_party/flash-linear-attention
 pip install -e . --no-build-isolation
 ```
 
-**Build fat wheel (SM90 + SM100 + SM103):**
+**Build a wheel with the SM90 extension (SM100/SM103 use CuTeDSL JIT):**
 
 ```bash
 CULA_BUILD_ALL_ARCHS=1 python -m build --wheel --no-isolation

REPO_LAYOUT.md

@@ -7,7 +7,7 @@ cuLA/
 ├── cula/                              # Python package (pip install -e .)
 │   ├── __init__.py
 │   ├── backends.py                    # Generic backend registry, probing, verification, dispatch
-│   ├── cudac.py                       # Lazy proxy for the per-architecture CUDA extension
+│   ├── cudac.py                       # Lazy proxy for the remaining SM90 CUDA extension
 │   ├── utils.py                       # Architecture, stream-buffer, and cu_seqlens helpers
 │   │
 │   ├── kda/                           # KDA public API, wrappers, autograd, routing, and Triton support kernels
@@ -22,8 +22,8 @@ cuLA/
 │   │   ├── hopper_fused_fwd_opt.py    # SM90 CUDA C++ optimized/CP wrapper
 │   │   ├── chunk.py                   # SM100 modular chunk API + autograd
 │   │   ├── chunk_fwd.py               # SM100 forward orchestration
-│   │   ├── chunk_intra.py             # C++ forward intra + Triton backward intra
-│   │   ├── chunk_bwd.py               # Triton + FLA + CuTeDSL + C++ backward orchestration
+│   │   ├── chunk_intra.py             # CuTeDSL forward intra + Triton backward intra
+│   │   ├── chunk_bwd.py               # Triton + FLA + CuTeDSL backward orchestration
 │   │   ├── wy_intra.py / wy_recompute.py # Triton WY preparation kernels
 │   │   └── gate_l2norm_fused.py / l2norm_qk_fused.py # SM90 Triton preprocessing kernels
 │   │
@@ -52,15 +52,12 @@ cuLA/
 │       │       └── prefill_kernel.py   # Fixed and packed recurrent prefill kernel
 │       └── experimental/               # [non-KDA] unwired prototypes
 │
-├── csrc/                               # CUDA C++ / CUTLASS kernels for SM90 and SM100/SM103
+├── csrc/                               # Remaining CUDA C++ / CUTLASS kernels for SM90
 │   ├── api/
 │   │   ├── pybind.cu                   # Shared named binding definitions
 │   │   ├── pybind_sm90.cu              # SM90-specific binding translation unit
-│   │   ├── pybind_sm100.cu             # SM100/SM103-specific binding translation unit
-│   │   ├── kda_sm90.cu                 # SM90 host API
-│   │   └── kda_sm100.cu                # SM100/SM103 host API
+│   │   └── kda_sm90.cu                 # SM90 host API
 │   ├── kda/sm90/                       # Hopper fully-fused prefill kernels
-│   ├── kda/sm100/                      # Blackwell modular chunk kernels
 │   ├── kerutils/include/               # Shared device and host helpers
 │   └── cutlass/                        # CUTLASS submodule
 │
@@ -80,7 +77,7 @@ cuLA/
 | `cula/kda/` | Python / Triton | Public KDA APIs, wrappers, autograd, routing, and supporting Triton kernels. `kda_prefill` dispatches SM90 calls to FlashKDA first and then the fully-fused CUDA backend; explicit backend exports remain available. See [`cula/kda/README.md`](cula/kda/README.md). |
 | `cula/backends.py` · `cula/kda/backends/` | Python | Generic backend registry plus KDA-specific availability probes, capability verifiers, priorities, and adapters. |
 | `cula/ops/kda/` | Python (CuTeDSL) | CuTeDSL KDA kernels organized into SM100 modular kernels, SM90 FlashKDA K1+K2 and intracard CP, decode, and experimental code. The fully-fused SM90 implementation lives under `csrc/`, not here. |
-| `csrc/kda/{sm90,sm100}/` | CUDA C++ | Hopper fully-fused prefill and Blackwell modular chunk kernels. |
-| `csrc/api/` · `cula/cudac.py` | CUDA C++ / Python | Per-architecture `_cudac_sm90` and `_cudac_sm100` extensions, exposed lazily through the `cula.cudac` compatibility proxy. |
+| `csrc/kda/sm90/` | CUDA C++ | Hopper fully-fused prefill kernels. SM100/SM103 KDA forward is implemented entirely in CuTeDSL under `cula/ops/kda/sm100/`. |
+| `csrc/api/` · `cula/cudac.py` | CUDA C++ / Python | The remaining `_cudac_sm90` extension and its lazy compatibility proxy. Blackwell public APIs do not load a native extension. |
 | `cula/ops/lightning/` · `cula/ops/experimental/` | Python (CuTeDSL) | `[non-KDA]` Lightning/linear-attention kernels and prototypes. Lightning prefill dispatches to the SM90 or SM100 backend from Q's device capability. |
 | `cula/ops/{inv,ptx}.py` · `cula/ops/sm100/ptx.py` | Python | Shared low-level helpers used across operators. |

benchmarks/bench_kda_sm100_intra_fused.py

@@ -279,7 +279,7 @@ def run_fla_with_gk():
         )
         return (gk, *fla)
 
-    def run_csrc_with_gk():
+    def run_cutedsl_with_gk():
         gk = kda_gate_chunk_cumsum(
             g=g,
             A_log=a_log,
@@ -288,7 +288,7 @@ def run_csrc_with_gk():
             chunk_size=BT,
             lower_bound=args.lower_bound,
         )
-        csrc = cula_chunk_kda_fwd_intra(
+        cutedsl = cula_chunk_kda_fwd_intra(
             q=q,
             k=k,
             v=k,
@@ -298,7 +298,7 @@ def run_csrc_with_gk():
             chunk_size=BT,
             safe_gate=True,
         )
-        return (gk, *csrc)
+        return (gk, *cutedsl)
 
     valid_mask = _valid_lower_mask(args.B, args.T, args.H, q.device)
 
@@ -336,7 +336,7 @@ def compare(cutedsl: tuple[torch.Tensor, ...], fla_with_gk: tuple[torch.Tensor,
 
     nt = args.B * (args.T // BT)
     detail = f"B={args.B} T={args.T} H={args.H} K={args.K} NT={nt} bias={not args.no_bias} variant={args.cutedsl_variant}"
-    run_ref = run_csrc_with_gk if args.baseline == "csrc" else run_fla_with_gk
+    run_ref = run_cutedsl_with_gk if args.baseline == "cutedsl" else run_fla_with_gk
     return BenchCase("equal", detail, run, run_ref, compare)
 
 
@@ -534,7 +534,7 @@ def run_fla_with_gk():
         )
         return (gk, *fla)
 
-    def run_csrc_with_gk():
+    def run_cutedsl_with_gk():
         gk = kda_gate_chunk_cumsum(
             g=g,
             A_log=a_log,
@@ -545,7 +545,7 @@ def run_csrc_with_gk():
             chunk_indices=chunk_indices,
             lower_bound=args.lower_bound,
         )
-        csrc = cula_chunk_kda_fwd_intra(
+        cutedsl = cula_chunk_kda_fwd_intra(
             q=q,
             k=k,
             v=k,
@@ -557,7 +557,7 @@ def run_csrc_with_gk():
             chunk_size=BT,
             safe_gate=True,
         )
-        return (gk, *csrc)
+        return (gk, *cutedsl)
 
     valid_mask = _valid_lower_mask_varlen(seq_lens, args.H, q.device)
 
@@ -601,7 +601,7 @@ def compare(cutedsl: tuple[torch.Tensor, ...], fla_with_gk: tuple[torch.Tensor,
         f"seqs={len(seq_lens)} total={total_t} H={args.H} K={args.K} NT={nt} "
         f"T_launch={t_launch} pure={pure} dist={args.dist} bias={not args.no_bias} seq_lens=[{seq_preview}]"
     )
-    run_ref = run_csrc_with_gk if args.baseline == "csrc" else run_fla_with_gk
+    run_ref = run_cutedsl_with_gk if args.baseline == "cutedsl" else run_fla_with_gk
     return BenchCase("varlen", detail, run, run_ref, compare)
 
 
@@ -713,7 +713,7 @@ def main() -> None:
     parser.add_argument(
         "--no-fla", action="store_true", help="Only time CuTeDSL; skip baseline timing and accuracy comparison."
     )
-    parser.add_argument("--baseline", choices=("csrc", "fla"), default="csrc")
+    parser.add_argument("--baseline", choices=("cutedsl", "fla"), default="cutedsl")
     parser.add_argument(
         "--with-recompute-wu",
         action="store_true",

benchmarks/bench_recompute_wu.py

@@ -34,7 +34,6 @@
 
 from fla.ops.kda.wy_fast import recompute_w_u_fwd as fla_recompute_w_u_fwd
 
-import cula.cudac as cula_cuda
 from benchmarks.utils import (
     SEED,
     exclusive_cumsum,
@@ -45,6 +44,7 @@
     triton_bench_fn,
 )
 from cula.kda.chunk_intra import chunk_kda_fwd_intra as cula_chunk_kda_fwd_intra
+from cula.ops.kda.sm100.recompute_wu import recompute_w_u_fwd as cula_recompute_w_u_fwd
 
 # Constant params
 B, H, D = 2, 64, 128
@@ -106,17 +106,19 @@ def run_fla_recompute_wu(k, v, beta, Akk, q, gk, cu_seqlens, chunk_indices, disa
 
 
 def run_cula_recompute_wu(k, v, beta, Akk, q, gk, cu_seqlens, chunk_indices, chunk_size, disable_recompute):
-    """cuLA recompute_w_u (handles both MHA and GVA; w/u/qg/kg allocated in HV head space)."""
-    B_flat, T, HV_out, Dv = v.shape
-    w = torch.empty_like(v)
-    u = torch.empty_like(v)
-    qg = torch.empty(B_flat, T, HV_out, Dv, device=q.device, dtype=q.dtype) if disable_recompute else None
-    kg = torch.empty_like(v) if gk is not None else None
-
-    cula_cuda.recompute_w_u_cuda(
-        k, v, beta, Akk, gk, cu_seqlens, chunk_indices, w, u, kg, chunk_size, q if disable_recompute else None, qg
+    """Run the pure CuTeDSL recompute-WU implementation."""
+    if chunk_size != BT:
+        raise ValueError(f"CuTeDSL benchmark expects chunk_size={BT}, got {chunk_size}")
+    return cula_recompute_w_u_fwd(
+        k,
+        v,
+        beta,
+        Akk,
+        gk,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        q=q if disable_recompute else None,
     )
-    return w, u, qg, kg
 
 
 # ==============================================================================