Repository navigation
refactor(mlx-cpp): organize mlxcel's MLX-side layer as mlxcelverse (backend overlays and mlxcel kernels) #1816
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:
-
Per-backend source overlays that replace or add MLX files before MLX is compiled:
src/lib/mlx-cpp/patches/mlx/backend/metal/(3 files) andsrc/lib/mlx-cpp/patches/mlx/backend/cuda/**(22 files), always copied;src/lib/mlx-cpp/patches-cuda/{dtype,fast,ops}.cpp, copied intomlx/only for CUDA builds;src/lib/mlx-cpp/patches-rocm/**from build(rocm): vendor the ROCm backend into mlxcelverse and add the rocm cargo feature #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.
-
mlxcel's own kernels, fusions and extension functions built on MLX's public custom-kernel APIs (
fast::metal_kernel,fast::cuda_kernel, andfast::hip_kernelon ROCm):src/lib/mlx-cpp/turbo/(20 files: paged attention, sampling, fused norm, fused RoPE+append, TurboQuant SDPA and others) and the kernel bodies insrc/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 inNOTICE).
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
- Move the directories with
git mvso history follows. Updatemlx_apply_source_overlaysso 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). - Update every path reference:
src/lib/mlxcel-core/build.rs(build_mlxsource dir,rerun-if-changedlines 240-260),.github/workflows/ci.ymlpath filters and the license-header gate comments,scripts/ci/check_kernel_dtype_keys.py,scripts/insert_apache_header.pyskip 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 examplesrc/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. - Write
src/lib/mlxcelverse/README.mdand a short provenance note per overlay directory. - Keep the build cache working: the moved CMake source directory must not trigger an MLX refetch loop, and
purge_stale_mlx_cachebehavior must stay the same.
Acceptance criteria
- For Metal, CUDA and ROCm builds, the set of files copied into
_deps/mlx-srcis 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.pyandmake verify-kernel-dtype-keys, and the CI path filters still trigger the same jobs for overlay and kernel changes. -
CONTRIBUTING.mdand 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_overlays18-82)src/lib/mlxcel-core/build.rs(rerun-if-changed240-260,build_mlx343).github/workflows/ci.yml(path filters around 104-160, license-header gate comments around 315)
Activity
refactor(mlxcelverse): organize MLX-side overlays and custom kernels …