refactor(kda): reorganize KDA backends into arch-first layout and add lazy imports

REPO_LAYOUT.md

@@ -1,101 +1,71 @@
 # Repository Layout
 
+Legend: `[exp]` experimental / unwired · `[non-KDA]` other operator.
+
 ```
 cuLA/
 ├── cula/                         # Python package (pip install -e .)
-│   ├── kda/                      # KDA (Kimi Delta Attention) operators
-│   │   ├── chunk.py              # End-to-end chunk KDA (fwd + bwd entry point)
-│   │   ├── chunk_fwd.py          # Chunk forward dispatch
-│   │   ├── chunk_intra.py        # Intra-chunk forward logic
-│   │   ├── blackwell_fused_fwd.py  # Fused KDA forward (SM100)
-│   │   └── hopper_fused_fwd.py     # Fused KDA forward (SM90)
-│   ├── lightning/                # Lightning Attention operators
-│   │   └── la_decode.py          # Single-token decode kernel (CuTe DSL)
-│   ├── ops/                      # CuTe DSL kernel implementations
-│   │   ├── chunk_delta_h_sm100.py      # Chunk delta-H kernel (SM100)
-│   │   ├── fwd_o_sm100.py              # Forward output kernel (SM100)
-│   │   ├── lightning_attn_sm100.py     # Lightning Attention prefill kernel (SM100)
-│   │   ├── linear_attn_sm100.py        # Generic linear attention kernel (SM100)
-│   │   ├── kda_fully_fused_sm100_wip.py # WIP fully fused KDA kernel (SM100)
-│   └── utils.py                  # Shared utilities
-│
-├── csrc/                         # CUDA C++ / CUTLASS kernels
-│   ├── api/                      # PyBind11 bindings
-│   │   ├── pybind.cu             # Python ↔ CUDA binding entry
-│   │   ├── kda_sm90.cu           # SM90 API wrappers
-│   │   └── kda_sm100.cu          # SM100 API wrappers
-│   ├── kda/
-│   │   ├── sm90/                 # Hopper KDA kernels (CUTLASS 3.x)
-│   │   │   ├── kda_fwd_sm90.cu
-│   │   │   ├── kda_fwd_sm90_safe_gate.cu
-│   │   │   ├── prefill_kernel.hpp
-│   │   │   ├── collective/       # CUTLASS collective mainloop
-│   │   │   ├── device/           # Device-level kernel wrappers
-│   │   │   ├── kernel/           # Kernel-level logic
-│   │   │   └── utils/            # SM90-specific helpers
-│   │   └── sm100/                # Blackwell KDA kernels (CUTLASS 3.x)
-│   │       ├── kda_fwd_sm100.cu
-│   │       ├── kda_fwd_common.cuh
-│   │       ├── kda_fwd_intra_kernel_sm100.hpp
-│   │       ├── kda_fwd_intra_mainloop_sm100.hpp # Chunk intra mainloop
-│   │       ├── kda_fwd_recomp_w_u_kernel_sm100.hpp 
-│   │       ├── kda_fwd_recomp_w_u_mainloop_sm100.hpp # Recompute W&U mainloop
-│   │       ├── kda_config.hpp
-│   │       ├── fwd_helpers.hpp
-│   │       ├── sm100_umma_ext.hpp
-│   │       └── tile_scheduler.hpp
-│   └── kerutils/
-│       └── include/              # Shared C++ header utilities
-│
-├── benchmarks/                   # Performance benchmarks
-│   ├── bench_kda.py              # KDA fixed + varlen benchmark
-│   ├── bench_lightning_attn.py   # Lightning Attention prefill + varlen
-│   ├── bench_la_decode_vs_fla.py # Decode: la_decode vs fla fused_recurrent
-│   ├── bench_kda_fused_fwd.py    # KDA fused forward benchmark
-│   ├── bench_kda_chunk_intra.py  # KDA chunk intra benchmark
-│   ├── bench_chunk_delta_h.py    # Chunk delta-H benchmark
-│   ├── bench_fwd_o.py            # Forward output benchmark
-│   ├── bench_linear_attn.py      # Linear attention benchmark
-│   ├── generate_benchmark_md.py  # Auto-generate BENCHMARK_GB200.md (Blackwell)
-│   ├── generate_benchmark_hopper_md.py  # Auto-generate BENCHMARK_H200.md (Hopper)
-│   └── utils.py                  # Benchmark utilities
-│
-├── tests/                        # Unit / integration tests
-│   ├── test_kda_compare_fla.py   # Modular KDA forward vs FLA Triton
-│   ├── test_kda.py               # Modular KDA forward vs naive reference
-│   ├── test_kda_fused_fwd.py     # Fused KDA forward tests
-│   ├── test_chunk_delta_h.py     # Chunk delta-H tests
-│   ├── test_fwd_o.py             # Forward output tests
-│   ├── test_compare_with_fla.py  # General FLA comparison
-│   ├── test_lightning_attn.py    # Lightning Attention tests
-│   └── test_la_decode.py         # Decode kernel tests
-│
-├── docs/                         # Design documents
-│   ├── chunk_delta_h_pipeline.md
-│   ├── fwd_o_pipeline.md
-│   └── lightning_attn_pipeline.md
+│   ├── __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,
+│   │   │                         #                  kda_decode, fused_sigmoid_gating_delta_rule_update
+│   │   ├── chunk.py              # chunk_kda + autograd — SM100 modular path (train + Blackwell prefill)
+│   │   ├── 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)
+│   │
+│   ├── lightning/                # [non-KDA] Lightning Attention operator (LinearAttentionChunkwiseDecay, lightning_attn_fwd, linear_attention_decode)
+│   │   └── __init__.py
+│   │
+│   └── ops/                      # backend kernels (CuTe DSL / TVM-FFI) + shared helpers
+│       ├── __init__.py           # exports kda_decode, fused_sigmoid_..., linear_attention_decode
+│       ├── inv.py / ptx.py       # shared low-level helpers
+│       ├── sm100/                # SM100 shared helper only
+│       │   └── 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
+│       │   ├── 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
+│       │   ├── decode/           #   single-token decode
+│       │   │   ├── cute.py       #     kda_decode / fused_sigmoid_gating_delta_rule_update (CuTe DSL)
+│       │   │   └── reference_fla.py
+│       │   └── experimental/sm100_fused/   # [exp] unwired fully-fused
+│       │       ├── kda_fully_fused_wip.py   #   KDAChunkwise (~6k lines)
+│       │       └── wrapper.py                #   flash_kda_prefill (dead path; raises on SM100 dispatch)
+│       │
+│       ├── lightning/            # [non-KDA] Lightning/linear attention kernels
+│       │   ├── prefill_sm100.py  #   Lightning Attn prefill (LinearAttentionChunkwiseDecay, lightning_attn_fwd[_varlen])
+│       │   └── decode.py         #   linear_attention_decode
+│       └── experimental/
+│           └── linear_attn_prototype.py     # [non-KDA] unwired normalized-linear-attn prototype
 │
-├── third_party/
-│   └── flash-linear-attention/   # FLA submodule (baseline)
+├── 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)
+│   └── kerutils/include/         # shared C++ headers (generic device helpers sm80/sm90/sm100, host)
 │
-├── BENCHMARK_GB200.md            # Auto-generated Blackwell benchmark results
-├── BENCHMARK_H200.md             # Auto-generated Hopper benchmark results
-├── README.md                     # Project overview
-├── setup.py                      # Build configuration
-├── pyproject.toml                # Project metadata
-└── LICENSE
+├── benchmarks/  tests/  docs/
+├── scripts/build_wheel.sh
+├── third_party/flash-linear-attention/      # FLA submodule (baseline + reused gate/CP ops)
+├── README.md  USAGE.md  REPO_LAYOUT.md  RECOMMENDED_CODING_STYLE.md
+└── setup.py  pyproject.toml  LICENSE
 ```
 
 ## Key Directories
 
 | Directory | Language | Description |
 |-----------|----------|-------------|
