Skip to content

refactor(mlx-cpp): organize mlxcel's MLX-side layer as mlxcelverse (backend overlays and mlxcel kernels) #1816

Closed

Description

@inureyes

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:

    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)

Activity

  1. added
    status:backlogIn the backlog, not yet ready
    type:refactorCode restructuring without changing functionality
    area:architectureArchitecture and code structure changes
    on Sep 11, 2026
  2. self-assigned this
    on Sep 12, 2026
  3. added a commit that references this issue on Oct 14, 2026
    096152c
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:architectureArchitecture and code structure changespriority:lowLow prioritystatus:doneCompletedtype:refactorCode restructuring without changing functionality

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions