Skip to content

perf(cuda): enable hardware NVFP4/MXFP4 conversion on Blackwell builds #1934

Description

@inureyes

Problem / Background

No build this project produces, released or tested, compiles MLX's hardware block-float conversion. Every release artifact and every CI job lists Blackwell targets in their plain form (100, 120, 121), and MLX's cvt.rn.satfinite.e2m1x2.f32 converters are compiled only when the architecture-specific form (100a, 103a, 120a, 121a) is present. NVFP4 and MXFP4 are first-class capabilities here, so "the hardware instruction is compiled out everywhere we build" is the defect, not the build-list inconsistency by itself. A second consequence is that a default local build on a Blackwell host resolves to 121a and therefore compiles different device code than the artifact users install, from the same source, with nothing in the build output naming the difference.

The fix is not a binary choice between the plain and the a form. A combined list carries both, and that is the recommended direction below.

Everything below was measured on a GB10 (DGX Spark, sm_121) host with CUDA 13.0 (nvcc V13.0.88) and CMake 3.28.3, against the pinned MLX tree.

Current Behavior

resolve_cuda_architectures (src/lib/mlxcel-core/build.rs:605-614) honors an explicit MLX_CUDA_ARCHITECTURES verbatim, otherwise auto-detects via nvidia-smi and passes the result through sm_arch_with_suffix (build.rs:658-663), which appends a to any SM >= 90. That rule was written for Hopper, where 90a gates qmm_sm90, and mirrors MLX's own CMake (mlx/backend/cuda/CMakeLists.txt:175-177) inside its auto-detect branch. The comment at build.rs:412-422 records exactly this. So a default local build here resolves to 121a.

What ships and what CI tests do not. .github/workflows/release.yml:541 sets "90a;100;121" and :764 sets "80;86;89;90a;100;120". .github/workflows/ci.yml sets "121" at lines 517, 888 and 1070 (and "70" at 1153 for the sm_70 compile check). Because build.rs forwards an explicit value verbatim, the a is never appended for any of them. The justification comment at release.yml:757-760 ("Blackwell (sm_100/sm_120) has no MLX gate") is stale against the pinned MLX: the __CUDA_ARCH_SPECIFIC__ guard below is a Blackwell gate.

