Repository navigation
update(rocm): skip the dequant LRU for GEMMs routed above the WMMA ceiling - #2236
Merged
Merged
Conversation
inureyes
added
status:review
…iling bf16 affine GEMMs at or above the fused WMMA kernel's row ceiling on RDNA 3.5 (128 rows by default) go to dequantize + hipBLASLt, and every dequantized weight went into the 8-entry, 256 MB LRU. A forward pass runs far more distinct projections than that, so prefill cycled the cache without a hit and its last entries, up to 256 MB of bf16 weights plus references to the quantized sources, stayed alive through decode. select_qmm_route now returns QmmRoute::DequantGemmAboveWmmaCeiling from the ceiling branch only. eval_gpu runs it exactly like DequantGemm but dequantizes into a temporary that the encoder frees after the GEMM, and quantized_matmul_runs_dequant_gemm accepts both enumerators. f16 and non-WMMA shapes keep the cache unchanged. The cache moved into one struct with relaxed atomic hit, miss, insert, eviction and bypass counters. rocm::dequant_cache_stats() (stubbed in no_rocm.cpp) reads them with the entry count and bytes, the bridge exposes them as mlxcel_core::rocm_qmm_cache::dequant_cache_stats(), and MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1 prints them on stderr at exit. tests/rocm_qmm_dequant_cache.rs runs one case per child process on gfx1151: bf16 above the default and an env ceiling bypasses the cache and matches dequantize + matmul bytewise, forced WMMA touches nothing, and f16 still misses, inserts and then hits. With the ceiling branch returning DequantGemm again both bf16 ceiling cases fail. Refs #2151
Adds a dated section to docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md and updates LOCAL_FIXES item 29 with what the ceiling route change measured on gfx1151 at pp2048/tg128, medians of three guarded runs on 8753883 plus one confirmation run per model on the rebased head. Gemma 3 4B 4-bit (bf16): the cache made 0 hits in 340 lookups and held 230 MiB after prefill; with the change it is bypassed 340 times and stays empty, active memory after prefill drops from 2.80 to 2.56 GB, MLX peak stays 4.52 GB, prefill and decode stay within run-to-run range, and the logit traces are byte-identical. Llama 3.1 8B 4-bit (f16) is unchanged and also gets 0 hits while holding 224 MiB; turning that default off is tracked in #2232. Refs #2151
inureyes
force-pushed
the
update/issue-2151-skip-dequant-lru
branch
from
October 8, 2026 11:38
b41cbfb to
79a8c9b
Compare
Bilingual report for the dequant LRU skip above the WMMA ceiling, with the full make verify-rocm result: 162 suites, 12191 passed, 0 failed, 399 ignored.
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
bf16 affine GEMMs that
select_qmm_routesends to dequantize + hipBLASLt because they are at or above the WMMA row ceiling (128 rows on RDNA 3.5) no longer go through the dequantized-weight LRU. A forward pass runs more distinct projections than the 8-entry, 256 MB cache holds, so it never hit, and its last entries stayed alive through decode.QmmRoute::DequantGemmAboveWmmaCeiling, returned only from the ceiling branch.QuantizedMatmul::eval_gpuruns the same dequantize and GEMM but into a temporary freed after the GEMM;quantized_matmul_runs_dequant_gemmaccepts both enumerators. f16 and non-WMMA shapes keep the cache unchanged.rocm::dequant_cache_stats()in the overlay (stubbed inno_rocm.cpp),mlxcel_core::rocm_qmm_cache::dequant_cache_stats()through the bridge, andMLX_ROCM_QMM_DEQUANT_CACHE_STATS=1to print them at exit.tests/rocm_qmm_dequant_cache.rs, LOCAL_FIXES item 29,docs/environment-variables.md, and a dated section indocs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md.Measurements (gfx1151, pp2048/tg128, medians of 3 on 8753883, each run under
rocm_gpu_guard.sh --idle-secs 60)Prefill and decode differences are inside the run-to-run range; no speedup is claimed. Before counts come from a counters-only build that routes like 8753883. One confirmation run per model on the rebased head (main ad84435) matched the after rows. The f16 cache also never hits, so #2232 tracks turning its default off.
Verification
All on gfx1151 (Radeon 8060S), ROCm 10.0.0 / HIP 7.15. The two test binaries and the overlay, dtype-key, port-dispatch and fmt gates were rerun on the final head (rebased onto main ae343d8):
cargo test --release --features rocm --test rocm_qmm_dequant_cache -- --test-threads=1: pass. With the ceiling branch returningDequantGemmagain, the two bf16 ceiling cases fail (misses=1 inserts=1 bypasses=0 entries=1on the first pass); the forced-WMMA and f16 cases still pass.cargo test --release --features rocm --test rocm_qmm_env -- --test-threads=1: pass.examples/logit_trace, Gemma 3 4B, 4 x 256-token chunks) before and after: byte-identical;scripts/compare_logit_traces.py --decided 2.0reports 0 / 1024 disagreements.make verify-rocm-overlay verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt,cargo test --test dead_doc_pointers,cargo clippy -p mlxcel-core --features rocm --lib --tests -- -D warnings,cargo clippy -p mlxcel --features rocm --test rocm_qmm_dequant_cache -- -D warnings(before the rebase): pass.make verify-rocmwithMLXCEL_ROCM_SMOKE_MODEL=models/mlx/Qwen3-0.6B-4bitat79a8c9bbon6dbfe7d7(the later commit only adds the technical report):[verify-rocm] OK, 162 cargo test suites, 12191 passed, 0 failed, 399 ignored; the ROCm smoke generated 32 tokens on the GPU. An earlier run on theae343d84base was also green (12190 passed, 0 failed) before main gained update(rocm): route top-k with top-p up to vocab 152064 on ROCm #2218, which touches the bridge files.MLXCEL_BRIDGE_ROCM_BACKEND); no CRITICAL or HIGH findings.Not verified: Metal and CUDA (not available on this host). The change is in the ROCm overlay and the bridge function is stubbed to zeros off ROCm (
#ifdef MLXCEL_BRIDGE_ROCM_BACKEND); the unit teststats_are_zero_off_rocmcovers that stub but has not been built or run on those backends.Closes #2151