Skip to content

chore: Reconcile auto-detect CUDA arch with the shipped build shape #1943

Description

@inureyes

Problem / Background

sm_arch_with_suffix (src/lib/mlxcel-core/build.rs:658-663) appends CUDA's architecture-specific a suffix to any SM >= 90 when MLX_CUDA_ARCHITECTURES is unset and the architecture is auto-detected from nvidia-smi (detect_cuda_arch, build.rs:618-648). On a GB10 that resolves the whole mlx target to 121a.

PR #1938 (merged 0bbfa95d, closing #1934) decided the opposite shape: the shipped workflow lists stay plain (release.yml:551 90a;100;121, release.yml:782 80;86;89;90a;100;120, ci.yml:1359 121), and MLX's hardware block-float converter is obtained by injecting --generate-code=arch=compute_121a,code=[sm_121a] onto the single source fp_quantize.cu via set_source_files_properties (src/lib/mlx-cpp/CMakeLists.txt:307-309). The measured reason, not repeated here, is in that PR: the whole-target form grows libmlx.a by 7.1 percent and gives every decode kernel a second architecture-specific image, while the per-source form costs 0.03 percent and leaves qmv.cu.o and fp_qmv.cu.o byte-identical.

So a developer building locally on this host without setting the variable now gets the form the project just decided against, and local and shipped builds differ in which machine code the decode kernels run rather than merely in architecture coverage. That is the same class of divergence that started #1934, pointing the other way.

The divergence is specific to SM >= 100. On Hopper, auto-detect produces 90a, which is also what the release lists carry, so those already agree.

Current Behavior

The injection selection loop (src/lib/mlx-cpp/CMakeLists.txt:277-294) only matches plain entries (^([0-9][0-9]+)$) and only for _sm GREATER_EQUAL 100, and an entry that already carries a is recorded in _already_arch_specific (CMakeLists.txt:271-276) so nothing is added for it, because asking nvcc for the same --generate-code twice is an error. A 121a auto-detect build therefore leaves _fp_quantize_arch_flags empty: no per-source injection fires at all, and the converter arrives from the whole-target suffix instead. The two shapes are mutually exclusive by construction, not additive.

Proposed Solution

This issue asks for a recorded decision plus the change that follows from it, not for a predetermined edit. PR #1938's body proposes dropping the suffix for SM >= 100 and leaving Hopper's 90a alone. The alternative of having auto-detect emit the same per-source shape the release build uses converges on exactly that for Blackwell: because of CMakeLists.txt:277-294 above, auto-detect emitting plain 121 is what makes the existing injection fire, with no CMake edit. Whichever framing is chosen, justify it in the PR body against the mechanism above rather than assuming it.

What is genuinely open is Hopper. The doc comment on sm_arch_with_suffix (build.rs:653-657) and the comment at build.rs:416-421 justify the suffix as unlocking Hopper features and gating MLX_CUDA_SM90A_ENABLED on the literal string 90a. In the pinned MLX tree (GIT_TAG 81ba1c6a, src/lib/mlx-cpp/CMakeLists.txt:171) that macro has no definition anywhere: the only occurrence is a lablup-authored comment in the patched mlx/backend/cuda/quantized/quantized.cpp:52. qmm_sm90.cu is in MLX's unconditional source list (mlx/backend/cuda/CMakeLists.txt:67) and carries no preprocessor conditional of its own. That does not settle the question, because something it includes may gate on the architecture-specific macro the way CUTLASS derives CUTLASS_ARCH_MMA_SM121A_ENABLED from it. Resolve it with evidence rather than with the comment: cross-compile the same tree at 90 and at 90a on this host and diff qmm_sm90.cu.o with cuobjdump. No Hopper device is needed. If the suffix buys nothing there either, the "90a" no-detection fallback literal at build.rs:609 follows the same decision.

Scope

In scope: src/lib/mlxcel-core/build.rs (the suffix rule, and the fallback literal if the Hopper decision moves it); scripts/ci/check_cuda_arch_lists.py and scripts/ci/check_cuda_arch_lists_test.sh (guard coverage for the auto-detect path); docs/installation.md:166-197 (the prose describing the rule).

Out of scope: the workflow architecture lists, which PR #1938 already settled; the per-source injection mechanism itself; any change to MLX's own CMake or to the patched files under src/lib/mlx-cpp/patches/.

Implementation Notes

  • Reuse: extend the existing cuda-arch-lists job (.github/workflows/ci.yml:341-377) rather than adding a new one. Model the new check on check_injection (scripts/ci/check_cuda_arch_lists.py:193-207): a static read of build.rs asserting the suffix rule excludes major >= 10, with failure text naming this issue the way the existing messages name perf(cuda): enable hardware NVFP4/MXFP4 conversion on Blackwell builds #1934. Add the matching case to scripts/ci/check_cuda_arch_lists_test.sh, which ci.yml:376-377 already runs.
  • Constraints: the guard job runs on ubuntu-latest with no CUDA toolkit and no GPU, so the check must be static text analysis. cuda-blockfloat (ci.yml:1339) cannot cover the auto-detect path without a second full CUDA build, and it pins MLX_CUDA_ARCHITECTURES explicitly at ci.yml:1359 so it never exercises auto-detection. build.rs has no #[cfg(test)] mod tests and cargo test does not run build-script tests, so do not add a unit test there.
  • Edge cases: a host reporting several distinct compute capabilities (the dedup path at build.rs:639-647), a mixed Hopper and Blackwell host, and nvidia-smi absent or unparseable, which falls through to the build.rs:609 literal. An explicitly set MLX_CUDA_ARCHITECTURES must stay honored verbatim, which is the documented escape hatch.
  • Error handling: no new runtime failure path. cuda_arch.rs already reports the compiled list at startup (src/lib/mlxcel-core/src/cuda_arch.rs:416, the ; compiled for [...] summary), and that line is the observable for what a given binary was built for.
  • Docs: docs/installation.md:166-172 states the auto-detect rule while 174-197 tells operators not to set a on Blackwell, so the document currently contradicts itself on this path. check_docs (check_cuda_arch_lists.py:210-222) only verifies that release lists appear verbatim, so nothing catches this prose drifting.

