refactor(mlx-cpp): organize mlxcel's MLX-side layer as mlxcelverse (backend overlays and mlxcel kernels)
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Related to #1801. Not a sub-issue of it: this touches the shipping Metal and CUDA paths, and the epic is scoped to ROCm. Start after #1802 merges so the reorganization moves all three backends at once.
## Context
**mlxcelverse** is the name for everything mlxcel builds on top of upstream MLX (defined in #1801). It already exists in practice, spread over two places and two kinds of code:
1. **Per-backend source overlays** that replace or add MLX files before MLX is compiled:
- `src/lib/mlx-cpp/patches/mlx/backend/metal/` (3 files) and `src/lib/mlx-cpp/patches/mlx/backend/cuda/**` (22 files), always copied;
- `src/lib/mlx-cpp/patches-cuda/{dtype,fast,ops}.cpp`, copied into `mlx/` only for CUDA builds;
- `src/lib/mlx-cpp/patches-rocm/**` from #1802, copied only for ROCm builds.
These follow the MLX pin: every bump needs a 3-way merge per file and a line-by-line comparison with upstream. Provenance is Apple MIT, plus the MIT NripeshN/mlx fork for the ROCm files.
2. **mlxcel's own kernels**, fusions and extension functions built on MLX's public custom-kernel APIs (`fast::metal_kernel`, `fast::cuda_kernel`, and `fast::hip_kernel` on ROCm): `src/lib/mlx-cpp/turbo/` (20 files: paged attention, sampling, fused norm, fused RoPE+append, TurboQuant SDPA and others) and the kernel bodies in `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`. These only need public-API compatibility across bumps and are Lablup code (some with third-party notices in `NOTICE`).
Because the two kinds sit side by side in `src/lib/mlx-cpp/` without a stated boundary, the maintenance rules (merge-and-review versus API compatibility) and the provenance rules (never stamp a Lablup header on upstream-derived files) are enforced by convention and by comments in `.github/workflows/ci.yml` rather than by layout.
## Scope
Rename and reorganize so the layout states the model, without changing any build output. Proposed shape (final names decided in the PR):
```
src/lib/mlxcelverse/
README.md what mlxcelverse is, the two kinds, the rules for each
CMakeLists.txt today's src/lib/mlx-cpp/CMakeLists.txt
overlays/
common/ files every GPU build copies today (if any remain after the split)
metal/ today's patches/mlx/backend/metal/**
cuda/ today's patches/mlx/backend/cuda/** and patches-cuda/**
rocm/ today's patches-rocm/** (with UPSTREAM, LOCAL_FIXES.md)
kernels/ today's turbo/**, and a home for ROCm ports from #1814
```
Whether the kernel bodies in `mlx_cxx_kernels.cpp` move into `kernels/` is decided in the PR; moving them is optional.
## Implementation plan
1. Move the directories with `git mv` so history follows. Update `mlx_apply_source_overlays` so each backend copies exactly the files it copies today (the always-copied Metal and CUDA backend overlays keep being copied for every build unless the PR shows a narrower rule is equivalent).
2. Update every path reference: `src/lib/mlxcel-core/build.rs` (`build_mlx` source dir, `rerun-if-changed` lines 240-260), `.github/workflows/ci.yml` path filters and the license-header gate comments, `scripts/ci/check_kernel_dtype_keys.py`, `scripts/insert_apache_header.py` skip rules if any, `CONTRIBUTING.md` (pin-bump procedure), `docs/architecture.md`, `docs/environment-variables.md`, and Rust sources that cite the paths in comments or tests (for example `src/lib/mlxcel-core/src/layers.rs`, `src/lib/mlxcel-core/src/cache/turbo/sparse_v.rs`, `src/lib/mlxcel-core/src/paged_v2/mod.rs`, `src/lib/mlxcel-core/src/autotune/ops/*.rs`). Historical benchmark pages and ADRs keep their old paths as written, with no rewrite of history.
3. Write `src/lib/mlxcelverse/README.md` and a short provenance note per overlay directory.
4. Keep the build cache working: the moved CMake source directory must not trigger an MLX refetch loop, and `purge_stale_mlx_cache` behavior must stay the same.
## Acceptance criteria
- [ ] For Metal, CUDA and ROCm builds, the set of files copied into `_deps/mlx-src` is identical before and after (compare the configure logs or a checksum listing of the patched tree).
- [ ] Greedy output for a fixed prompt is byte-identical before and after on one Metal host, one CUDA host and the ROCm host.
- [ ] All existing gates are green, including `insert_apache_header.py --check`, `check_cross_repo_refs.py` and `make verify-kernel-dtype-keys`, and the CI path filters still trigger the same jobs for overlay and kernel changes.
- [ ] `CONTRIBUTING.md` and the README describe the two kinds and their maintenance rules.
## Coordination
This moves shared build surfaces that several ongoing efforts touch (CUDA overlay work, Metal kernels, the ROCm epic). Announce the start and land it as a single mechanical PR with no functional change, so in-flight branches can rebase onto a pure rename.
## References
- `src/lib/mlx-cpp/CMakeLists.txt` (`mlx_apply_source_overlays` 18-82)
- `src/lib/mlxcel-core/build.rs` (`rerun-if-changed` 240-260, `build_mlx` 343)
- `.github/workflows/ci.yml` (path filters around 104-160, license-header gate comments around 315)
Contributor guide
Assessment
This issue has not been assessed yet.