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

[cherhh]

📌 Description

Reorganization

Restructure the flat cula/ops/*.py layout into an arch-first hierarchy:

Before	After
cula/ops/kda_decode.py	cula/ops/kda/decode/cute.py
cula/ops/fwd_o_sm100.py	cula/ops/kda/sm100/fwd_o.py
cula/ops/sm100/chunk_delta_h.py	cula/ops/kda/sm100/delta_h.py
cula/ops/sm100/bwd_wy_dqkg.py	cula/ops/kda/sm100/bwd_wy_dqkg.py
cula/ops/sm100/cp/*	cula/ops/kda/sm100/cp/*
cula/ops/intrinsics_sm100.py	cula/ops/sm100/ptx.py
cula/ops/ptx_umma_ext.py	(consolidated into cula/ops/sm100/ptx.py)
cula/ops/kda_fully_fused_sm100_wip.py	cula/ops/kda/experimental/sm100_fused/
cula/ops/la_decode.py	cula/ops/lightning/decode.py
cula/ops/prefill_sm100.py	cula/ops/lightning/prefill_sm100.py

New cula/ops/ layout: see updated REPO_LAYOUT.md.

Lazy imports (PEP 562)

Add __getattr__ in cula/__init__.py, cula/kda/__init__.py, and cula/ops/__init__.py so that importing the top-level package does not eagerly import CuTeDSL/CUDA kernel modules.

Intracard-fwd_h cleanup

Separate policy from execution in the SM100 intra-card CP path:

• Policy decisions are now handled entirely by sm100_intracard_cp_decision in policy.py
• intracard_fwd_h is a pure executor: split and run, or raise NotSplittableError
• The caller owns the fallback logic based on the policy result (forced CP → re-raise, auto → fall through to serial)

What is NOT changed

• SM90 prefill kernel path: still uses the existing C++ kernel under csrc/kda/sm90/
• SM100 kernel behavior: unchanged (pure move/rename)

🔍 Related Issues

N/A

🧪 Tests

• Pre-commit (clang-format + ruff + ruff-format) — all passing

Reviewer Notes

[gemini-code-assist]

Code Review

This pull request reorganizes the repository layout by migrating KDA backend kernels to an arch-first structure (e.g., grouping SM100 Blackwell kernels, decode, and experimental paths) and refactoring shared PTX/MLIR helpers. It also introduces a centralized context-parallel dispatch policy in policy.py for SM100. Regarding the changes, a performance improvement opportunity was identified in chunk_gated_delta_rule_fwd_h where cu_seqlens_cpu can be materialized twice when it is not provided, leading to redundant device-to-host synchronizations. Materializing it once at the beginning of the function would resolve this.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[fkuner]

The new cula/ops layout is a good step away from the old flat structure, but it still mixes multiple organization axes: algorithm (kda/, lightning/), architecture (sm100/), and maturity (experimental/). For example, ops/kda/sm100/ is arch-first under KDA, while ops/kda/decode/ is function-first, and shared SM100 helpers live at ops/sm100/. This makes it unclear where future SM90/SM100 decode, prefill, or CP code should go, so I think we should discuss and settle on one placement rule before merging.

I suggest making ops an internal backend namespace with a simpler rule: organize by algorithm first, and only encode architecture where a specific stage needs it. For the current code size, a pragmatic layout could be:

cula/ops/
├── common/
│   ├── inv.py
│   ├── ptx.py
│   └── ptx_sm100.py
├── kda/
│   ├── chunk_sm100/
│   │   ├── delta_h.py
│   │   ├── fwd_o.py
│   │   ├── bwd_wy_dqkg.py
│   │   └── cp/
│   │       ├── policy.py
│   │       ├── chunk_delta_h.py
│   │       ├── pre_scan.py
│   │       └── merge.py
│   ├── decode/
│   │   ├── cute.py
│   │   └── reference_fla.py
│   └── prefill/
│       ├── hopper.py
│       └── fused_wip_sm100/
│           ├── kernel.py
│           └── wrapper.py
├── lightning/
│   ├── prefill_sm100.py
│   └── decode.py
└── experimental/
    └── linear_attn_prototype.py

This keeps the hierarchy shallow while giving contributors a clear placement rule: KDA backends under ops/kda, Lightning backends under ops/lightning, cross-algorithm low-level helpers under ops/common, and stage-specific policy code next to the backend it controls, e.g. chunk_sm100/cp/policy.py rather than ops/kda/policy.py.

This is not a blocker for the refactor itself, but I think it is worth discussing so we can agree on a clear convention before this layout becomes the pattern for future backends.

[fkuner]

Nit: git diff --check reports a blank line at EOF in REPO_LAYOUT.md and trailing whitespace in docs/chunk_delta_h_pipeline.md. Not functionally important, but worth cleaning up to keep whitespace checks clean.

REPO_LAYOUT.md:72: new blank line at EOF.
docs/chunk_delta_h_pipeline.md:3: trailing whitespace.
+> 文件: cula/ops/kda/sm100/delta_h.py

[cherhh]

Nit: git diff --check reports a blank line at EOF in REPO_LAYOUT.md and trailing whitespace in docs/chunk_delta_h_pipeline.md. Not functionally important, but worth cleaning up to keep whitespace checks clean.

REPO_LAYOUT.md:72: new blank line at EOF. docs/chunk_delta_h_pipeline.md:3: trailing whitespace. +> 文件: cula/ops/kda/sm100/delta_h.py

Fixed.

[cherhh]

The new cula/ops layout is a good step away from the old flat structure, but it still mixes multiple organization axes: algorithm (kda/, lightning/), architecture (sm100/), and maturity (experimental/). For example, ops/kda/sm100/ is arch-first under KDA, while ops/kda/decode/ is function-first, and shared SM100 helpers live at ops/sm100/. This makes it unclear where future SM90/SM100 decode, prefill, or CP code should go, so I think we should discuss and settle on one placement rule before merging.

I suggest making ops an internal backend namespace with a simpler rule: organize by algorithm first, and only encode architecture where a specific stage needs it. For the current code size, a pragmatic layout could be:

cula/ops/
├── common/
│   ├── inv.py
│   ├── ptx.py
│   └── ptx_sm100.py
├── kda/
│   ├── chunk_sm100/
│   │   ├── delta_h.py
│   │   ├── fwd_o.py
│   │   ├── bwd_wy_dqkg.py
│   │   └── cp/
│   │       ├── policy.py
│   │       ├── chunk_delta_h.py
│   │       ├── pre_scan.py
│   │       └── merge.py
│   ├── decode/
│   │   ├── cute.py
│   │   └── reference_fla.py
│   └── prefill/
│       ├── hopper.py
│       └── fused_wip_sm100/
│           ├── kernel.py
│           └── wrapper.py
├── lightning/
│   ├── prefill_sm100.py
│   └── decode.py
└── experimental/
    └── linear_attn_prototype.py

This keeps the hierarchy shallow while giving contributors a clear placement rule: KDA backends under ops/kda, Lightning backends under ops/lightning, cross-algorithm low-level helpers under ops/common, and stage-specific policy code next to the backend it controls, e.g. chunk_sm100/cp/policy.py rather than ops/kda/policy.py.

This is not a blocker for the refactor itself, but I think it is worth discussing so we can agree on a clear convention before this layout becomes the pattern for future backends.

Agreed the layout needs one clear placement rule. Since this isn't a blocker: I'd propose we merge this PR and #103 first, then do the layout move in a dedicated follow-up PR. Moving files now would force a large rebase of #103, while a follow-up can apply the agreed convention to everything (including the new SM90 code) in one shot. I'll open that PR and we can settle the exact tree there(maybe similar to what is shown below).

  cula/ops/
  ├── common/                  # inv.py, ptx.py, ptx_sm100.py
  ├── kda/
  │   ├── cp_mode.py           # CP vocabulary shared by sm90 + sm100
  │   ├── chunk/
  │   │   ├── sm90/            # fwd, k1, k2 + cp/{plan,driver,pre_scan,merge}
  │   │   ├── sm100/           # policy, delta_h, fwd_o, bwd_wy_dqkg + cp/...
  │   │   └── sm100_fused_wip/ # experimental fused backend
  │   └── decode/
  ├── lightning/
  └── experimental/            # cross-algorithm prototypes only

[fkuner]

/LGTM

[zheyang0825]

LGTM