Repository navigation
perf(rocm): batch sorted mxfp4 gather_qmm prefill by expert - #2164
Merged
Merged
Conversation
gpt-oss-20b spent 99.8% of a 512-token prefill on gfx1151 in the per-row non-affine gather_qmv_kernel (about 940 ms per call), because GatherQMM::eval_gpu sent only affine weights to the expert-batched kernel, which reads each expert's weights once per four routed rows. The expert-batched kernel gains an AFFINE=false arm for mxfp4 (group size 32, E8M0 scales, no biases): each e2m1 nibble is decoded without a branch through fp16 bits and each word's dot product is rescaled by 2^14 with the group's scale. bf16 mxfp4 reaches it under the existing gate and MLX_ROCM_GATHER_QMV_EXPERT_BATCHED switch; that is one new instantiation, and qmm.hip compile time is unchanged. LOCAL_FIXES item 31 records it. tests/rocm_gather_qmm_expert_batched.rs adds mxfp4 cases at gpt-oss-20b's expert shape and a partial-column shape, flat and broadcast layouts, against the unsorted path and a dequantized f32 reference; they fail with the nibble sign dropped. Measured on gfx1151 (guarded, alternated, six rounds): prefill 7.53 to 526.3 tok/s median, decode 8.45 and 8.49 tok/s (unchanged), and no decided-position logit_trace disagreement against the per-row path. Results: docs/benchmark_results/rocm-moe-prefill-mxfp4-gfx1151-2026-10-06.md. Closes #2106
inureyes
added
status:review
4 tasks done
…rnel Review follow-up for #2164. The mxfp4 cases passed even if the gate stopped sending mxfp4 to the expert-batched kernel, since the per-row fallback is also correct. At gpt-oss-20b's K = 2880 the test now also runs the sorted call with the kernel switched off and requires the bytes to differ; at K = 512 both kernels round to identical bf16 outputs, so the narrow shape skips the check. The two unsafe FFI calls the test rewrote get SAFETY comments, and LOCAL_FIXES items 9 and 31 note the mxfp4 scope and the fork policy. Refs #2106
This was referenced Oct 6, 2026
inureyes
added a commit
that referenced
this pull request
Oct 6, 2026
## Summary The shuffle-based HIP ports have only run on 32-lane wavefronts, and their `#error` guards on `__AMDGCN_WAVEFRONT_SIZE` are inert because AMD clang 23 defines neither macro spelling (rechecked here with `hipcc -E -dM` for gfx1151, gfx942 and gfx90a). This PR holds those ports to 32-lane devices at port selection, so a wave64 (CDNA) device takes the graph fallbacks instead of running unvalidated kernels. Nothing changes on gfx1151. ## What changed - `port_for` (`kernel_port.cpp`) returns a ROCm port only when `rocm_port_allowed(ports.rocm_any_wave_size, rocm_port_warp_size())`, so `has_kernel_port`, every `*_available()` predicate and `select_kernel_port` agree. The refusal message names the width; stderr notes a non-32 wavefront once per process. - The width is the hardware value from a new overlay call, `mlx::core::rocm::device_warp_size()` (`hipDeviceAttributeWarpSize`, ignores `MLX_ROCM_FORCE_WARP_SIZE`, 0 on error), recorded as LOCAL_FIXES item 32 (new "Additions to the fork's API" section; #2164 took 31) and stubbed in `no_rocm.cpp`. It clears only an HIP error it raised itself. - `KernelPorts::rocm_any_wave_size` (default false) is set on the five tables whose HIP body has no lane-level operation: RoPE append, paged merge, Gumbel, rejection, and xIELU. xIELU is not in the issue's list, but checking the current code shows its body is purely elementwise; its inert guard is dropped, as the sampler and RoPE bodies already have none. The other ten ROCm tables (bitlinear, ssm, mamba1_scan, moe_gateup, moe_down, moe_fc1_relu2, add3_layer_norm, fused_norm, paged_attention, paged_v2_partial) stay wave32-only. #2107's fused norm is wave32-only (shuffles) and its RoPE port is any-wave. - `bitlinear_kernel_available` now reads its port table (it returned "any GPU backend" before), so BitNet is refused at load on wave64 rather than failing mid-request. - `verify-kernel-port-dispatch`: rule 5 pins the any-wave tables with their HIP source, follows each table's `.rocm` getter to its holder and requires exactly that source (no header argument, one raw literal), and fails on lane intrinsics, AMDGCN lane builtins, wavefront macros or 32-lane index arithmetic in it; rule 6 keeps the test seam `set_rocm_port_warp_size_for_tests` out of non-test code, with an exact use count in its definition files. `scripts/ci/check_kernel_port_dispatch_test.sh` (run by the make target and CI job) proves each rule fails on a mutated copy (13 cases). - `mamba1_selective_scan` is now `-> Result<()>`, since `select_kernel_port` can refuse there; Mamba, Jamba and its parity test `expect` the gated launch. BitLinear's refusal names the BitNet load-time refusal instead of a graph fallback it lacks. - Parity tests that assert a ROCm port exists (fused MoE, relu2, SSM, fused norm, Mamba1, add3) skip on a wave64 device through `test_support::kernel_ports::skip_for_wave32_only_rocm_port` instead of failing. - `build.rs` now tracks `kernel_port.*` and `gpu_backend.*`; edits to them did not rebuild the bridge before (found when a negative check passed on a stale library). - Guard comments now name the host-side check as the protection; `docs/installation.md` gains a "Wave64 devices" row; CONTRIBUTING and the paged-test skip message mention the hold. ## Verification (gfx1151) - `cargo test --features rocm -p mlxcel-core --lib kernel_port`: 2 passed (`rocm_port_allowed` truth table, runs on every backend). - `tests/rocm_wave_size.rs`: `device_warp_size() == 32` on gfx1151; with `MLX_ROCM_FORCE_WARP_SIZE=64` in a child (the child's device log confirms the override applied) it still reads 32; all 16 port predicates are true. - `tests/rocm_wave64_port_refusal.rs`, with the reported width forced to 64 through the test seam: all 11 wave32-only predicates false, the five any-wave ones true; `bitlinear_matmul`, `fused_add_rms_norm` and `fused_moe_expert_kernel` (moe_gateup) refuse naming the width; `layers::fused_add_rms_norm` (fusion on) and `residual_add3_layer_norm` return the graph pair's bytes. With the hold disabled in `port_for`, this test fails listing all 11 tables. - Existing port tests (debug profile): fused norm, RoPE, MoE, relu2, SSM, Mamba1, Gumbel, rejection, fixed-key, paged v2, MLA, residual add3: 357 passed; 2 `cache::paged_detach` tests fail on a `debug_assert` in the dev profile only (the gate's `test-fast` profile has no debug assertions; unrelated to ports). Apertus, BitNet and Cohere2 tests in the mlxcel lib: 26 passed. `cpu_device_custom_kernel_gates`: passed. - `make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt verify-rocm-overlay`: pass. - `make verify-rocm` on the final head (merged with origin/main 08e2269): OK, 151 test suites, 11969 passed, 0 failed, 381 ignored; clippy, smoke, overlay and script gates pass. The targeted tests above were re-run after the review fixes (29 mlxcel-core, 30 Mamba/Jamba, the three integration files): all pass. - Meta-Llama-3.1-8B-Instruct-4bit, release binary under `scripts/rocm_gpu_guard.sh` (clean attempt): greedy 64 tokens with `MLXCEL_FUSED_ADD_RMSNORM=0` and `=1` give identical text (33.88 and 33.90 tok/s, single run, not a performance claim); `rocprofv3 --kernel-trace` of a sampled run shows `custom_kernel_mlxcel_fused_add_rms_norm` (wave32-only, 544 launches) and `custom_kernel_mlxcel_gumbel_max_sample` (any-wave) still selected, and no wavefront notice on stderr. Not verified: no wave64 device is available, so the refusal is proven only through the test seam. Metal and CUDA are not available on this host; the change touches the shared `KernelPorts` struct and `port_for`, whose Metal and CUDA cases are unchanged, and `bitlinear_kernel_available`, which answers the same on Metal and CUDA (both tables have ports). Closes #2147
This was referenced Oct 7, 2026
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.
gpt-oss-20b spent 99.8% of its 512-token prefill on gfx1151 in the per-row non-affine
gather_qmv_kernel(about 940 ms per call), becauseGatherQMM::eval_gpusent only affine weights to the expert-batched kernel. This gives that kernel an mxfp4 arm.gather_qmv_expert_batched_kernel(patches-rocm/.../quantized/qmm.hip) takesAFFINE=falsefor mxfp4 (group size 32, E8M0 scales, no biases). Each e2m1 nibble is decoded without a branch through fp16 bits (value times 2^-14), and each word's dot product is rescaled by 2^14 with the group's scale.M == 1,B >= 64,E <= 64,B / E >= 4) and the sameMLX_ROCM_GATHER_QMV_EXPERT_BATCHEDswitch, on by default per the perf(rocm): fix expert-batched gather_qmm for bf16 MoE prefill #2066 rule. One new instantiation; f16 mxfp4 and mxfp8 keep the per-row kernel.qmm.hipcompile time is unchanged (37.9/37.8 s against 38.1/37.9 s).LOCAL_FIXES.mditem 31 records it (not marked for upstreaming, per the 2026-10-06 fork policy); item 10 points to it.docs/benchmark_results/rocm-moe-prefill-mxfp4-gfx1151-2026-10-06.md.Verification (gfx1151)
cargo test --release --features rocm --test rocm_gather_qmm_expert_batched -- --test-threads=1: 2 passed (48 affine cases unchanged, 8 new mxfp4 cases at gpt-oss-20b's expert shape and a partial-column shape, flat and broadcast layouts). The mxfp4 relative errors match the unsorted path to four digits (1.6e-3 to 1.7e-3). With the nibble sign bit dropped the first mxfp4 case fails at 1.377. At K = 2880 the test also requires the sorted call to differ from the same call with the kernel switched off, so it fails if the gate stops reaching the kernel.mlxcel-bench-decode --prompt-tokens 512 -n 128 --warmup-tokens 20 --ignore-eos, kernel off/on alternated, six rounds, all underscripts/rocm_gpu_guard.sh: prefill median 7.53 to 526.3 tok/s (69.9x); decode median 8.45 and 8.49 tok/s, ranges overlapping, unchanged (a 16-token control where neither arm reaches the kernel shows about 1% host spread).logit_tracew8andw256, kernel on against kernel off (no gpt-oss traces exist): 0 / 88 and 0 / 81 decided-position mismatches at--decided 2.0.make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt verify-rocm-overlayandcargo test --test dead_doc_pointerspass.make verify-rocm(withMLXCEL_ROCM_SMOKE_MODEL=models/mlx/Qwen3-0.6B-4bit) on 826f9d6 plus the review commit's test change: every step through clippy and the smoke passed, andverify-test-rocmgave 11962 passed, 1 failed. The failure was a first version of the dispatch check above, which also ran at K = 512 where both kernels give identical bf16 bytes; it now runs at K = 2880 only, and the file passes (2/2) with clippy clean on the final commit.Not verified: Metal and CUDA (not available on this host). The change is in the ROCm overlay only; no Metal or CUDA file is touched.
Review: pr-reviewer and pr-security-checker found no CRITICAL, HIGH or MEDIUM issues; the LOW items (dispatch check, SAFETY comments, LOCAL_FIXES wording) are in e87ffa7. Left as is: a sorted index value of E or more leaves its output rows unwritten in this kernel, as it already did for affine; routing never produces one.
Follow-up: decode is now dominated by the same per-row mxfp4 kernel on the unsorted
B = 4path (about 100 ms of a 115 ms step), outside this issue's scope.Closes #2106