perf(cuda): enable hardware NVFP4/MXFP4 conversion on Blackwell builds - #1938
Merged
Merged
Conversation
inureyes
added
status:review
10 tasks
MLX compiles its hardware NVFP4/MXFP4 converters only under an architecture-specific target (`__CUDA_ARCH_SPECIFIC__`, spelled `a` in CMake), so every plain Blackwell entry this repository ships or tests compiles the `cvt.rn.satfinite.e2m1x2.f32` path out and silently takes the scalar CUTLASS fallback. The difference between `90a;100;121` and `90a;100a-real;100;121a-real;121` is two tokens in a YAML env block and nothing in the build output names it, so review cannot see it. `scripts/ci/check_cuda_arch_lists.py` parses every `MLX_CUDA_ARCHITECTURES` assignment under `.github/workflows/`, rejects any entry that is not a CMake architecture spelling, rejects the family-specific `f` suffix by name (it satisfies the dispatcher gate in `nvfp4_quantize.cuh` but not the converter gate, so nvcc fails with three errors in that header), and requires every plain Blackwell entry to have a cubin-emitting `a` sibling in the same list. Both surviving forms pass: the combined `121a-real;121` and the bare `121a`. It also requires every list in `release.yml` to appear verbatim in `docs/installation.md`, which is the drift that kept this invisible. `scripts/ci/check_cuda_arch_lists_test.sh` drives nine throwaway roots through the checker and asserts the exit status and the operator-facing line for each. The checker fails against the current workflows, which is the point; the lists move in the next commit. Refs #1934
Every architecture list this repository shipped or tested named Blackwell plainly (`100`, `120`, `121`), and MLX compiles its hardware NVFP4/MXFP4 converters only under an architecture-specific target, so `cvt.rn.satfinite.e2m1x2.f32` was compiled out of every release artifact and every CI job. Block-float quantization silently took a scalar CUTLASS conversion sequence instead, stochastic-rounding quantization was statically unavailable, and the tests named for MXFP4 and NVFP4 had never exercised the path they cover. Each Blackwell target is now named twice, `<sm>a-real` beside plain `<sm>`. The architecture-specific entry contributes the cubin that carries the hardware instruction; the plain entry beside it keeps the forward-JIT-capable `compute_<sm>` PTX that a bare `<sm>a` would discard. Hopper stays `90a`, because those converters also gate on compute capability 10.0 and 9.0 cannot reach them however it is spelled. `cuda-arch-lists` runs the new guard on every PR with no toolchain. `cuda-blockfloat` builds the shipped list on GB10, asserts with `cuobjdump` that the `sm_121a` cubin and the `F2FP.SATFINITE.E2M1` instruction really reached `libmlx.a`, and then runs the block-float quantization tests, which is the half that was missing: nothing tested the converter because nothing compiled it. The assertion is on SASS and not on PTX, because under the combined form the emitted PTX comes from the plain `compute_121` pass and takes the fallback arm, so grepping PTX for the converter returns zero on a correct build. `docs/installation.md` gains the Blackwell rule, the archive-reading commands that distinguish the two builds, and the PTX trap. `cuda_arch_tests.rs` pins the tradeoff rather than only the constants: the `a-real` half covers its exact target and nothing else, the plain half is what still serves a later Blackwell minor by PTX, and a bare `121a` serves nothing there. Refs #1934
`cuobjdump --dump-sass` on the 178 MB `libmlx.a` takes about ten minutes; extracting `fp_quantize.cu.o` first and disassembling that takes under a second. Measured on the plain `121` archive both report zero `F2FP.SATFINITE.E2M1`, and the single object reproduces the whole 128,920-line sm_121 disassembly the issue quotes, so the cheap form carries the same verdict. The archive-wide spelling stays in `docs/installation.md`, which is where an operator checking one build by hand will read it. The step also prints the `F2FP.SATFINITE.E4M3` count beside it. That converter is not gated on the architecture-specific target, so it is present either way, and printing both makes "nothing compiled at all" read differently from "the fp4 gate is off". The temporary disassembly now lives under `RUNNER_TEMP` rather than a fixed `/tmp` path, which two jobs on this self-hosted runner could otherwise collide on. Refs #1934
Verified on this GB10 by compiling a marker kernel under both `-gencode` sets and running it: the combined form prints `ARCH-SPECIFIC (__CUDA_ARCH__=1210, __CUDA_ARCH_SPECIFIC__=1210)` and the plain form prints `generic`. So on a device the `a` entry matches, the architecture-specific image is what executes and the plain entry is purely the forward-compatibility fallback. Worth stating, because a reader can otherwise reasonably assume naming a target twice means the hardware path is merely available rather than selected. Refs #1934
`shapeless_compile_audit_harness` carries `#[ignore = "run with scripts/audit_shapeless_compile.sh on a reachable GPU"]`, so naming it in the job's filter contributes a skipped line and no coverage. `global_scale` alone selects the six `ffi_tests` cases that drive `quantize_weights_with_mode` with `mxfp4` and native NVFP4 sidecars, which is every test in this crate that reaches MLX's block-float quantization kernels. Verified against the built test binary: the filter matches seven tests, six run and one is the ignored harness. Refs #1934
An earlier revision of this branch raised the ceiling from 480 to 600 minutes on the reasoning that two more code-generation passes need proportionally more time. They do, and the existing ceiling already covers it: a cold six-pass build was observed at about three hours, so eight passes is about four against an eight-hour ceiling. Changing a timeout that does not need to change only adds something for review to weigh, so the number stays and the arithmetic is now written down beside it. Refs #1934
The first run of `cuda-blockfloat` failed at the verification step with "cuobjdump not found". Everything before it was correct: the job env carried `121a-real;121`, and the build step compiled the tests for about fifteen minutes, so the list reached build.rs and CMake on the runner. Only the tool lookup failed, because neither `cuobjdump` nor `nvcc` is on PATH on GB10; the toolkit lives at `/usr/local/cuda/bin`, and the lookup tried PATH first and then nvcc's directory, which is also PATH-resolved, so both arms missed. Both jobs now resolve the toolkit the way `src/lib/mlxcel-core/build.rs` does, `CUDA_HOME` first and then the conventional location, with PATH as a last resort rather than as the primary, and print the resolved path. That also means the gate reads the same toolkit that compiled the archive. Not-found stays a hard failure: skipping would make the gate pass while checking nothing. `cuda-sm70-compile` carried the identical lookup, latent only because its toolkit probe skips the later steps on CUDA 13, and its nvcc probe already had the absolute fallback its cuobjdump lookup lacked; both halves are now consistent. The step also asserts that the archive it picked was configured with the list the job pins. `find | head -1` chooses among however many build directories a checkout holds, one per feature set and architecture list, so without this the gate could measure a neighbour and pass. Verified by running the step verbatim out of the workflow, with PATH stripped of the toolkit, against real archives: it passes on the `121a-real;121` build reporting 216 fp4 conversions, 98 `sm_121` and 98 `sm_121a` cubins and 98 plain PTX images, and fails on the `121` build twice over, once on the configured-list assertion and once, with that relaxed, on the SASS count reaching zero. Refs #1934
inureyes
force-pushed
the
update/issue-1934-blackwell-arch-specific
branch
from
September 21, 2026 01:57
fff35fb to
01fe5bb
Compare
MLX compiles its hardware NVFP4/MXFP4 converters only under an architecture-specific target, so every artifact this project shipped compiled them out. The fix this branch carried until now answered that from the architecture list, naming Blackwell `<sm>a-real` beside plain `<sm>`. That works and it charges every translation unit for a converter exactly one of them reaches: CUTLASS derives `CUTLASS_ARCH_MMA_SM121A_ENABLED` from the same macro, so the decode kernels get a second image too, and on a matching device that is the image the driver loads. Measured on GB10, 12.7 MB of extra archive and a different `qmv` on every Blackwell machine. `fp_quantize.cu` is the only translation unit that can reach the converters: it is the only file that includes `fp_quantize.cuh`, which is the only file that includes `nvfp4_quantize.cuh`. All five kernels live in `fp_quantize.cuh`, the gated functions are `__device__ __forceinline__` so they never leave the TU that instantiates them, the object exports only ordinary host entry points, and MLX compiles whole-program with no `-rdc`, so no device symbol has to resolve across objects. `src/lib/mlx-cpp/CMakeLists.txt` therefore adds the extra architecture to that one source through `set_source_files_properties(... TARGET_DIRECTORY mlx ...)`, derived from whatever Blackwell entries the list already names, and fails the configure if the file ever moves. Measured on GB10 against the same tree: the archive grows 49,536 bytes instead of 12,721,576, `fp_quantize.cu.o` carries 216 `F2FP.SATFINITE.E2M1` instructions where the plain build has none, and `qmv.cu.o` and `fp_qmv.cu.o` disassemble to the same bytes as the plain build with no second image at all. So the architecture lists go back to plain everywhere, and an `a` in one is now the regression. The guard inverts with them: it still rejects `121f`, which does not compile, and now rejects a target-level Blackwell `a`, and additionally requires the per-source injection to still be present, because without it plain lists mean the converters are silently compiled out again. Refs #1934
A sequential A/B on the qwen3.5-4b DFlash pairing reported the architecture-specific build about 2.6% slower at draft width 4 with disjoint ranges. Two arms measured one after the other are only as comparable as the host is still between them, so this re-measures them interleaved: one server start, one request, teardown, next binary, with the opening arm rotating every round. It does not reproduce. Across five rounds at width 4 and five at classic, every range overlaps, the paired within-round deltas change sign, and the architecture-specific arm is nominally faster at width 4 (ratio of means 1.0052). The third arm is what makes that conclusive rather than merely negative. It compiles the converter per translation unit, so its decode kernels are the same bytes as the plain build's, verified by disassembly hash, with no second image for the driver to prefer. It measures 1.11% slower at width 4 and 1.48% slower at classic with all five paired deltas negative. Five same-sign deltas at about 1.5% are therefore reachable on this host with no code difference at all, which is the floor the 2.6% has to be read against. The driver imports `run_arm` from #1797's served harness rather than reimplementing the host gate, the NVRM trip wire and the DFlash diagnostics parsing, so the conditions are the same code that produced the records this compares against. Conditions were clean throughout: 30 arms, no errors, zero NVRM delta, no CI job overlapping any arm, load1 between 0.01 and 0.47. Refs #1934
Two edges in the per-source injection, both found by driving its list parsing directly rather than by building. A list that already names a Blackwell capability architecture-specific, `121a-real;121` for instance, would have had `--generate-code=arch=compute_121a` appended a second time, and nvcc rejects the duplicate rather than ignoring it. The workflow guard rejects that spelling, but an operator following the pre-#1934 advice to build it by hand would have hit it. The injection now skips any capability the list already covers. `IN_LIST` needs policy CMP0057, so it works or does not depending on the `cmake_minimum_required` of whoever includes this file. `list(FIND)` has no such dependency. This file's own project declares 3.16, which would have set the policy, but the construct should not depend on that. Driven against every list this repository builds plus the rejected forms: `90a;100;121` and `80;86;89;90a;100;120` each get exactly their two Blackwell entries, `121` gets one, `70` and `90a` get none, and the three architecture-specific spellings get none. Refs #1934
The two job comments still described the gates as catching a plain-only Blackwell list, which is now the correct spelling. They catch the opposite: a target-level `a` entry, and a missing per-source injection. The block-float job's comment also promised the SASS assertion would survive a change of mind between the two architecture-specific forms, which is no longer the choice being made. Refs #1934
Fifteen runs across three binaries produced one distinct completion from the 158-token prompt. It is worth stating because it rules out the arms having done different amounts of work: what the throughput numbers compare is kernel speed, not a shorter or longer generation. Refs #1934
This was referenced Sep 21, 2026
inureyes
added a commit
that referenced
this pull request
Sep 21, 2026
#1946) * chore(build): stop auto-detect at the shipped Blackwell spelling With MLX_CUDA_ARCHITECTURES unset, `sm_arch_with_suffix` appended CUDA's architecture-specific `a` suffix to every SM from 90 up, so a default `cargo build --features cuda` on a GB10 resolved the whole MLX target to `121a`. That is the whole-target form PR #1938 decided against: every translation unit compiles under `__CUDA_ARCH_SPECIFIC__`, CUTLASS keys `CUTLASS_ARCH_MMA_SM121A_ENABLED` on the same macro, and the per-source injection in `src/lib/mlx-cpp/CMakeLists.txt` never fires at all, because an entry already carrying `a` is recorded as architecture-specific and skipped rather than given a duplicate `--generate-code`. A local build and a shipped build then differ in the machine code their decode kernels run, not merely in architecture coverage. Both functions that spell a list from a detected capability now stop suffixing at a shared `FIRST_PLAIN_SM` boundary: `sm_arch_with_suffix` in the build script, and `CudaArchMismatch::suggested_architecture`, whose startup message was telling an operator on a Blackwell card to rebuild with `MLX_CUDA_ARCHITECTURES=121a`, the one value this repository rejects everywhere else. Hopper keeps `90a` because `release.yml` ships `90a`. The stale rationale for it is replaced rather than repeated: upstream MLX commit `44540d12` moved `qmm_sm80`, `qmm_sm90` and `gather_gemm` to runtime NVRTC compilation and deleted the `MLX_CUDA_SM90A_ENABLED` definition the old comments cited, and `jit_module.cpp` now derives the NVRTC `--gpu-architecture` from the running device, appending `a` itself from compute capability 9 up. `check_cuda_arch_lists.py` grows rule 6 for the auto-detect path, which rule 3 cannot see because no workflow writes that list down, and `check_cuda_arch_lists_test.sh` grows three cases for it. Refs #1943, #1934, PR #1938. * docs(installation): describe one CUDA architecture rule, not two The section stated the auto-detect rule as suffixing every SM from 90 up, citing a Hopper quantized-kernel gate, and then told operators not to set `a` on Blackwell and that building `121a` is a step backwards. Both passages described the same code path in opposite terms, and `check_docs` in the guard script only verifies that the release lists appear verbatim, so nothing caught it. It now states the rule once, the way the build script and the release workflow both spell it, and explains what `121a` costs and why the per-source injection makes it unnecessary. The Hopper paragraph records what is actually true at the current MLX pin: upstream `44540d12` moved the Hopper quantized kernel to runtime NVRTC compilation and removed the `MLX_CUDA_SM90A_ENABLED` definition this section used to cite, `90a` is kept because both release lists ship it, and `CUTLASS_ARCH_MMA_SM90A_ENABLED` still keys on `__CUDA_ARCH_FEAT_SM90_ALL` so a later pin can make it matter again. Refs #1943. * chore: state the Hopper cross-compile result precisely The first pass said the quantized translation units produce identical SASS at `90` and `90a`. Measured across all five, that was too strong in one direction and too weak in another. `qmm_sm90.cu`, `qmm.cu` and `qmm_sm80.cu` emit no device function at either spelling, because upstream `44540d12` moved those kernels to runtime NVRTC compilation; their objects carry 21 lines of fatbin header and zero kernel symbols. `qmv.cu` (2898 kernels) and `fp_qmv.cu` (54) do emit device code, and it is instruction-identical, but their cubins are not byte-identical: the `90a` build adds `EF_CUDA_ACCELERATORS` to the `.headerflags` line, which is the ELF flag marking a cubin architecture-specific. So the suffix changes no instruction at this pin, which is what the decision rests on, and the text now says that rather than implying the objects match byte for byte. Also drops the two remaining statements of the pre-#1943 rule that the new guard does not read: the inline comment in `detect_cuda_arch`, and the `90a is load-bearing` rationale on the x86_64 release job, which cited the same deleted macro. The architecture lists themselves are untouched. Refs #1943. * docs(build): say why the no-detection fallback stays 90a The issue tied the `90a` fallback literal to the Hopper decision, and the cross-compile shows the suffix buys no instruction at this pin, so leaving the literal in place needed a reason rather than silence. The reason is the same one that keeps Hopper suffixed at all: `release.yml` ships `90a`, and a fallback spelled `90` would hand a host without `nvidia-smi` a build shape no published archive uses, which is the divergence this issue closed on Blackwell. The doc comment on `resolve_cuda_architectures` said only that `90a` is the last resort. Refs #1943. * docs(core): stop the compiled-list example showing a rejected shape The doc comment on `compiled_cuda_architectures` illustrated the string with `"80;86;89;90a;100a-real;100;120a-real;120"`, which names Blackwell architecture-specific twice. A doc comment on a parser is where a reader copies a list from, and after this branch the project's own guard rejects that shape, so the file's example and the file's rule disagreed. It now shows the real x86_64 release list, which exercises the same parser features that mattered, pre-Blackwell plain entries and a Hopper `a` entry, and cannot drift into recommending something the build refuses. The `-real` and `-virtual` forms are still described on `parse_cuda_arch_entry` and round-tripped in `entry_display_round_trips_through_the_parser`, which is where an illustration of the grammar belongs. The same pass corrects the header on the release-list constants in the test module, which still justified Hopper's `90a` by the quantized-kernel gate that upstream deleted. Refs #1943. * chore(ci): correct the reasons that named the pre-#1943 auto-detect shape Four comments justified themselves by what auto-detection used to produce. Three CUDA jobs in `ci.yml` said their pin exists because auto-detection "yields `121a` here"; after this branch it yields `121`. Each now states the reason that survives, that a gate should name the list it compiles rather than infer it from whichever GPU the runner has, so the pin keeps a rationale instead of becoming a bare value someone deletes as redundant. `scripts/bench_block_width.sh` is the one that inverted rather than going stale. It told a reader that a row measured on an auto-detected build is not comparable with records pinned at plain `121`, which after this branch is backwards, since an auto-detected GB10 build now resolves to exactly `121`. A wrong comparability rule in a benchmark harness produces discarded runs or published numbers someone trusts, so it now says which rows are actually suspect: those recorded between #1934 and #1943 on an unpinned build. The aarch64 release job carried the same deleted-macro rationale the x86_64 one did, missed in the earlier pass because only one of the two was grepped. Refs #1943. * chore(bench): correct the reused width harness on auto-detected builds The `draft-block-width-gb10-2026-09-20/harness/` directory is a live tool despite its dated path: the post-#1939 width sweep and the #1943 architecture A/B both re-used it, and `draft-block-width-post-1939-gb10-2026-09-21/` ships no `harness/` of its own for exactly that reason. Its four files told a reader that `build.rs` auto-detects `121a` here, so a row from an unpinned build is not comparable with the pinned records. After this branch that is backwards: auto-detection resolves to plain `121`, the same list those records pin. Each now keeps its pin and its reason, states that the pin no longer corrects a shape but still names what was compiled, and identifies the rows that are genuinely suspect, those recorded between #1934 and #1943 on an unpinned build, rather than implying every auto-detected build is. The record prose under `docs/benchmark_results/` and the `TECHNICAL_REPORTS` entry are left as written; they document what was true when they were measured. Refs #1943.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
No build this project produced compiled MLX's hardware NVFP4/MXFP4 converters. The converters in
mlx/backend/cuda/quantized/nvfp4_quantize.cuhare gated on__CUDA_ARCH_SPECIFIC__, every architecture list here names Blackwell plainly, socvt.rn.satfinite.e2m1x2.f32was compiled out of every release artifact and every CI job: block-float quantization silently took a scalar CUTLASS fallback and stochastic-rounding quantization was statically unavailable.The fix adds the architecture-specific image to
fp_quantize.cualone, insrc/lib/mlx-cpp/CMakeLists.txt, rather than to the wholemlxtarget. The architecture lists are unchanged.Why not just put
121ain the listThat was this PR's first implementation, and it works. It also charges every translation unit for a converter exactly one of them reaches, because CUTLASS derives
CUTLASS_ARCH_MMA_SM121A_ENABLEDfrom the same macro. Measured on GB10, same tree, three builds differing only in how MLX was compiled:121121a-real;121libmlx.amlxcel-serverF2FP.SATFINITE.E2M1infp_quantize.cu.oqmv.cu.o.target sm_121)On a device the suffix matches, the architecture-specific image is the one the driver loads, not an alternative it can ignore: a marker kernel compiled under both
-gencodesets printsARCH-SPECIFIC (__CUDA_ARCH__=1210, __CUDA_ARCH_SPECIFIC__=1210)for the list form andgenericfor plain. So the list form does not merely add code, it replaces the code every kernel on every Blackwell host runs.The decode regression: measured, and it does not hold
A sequential A/B reported the list form about 2.6% slower at draft width 4 with disjoint ranges. Re-measured interleaved at round granularity, one server start and one request per arm with the opening arm rotating, it does not reproduce. Full record and harness in
docs/benchmark_results/data/blackwell-arch-specific-gb10-2026-09-21/.121a-real;121121a-real;121Plain measured 65.96-67.81 at width 4 and 55.72-59.39 at classic. Every range overlaps.
The sequential result was drift, and this PR's own arm is what proves it. Its decode kernels are the same bytes as plain, verified by disassembly hash, with no second image to prefer, and it still measures 1.11% and 1.48% slower with all five paired deltas negative at classic. Five same-sign deltas at ~1.5% are reachable here with zero code difference, so the 2.6% sits inside this method's floor.
So why change the approach at all
Not for throughput, which is the honest answer. For artifact size, 257x less of it, and for blast radius: a converter fix should not replace the machine code of
qmvon every sm_100, sm_120 and sm_121 user's machine, even when the one pairing measured here does not slow down. The list form also needs six workflow edits kept in sync; this needs none.Against that, the list form is the more conventional mechanism and would automatically pick up any future MLX code gated on the same macro.
cuda-blockfloatnow asserts that the set of MLX files using either arch-specific macro is exactlynvfp4_quantize.cuh, so a pin bump that grows that surface fails loudly instead of silently losing coverage.Why this is safe per translation unit
fp_quantize.cuis the only file that includesfp_quantize.cuh, which is the only file that includesnvfp4_quantize.cuh. All five__global__kernels are infp_quantize.cuh; the gated functions are__device__ __forceinline__and never leave the TU that instantiates them; the object exports only ordinary host entry points (mlx::core::fp_quantizeand siblings); and MLX compiles whole-program, no-rdc, so no device symbol resolves across objects. CMake fails the configure if the file ever moves.Gates
cuda-blockfloatgreen on CI, printingF2FP.SATFINITE.E2M1 occurrences in fp_quantize.cu.o: 216,qmv.cu.o images: sm_121,fp_qmv.cu.o images: sm_121,architecture-specific cubin images in the whole archive: 1,forward-JIT-capable plain PTX images (.target sm_121): 98,MLX files gated on an architecture-specific target: backend/cuda/quantized/nvfp4_quantize.cuh, then 6 block-float tests passing on the runner121a-real;121build (qmv.cu.o carries an architecture-specific image), which is the build this PR previously proposedcheck_cuda_arch_lists.pyclean; its self-test covers 10 cases including the target-levelaform,121f, a missing per-source injection, and doc driftcargo test -p mlxcel-core --release --features cuda --lib -- --test-threads=1 global_scale: 6 passed;cuda_arch: 27 passedcargo clippy -p mlxcel --lib --tests -- -D warningsat default features andcargo fmt --checkclean121a-real;121builds of this treecargo check --features cuda,xla-iree --all-targetscannot run here:mlxcel-xla's build script aborts withoutIREE_DIST, which CI provisions as a job stepOpen question, not decided here
sm_arch_with_suffix(build.rs:658-663) still appends a bareaon auto-detect, so a default local build on a Blackwell host is121a: the whole-target form this PR argues against. With the per-source injection that suffix now buys nothing it does not already get, so I would propose dropping it for SM >= 100 and letting the injection do the work, which also makes a local build match a release build. That changesbuild.rsbehavior and the issue put it out of scope, so it is left for review.Closes #1934