There are two gates in mlx/backend/cuda/quantized/nvfp4_quantize.cuh, and they are not independent. Line 26-27, (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && defined(__CUDA_ARCH_SPECIFIC__), defines the PTX converters scale_cvt_bf16x4_to_fp4x4_rn and siblings (the cvt.rn.satfinite.e2m1x2.f32 asm is at lines 57, 59, 130, 131, 204, 205). Line 315-316, (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000), selects scale_cvt_Tx4_to_fp4x4_fast over scale_cvt_Tx4_to_fp4x4_fallback and carries static_assert(!USE_SR, "Stochastic rounding (USE_SR=true) requires CUDA >= 12.8 and compute capability >= 1000."), so stochastic-rounding quantization is statically unavailable in every build we ship or test.

Compiling marker kernels under each target and reading symbols back with cuobjdump -symbols:

target converters defined (gate 1) dispatcher takes fast path (gate 2)
sm_80, sm_90a no no
sm_100, sm_120, sm_121 no no
sm_121f no yes
sm_100a, sm_103a, sm_120a, sm_121a yes yes

sm_90a has __CUDA_ARCH_SPECIFIC__ but fails __CUDA_ARCH__ >= 1000 (Hopper is 900). sm_121f satisfies only the dispatcher gate, so the fast path calls converters that do not exist. Compiling the real translation unit mlx/backend/cuda/quantized/fp_quantize.cu with its own flags from compile_commands.json, varying only -arch:

target result
sm_121 compiles, 128,920 SASS lines, 229 convert-class ops
sm_121f fails, 3 errors in nvfp4_quantize.cuh
sm_121a compiles, 41,745 SASS lines, 673 convert-class ops

The 3.1x code-size reduction is a scalar CUTLASS conversion sequence collapsing into one hardware instruction. Read it as a static code-size measurement; no runtime benchmark has been taken.

MXFP4 is governed by the same gate. fp_quantize.cuh includes nvfp4_quantize.cuh at line 6 and calls scale_cvt_Tx4_to_fp4x4 at lines 211 and 508, inside templates parameterized on <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>. The bits == 4 arm is shared by MXFP4 (use_mx_scale = true) and NVFP4 (use_mx_scale = false), so one change fixes both. mlxcel already exercises MXFP4 through quantize_weights_with_mode (src/lib/mlxcel-core/src/ffi_tests.rs:2379, cases at 2674-3107). The decode kernels are untouched: qmv, fp_qmv and qmm_sm80 use neither macro, and MLX_CUDA_SM90A_ENABLED keys on the literal string "90a".

Scope honestly stated: on the current default path the gated converter is not called. nvfp4_repack_strategy (src/models/sanitize.rs:533-547) returns DirectTranscode unless dense repack is forced, and that branch continues at sanitize.rs:733 before reaching the mlx::core::quantize call at sanitize.rs:851-857. The argument is not a measured regression on today's default. It is that the capability is compiled out of every build, tests included, so any use of it silently gets the fallback, and the DirectTranscode default was chosen without the hardware path ever having been available to compare against.

Proposed Solution

The choice is three-way, not two-way. Measured today by compiling a translation unit carrying MLX's own cvt.rn.satfinite.e2m1x2.f32 asm block under the architecture gate, then reading the resulting fatbin:

form SASS images PTX hardware converter
121 (today) sm_121 .target sm_121 no
121a sm_121a .target sm_121a yes
121a-real;121 sm_121a, sm_121 .target sm_121 yes

The combined form is the recommended default, for review to weigh rather than a decided outcome. It carries the hardware converter as sm_121a SASS, keeps plain sm_121 SASS, and keeps forward-JIT-capable plain PTX, so it does not pay the forward-compatibility cost that 121a alone does. CMake 3.28.3 expands CUDA_ARCHITECTURES "121a-real;121" to exactly --generate-code=arch=compute_121a,code=[sm_121a] and --generate-code=arch=compute_121,code=[compute_121,sm_121] (verified), and build.rs:427 passes MLX_CUDA_ARCHITECTURES straight into that property, so the env value needs no new plumbing. The remaining cost is archive size, since the sm_121a cubin is additional; that, not forward compatibility, is the thing to measure if size matters.

Concretely, for review: release.yml:541 becomes "90a;100a-real;100;121a-real;121", release.yml:764 becomes "80;86;89;90a;100a-real;100;120a-real;120", and ci.yml:517, :888, :1070 become "121a-real;121". Leave 90a, the pre-Hopper entries, and the sm_70 job at :1153 alone; 90a gains nothing here because Hopper fails the __CUDA_ARCH__ >= 1000 gate, and the x86_64 list already accepts the plain-a tradeoff there. Add 103a-real;103 only if sm_103 is ever targeted. If review prefers the simpler a-only form, the x86_64 90a precedent supports it, but the forward-JIT loss is then real and should be recorded as accepted.

121f must not be used in any form: it fails to compile, per the table above.

Scope

In scope: .github/workflows/release.yml (both matrices and the stale rationale comment at :757-760), .github/workflows/ci.yml (the three 121 jobs), docs/installation.md:161-211 (which recommends MLX_CUDA_ARCHITECTURES=121 at :179-180 and names both release lists at :200-205), and a regression guard.

Out of scope: changing the DirectTranscode default (record the decision, see below), any runtime throughput work, and build.rs's sm_arch_with_suffix rule if review keeps it implicit.

Implementation Notes

  • Reuse: the binary already records its list. mlxcel_core::hardware::compiled_cuda_architectures() (src/lib/mlxcel-core/src/cuda_arch.rs:57) reads MLXCEL_CUDA_ARCHITECTURES, and MLXCEL_TRACE_ARCH prints it next to the Detected N GPU(s) startup line. Quote that in benchmark records rather than adding new plumbing; most existing records under docs/benchmark_results/ do not name the architecture list, which is why this went unnoticed.
  • Correcting an earlier claim: the plain form is not PTX-free. flags.make for the release-profile MLX build shows --generate-code=arch=compute_121,code=[compute_121,sm_121], and cuobjdump on target/release/build/mlxcel-core-*/out/build/lib/libmlx.a reports 98 sm_121 cubins and PTX with .target sm_121. This matters because it is the forward-JIT fallback the combined form preserves and the a-only form would discard.
  • Do not gate on the PTX spelling: under the recommended combined form the emitted PTX comes from the plain compute_121 pass and takes the fallback arm, so cuobjdump --dump-ptx | grep 'cvt.rn.satfinite.e2m1x2' returns 0 even on a correct build (measured). Gate on SASS instead, which holds for both the a-only and the combined form.
  • Edge cases: a fat-binary list mixing plain, a and -real entries is legal and already partly shipped (90a beside 100); the sm_70 job must keep CUDA 12.x, since CUDA 13 rejects compute_70; build.rs's auto-detect still yields 121a alone, so a default local build and a release build stay different unless review also changes sm_arch_with_suffix.
  • Error handling: a wrong list already fails loudly at startup with the runtime mismatch check from chore(core): expose CUDA compute capability to the mlxcel runtime, build, and diagnostics #1537, naming both the compiled list and the detected capability, so a mistake here does not degrade into an opaque CUDA load failure.

Acceptance Criteria

  • release.yml:541 and :764 and the three ci.yml Blackwell jobs carry an architecture-specific Blackwell entry superseded, see the refresh below: the lists stay plain and the architecture-specific image goes to fp_quantize.cu alone; the stale rationale comment at release.yml:757-760 and docs/installation.md:161-211 are updated to match, including which of the three forms was chosen and why
  • cuobjdump --dump-sass <libmlx.a> | grep -c 'F2FP.SATFINITE.E2M1' is greater than 0 on a build with the new list; it is 0 on today's 121 build (measured), so this check fails on current code and needs no GPU. This is the gate because it holds for both the a-only and the combined form
  • cuobjdump --list-elf <libmlx.a> | grep -c 'sm_121a' is greater than 0 on the new build; it is 0 today, with 98 sm_121 (measured)
  • (the combined form was not adopted; recorded anyway) cuobjdump --dump-ptx <libmlx.a> | grep -c '.target sm_121$' is still greater than 0, confirming forward-JIT-capable plain PTX was kept, and the archive size delta against the 121 build is recorded
  • At least one CI job builds with the chosen form and runs the existing MXFP4/NVFP4 tests in src/lib/mlxcel-core/src/ffi_tests.rs, so the hardware path is compiled and tested rather than only compiled
  • A guard fails the build or CI if any Blackwell entry in release.yml or ci.yml regresses to plain-only, with 121f rejected explicitly and the compile failure cited as the reason
  • Decode path confirmed unchanged: qmv, fp_qmv and qmm_sm80 use neither macro, and greedy generation on a quantized checkpoint is token-identical before and after
  • A decision on whether DirectTranscode remains the right default is recorded on this issue, now that DenseNative can reach the hardware converter
  • Either the dense-repack load path is measured under both lists, or the issue records that it was not measured because the path is opt-in behind MLXCEL_NVFP4_DENSE_REPACK
  • Integrated in the real build flow: the shipped artifacts and CI jobs carry the new list, not a local override

Verification

MLX_CUDA_ARCHITECTURES="121a-real;121" cargo build --release --features cuda
A=$(ls -d target/release/build/mlxcel-core-*/out/build/lib/libmlx.a | head -1)
cuobjdump --dump-sass "$A" | grep -c 'F2FP.SATFINITE.E2M1'   # expect > 0 (0 before)
cuobjdump --list-elf  "$A" | grep -c 'sm_121a'               # expect > 0 (0 before)
cuobjdump --dump-ptx  "$A" | grep -c '.target sm_121$'       # expect > 0 (plain PTX kept)
ls -l "$A"                                                    # record the size delta
cargo test -p mlxcel-core --release --features cuda -- --test-threads=1 mxfp4 nvfp4
MLXCEL_TRACE_ARCH=1 ./target/release/mlxcel generate --model <quantized-model> --prompt "hi" --max-tokens 8

A pass is: the SASS and sm_121a counts move from 0, the plain PTX count stays above 0, the block-float tests pass, and the MLXCEL_TRACE_ARCH line reports the new list while the selected quantized-matmul path is unchanged from the 121 build. CUDA runs on this host require --test-threads=1.

Technical Considerations

Open questions for review. First, with the hardware converter available, is DirectTranscode still the right nvfp4_repack_strategy default, or should DenseNative be reconsidered? The current default predates the hardware path being compilable at all. Second, should sm_arch_with_suffix (build.rs:658-663) keep its blanket >= 90 rule, or emit the same combined form the release lists use, so that a default local build and a shipped artifact stop diverging?

Related: #1537 (compiled-architecture plumbing and the runtime mismatch check), #705 (the direct-transcode default), #693 (the dense repack fallback), #637 (Blackwell quantized GEMM).


Refresh 2026-09-20 (PR #1938). All criteria met as written at that point, with the combined 121a-real;121 form in the release and CI lists.

Refresh 2026-09-21 (PR #1938), and the first criterion is deliberately not met any more. Measuring the combined form found what the criterion did not anticipate: naming Blackwell architecture-specific in the list compiles every translation unit that way, because CUTLASS keys CUTLASS_ARCH_MMA_SM121A_ENABLED on the same macro, and on a matching device that image is the one the driver loads. It costs 12.7 MB of archive and replaces the machine code of qmv on every Blackwell host, to fix a converter no decode kernel calls.

fp_quantize.cu is the only translation unit that can reach the converters, so the PR now gives the architecture-specific image to that one source and leaves the lists plain. Same 216 F2FP.SATFINITE.E2M1 instructions, one architecture-specific cubin in the archive instead of 98, qmv.cu.o and fp_qmv.cu.o byte-identical to a plain build, +49,536 bytes instead of +12,721,576.

A reported 2.6% decode regression from the combined form did not survive an interleaved re-measurement and is recorded as drift; the record is in docs/benchmark_results/data/blackwell-arch-specific-gb10-2026-09-21/. The change of approach rests on artifact size and blast radius, not on throughput.

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:highHigh prioritystatus:doneCompletedtype:performancePerformance improvements

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions