TheRock/rocm-libraries/dnn-providers/hip-kernel-provider/ROCKE_LAYOUT_RESTRUCTURE.md

# rocKE Layout Restructure: Final Structure (SDPA carve)

**Branch:** `users/bharriso/rocke-layout-restructure` (off latest `develop`,
includes the merged `rocke-client` skeleton #8928).
**Scope (locked):** wrap existing `rocKE/` under a `platform/` parent (untouched
internally), and extract only the SDPA/MHA vertical into a new `library/`.
Everything non-attention stays in platform. C++ engine and byte-identity parity
stay where they are (dormant until JIT: status, not a move).

Paths under `dnn-providers/hip-kernel-provider/` (**BASE**), fixed location for now.

> **Update (AICK-1737):** we decided to release rocKE as kernel bundles under
> `hip-kernel-provider`, so the `library/api/` C++ client has been **removed**.
> This doc drops the `api/` layer; the rest of the carve is unaffected.

---

## 1. Top-level final tree

```
BASE/
├── CMakeLists.txt                      # add_subdirectory(rocke)  (was rocke-client + rocKE)
└── rocke/                              # project parent
    ├── CMakeLists.txt                  # add_subdirectory(platform); library python via rocke-pyenv/rocke-wheels
    ├── BUILDING.md                     # build/dev guide (editable installs, PYTHONPATH, rocke.assets)
    ├── platform/                       # former rocKE/, reparented verbatim (only attention-unwiring edits)
    │   ├── CMakeLists.txt              # former rocKE/CMakeLists.txt
    │   ├── pyproject.toml  requirements.txt
    │   ├── AGENTS.md  BUILD.md  README.md
    │   ├── Cpp/                        # C++ engine + bindings (dormant until JIT; unchanged)
    │   ├── Python/
    │   │   └── rocke/                  # the SDK package (import name `rocke`, installable)
    │   │       ├── __init__.py         # EDIT: drop attention re-exports only
    │   │       ├── __main__.py
    │   │       ├── core/               # ir, lower_llvm/hip/cktile, passes, ir_serialize, verify, arch, isa, backend
    │   │       ├── helpers/            # all helpers stay (generic toolkit, incl. attention-infra primitives)
    │   │       ├── runtime/            # comgr, hip_module, launcher
    │   │       ├── analysis/  benchmark/
    │   │       ├── instances/          # all kernels EXCEPT attention  (EDIT __init__s: drop attention)
    │   │       ├── dispatch/           # all families EXCEPT attention  (EDIT __init__s: drop attention)
    │   │       ├── heuristics/         # stays (1 lazy attention import repointed to library)
    │   │       ├── examples/           # all examples EXCEPT gfx*/attention
    │   │       └── sweep.py sweep_bench.py run_manifest.py torch_backend.py
    │   ├── tests/                      # C++ tests + platform pytest (minus attention python tests)
    │   ├── tools/  cmake/  dsl_docs/
    └── library/                        # the SDPA/MHA product: build-time-only Python (relies on platform `rocke.*`)
        ├── pyproject.toml              # build-time metadata (library python is NOT installed into the product)
        ├── kernels/                    # was platform instances attention (the SDPA kernel defs)
        ├── builders/                   # was platform examples/gfx*/attention (build-time drivers)
        ├── dispatch/                   # was dispatch/families/attention (1-of-N SDPA select)
        ├── tests/                      # the python attention tests
        └── benchmarks/                 # per-arch attention benchmark drivers + shape data
# removed: BASE/rocKE/ and BASE/rocke-client/  (their content now under rocke/)
# removed (AICK-1737): BASE/rocke/library/api/  and  BASE/rocke/library/CMakeLists.txt
```

## 2. platform: what it is
The entire former `rocKE/`, reparented as `rocke/platform/`, internally unchanged
except the attention-unwiring edits (§4). It remains:
- the installable **SDK** (`rocke` package at `platform/Python/rocke`, via
  `platform/pyproject.toml`), usable standalone by others to build kernels;
- the home of `core`, all `helpers` (the generic toolkit, including
  attention-infra primitives `mfma_attention`, `mfma_attention_bwd`, `qk_scale`,
  `attention`, `rotary`; reusable building blocks), `runtime`, `analysis`,
  `benchmark`, and all non-attention `instances`/`dispatch`/`examples`,
  `heuristics`, the C++ engine (`Cpp/`), tools, docs, tests.

## 3. library: the SDPA/MHA carve (relies on platform helpers)
Build-time Python (NOT installed into the product). Self-contained SDPA kernels
that import the platform SDK as `rocke.*`.

| `library/` dir | Source (from platform) | Contents |
|---|---|---|
| `kernels/` | `instances/common/` attention + arch attention | `attention_unified`, `attention_arch`, `_fmha_common`, `_fmha_warp_body`, `fmha_{appendkv,arch,bwd,fwd_fp8,head_grouping,mfma,paged_prefill,splitkv_decode,varlen}`, `sage_attention`, `sparse_attention`; `gfx1151/wmma_fmha_fwd`; `gfx1250/{_wmma_attention_common,attention_tiled_2d,attention_tiled_3d,wmma_attention_fwd}`; `gfx942/{attention_tiled_2d,attention_tiled_3d}`; `gfx950/{attention_tiled_2d,attention_tiled_2d_fastkv_regp,attention_tiled_3d}`; plus the inline arch-router builders `build_unified_attention_2d_tiled` / `build_unified_attention_3d_tiled` lifted out of `instances/__init__.py` |
| `builders/` | `examples/gfx{942,950,1151,1250}/attention/`, `examples/common/fmha_fwd_verify_hip.py`, plus the attention halves split out of mixed `examples/common` drivers (`parity_fmha_extended` from `parity_extended_kernels`, `hip_lowering_attention_parity` from `hip_lowering_parity`), the `dsl_probe`/`gen_sweep`/`stage1_benchmark` attention slices (`dsl_probe_attention_demos`, `gen_sdpa_sweep_data`, `benchmark_rocke_unified_attention`), and `gfx950/qwen3_30b_a3b/` decode drivers | the SDPA verify/parity/bench drivers (build-time) |
| `dispatch/` | `dispatch/families/attention.py` | `ATTENTION_REGISTRY`, `AttentionRequest`, `dispatch_attention`: the 1-of-N SDPA select |
| `tests/` | `tests/instances/test_gfx1250_attention.py` (+ any other python attention tests) | python attention tests (C++ `tiled_attention_2d_reentrancy.cpp` stays in platform; it tests the C++ engine) |

**Naming:** `kernels` (was *instances*) and `builders` (was *examples*) are
product nouns; platform keeps its DSL terms `instances`/`examples`.

## 4. The only platform edits (attention-unwiring)
Removing attention from the `instances` package requires editing its re-export
sites (all mechanical "drop/relocate attention"):
- `instances/__init__.py`: strip the attention re-export block (`UnifiedAttention*`, `build_unified_attention_*`) and **move** the inline `build_unified_attention_2d_tiled`/`_3d_tiled` arch-routers to `library/kernels/`.
- `instances/gfx942/__init__.py`, `instances/gfx950/__init__.py`, `instances/gfx1250/__init__.py`: drop attention re-exports.
- `dispatch/__init__.py`, `dispatch/families/__init__.py`: drop `dispatch_attention` / `ATTENTION_REGISTRY`.
- `heuristics/gen_sweep_data.py`: drop the attention sweep-generation path (moved to `library/builders/common/gen_sdpa_sweep_data.py`); the platform copy only *names* that tool in its help text, so it holds no `import` of the library and platform stays fully standalone.
- `examples/run_all.py` REGISTRY + `examples/_goldens/fmha_fwd_hip_mha.json`: remove the `fmha_fwd_hip_mha` (`family="attention"`) entry; its driver `examples/common/fmha_fwd_verify_hip.py` moves to `library/builders/`. Audit `_goldens/` for any other attention entries.
- `rocke/__init__.py`: drop any attention re-export it carries.
- `helpers/__init__.py`: untouched (all helpers stay).
Non-attention platform code is otherwise unchanged.

## 5. Imports
- Platform package name `rocke` is preserved, so non-attention code needs no
  import changes.
- **Library** python is build-time; its source root is **`rocke/library/`**
  (NOT `rocke/`, which would put `platform/` on `sys.path` and shadow the stdlib
  `platform` module). So `kernels`, `builders`, `dispatch` are top-level packages
  under that root (not a `library.*` package). Caveat: these are generic
  top-level names, acceptable for a build-time-only, controlled-PYTHONPATH
  context; wrap them in one named package later if needed.
- Modules import the SDK as `rocke.*` (e.g. `from rocke.helpers.mfma_attention
  import ...`, `from rocke.core.ir import ...`, `from rocke.dispatch.core import ...`).
- **Codemod for the moved files:** relative imports that pointed at platform
  (`from ..core`, `from ...helpers`, dispatch's `from ..core` to `rocke.dispatch.core`,
  etc.) become absolute `from rocke.<...>`; refs to the moved kernels become `from
  kernels.<...>` (e.g. `dispatch/attention` becomes `from kernels.common.attention_unified`);
  relative imports within a moved subpackage stay relative.
- **`rocke.examples.*.attention*` and `rocke.examples.common.fmha_fwd_verify_hip`
  become `builders.<...>`** (statements, string literals, `_goldens` JSON `module`,
  REGISTRY rows, `-m` docs), full-occurrence discipline, scoped to attention only.
- Build-time PYTHONPATH (conftest + a `cmake -E env` helper): `platform/Python`
  (for `rocke`) + `rocke/library/` (for `kernels`/`builders`/`dispatch`).
- **Editable-install migration:** per-script `sys.path`/`parents[N]`/`Path(__file__)`
  hacks are replaced by editable installs plus a `rocke.assets` accessor module
  (`platform_root`/`dsl_docs_dir`/`shape_utils_dir`, env overrides
  `ROCKE_PLATFORM_ROOT`/`ROCKE_DSL_DOCS`); the only residual path handling lives
  in the test `conftest.py` bootstraps and `assets.py` itself.
- One-way rule: `library` imports `platform` only; platform never imports
  `library` (verified: zero `kernels`/`builders`/`dispatch` imports anywhere
  under `platform/`, static and at runtime).

## 6. CMake
- `BASE/CMakeLists.txt`: `add_subdirectory(rocke)` replaces the `rocke-client` +
  `rocKE` adds; keep the `ROCKE_INSTALL_*` staging vars set before it and
  `add_subdirectory(src)` after (CTest GLOBAL-staging order preserved).
- `rocke/CMakeLists.txt`: `add_subdirectory(platform)`. The library python
  (`kernels`/`builders`/`dispatch`) is not a CMake subdirectory; it is
  editable-installed and wheeled by the `rocke-pyenv` / `rocke-wheels` blocks in
  this file, which reference `library/pyproject.toml` by path.
- `rocke/platform/CMakeLists.txt` = former `rocKE/CMakeLists.txt`, unchanged
  (still installs the `rocke` SDK package, now without attention, builds C++
  tests, stages CTest).
- ~~`rocke/library/CMakeLists.txt` + `api/CMakeLists.txt`~~: removed (AICK-1737)
  along with the `library/api/` client. Library is now Python-only.

## 7. Tests
- **platform/tests/**: C++ tests (incl. `tiled_attention_2d_reentrancy.cpp`),
  non-attention python tests; re-rooted `conftest.py` (platform root).
- **library/tests/**: the python attention tests (`test_gfx1250_attention.py`
  + any others importing the moved kernels/builders) with a conftest inserting
  `platform/Python` + `rocke/`.

## 8. Verification
- **Platform standalone:** clean venv, `pip install rocke/platform`, then
  `python -c "import rocke, rocke.core, rocke.helpers, rocke.instances, rocke.dispatch; from rocke import lower_kernel_to_llvm"` succeeds without library present. (Proves attention removal is clean and platform stays self-contained.)
- **Platform has no attention *kernels*:** `instances`/`dispatch` carry no
  attention modules and `dispatch.dispatch_attention` is no longer importable,
  but platform retains the attention-infra *helpers* (`helpers/attention`,
  `mfma_attention`, `qk_scale`, ...) and `instances/gfx1250/qwen3_kv_cache` by design.
- **Library on platform:** with both roots on the path (`platform/Python` +
  `rocke/library/`),
  `python -c "import kernels.common.attention_unified, dispatch, builders.gfx1250.attention..."` and `dispatch.attention.dispatch_attention` resolve.
- **Functional wiring (not just imports):** after the move, actually *run* the
  SDPA path end-to-end: invoke `dispatch.attention.dispatch_attention` to select
  and build an attention kernel, and run one `builders` driver, to prove the
  moved kernels+dispatch+builders are wired and produce output, not merely import.
- **Zero-leftover:** `search 'rocke\.examples.*attention'` and stale
  `rocke.instances.common.attention*` references resolve to none.
- **pytest parity:** identical pass/skip/fail for the moved attention tests
  (now under library) and the remaining platform tests, pre vs post.
- **CMake both modes** (provider opt-in + default OFF) configure/build/install;
  installed CTest lists the platform engine tests staged into the provider bucket.
- **Move integrity:** `git diff -M` shows renames; `pre-commit` clean.

## 9. Notes / follow-ups (out of scope here)
- **AOT** pre-built kernel list (Python builders to platform lowering to
  artifacts) is a follow-up, now tracked under AICK-1704 (migration into
  `hip-kernel-provider`, off POC / PR #9533).
- Optional later: rename `platform/Python` to `platform/python`, `Cpp` to `cpp`.
```