-| `cula/ops/` | Python (CuTe DSL) | Warp-specialized GPU kernels written in CuTe DSL — compiled to CUDA at import time |
-| `cula/kda/` | Python | KDA operator dispatch — selects SM90 or SM100 path, handles chunking and autograd |
-| `cula/lightning/` | Python (CuTe DSL) | Lightning Attention decode kernel |
-| `csrc/kda/sm90/` | CUDA C++ | Hopper KDA kernels using CUTLASS 3.x collective API |
-| `csrc/kda/sm100/` | CUDA C++ | Blackwell KDA kernels using CUTLASS 3.x + UMMA extensions |
-| `csrc/api/` | CUDA C++ | PyBind11 entry points exposing C++ kernels to Python |
-| `benchmarks/` | Python | Performance benchmarks vs FLA Triton baselines |
-| `tests/` | Python | Correctness tests (pytest) |
-| `docs/` | Markdown | Internal pipeline design notes |
+| `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/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

@@ -137,7 +137,7 @@ import os
 os.environ["CULA_INTRACARD_CP"] = "1"
 
 import torch
-from cula.ops.chunk_delta_h import chunk_gated_delta_rule_fwd_h
+from cula.ops.kda.sm100.delta_h import chunk_gated_delta_rule_fwd_h
 
 B, T, H, K, V = 1, 65536, 8, 128, 128
 device = 'cuda'

benchmarks/bench_chunk_delta_h.py

@@ -50,7 +50,7 @@
 from benchmarks.utils import benchmark_cuda_mode_fn, relative_rms_error_max_mean_abs
 
 # ─── CuTe DSL wrapper (TVM-FFI compile cache) ───
-_delta_h_mod = importlib.import_module("cula.ops.chunk_delta_h_sm100")
+_delta_h_mod = importlib.import_module("cula.ops.kda.sm100.delta_h")
 chunk_gated_delta_rule_fwd_h = _delta_h_mod.chunk_gated_delta_rule_fwd_h
 
 # ─── FLA baseline imports ───

benchmarks/bench_fwd_o.py

@@ -47,7 +47,7 @@
 sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
 
 # ─── CuTe DSL wrapper (TVM-FFI compile cache) ───
-_fwd_o_mod = importlib.import_module("cula.ops.fwd_o_sm100")
+_fwd_o_mod = importlib.import_module("cula.ops.kda.sm100.fwd_o")
 chunk_gla_fwd_o = _fwd_o_mod.chunk_gla_fwd_o
 build_chunk_indices = _fwd_o_mod.build_chunk_indices
 

benchmarks/bench_intracard_cp.py

@@ -47,7 +47,7 @@
     set_seed,
 )
 from cula.kda.chunk_fwd import chunk_kda_fwd
-from cula.ops.cp.chunk_delta_h import (
+from cula.ops.kda.sm100.cp.chunk_delta_h import (
     compute_subseq_len,
     prepare_subseq_cu_seqlens,
     should_use_intracard_cp,

benchmarks/bench_kda_bwd_wy_dqkg_sm100.py

@@ -50,7 +50,7 @@
     relative_rms_error_rel_max_mean_abs,
     set_seed,
 )
-from cula.ops.chunk_wy_dqkg_sm100 import chunk_kda_bwd_wy_dqkg_fused as cula_chunk_kda_bwd_wy_dqkg_fused
+from cula.ops.kda.sm100.bwd_wy_dqkg import chunk_kda_bwd_wy_dqkg_fused as cula_chunk_kda_bwd_wy_dqkg_fused
 
 torch.backends.cuda.matmul.allow_tf32 = True
 

benchmarks/bench_kda_decode.py

@@ -55,7 +55,7 @@
 
 from benchmarks.utils import benchmark_cuda_fn, relative_rms_error_rel_max
 from cula.kda import fused_sigmoid_gating_delta_rule_update as cula_fused
-from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as fla_fused
+from cula.ops.kda.decode.reference_fla import fused_sigmoid_gating_delta_rule_update as fla_fused
 
 
 # ──────────────────────────────────────────────────────────────────────

benchmarks/bench_kda_fused_fwd.py

@@ -18,7 +18,7 @@
 
 Automatically selects the cuLA fully-fused implementation based on the current
 GPU architecture:
-    - sm100 (Blackwell) → cula.kda.blackwell_fused_fwd.flash_kda_prefill
+    - sm100 (Blackwell) → cula.ops.kda.experimental.sm100_fused.wrapper.flash_kda_prefill
   - sm90  (Hopper)    → cula.kda.hopper_fused_fwd.cula_kda_prefill
 
 Compares:

benchmarks/bench_la_decode_vs_fla.py

@@ -56,7 +56,7 @@
 from fla.ops.common.fused_recurrent import fused_recurrent_fwd, fused_recurrent_fwd_kernel
 
 from benchmarks.utils import benchmark_cuda_fn, relative_rms_error
-from cula.ops.la_decode import _get_compiled_kernel, linear_attention_decode
+from cula.ops.lightning.decode import _get_compiled_kernel, linear_attention_decode
 from cula.utils import USE_FAST_MATH
 
 

benchmarks/bench_lightning_attn.py

@@ -53,7 +53,7 @@
 from fla.ops.simple_gla.chunk import chunk_simple_gla_fwd
 
 from benchmarks.utils import gen_random, gen_skewed, gen_uniform, relative_rms_error, time_cuda_fn
-from cula.ops.lightning_attn_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen
+from cula.ops.lightning.prefill_sm100 import lightning_attn_fwd, lightning_attn_fwd_varlen
 
 # =============================================================================
 # Constants

benchmarks/bench_linear_attn.py

@@ -25,7 +25,7 @@
 # from fla.ops.linear_attn.naive import naive_recurrent_linear_attn
 from fla.utils import assert_close, device
 
-from cula.ops.linear_attn_sm100 import LinearAttentionChunkwise
+from cula.ops.experimental.linear_attn_prototype import LinearAttentionChunkwise
 
 os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1"))  # Enable fast ops in FLA for fair comparison
 

cula/__init__.py

@@ -16,9 +16,3 @@
     from cula._version import version as __version__
 except ImportError:
     __version__ = "0.1.0"
-
-from cula.ops.lightning_attn_sm100 import LinearAttentionChunkwiseDecay
-
-__all__ = [
-    "LinearAttentionChunkwiseDecay",
-]

cula/kda/__init__.py

@@ -12,15 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cula.kda.blackwell_fused_fwd import flash_kda_prefill as kda_prefill_blackwell
-from cula.kda.chunk import chunk_kda
-from cula.kda.hopper_fused_fwd import cula_kda_prefill as kda_prefill_hopper
-from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode
+"""Public KDA API exports for chunk, prefill, and decode"""
 
 __all__ = [
     "chunk_kda",
-    "kda_prefill_blackwell",
     "kda_decode",
     "fused_sigmoid_gating_delta_rule_update",
     "kda_prefill_hopper",
 ]
+
+_LAZY = {
+    "chunk_kda": ("cula.kda.chunk", "chunk_kda"),
+    "kda_prefill_hopper": ("cula.kda.hopper_fused_fwd", "cula_kda_prefill"),
+    "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"),
+    "fused_sigmoid_gating_delta_rule_update": (
+        "cula.ops.kda.decode.cute",
+        "fused_sigmoid_gating_delta_rule_update",
+    ),
+}
+
+
+def __getattr__(name):
+    target = _LAZY.get(name)
+    if target is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    import importlib
+
+    return getattr(importlib.import_module(target[0]), target[1])
+
+
+def __dir__():
+    return sorted(__all__)

cula/kda/chunk.py

@@ -15,6 +15,10 @@
 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
 # Related files are modified and supported by the Moonshot AI Team
 
+"""SM100 modular chunk KDA public API and autograd wrapper"""
+
+from typing import Literal
+
 import torch
 from fla.modules.l2norm import l2norm_bwd, l2norm_fwd
 from fla.ops.cp import FLACPContext
@@ -23,6 +27,7 @@
 
 from cula.kda.chunk_bwd import chunk_kda_bwd
 from cula.kda.chunk_fwd import chunk_kda_fwd
+from cula.ops.kda.policy import IntracardCPMode, resolve_intracard_cp_mode
 
 
 class ChunkKDAFunction(torch.autograd.Function):
@@ -50,6 +55,7 @@ def forward(
         disable_recompute: bool = False,
         return_intermediate_states: bool = False,
         cp_context: FLACPContext | None = None,
+        use_intracard_cp: IntracardCPMode | None = None,
     ):
         chunk_size = 64
 
@@ -85,6 +91,7 @@ def forward(
             disable_recompute=disable_recompute,
             return_intermediate_states=return_intermediate_states,
             cp_context=cp_context,
+            use_intracard_cp=use_intracard_cp,
         )
 
         if return_intermediate_states:
@@ -211,6 +218,7 @@ def backward(
             None,
             None,
             None,
+            None,
         )
 
 
@@ -233,6 +241,7 @@ def chunk_kda(
     disable_recompute: bool = False,
     return_intermediate_states: bool = False,
     cp_context: FLACPContext = None,
+    use_intracard_cp: Literal["auto"] | bool | None = None,
     **kwargs,
 ):
     r"""
@@ -345,7 +354,15 @@ def chunk_kda(
         )
     """
 
+    # just for backward compatibility, resolve the deprecated `use_cp` argument
+    # TODO: maybe we can remove this in the future
+    use_cp_alias = kwargs.pop("use_cp", None)
+    use_intracard_cp = resolve_intracard_cp_mode(use_intracard_cp, use_cp_alias)
+
     if cp_context is not None:
+        if use_intracard_cp is True:
+            raise ValueError("use_intracard_cp=True cannot be combined with FLA cp_context.")
+        use_intracard_cp = False
         assert initial_state is None, "Initial state is not supported for CP"
         assert output_final_state is False, "Output final state is not supported for CP"
         assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP"
@@ -354,6 +371,11 @@ def chunk_kda(
         if cp_context.cu_seqlens_cpu is not None:
             cu_seqlens_cpu = cp_context.cu_seqlens_cpu
 
+    if return_intermediate_states:
+        if use_intracard_cp is True:
+            raise ValueError("use_intracard_cp=True is not supported with return_intermediate_states=True.")
+        use_intracard_cp = False
+
     if cu_seqlens is not None:
         if q.shape[0] != 1:
             raise ValueError(
@@ -414,4 +436,5 @@ def chunk_kda(
         disable_recompute,
         return_intermediate_states,
         cp_context,
+        use_intracard_cp,
     )