test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate

[inureyes]

Part of #1801. Phase 3. Depends on #1803, #1805, #1806. Rows for mxfp8 and mxfp4 checkpoints are added when #1807 and #1808 land. (#1803 and #1805 are closed; #1806 is the only open dependency and is the last failing item in the ROCm gate.)

Context

The spike validated MLX ops and mlx-lm generation, not mlxcel's own model code. mlxcel judges numerics with a teacher-forced logit trace and the decided-position mismatch rate (docs/benchmarks.md:408, "Judging a change that moves the numbers"), because byte equality is not expected across backends and perplexity can hide greedy flips. Real-model integration tests skip silently when a checkpoint is missing (tests/common/mod.rs:36, repo_model_dir searches models/, models/mlx/, models/mlx-big/) unless MLXCEL_REQUIRE_MODELS=1.

There is no ROCm test gate. (Obsolete: make verify-test-rocm landed in #1826 and the aggregate make verify-rocm in #2008; see plan item 1.) make verify uses --features metal,accelerate and is not usable on Linux; the CUDA analogue is make verify-test-cuda (Makefile:775).

Scope

A reproducible correctness matrix for ROCm against a Metal reference, and a local test gate for ROCm.

Status at a glance (2026-09-29, origin/main d8d34e2b)

• Done: the gate target, the aggregate, the first twelve matrix pairs and their document (test(rocm): add the gfx1151 correctness matrix and its gate #1826), and the server check (test(rocm): verify mlxcel-server chat completions on gfx1151 #1831, f9ece5d9).
• Open: the gate is not green (one failing target, the NVFP4 abort of feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806), and the sliding-window, SSM-hybrid, VLM and Qwen2.5-7B-Instruct-4bit rows have not been run. Those rows were blocked on refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803 and feat(rocm): memory estimation, device info and hardware detection on AMD UMA hosts #1805, which are both closed, so they are now runnable.

Implementation plan

1. Gate. make verify-test-rocm = cargo test --workspace --profile test-fast --features rocm --no-fail-fast -- --test-threads=1 (test binaries must not share the GPU; see the note on verify-test-cuda). Done: the target is at Makefile:780 (added by test(rocm): add the gfx1151 correctness matrix and its gate #1826, 1206f863) and is the last step of the aggregate make verify-rocm at Makefile:804 (test(rocm): add a local verify-rocm gate with a shared build-and-generate smoke #2008), which also runs crate versions, kernel dtype keys, kernel port dispatch, llama-compat manifest, cargo fmt --check, ROCm clippy and verify-rocm-smoke. What remains is getting it green and recording the pass count; a run reporting far fewer tests than the workspace total (about 11k) is a failure.
   • Current result (make verify-rocm minus the smoke on gfx1151, d8d34e2b): exactly one failing target, -p mlxcel-core --lib, and exactly one terminate called, both from ffi_tests::compiled_qgelu_mlp_global_scale_native_nvfp4_prefill_matches_reference, the NVFP4 abort tracked by feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806. Every other step passes.
   • The abort hides later tests. It terminates the mlxcel-core lib test binary, so no test sorting alphabetically after that one reaches the runner at all. A test missing from the failure list is therefore not a pass. Until feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 is fixed, any evidence this matrix needs from a later-sorting mlxcel-core test has to come from running that test directly, and the pass count in the matrix document must say which tests were run that way.
   • New coverage: rms_norm_small_axis_tests (MLX's own fast::rms_norm dispatch, not an mlxcel port) used to be skipped on ROCm by a Metal-or-CUDA gate; since refactor: finish the kernel port standardization #2029 it runs on ROCm and both sweeps pass (run directly, because it sorts after the abort).
   • verify-rocm-smoke did not run on d8d34e2b because its checkpoint had been under /tmp, which the host clears on reboot; set MLXCEL_ROCM_SMOKE_MODEL to a local checkpoint to run it.
2. Model matrix (affine 4-bit first, one per path mlxcel distinguishes). Rows marked done are in docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md (test(rocm): add the gfx1151 correctness matrix and its gate #1826) at widths 1, 8 and 256:
   • dense: Meta-Llama-3.1-8B-Instruct-4bit (done, test(rocm): add the gfx1151 correctness matrix and its gate #1826), Qwen2.5-7B-Instruct-4bit (not run)
   • small smoke: Qwen3-0.6B-4bit (done, test(rocm): add the gfx1151 correctness matrix and its gate #1826)
   • sliding window: gemma-3-4b-it-4bit (not run)
   • MoE: Qwen3-30B-A3B-4bit, Mixtral-8x7B-Instruct-v0.1-4bit (both done, test(rocm): add the gfx1151 correctness matrix and its gate #1826, with MLXCEL_FUSED_MOE=0 on both sides, which refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803 (PR refactor(core): route custom kernels by GPU backend kind #1877) made unnecessary on ROCm per PR ci(rocm): build, link and generate on a self-hosted gfx1151 runner #1989)
   • SSM hybrid: granite-4.0-h-tiny-4bit, NVIDIA-Nemotron-3-Nano-30B-A3B-4bit (not run; SSM and gated-delta kernels are Metal-only, so ROCm runs the graph paths). See the verification note below: the Nemotron-H row settles one claim that could only be traced, not executed.
   • one VLM on the default path (not run)
   • Mixtral-8x7B-Instruct-v0.1-4bit (E=8) doubles as the guard for the ROCm expert-batched gather path that the overlay keeps opt-in: its width-256 trace goes through sorted gather_qmm with at most 64 experts.
3. Reference. Metal logit traces from an M1 Ultra host (Apple GPU generation 13, so the Metal path without NAX; record this, since M5-class hosts gated by hw.has_neural_accelerator && hw.macos_supports_na take different kernels) with the same checkpoint, corpus and arguments, built from the same mlxcel commit. Corpus tests/fixtures/wikitext2_excerpt.txt; logit_trace MODEL_DIR TEXT_FILE CHUNK_TOKENS MAX_CHUNKS TOPK PREFILL at width 1 (1 128 8 0), width 8 after 512 tokens of context (8 80 8 512) and width 256 (256 2 8 0). Note the known logit_trace defect that chunk 0 gets no BOS (fix(examples): anchor logit_trace chunk 0 with the BOS token #1785); it affects both sides equally. Done for the first four models: benchmarks/logit_traces/metal_m1u_bec64748/ and benchmarks/logit_traces/rocm_gfx1151_bec64748/, both at bec64748. The new rows need a fresh Metal trace set at whatever commit they are run on.
4. Output A/B. scripts/ab_output_equality.sh with --out-dir set explicitly (its default uses BSD-style mktemp -d -t, which GNU mktemp may reject; still true at scripts/ab_output_equality.sh:129) for greedy smoke on each model.
5. Server. Start mlxcel-server with one dense and one MoE model and exercise /v1/chat/completions (streaming and non-streaming). Done in test(rocm): verify mlxcel-server chat completions on gfx1151 #1831 (f9ece5d9): scripts/server_chat_smoke.sh passes for Meta-Llama-3.1-8B-Instruct-4bit and Qwen3-30B-A3B-4bit on the Radeon 8060S, streaming and non-streaming agree, and neither server log has an error, panic or terminate. The MoE run then needed MLXCEL_FUSED_MOE=0; since fix(rocm): guard the four remaining custom-kernel launchers #2018 and refactor(core): choose custom-kernel ports through one helper #2026 the fused MoE launcher refuses with a catchable error and the caller takes the graph path instead of aborting.
6. Store the matrix, commands and thresholds in docs/benchmark_results/ next to the benchmark page from perf(bench): ROCm support in the benchmark harness and a published gfx1151 baseline #1810. Done for the first run (docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md); extend it, or add a dated sibling, for the remaining rows.

Verification note: the Nemotron-H row settles the fused-MoE gate

This is not extra work, only a claim to record against a row already in the matrix. PR #2029 (merged as d8d34e2b) added && mlxcel_core::custom_kernels_available() to use_fused in src/models/nemotron_h.rs:1359, so that a backend with no fused-MoE port takes forward_nonfused instead of reaching the launcher. The comment above it (src/models/nemotron_h.rs:1349-1357) records it as defensive, and the same PR made the bridge's fused_moe_forward return Result (src/lib/mlxcel-core/src/lib.rs:1782), so a refusal is now catchable rather than a std::terminate.

That added term was traced, not executed. Unlike fused_xielu, whose C++ entry point returns an elementwise fallback before it resolves a port at all (src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:463, the if (!mlx::core::metal::is_available()) early return), fused_moe_forward (mlx_cxx_kernels.cpp:1356) has no such fallback, so on a portless backend the pre-change code reached run_fused_moe_two_kernel (:2306). No Nemotron-H checkpoint was available on the ROCm validation host, so whether it actually aborted there is unconfirmed.

Running the Nemotron-H row on ROCm settles it. If the pre-change behavior was an abort, the guard is a confirmed fix; if generation was already fine, the added term is merely defensive and harmless. Either way the term is safe: on Metal and CUDA custom_kernels_available() is true so nothing changes, and on ROCm it selects forward_nonfused, which is the same arm the non-quantized path already takes. Record which of the two it turned out to be in the matrix document (reproducing the pre-change behavior means building d8d34e2b^), and note it on #1814 if the port order should change as a result.

Judging cross-backend results

Metal and ROCm are different accelerators with different reduction orders, accumulation dtypes, kernel choices per width and transcendental implementations, so byte-identical logits or free-running outputs are not the criterion.

• Metric. Decided-position mismatch from scripts/compare_logit_traces.py (default --decided 2.0) at each of the three widths. Report top-1 mismatch, where the mismatched token ranks in the reference, perplexity difference and any NaN alongside it, as context rather than as the gate.
• Threshold. Not zero. Set it from a measured noise floor: the spread between backends that are already trusted on the same checkpoint, corpus and arguments (for example Metal on two GPU generations, and CUDA against Metal). Which combinations to measure is decided after the first ROCm results are in. For scale, the 81ba1c6a pin bump alone moved top-1 by up to 15 of 512 positions on unchanged Metal hardware while decided-position mismatch stayed at 0. (First results, test(rocm): add the gfx1151 correctness matrix and its gate #1826: zero disagreements on decided positions across all twelve pairs, and the largest reference gap at which the backends ever disagree is 1.125 logits, which compare_logit_traces.py now prints itself. No noise floor has been measured yet.)
• Align known systematic differences first, so only backend numerics remain: quantized checkpoints keep bf16 scales on both sides; unquantized checkpoints run on Metal with MLXCEL_KEEP_BF16=1 (Metal otherwise converts bf16 weights to f16 at load); MoE models are traced with fused MoE on and off on Metal and compared with whichever path ROCm takes after refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803; both sides use the same mlxcel commit.
• Expected sources of difference that should not by themselves block: GEMV/GEMM kernel selection per width M, reduction order and accumulation dtype, exp/sigmoid implementations (for example upstream's switch to precise::exp, Use precise::exp in Sigmoid so compiled and eager sigmoid agree ml-explore/mlx#4461), and MoE router logits whose last-digit differences select a different expert.
• Free-running greedy output is a smoke check only (fluency, no NaN, no empty output); it diverges after a few tokens even between healthy backends.

Acceptance criteria

• make verify-test-rocm exists and passes on gfx1151 with MLXCEL_REQUIRE_MODELS=1 for the checkpoints present. (Exists, Makefile:780, test(rocm): add the gfx1151 correctness matrix and its gate #1826. Does not pass: one failing target, the feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 NVFP4 abort, which also stops every later-sorting mlxcel-core test from running.)
• Every model in the matrix generates on ROCm; decided-position mismatch rates against the Metal reference are reported per model and width, and fall within the measured noise floor described above, or each outlier has a filed follow-up. (Four of nine models reported, all at zero decided-position disagreements; noise floor not measured.)
• The Nemotron-H row records whether the pre-guard code aborted on ROCm, so the use_fused gate is documented as a fix or as defensive.
• mlxcel-server serves chat completions on ROCm for one dense and one MoE model. (f9ece5d9, test(rocm): verify mlxcel-server chat completions on gfx1151 #1831)
• The matrix document lists mlxcel commit, MLX pin, source fork commit of the mlxcelverse ROCm overlay, ROCm/HIP version, gfx target and host. (All present except the fork commit, which is only implied by src/lib/mlx-cpp/patches-rocm/UPSTREAM at bec64748, 75915908; add it explicitly when the document is next extended.)

Validation

make verify-rocm
make verify-test-rocm
cargo build --release --features rocm --example logit_trace
./target/release/examples/logit_trace models/mlx/Meta-Llama-3.1-8B-Instruct-4bit corpus.txt > rocm.tsv
python3 scripts/compare_logit_traces.py metal.tsv rocm.tsv
scripts/server_chat_smoke.sh

References

• docs/benchmarks.md, examples/logit_trace.rs, scripts/compare_logit_traces.py, scripts/ab_output_equality.sh, scripts/server_chat_smoke.sh
• tests/common/mod.rs:36 (repo_model_dir), Makefile:775 (verify-test-cuda), Makefile:780 (verify-test-rocm), Makefile:804 (verify-rocm)
• docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md, TECHNICAL_REPORTS/2029-finish-kernel-port-standardization-20260929.en.md

Refresh log

2026-09-29

• Replaced the "in flight" Nemotron-H citations (branch fix/correct-sdpa-exemption-rationale, working-tree line) with the merged location, src/models/nemotron_h.rs:1359 from refactor: finish the kernel port standardization #2029 (d8d34e2b), and moved the kernel line numbers to mlx_cxx_kernels.cpp:463, :1356 and :2306; Makefile lines moved to :775, :780 and :804.
• Recorded the current gate result: one failing target, the feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 NVFP4 abort, which terminates the mlxcel-core test binary so later-sorting tests never run and cannot be counted as passes.
• Checked off the server criterion (f9ece5d9, test(rocm): verify mlxcel-server chat completions on gfx1151 #1831) and marked the matrix rows and plan items that test(rocm): add the gfx1151 correctness matrix and its gate #1826 delivered; the remaining rows are unblocked now that refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803 and feat(rocm): memory estimation, device info and hardware detection on AMD UMA hosts #1805 are closed.
• Added the rms_norm_small_axis_tests coverage from refactor: finish the kernel port standardization #2029, and noted the matrix document's missing explicit fork commit, which kept that criterion unchecked.

[inureyes]

Both deliverables landed in #1826, so closing.

The matrix is docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md, with the traces it was computed from under benchmarks/logit_traces/: metal_m1u_bec64748/ for the Metal reference and rocm_gfx1151_bec64748/ for ROCm, both at mlxcel bec64748 with MLX pin 81ba1c6a, the same corpus and arguments, and the same checkpoint revisions. Twelve model and width pairs, zero disagreements on decided positions, and the threshold-free form of the result: the largest reference gap at which the two backends ever disagree is 1.125 logits, so any --decided at or above 1.2 gives the same answer. scripts/compare_logit_traces.py now prints that number itself, which is what lets a CUDA arm reuse the wording without recomputing it by hand.

The gate is make verify-test-rocm. It is not green, and the document says so with the counts: the remaining failures and aborts are missing primitives (#1825), the cuda_kernel routing path (#1803), backend identification (#1805) and nvfp4 group 16 (#1806). None of them is a wrong-numbers defect, and each has its own issue, so this issue does not need to stay open for them.

Worth recording where the value came from: the matrix itself found no defect, while make verify-test-rocm found three real ones (#1822, #1823, #1824) plus the rpath gap that kept two crates' test binaries from starting. A logit matrix checks the decode path; the gate checks the rest.