Acceptance Criteria

  • A local build on this host with MLX_CUDA_ARCHITECTURES unset configures MLX with the same architecture shape a shipped build uses: the MLX_CUDA_ARCHITECTURES: entry in that build directory's CMakeCache.txt reads 121, not 121a. Read the cache file rather than the configure output, which cargo swallows unless the build runs with -vv.
  • cuobjdump confirms the decode kernels are unchanged between that auto-detect build and an explicit MLX_CUDA_ARCHITECTURES=121 build: qmv.cu.o and fp_qmv.cu.o extracted with ar x list only sm_121 under --list-elf, their --dump-sass SHA-256 hashes match across the two builds, and each archive carries exactly one architecture-specific image. This is the property the change exists to restore, and it needs no GPU.
  • scripts/ci/check_cuda_arch_lists.py fails on a build.rs whose suffix rule still produces an a for major >= 10, and scripts/ci/check_cuda_arch_lists_test.sh gains a case covering that, so the auto-detect path is guarded the way the workflow lists and the CMake injection already are.
  • The Hopper decision is recorded in the PR body with the cuobjdump evidence described above, and build.rs:653-657 plus build.rs:416-421 are rewritten to state what is actually true of the pinned MLX rather than citing MLX_CUDA_SM90A_ENABLED, whichever way the decision goes.
  • docs/installation.md:166-197 describes one rule instead of two, and no longer tells operators to avoid on Blackwell what the build script does for them.
  • The change is in the real build path, not behind a new opt-in flag: an unset MLX_CUDA_ARCHITECTURES is what a default cargo build --features cuda on this host takes.

Verification

python3 scripts/ci/check_cuda_arch_lists.py
bash scripts/ci/check_cuda_arch_lists_test.sh
cargo fmt --check
cargo clippy -p mlxcel --lib --tests -- -D warnings
cargo test -p mlxcel-core --release --features cuda --lib -- --test-threads=1 cuda_arch

# Two arms, each in its own target directory. `build.rs:293` declares
# `cargo:rerun-if-env-changed=MLX_CUDA_ARCHITECTURES`, and the rebuild reuses the same OUT_DIR,
# so two arms sharing one target directory overwrite each other's `libmlx.a` and there is
# nothing left to compare. Run the first on a Blackwell host with nvidia-smi present.
AUTO=$PWD/target-arch-auto
PLAIN=$PWD/target-arch-121
env -u MLX_CUDA_ARCHITECTURES CARGO_TARGET_DIR="$AUTO" cargo build -p mlxcel-core --release --features cuda
MLX_CUDA_ARCHITECTURES=121 CARGO_TARGET_DIR="$PLAIN" cargo build -p mlxcel-core --release --features cuda

for ARM in "$AUTO" "$PLAIN"; do
  LIB=$(find "$ARM/release/build" -path '*/out/build/lib/libmlx.a' | head -1)
  sed -n 's/^MLX_CUDA_ARCHITECTURES:[^=]*=//p' "$(dirname "$LIB")/../CMakeCache.txt"   # both expect: 121
  WORK=$(mktemp -d) && ( cd "$WORK" && ar x "$LIB" qmv.cu.o fp_qmv.cu.o
    cuobjdump --list-elf qmv.cu.o | grep -oE 'sm_[0-9]+[af]?' | sort -u   # expect: sm_121 only
    cuobjdump --dump-sass qmv.cu.o | sha256sum                            # expect: equal across arms
    cuobjdump --dump-sass fp_qmv.cu.o | sha256sum )                       # expect: equal across arms
  cuobjdump --list-elf "$LIB" | grep -cE 'sm_[0-9]+a'                     # expect: 1
done

Run these from the repository root. The guard scripts and the cuda_arch unit tests run anywhere; the rest needs a CUDA toolkit. cuobjdump ships beside nvcc and is often not on PATH, so reuse the resolution ci.yml:1381-1397 already does instead of assuming it.

Technical Considerations

This is a build-shape correctness issue, not a performance one. A throughput A/B on the whole-target form was run and refuted: the apparent regression sits inside the measurement method's own floor, shown by a negative-control arm whose decode kernels are the same bytes and which still measured about 1.5 percent slow with all five paired deltas the same sign. The record, harness and raw arms are under docs/benchmark_results/data/blackwell-arch-specific-gb10-2026-09-21/. Do not re-run that A/B as part of this work, and do not justify the change on speed.

Refs #1934, PR #1938.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificpriority:mediumMedium prioritystatus:doneCompletedtype:choreMaintenance tasks (build, CI, etc.)

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions