diff --git a/TECHNICAL_REPORTS/2236-rocm-dequant-lru-skip-above-wmma-ceiling-20261008.en.md b/TECHNICAL_REPORTS/2236-rocm-dequant-lru-skip-above-wmma-ceiling-20261008.en.md new file mode 100644 index 000000000..62ffa9b15 --- /dev/null +++ b/TECHNICAL_REPORTS/2236-rocm-dequant-lru-skip-above-wmma-ceiling-20261008.en.md @@ -0,0 +1,156 @@ +# Technical Report: PR #2236 - Skip the Dequant LRU for GEMMs Routed Above the WMMA Ceiling + +**Date**: 2026-10-08 + +**Status**: Implemented and verified on the gfx1151 host (ROCm/HIP); branch commits `8964f610` (code) and `79a8c9bb` (docs), rebased onto origin/main `6dbfe7d7`, merged after the full `make verify-rocm` gate. Closes #2151. Follow-up: #2232. + +**Languages**: HIP/C++ (ROCm overlay `qmm.hip`, `rocm.h`, `no_rocm.cpp`; mlxcel-core bridge `mlx_cxx_bridge.cpp`, `mlx_cxx_bridge.h`), Rust (mlxcel-core `lib.rs`, new `rocm_qmm_cache.rs`; new `tests/rocm_qmm_dequant_cache.rs`), Markdown (`LOCAL_FIXES.md`, `docs/environment-variables.md`, `docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md`) + +**Risk Level**: Low. The ceiling-routed GEMM runs the same dequantize kernel and the same hipBLASLt GEMM as before; only the cache lookup and insert are skipped, and the logit traces are byte-identical. The f16 and non-WMMA callers of the cached route keep the cache unchanged. The new counters are relaxed atomics bumped once per GEMM on the dequantize route, and the cache refactor moves the existing statics into one struct without changing the eviction rule. + +## Executive Summary + +Since #2085, a bf16 affine GEMM of 128 rows or more on RDNA 3.5 leaves the fused WMMA kernel for dequantize plus hipBLASLt. That route stored every dequantized weight in an LRU of 8 matrices or 256 MB. A forward pass runs far more distinct projections than 8, so prefill cycled the cache without a single hit, and the last entries (up to 256 MB of bf16 weights, plus references to the quantized sources) stayed alive through decode. + +PR #2236 adds `QmmRoute::DequantGemmAboveWmmaCeiling`, returned only from the ceiling branch of `select_qmm_route`. `QuantizedMatmul::eval_gpu` runs it like `DequantGemm` but dequantizes into a temporary that the command encoder frees after the GEMM. The cache now lives in one struct with hit, miss, insert, eviction and bypass counters, readable through `rocm::dequant_cache_stats()`, the bridge, and `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1`. + +On gfx1151 with Gemma 3 4B 4-bit (bf16) at pp2048/tg128, the cache made 0 hits in 340 lookups and held 6 entries (230 MiB) at exit before the change. After it, the GEMMs bypass the cache 340 times, the cache stays empty, and active memory after prefill drops from 2.80 GB to 2.56 GB. MLX peak memory stays at 4.52 GB and prefill and decode stay within run-to-run noise. Llama 3.1 8B 4-bit (f16 scales) is unchanged and also gets 0 hits in 320 lookups while holding 224 MiB; turning that default off is #2232. + +## 1. Problem Statement + +### 1.1 The ceiling route and its cache + +`select_qmm_route` in `src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip` decides how a quantized matmul runs. For a WMMA-eligible bf16 shape, it computes `hipblaslt_is_faster = env != 1 && dequant && M >= wmma_qmm_max_m(d) && !fp8()`, where `wmma_qmm_max_m` is 128 on RDNA 3.5. Before this PR that branch returned `QmmRoute::DequantGemm`, the same enumerator that f16 checkpoints and non-WMMA shapes reach through `if (dequant)`. + +In `QuantizedMatmul::eval_gpu`, `DequantGemm` looked up a process-wide LRU keyed by the weight, scales and bias pointers (`DequantCacheKey`). On a miss it dequantized the weight and inserted it whenever `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE` (default 8) and `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES` (default 256 MB) were both non-zero. Each entry also held references to the quantized source arrays. + +### 1.2 Why the cache never hits + +A prefill touches each projection once per forward pass, and a model has many more distinct projections than the cache holds (a Llama-class 8B has 32 layers of 7 projections). Every lookup is a miss, every insert evicts an older entry, and nothing is looked up again before it is evicted. The miss is structural: no key ordering or hashing change fixes it. Decode at batch 1 runs one row, takes qmv, and never reaches the cache, so the entries left at the end of prefill are never used. + +`LOCAL_FIXES.md` item 29 and the #2085 technical report already recorded that the cache gets no hits in prefill and keeps up to 256 MB alive afterwards. The memory cost had not been measured. + +### 1.3 The measured cost + +With counters added (a counters-only build that routes and caches as `87538835` does), Gemma 3 4B 4-bit at pp2048/tg128 made 340 lookups: 0 hits, 340 misses, 340 inserts, 334 evictions, and 6 entries (241,172,480 bytes, 230 MiB) still cached at exit. Those 230 MiB stayed allocated from the end of prefill through decode. + +## 2. Change Summary + +| Area | Change | +|---|---| +| `patches-rocm/.../quantized/qmm.hip` | New `QmmRoute::DequantGemmAboveWmmaCeiling`, returned from the ceiling branch only. `eval_gpu` accepts it in the dequantize + GEMM arm; `use_cache` is true only for `DequantGemm`, so the new route dequantizes into a temporary and bumps `bypasses`. `quantized_matmul_runs_dequant_gemm` accepts both enumerators. The cache statics move into `struct DequantCache` with `evict_to()`; `DequantCacheCounters` holds relaxed atomic hit, miss, insert, eviction and bypass counters; `print_dequant_cache_stats()` is registered with `std::atexit` when `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1` | +| `patches-rocm/.../rocm.h`, `no_rocm.cpp` | `struct DequantCacheStats` and `MLX_API dequant_cache_stats()`; the stub returns zeros | +| `mlxcel-core/cpp/mlx_cxx_bridge.{cpp,h}` | `rocm_dequant_cache_stats()` returns the seven values as a `rust::Vec`, seven zeros without `MLXCEL_BRIDGE_ROCM_BACKEND` | +| `mlxcel-core/src/lib.rs`, `rocm_qmm_cache.rs` (new) | FFI declaration and `pub mod rocm_qmm_cache` with `DequantCacheStats` and `dequant_cache_stats()`; unit test `stats_are_zero_off_rocm` | +| `tests/rocm_qmm_dequant_cache.rs` (new) | Four cases, one child process each, on ROCm | +| `LOCAL_FIXES.md` item 29, `docs/environment-variables.md`, `docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md` | New behavior, the stats variable, and the dated measurement section | + +Two commits: the code change (`8964f610`) and the measurements (`79a8c9bb`). 11 files, 812 insertions and 53 deletions. + +## 3. Design + +### 3.1 A separate enumerator, not a flag on the call + +The ceiling branch is the only place that knows the GEMM came from a WMMA shape sent away by row count. A new enumerator carries that fact to `eval_gpu` without a second set of conditions there. Everything the two routes share stays shared: the same arm handles batch-shape checks, the dequantize kernel, the hipBLASLt GEMM and its rocBLAS fallback. Only one condition differs: + +```cpp +const bool use_cache = route == QmmRoute::DequantGemm && cache_cap > 0 && + cache_max_bytes > 0; +``` + +With `use_cache` false, the weight takes the existing cache-off path: `w_dequant` gets a fresh allocation, and `enc.add_temporary(w_dequant)` keeps it alive until the GEMM completes. This is the path `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE=0` already used. + +`quantized_matmul_runs_dequant_gemm`, which mlxcel's dense prefill path uses to know whether `quantized_matmul` and `dequantize` + `matmul` give the same bytes, is the other reader of the route. It now accepts both enumerators, so its answer for ceiling-routed shapes is unchanged. + +### 3.2 Scope: only the ceiling route + +The issue asked to leave the other `DequantGemm` callers alone and measure them first. f16 checkpoints and non-WMMA shapes still use the cache in this PR. The measurement below shows the f16 cache also never hits, but changing its default is a separate decision with its own case to check (#2156, f16 batch-4 decode on `DequantGemm`, where a decode loop reusing the same projections could hit), so it went to #2232. + +Edge cases follow from the route logic: `MLX_ROCM_WMMA_QMM=1` sets `env == 1`, which never takes the ceiling branch, so the fused kernel runs and neither the cache nor the bypass counter moves. An `MLX_ROCM_WMMA_QMM_MAX_M` override changes the ceiling and still produces the new enumerator. The fp8 route is unaffected. + +### 3.3 Counters and the stats hook + +The issue's acceptance criteria needed a way to assert the cache is empty after a bf16 prefill. The cache's `static` locals inside `eval_gpu` could not be read from outside, so they moved into `struct DequantCache` (mutex, LRU list, byte count, entry map) behind `dequant_cache()`, and the two copies of the eviction loop became one `evict_to()` method. Eviction behavior is the same; `evict_to()` also bumps the eviction counter. + +The counters are `std::atomic` with relaxed ordering. They are bumped once per GEMM on the dequantize route and never read on the hot path. `read_dequant_cache_stats()` loads them and reads the entry count and bytes under the cache mutex. Three readers sit on top: + +- `rocm::dequant_cache_stats()` in `rocm.h`, stubbed in `no_rocm.cpp`. +- `mlxcel_core::rocm_qmm_cache::dequant_cache_stats()` through the bridge, which returns zeros off ROCm. +- `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1`, which registers an `std::atexit` handler that prints one `[ROCm] qmm dequant cache: ...` line on stderr. The handler is registered inside `dequant_cache()` after the static cache is constructed, so it runs before the cache is destroyed. The ceiling route calls `dequant_cache()` before counting the bypass, so the report is registered even when the cache itself is never used. + +### 3.4 The test + +`tests/rocm_qmm_dequant_cache.rs` runs each case in a fresh child process, because the cache, its counters and the environment knobs are process-wide and read once. The parent test runs the children one after another so they never share the GPU. Each child sets `MLX_ROCM_QMM_DEQUANT_M_THRESHOLD=1` so the dequantize route is eligible at any row count, runs one quantized GEMM twice, compares each result with an f32 `dequantize` + `matmul` on the CPU stream (relative tolerance 2e-2), and checks the counters after each pass: + +- `bf16_above_default_ceiling` (256 rows, 4096x4096): 1 then 2 bypasses, an empty cache, and output bytes equal to `dequantize` + `matmul` on the GPU. +- `bf16_above_env_ceiling` (64 rows, `MLX_ROCM_WMMA_QMM_MAX_M=32`): the same with an environment ceiling. +- `bf16_forced_wmma` (`MLX_ROCM_WMMA_QMM=1`): no counter moves. +- `f16_cached` (128 rows): a miss and an insert on the first pass, a hit on the second, one entry of `n * k * 2` bytes. + +With the ceiling branch returning `DequantGemm` again, the two bf16 ceiling cases fail on their first pass (`misses=1 inserts=1 bypasses=0 entries=1`), and the forced-WMMA and f16 cases still pass: 2 of 4 cases fail. + +## 4. Production Impact + +Measured on gfx1151 (Radeon 8060S), `scripts/bench_decode.sh --prompt-tokens 2048 --max-tokens 128`, before on main `87538835` and after on that commit plus the change, medians of 3 (range), each run under `scripts/rocm_gpu_guard.sh --idle-secs 60` with arms alternated: + +| Model | Scales | Arm | Prefill tok/s | Decode tok/s | MLX peak | Active after prefill | Active at end | +|---|---|---|---:|---:|---:|---:|---:| +| gemma-3-4b-it-4bit | bf16 | before | 2872 (2819-2910) | 62.75 (62.68-62.76) | 4.52 GB | 2.80 GB | 2.97 GB | +| gemma-3-4b-it-4bit | bf16 | after | 2868 (2839-2945) | 62.71 (62.55-62.73) | 4.52 GB | 2.56 GB | 2.73 GB | +| Meta-Llama-3.1-8B-Instruct-4bit | f16 | before | 1136 (1106-1136) | 36.63 (32.36-36.84) | 6.92 GB | 4.75 GB | 4.75 GB | +| Meta-Llama-3.1-8B-Instruct-4bit | f16 | after | 1132 (1048-1143) | 36.67 (36.58-36.92) | 6.92 GB | 4.75 GB | 4.75 GB | + +Cache counters (warmup and measured prefill together; before from the counters-only build): + +| Model | Arm | Hits | Misses | Inserts | Evictions | Bypassed | Entries at exit | Bytes at exit | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| gemma-3-4b-it-4bit | before | 0 | 340 | 340 | 334 | 0 | 6 | 241,172,480 | +| gemma-3-4b-it-4bit | after | 0 | 0 | 0 | 0 | 340 | 0 | 0 | +| Meta-Llama-3.1-8B-Instruct-4bit | before | 0 | 320 | 320 | 318 | 0 | 2 | 234,881,024 | +| Meta-Llama-3.1-8B-Instruct-4bit | after | 0 | 320 | 320 | 318 | 0 | 2 | 234,881,024 | + +For bf16 checkpoints on RDNA 3.5, active memory after prefill and at the end of the run drops by 0.24 GB, the size of the cache's leftover entries. MLX peak memory does not move, because the peak comes from the per-call dequantized copies inside prefill, which the cache never removed (item 28 bounds those). Prefill and decode are within the run-to-run range; no speedup is claimed. Llama 3.1 8B is the f16 control: its GEMMs never take the ceiling route, so it reads the same in both arms. + +One confirmation run per model on the change rebased onto main `ad844354` matched the after rows: Gemma 3 4B 2885 tok/s prefill, 62.14 tok/s decode, 4.52 GB peak, 2.56 GB active after prefill, 340 bypasses and an empty cache; Llama 3.1 8B 1119 tok/s prefill, 36.66 tok/s decode, 6.92 GB peak, 4.75 GB active, 320 misses, 0 hits and 2 entries. + +Metal and CUDA are unaffected: the change is in the ROCm overlay, and the bridge function returns zeros off ROCm. + +## 5. Documentation + +- **`LOCAL_FIXES.md` item 29.** The sentence that said the LRU gets no hits in prefill but keeps its last entries alive now describes `QmmRoute::DequantGemmAboveWmmaCeiling`, the Gemma 3 4B numbers (0 hits in 340 lookups and 230 MiB cached before; 340 bypasses, an empty cache, 2.80 to 2.56 GB active after prefill, peak unchanged after), the counters, `dequant_cache_stats()`, `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1`, and the new test. +- **`docs/environment-variables.md`.** The `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE` text now says bf16 GEMMs sent above `MLX_ROCM_WMMA_QMM_MAX_M` never use the cache, and `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1` is documented. +- **`docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md`.** A new section, "2026-10-07: the dequantized-weight cache above the WMMA ceiling (issue #2151)", with the method, both tables, the logit check and the confirmation run. It notes that the third Gemma before run was taken on 2026-10-08 after the first failed in the harness, and that Gemma's end-of-run active memory on the rebased head (2.56 GB against 2.73 GB on `87538835`) moved with main and was not isolated. + +## 6. Verification + +On gfx1151 (Radeon 8060S), ROCm 10.0.0 / HIP 7.15: + +- **New test.** `cargo test --release --features rocm --test rocm_qmm_dequant_cache -- --test-threads=1`: pass. With the ceiling branch returning `DequantGemm` again, 2 of 4 cases fail (the two bf16 ceiling cases, on their first pass); the forced-WMMA and f16 cases still pass. +- **Route env test.** `cargo test --release --features rocm --test rocm_qmm_env -- --test-threads=1`: pass. +- **Logits.** `examples/logit_trace` on gemma-3-4b-it-4bit (README.md as the corpus, 4 chunks of 256 tokens, so every projection runs at 256 rows and takes the ceiling route), before and after: byte-identical traces; `scripts/compare_logit_traces.py --decided 2.0` reports 0 of 1024 disagreements and an identical perplexity. +- **Targeted gates.** `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` and `cargo clippy -p mlxcel --features rocm --test rocm_qmm_dequant_cache -- -D warnings`: pass. +- **Full gate.** `make verify-rocm` with `MLXCEL_ROCM_SMOKE_MODEL=models/mlx/Qwen3-0.6B-4bit` at head `79a8c9bb` on `6dbfe7d7`: `[verify-rocm] OK`, 162 cargo test suites with 12191 passed, 0 failed and 399 ignored (8906 passed and 165 ignored in the mlxcel-core lib suite), and the ROCm smoke generated 32 tokens on the GPU. An earlier run on the `ae343d84` base was also green (12190 passed, 0 failed); main then gained #2218, which touches the bridge files, so the branch was rebased and the gate re-run. + +Not verified: Metal and CUDA, which are not available on this host. The bridge stub is covered by the unit test `stats_are_zero_off_rocm`, which has not been built or run on those backends. + +## 7. Technical Decisions + +- **Skip the cache only for the ceiling route.** That is the route the issue measured and the one where the cache's miss pattern was known. The f16 route also showed 0 hits, but #2156's decode-like case could hit, so its default change is tracked in #2232 rather than folded in here. +- **New enumerator over a global cache-size change.** Shrinking `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES` globally would keep the useless inserts and still pin memory; the issue rejected it. Changing the key was rejected too, since the miss is structural. +- **Reuse the cache-off path.** The temporary allocation and `add_temporary` already existed for `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE=0`, so the new route adds no new memory-management code. +- **Counters always on.** One relaxed atomic increment per GEMM is negligible next to a dequantize and a GEMM, and having them always present means a benchmark can read them with one environment variable and no rebuild. +- **One child process per test case.** The cache, the counters and the knobs are process-wide and read once, so each case needs a fresh process to start from zero. + +## 8. Residual Risks and Follow-ups + +- **#2232: the f16 cache.** Llama 3.1 8B (f16 scales) made 320 lookups with 0 hits, 320 inserts and 318 evictions, and kept 2 entries (224 MiB) alive after the run, identical in all three runs. #2232 proposes turning `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE` off by default on ROCm, after checking #2156's decode-like case with the counters. +- **Ceiling measured on gfx1151 only.** The change follows wherever the ceiling route is taken, but the route and these numbers were measured on gfx1151. +- **End-of-run memory difference between bases.** Gemma's active memory at the end of the run was 2.56 GB on the rebased head against 2.73 GB on `87538835` plus the change. Main moved between the two bases and the difference was not isolated. +- **Off-ROCm stub not exercised.** `stats_are_zero_off_rocm` covers the stub but has not run on Metal or CUDA. + +## 9. Learning Points + +- **A cache needs a hit rate before it earns its memory.** The LRU was added for reuse that a forward pass at these row counts never gives. Counters made that visible in one run: 0 hits in 340 lookups. +- **Peak and resident memory are different numbers.** Skipping the cache lowered active memory after prefill by 0.24 GB and left the MLX peak unchanged, because the peak comes from per-call transients the cache never removed. Measuring only peak would have shown no effect. +- **Carry routing facts in the route.** Adding an enumerator at the point that knows why the GEMM was rerouted kept `eval_gpu` to a one-line difference and kept the other reader of the route, `quantized_matmul_runs_dequant_gemm`, correct with a one-line change. +- **Make the before state measurable on the old routing.** The before counts came from a counters-only build that routes like `87538835`, which separated the effect of the counters from the effect of the route change. diff --git a/TECHNICAL_REPORTS/2236-rocm-dequant-lru-skip-above-wmma-ceiling-20261008.ko.md b/TECHNICAL_REPORTS/2236-rocm-dequant-lru-skip-above-wmma-ceiling-20261008.ko.md new file mode 100644 index 000000000..afae744e2 --- /dev/null +++ b/TECHNICAL_REPORTS/2236-rocm-dequant-lru-skip-above-wmma-ceiling-20261008.ko.md @@ -0,0 +1,156 @@ +# 기술 보고서: PR #2236 - WMMA 상한 위로 라우팅된 GEMM은 Dequant LRU를 건너뜀 + +**날짜**: 2026-10-08 + +**상태**: gfx1151 호스트(ROCm/HIP)에서 구현 및 검증 완료. 브랜치 커밋 `8964f610`(코드)와 `79a8c9bb`(문서), origin/main `6dbfe7d7` 위로 리베이스, 전체 `make verify-rocm` 게이트 통과 후 머지. #2151을 닫음. 후속: #2232. + +**언어**: HIP/C++(ROCm 오버레이 `qmm.hip`, `rocm.h`, `no_rocm.cpp`; mlxcel-core 브리지 `mlx_cxx_bridge.cpp`, `mlx_cxx_bridge.h`), Rust(mlxcel-core `lib.rs`, 신규 `rocm_qmm_cache.rs`; 신규 `tests/rocm_qmm_dequant_cache.rs`), Markdown(`LOCAL_FIXES.md`, `docs/environment-variables.md`, `docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md`) + +**위험도**: 낮음. 상한으로 라우팅된 GEMM은 이전과 같은 dequantize 커널과 같은 hipBLASLt GEMM을 실행합니다. 캐시 조회와 삽입만 건너뛰며, logit trace는 바이트 단위로 동일합니다. 캐시 경로를 쓰는 f16과 비 WMMA 호출자는 캐시를 그대로 씁니다. 새 카운터는 dequantize 경로의 GEMM마다 한 번 올리는 relaxed atomic이고, 캐시 리팩터링은 기존 static들을 구조체 하나로 옮길 뿐 축출 규칙은 바꾸지 않습니다. + +## 요약 + +#2085 이후 RDNA 3.5에서 128행 이상의 bf16 affine GEMM은 fused WMMA 커널을 떠나 dequantize + hipBLASLt로 갑니다. 이 경로는 dequantize한 가중치를 모두 8개 행렬 또는 256 MB 크기의 LRU에 저장했습니다. forward pass 한 번이 실행하는 서로 다른 projection은 8개보다 훨씬 많으므로, prefill은 한 번도 hit 없이 캐시를 순환시켰고, 마지막 항목들(bf16 가중치 최대 256 MB와 양자화 원본에 대한 참조)은 decode 내내 살아 있었습니다. + +PR #2236은 `select_qmm_route`의 상한 분기에서만 반환되는 `QmmRoute::DequantGemmAboveWmmaCeiling`을 추가합니다. `QuantizedMatmul::eval_gpu`는 이 경로를 `DequantGemm`처럼 실행하되, command encoder가 GEMM 뒤에 해제하는 임시 버퍼로 dequantize합니다. 캐시는 이제 hit, miss, insert, eviction, bypass 카운터를 가진 구조체 하나에 있고, `rocm::dequant_cache_stats()`, 브리지, `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1`로 읽을 수 있습니다. + +gfx1151에서 Gemma 3 4B 4-bit(bf16)를 pp2048/tg128로 돌렸을 때, 변경 전 캐시는 조회 340번에 hit 0번이었고 종료 시 항목 6개(230 MiB)를 들고 있었습니다. 변경 후 GEMM은 캐시를 340번 bypass하고, 캐시는 비어 있으며, prefill 후 active memory는 2.80 GB에서 2.56 GB로 줄었습니다. MLX peak memory는 4.52 GB 그대로이고 prefill과 decode는 실행 간 노이즈 범위 안입니다. Llama 3.1 8B 4-bit(f16 scales)는 변하지 않았고, 역시 조회 320번에 hit 0번이면서 224 MiB를 들고 있습니다. 그 기본값을 끄는 작업은 #2232입니다. + +## 1. 문제 정의 + +### 1.1 상한 경로와 그 캐시 + +`src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip`의 `select_qmm_route`는 양자화 matmul의 실행 방식을 정합니다. WMMA 대상인 bf16 shape에 대해 `hipblaslt_is_faster = env != 1 && dequant && M >= wmma_qmm_max_m(d) && !fp8()`를 계산하며, `wmma_qmm_max_m`은 RDNA 3.5에서 128입니다. 이 PR 전에는 이 분기가 `QmmRoute::DequantGemm`을 반환했는데, 이는 f16 체크포인트와 비 WMMA shape가 `if (dequant)`를 통해 도달하는 것과 같은 열거자입니다. + +`QuantizedMatmul::eval_gpu`에서 `DequantGemm`은 가중치, scales, bias 포인터로 키를 만든(`DequantCacheKey`) 프로세스 전역 LRU를 조회했습니다. miss가 나면 가중치를 dequantize하고, `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE`(기본 8)와 `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`(기본 256 MB)가 모두 0이 아니면 삽입했습니다. 각 항목은 양자화 원본 배열에 대한 참조도 들고 있었습니다. + +### 1.2 캐시가 hit하지 않는 이유 + +prefill은 forward pass마다 각 projection을 한 번씩 건드리고, 모델에는 캐시가 담을 수 있는 것보다 훨씬 많은 projection이 있습니다(Llama 계열 8B는 32개 레이어에 각 7개). 모든 조회는 miss이고, 모든 삽입은 오래된 항목을 축출하며, 축출되기 전에 다시 조회되는 항목은 없습니다. miss는 구조적이므로 키 순서나 해시를 바꿔도 해결되지 않습니다. 배치 1의 decode는 한 행을 실행해 qmv를 타고 캐시에 도달하지 않으므로, prefill 끝에 남은 항목은 쓰이지 않습니다. + +`LOCAL_FIXES.md` 항목 29와 #2085 기술 보고서는 캐시가 prefill에서 hit하지 않으면서 이후에도 최대 256 MB를 살려 둔다고 이미 기록했습니다. 메모리 비용은 측정되지 않았습니다. + +### 1.3 측정된 비용 + +카운터를 추가한 빌드(`87538835`와 똑같이 라우팅하고 캐시하는 카운터 전용 빌드)에서 Gemma 3 4B 4-bit는 pp2048/tg128에서 조회 340번을 했습니다. hit 0, miss 340, insert 340, eviction 334이고, 종료 시 항목 6개(241,172,480 바이트, 230 MiB)가 캐시에 남아 있었습니다. 이 230 MiB는 prefill 끝부터 decode 내내 할당된 채로 있었습니다. + +## 2. 변경 요약 + +| 영역 | 변경 | +|---|---| +| `patches-rocm/.../quantized/qmm.hip` | 상한 분기에서만 반환되는 `QmmRoute::DequantGemmAboveWmmaCeiling` 신설. `eval_gpu`는 dequantize + GEMM 분기에서 이를 받음. `use_cache`는 `DequantGemm`일 때만 true이므로 새 경로는 임시 버퍼로 dequantize하고 `bypasses`를 올림. `quantized_matmul_runs_dequant_gemm`은 두 열거자를 모두 받음. 캐시 static들은 `evict_to()`를 가진 `struct DequantCache`로 이동. `DequantCacheCounters`가 relaxed atomic hit, miss, insert, eviction, bypass 카운터를 가짐. `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1`이면 `print_dequant_cache_stats()`를 `std::atexit`에 등록 | +| `patches-rocm/.../rocm.h`, `no_rocm.cpp` | `struct DequantCacheStats`와 `MLX_API dequant_cache_stats()`. stub은 0을 반환 | +| `mlxcel-core/cpp/mlx_cxx_bridge.{cpp,h}` | `rocm_dequant_cache_stats()`가 일곱 값을 `rust::Vec`로 반환. `MLXCEL_BRIDGE_ROCM_BACKEND`가 없으면 0 일곱 개 | +| `mlxcel-core/src/lib.rs`, `rocm_qmm_cache.rs`(신규) | FFI 선언과 `DequantCacheStats`, `dequant_cache_stats()`를 담은 `pub mod rocm_qmm_cache`. 단위 테스트 `stats_are_zero_off_rocm` | +| `tests/rocm_qmm_dequant_cache.rs`(신규) | ROCm에서 케이스 4개를 각각 자식 프로세스 하나로 실행 | +| `LOCAL_FIXES.md` 항목 29, `docs/environment-variables.md`, `docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md` | 새 동작, stats 변수, 날짜가 붙은 측정 섹션 | + +커밋은 두 개입니다. 코드 변경(`8964f610`)과 측정 기록(`79a8c9bb`)입니다. 파일 11개, 812줄 추가, 53줄 삭제. + +## 3. 설계 + +### 3.1 호출 플래그가 아니라 별도 열거자 + +GEMM이 WMMA shape였다가 행 수 때문에 다른 경로로 보내졌다는 사실을 아는 곳은 상한 분기뿐입니다. 새 열거자는 `eval_gpu`에 두 번째 조건 묶음을 두지 않고 그 사실을 전달합니다. 두 경로가 공유하는 것은 그대로 공유합니다. 같은 분기가 배치 shape 검사, dequantize 커널, hipBLASLt GEMM과 그 rocBLAS 폴백을 처리합니다. 다른 조건은 하나뿐입니다. + +```cpp +const bool use_cache = route == QmmRoute::DequantGemm && cache_cap > 0 && + cache_max_bytes > 0; +``` + +`use_cache`가 false이면 가중치는 기존의 캐시 꺼짐 경로를 탑니다. `w_dequant`가 새로 할당되고, `enc.add_temporary(w_dequant)`가 GEMM이 끝날 때까지 이를 살려 둡니다. `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE=0`이 이미 쓰던 경로입니다. + +경로의 다른 독자는 `quantized_matmul_runs_dequant_gemm`이며, mlxcel의 dense prefill 경로는 이 함수로 `quantized_matmul`과 `dequantize` + `matmul`이 같은 바이트를 내는지 판단합니다. 이제 두 열거자를 모두 받으므로, 상한으로 라우팅된 shape에 대한 답은 바뀌지 않습니다. + +### 3.2 범위: 상한 경로만 + +이슈는 다른 `DequantGemm` 호출자는 그대로 두고 먼저 측정하라고 했습니다. f16 체크포인트와 비 WMMA shape는 이 PR에서도 캐시를 씁니다. 아래 측정에서 f16 캐시도 hit하지 않는 것으로 나왔지만, 그 기본값을 바꾸는 것은 따로 확인할 경우가 있는 별개의 결정이므로(#2156, `DequantGemm`을 타는 f16 배치 4 decode로, 같은 projection을 매 스텝 재사용하는 decode 루프는 hit할 수 있음) #2232로 넘겼습니다. + +경계 사례는 라우팅 로직에서 바로 나옵니다. `MLX_ROCM_WMMA_QMM=1`은 `env == 1`이 되어 상한 분기를 타지 않으므로, fused 커널이 실행되고 캐시도 bypass 카운터도 움직이지 않습니다. `MLX_ROCM_WMMA_QMM_MAX_M` 오버라이드는 상한을 바꾸고, 여전히 새 열거자를 만듭니다. fp8 경로는 영향이 없습니다. + +### 3.3 카운터와 stats 훅 + +이슈의 완료 조건은 bf16 prefill 후 캐시가 비어 있음을 단언할 방법이 필요했습니다. `eval_gpu` 안의 캐시 `static` 지역 변수는 밖에서 읽을 수 없었으므로, `dequant_cache()` 뒤의 `struct DequantCache`(뮤텍스, LRU 리스트, 바이트 수, 항목 맵)로 옮겼고, 두 벌이던 축출 루프는 `evict_to()` 메서드 하나가 되었습니다. 축출 동작은 같으며, `evict_to()`가 eviction 카운터도 올립니다. + +카운터는 relaxed ordering의 `std::atomic`입니다. dequantize 경로의 GEMM마다 한 번 올리고, hot path에서는 읽지 않습니다. `read_dequant_cache_stats()`가 카운터를 읽고, 캐시 뮤텍스 아래에서 항목 수와 바이트 수를 읽습니다. 그 위에 독자가 셋 있습니다. + +- `rocm.h`의 `rocm::dequant_cache_stats()`. `no_rocm.cpp`에 stub이 있습니다. +- 브리지를 거치는 `mlxcel_core::rocm_qmm_cache::dequant_cache_stats()`. ROCm이 아니면 0을 반환합니다. +- `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1`. stderr에 `[ROCm] qmm dequant cache: ...` 한 줄을 출력하는 `std::atexit` 핸들러를 등록합니다. 핸들러는 static 캐시가 생성된 뒤 `dequant_cache()` 안에서 등록되므로 캐시가 소멸되기 전에 실행됩니다. 상한 경로는 bypass를 세기 전에 `dequant_cache()`를 호출하므로, 캐시 자체를 한 번도 쓰지 않아도 보고가 등록됩니다. + +### 3.4 테스트 + +캐시, 카운터, 환경 변수 노브가 프로세스 전역이고 한 번만 읽히므로, `tests/rocm_qmm_dequant_cache.rs`는 각 케이스를 새 자식 프로세스에서 실행합니다. 부모 테스트는 자식을 하나씩 차례로 실행해 GPU를 공유하지 않게 합니다. 각 자식은 `MLX_ROCM_QMM_DEQUANT_M_THRESHOLD=1`로 어떤 행 수에서든 dequantize 경로를 대상으로 만들고, 양자화 GEMM 하나를 두 번 실행해 각 결과를 CPU 스트림의 f32 `dequantize` + `matmul`과 비교하며(상대 허용 오차 2e-2), 각 패스 뒤 카운터를 확인합니다. + +- `bf16_above_default_ceiling`(256행, 4096x4096): bypass 1번, 이어서 2번, 빈 캐시, GPU의 `dequantize` + `matmul`과 같은 출력 바이트. +- `bf16_above_env_ceiling`(64행, `MLX_ROCM_WMMA_QMM_MAX_M=32`): 환경 변수 상한으로 같은 확인. +- `bf16_forced_wmma`(`MLX_ROCM_WMMA_QMM=1`): 어떤 카운터도 움직이지 않음. +- `f16_cached`(128행): 첫 패스에 miss와 insert, 두 번째 패스에 hit, `n * k * 2` 바이트 항목 하나. + +상한 분기가 다시 `DequantGemm`을 반환하게 하면 두 bf16 상한 케이스가 첫 패스에서 실패하고(`misses=1 inserts=1 bypasses=0 entries=1`), forced-WMMA와 f16 케이스는 여전히 통과합니다. 4개 중 2개가 실패합니다. + +## 4. 프로덕션 영향 + +gfx1151(Radeon 8060S)에서 `scripts/bench_decode.sh --prompt-tokens 2048 --max-tokens 128`로 측정했습니다. 변경 전은 main `87538835`, 변경 후는 그 커밋에 변경을 더한 것이며, 3회 중앙값(범위)이고, 각 실행은 `scripts/rocm_gpu_guard.sh --idle-secs 60` 아래에서 두 쪽을 번갈아 돌렸습니다. + +| 모델 | Scales | 구분 | Prefill tok/s | Decode tok/s | MLX peak | Prefill 후 active | 종료 시 active | +|---|---|---|---:|---:|---:|---:|---:| +| gemma-3-4b-it-4bit | bf16 | 변경 전 | 2872 (2819-2910) | 62.75 (62.68-62.76) | 4.52 GB | 2.80 GB | 2.97 GB | +| gemma-3-4b-it-4bit | bf16 | 변경 후 | 2868 (2839-2945) | 62.71 (62.55-62.73) | 4.52 GB | 2.56 GB | 2.73 GB | +| Meta-Llama-3.1-8B-Instruct-4bit | f16 | 변경 전 | 1136 (1106-1136) | 36.63 (32.36-36.84) | 6.92 GB | 4.75 GB | 4.75 GB | +| Meta-Llama-3.1-8B-Instruct-4bit | f16 | 변경 후 | 1132 (1048-1143) | 36.67 (36.58-36.92) | 6.92 GB | 4.75 GB | 4.75 GB | + +캐시 카운터(warmup과 측정 prefill 합계, 변경 전은 카운터 전용 빌드): + +| 모델 | 구분 | Hits | Misses | Inserts | Evictions | Bypassed | 종료 시 항목 | 종료 시 바이트 | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| gemma-3-4b-it-4bit | 변경 전 | 0 | 340 | 340 | 334 | 0 | 6 | 241,172,480 | +| gemma-3-4b-it-4bit | 변경 후 | 0 | 0 | 0 | 0 | 340 | 0 | 0 | +| Meta-Llama-3.1-8B-Instruct-4bit | 변경 전 | 0 | 320 | 320 | 318 | 0 | 2 | 234,881,024 | +| Meta-Llama-3.1-8B-Instruct-4bit | 변경 후 | 0 | 320 | 320 | 318 | 0 | 2 | 234,881,024 | + +RDNA 3.5의 bf16 체크포인트에서는 prefill 후와 실행 종료 시 active memory가 캐시에 남던 항목 크기만큼인 0.24 GB 줄어듭니다. MLX peak memory는 변하지 않는데, peak는 prefill 안에서 호출마다 만드는 dequantize 사본에서 오고 캐시는 그것을 없앤 적이 없기 때문입니다(그 사본은 항목 28이 제한합니다). prefill과 decode는 실행 간 범위 안이며, 속도 향상은 주장하지 않습니다. Llama 3.1 8B는 f16 대조군입니다. 이 모델의 GEMM은 상한 경로를 타지 않으므로 두 쪽이 같게 나옵니다. + +main `ad844354` 위로 리베이스한 변경에서 모델당 한 번씩 돌린 확인 실행은 변경 후 행과 일치했습니다. Gemma 3 4B는 prefill 2885 tok/s, decode 62.14 tok/s, peak 4.52 GB, prefill 후 active 2.56 GB, bypass 340번, 빈 캐시였고, Llama 3.1 8B는 prefill 1119 tok/s, decode 36.66 tok/s, peak 6.92 GB, active 4.75 GB, miss 320번, hit 0번, 항목 2개였습니다. + +Metal과 CUDA는 영향이 없습니다. 변경은 ROCm 오버레이에 있고, 브리지 함수는 ROCm이 아니면 0을 반환합니다. + +## 5. 문서 + +- **`LOCAL_FIXES.md` 항목 29.** LRU가 prefill에서 hit하지 않으면서 마지막 항목을 살려 둔다고 적은 문장이 이제 `QmmRoute::DequantGemmAboveWmmaCeiling`, Gemma 3 4B 수치(변경 전 조회 340번에 hit 0번, 230 MiB 캐시됨; 변경 후 bypass 340번, 빈 캐시, prefill 후 active 2.80에서 2.56 GB, peak 변화 없음), 카운터, `dequant_cache_stats()`, `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1`, 새 테스트를 설명합니다. +- **`docs/environment-variables.md`.** `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE` 설명에 `MLX_ROCM_WMMA_QMM_MAX_M` 위로 보내진 bf16 GEMM은 캐시를 쓰지 않는다는 내용이 추가되었고, `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1`이 문서화되었습니다. +- **`docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md`.** "2026-10-07: the dequantized-weight cache above the WMMA ceiling (issue #2151)" 섹션이 추가되어 방법, 두 표, logit 확인, 확인 실행을 담습니다. 세 번째 Gemma 변경 전 실행은 첫 실행이 하네스에서 실패해 2026-10-08에 했다는 점, 그리고 리베이스된 헤드에서 Gemma의 종료 시 active memory(2.56 GB, `87538835`에서는 2.73 GB)가 main의 변화와 함께 움직였고 그 차이를 분리하지 않았다는 점도 적혀 있습니다. + +## 6. 검증 + +gfx1151(Radeon 8060S), ROCm 10.0.0 / HIP 7.15에서 실행했습니다. + +- **새 테스트.** `cargo test --release --features rocm --test rocm_qmm_dequant_cache -- --test-threads=1`: 통과. 상한 분기가 다시 `DequantGemm`을 반환하게 하면 4개 케이스 중 2개(두 bf16 상한 케이스, 첫 패스에서)가 실패하고, forced-WMMA와 f16 케이스는 여전히 통과합니다. +- **라우팅 환경 변수 테스트.** `cargo test --release --features rocm --test rocm_qmm_env -- --test-threads=1`: 통과. +- **Logits.** gemma-3-4b-it-4bit에서 `examples/logit_trace`(README.md를 코퍼스로, 256토큰 청크 4개이므로 모든 projection이 256행으로 실행되어 상한 경로를 탐)를 변경 전후로 실행: trace가 바이트 단위로 동일. `scripts/compare_logit_traces.py --decided 2.0`은 1024개 중 0개 불일치와 동일한 perplexity를 보고합니다. +- **대상 게이트.** `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`: 통과. +- **전체 게이트.** `6dbfe7d7` 위의 head `79a8c9bb`에서 `MLXCEL_ROCM_SMOKE_MODEL=models/mlx/Qwen3-0.6B-4bit`로 실행한 `make verify-rocm`: `[verify-rocm] OK`, cargo 테스트 스위트 162개에서 12191개 통과, 0개 실패, 399개 무시(mlxcel-core lib 스위트에서 8906개 통과, 165개 무시), ROCm 스모크는 GPU에서 32개 토큰을 생성했습니다. `ae343d84` 기준의 이전 실행도 통과했습니다(12190개 통과, 0개 실패). 이후 main에 bridge 파일을 건드리는 #2218이 들어와 브랜치를 리베이스하고 게이트를 다시 실행했습니다. + +검증하지 않은 것: 이 호스트에 없는 Metal과 CUDA. 브리지 stub은 단위 테스트 `stats_are_zero_off_rocm`이 다루지만, 그 백엔드에서 빌드하거나 실행하지 않았습니다. + +## 7. 기술적 선택과 그 이유 + +- **상한 경로에서만 캐시를 건너뜀.** 이슈가 측정한 경로이고 캐시의 miss 패턴이 알려진 경로입니다. f16 경로도 hit 0번이었지만 #2156의 decode 유사 경우는 hit할 수 있으므로, 그 기본값 변경은 여기에 합치지 않고 #2232에서 다룹니다. +- **전역 캐시 크기 변경 대신 새 열거자.** `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`를 전역으로 줄이면 쓸모없는 삽입은 그대로이고 메모리도 여전히 묶이므로, 이슈에서 기각했습니다. miss가 구조적이므로 키를 바꾸는 방안도 기각했습니다. +- **캐시 꺼짐 경로 재사용.** 임시 할당과 `add_temporary`는 `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE=0`용으로 이미 있었으므로, 새 경로는 메모리 관리 코드를 새로 추가하지 않습니다. +- **카운터는 항상 켜 둠.** GEMM마다 relaxed atomic 증가 한 번은 dequantize와 GEMM에 비하면 무시할 수 있고, 항상 있으므로 벤치마크가 재빌드 없이 환경 변수 하나로 읽을 수 있습니다. +- **테스트 케이스마다 자식 프로세스 하나.** 캐시, 카운터, 노브가 프로세스 전역이고 한 번만 읽히므로, 각 케이스가 0에서 시작하려면 새 프로세스가 필요합니다. + +## 8. 남은 위험과 후속 작업 + +- **#2232: f16 캐시.** Llama 3.1 8B(f16 scales)는 조회 320번에 hit 0번, insert 320번, eviction 318번이었고, 실행 후 항목 2개(224 MiB)를 살려 두었으며, 세 번 실행 모두 같았습니다. #2232는 #2156의 decode 유사 경우를 카운터로 확인한 뒤 ROCm에서 `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE`의 기본값을 끄자고 제안합니다. +- **상한은 gfx1151에서만 측정됨.** 변경은 상한 경로를 타는 모든 곳에 적용되지만, 그 경로와 이 수치는 gfx1151에서 측정했습니다. +- **베이스 간 종료 시 메모리 차이.** Gemma의 실행 종료 시 active memory는 리베이스된 헤드에서 2.56 GB, `87538835`에 변경을 더한 것에서 2.73 GB였습니다. 두 베이스 사이에 main이 바뀌었고 그 차이는 분리하지 않았습니다. +- **ROCm 외 stub은 실행되지 않음.** `stats_are_zero_off_rocm`이 stub을 다루지만 Metal이나 CUDA에서 실행된 적은 없습니다. + +## 9. 학습 포인트 + +- **캐시는 hit율이 있어야 메모리 값을 합니다.** 이 LRU는 이 행 수의 forward pass가 결코 주지 않는 재사용을 위해 추가되었습니다. 카운터가 실행 한 번으로 이를 드러냈습니다. 조회 340번에 hit 0번입니다. +- **peak와 상주 메모리는 다른 수치입니다.** 캐시를 건너뛰자 prefill 후 active memory는 0.24 GB 줄었지만 MLX peak는 변하지 않았습니다. peak는 캐시가 없앤 적 없는 호출당 임시 버퍼에서 오기 때문입니다. peak만 측정했다면 효과가 없는 것으로 보였을 것입니다. +- **라우팅에 관한 사실은 경로에 담습니다.** GEMM이 왜 다른 경로로 갔는지 아는 지점에 열거자를 추가하니, `eval_gpu`의 차이는 한 줄로 끝났고 경로의 다른 독자인 `quantized_matmul_runs_dequant_gemm`도 한 줄 변경으로 올바르게 유지되었습니다. +- **변경 전 상태를 이전 라우팅에서 측정 가능하게 만듭니다.** 변경 전 카운트는 `87538835`처럼 라우팅하는 카운터 전용 빌드에서 나왔고, 이로써 카운터의 효과와 라우팅 변경의 효과를 분리했습니다. diff --git a/docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md b/docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md index 21c9090ee..fa99943de 100644 --- a/docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md +++ b/docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md @@ -46,3 +46,31 @@ Decode is unchanged in every cell (M is 1, which never took the WMMA kernel). ML ## The dense prefill path on ROCm `MLXCEL_PREFILL_DEQUANT_MIN_M=1024` against unset, same binary, pp2048, three ABBA runs: gemma-3-4b-it-4bit 2871 against 2852 tok/s, Qwen3-0.6B-4bit 4174 against 4185. With the route change the dense path is eligible on ROCm only where `quantized_matmul` already runs the same GEMM, so turning it on changes neither bytes nor speed there. + +## 2026-10-07: the dequantized-weight cache above the WMMA ceiling (issue #2151) + +The route above keeps every dequantized weight in an LRU of 8 matrices or 256 MB (`MLX_ROCM_QMM_DEQUANT_CACHE_SIZE`, `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`). A forward pass runs more distinct projections than that, so the cache never hits, and its last entries stay alive into decode. GEMMs routed above the ceiling (`QmmRoute::DequantGemmAboveWmmaCeiling`) now dequantize into a temporary and leave the cache alone; f16 and the other shapes on the same GEMM keep it. `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1` prints the cache counters at exit. + +Same host. Before is main at `87538835`; after is that commit plus the change. The before and after medians were taken on 2026-10-07, except the third Gemma before run, which was taken on 2026-10-08 (the first failed in the harness, not the bench). `scripts/bench_decode.sh --prompt-tokens 2048 --max-tokens 128` (20-token same-process warmup, then the measured pass), each run in its own `scripts/rocm_gpu_guard.sh --idle-secs 60` with before and after alternated, every run CLEAN. Medians of three runs (range). "Active after prefill" is the `[Memory] after warmup (prefill + warmup tokens)` line, read once the warmup prefill and its decode tokens are done; "active at end" is the `after measured pass` line. + +| Model | Scales | Arm | Prefill tok/s | Decode tok/s | MLX peak | Active after prefill | Active at end | +|---|---|---|---:|---:|---:|---:|---:| +| gemma-3-4b-it-4bit | bf16 | before | 2872 (2819-2910) | 62.75 (62.68-62.76) | 4.52 GB | 2.80 GB | 2.97 GB | +| gemma-3-4b-it-4bit | bf16 | after | 2868 (2839-2945) | 62.71 (62.55-62.73) | 4.52 GB | 2.56 GB | 2.73 GB | +| Meta-Llama-3.1-8B-Instruct-4bit | f16 | before | 1136 (1106-1136) | 36.63 (32.36-36.84) | 6.92 GB | 4.75 GB | 4.75 GB | +| Meta-Llama-3.1-8B-Instruct-4bit | f16 | after | 1132 (1048-1143) | 36.67 (36.58-36.92) | 6.92 GB | 4.75 GB | 4.75 GB | + +Cache counters for the same run (warmup and measured prefill together). The before counts come from a build of the counters alone, which routes and caches exactly as `87538835` does; one run each, since the counts are deterministic and matched across the three after runs. + +| Model | Arm | Hits | Misses | Inserts | Evictions | Bypassed | Entries at exit | Bytes at exit | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| gemma-3-4b-it-4bit | before (counters build) | 0 | 340 | 340 | 334 | 0 | 6 | 241,172,480 | +| gemma-3-4b-it-4bit | after | 0 | 0 | 0 | 0 | 340 | 0 | 0 | +| Meta-Llama-3.1-8B-Instruct-4bit | before (counters build) | 0 | 320 | 320 | 318 | 0 | 2 | 234,881,024 | +| Meta-Llama-3.1-8B-Instruct-4bit | after | 0 | 320 | 320 | 318 | 0 | 2 | 234,881,024 | + +For Gemma 3 4B the cache held 230 MiB of bf16 weights after prefill and never hit; skipping it lowers active memory after prefill and at the end of the run by 0.24 GB. MLX peak memory does not move, because the peak comes from the per-call dequantized copies inside prefill, which the cache never removed. Prefill and decode differ by less than the run-to-run range. Llama 3.1 8B is the f16 control and reads the same in both arms: its GEMMs never take the ceiling route, and its cache also gets zero hits while holding 224 MiB after prefill; turning its default off is lablup/mlxcel#2232. Decode at batch 1 runs one row, which takes qmv and never reaches the cache. + +Logits: `examples/logit_trace` on gemma-3-4b-it-4bit (README.md as the corpus, 4 chunks of 256 tokens, so every projection runs at 256 rows and takes the ceiling route), before and after, gives byte-identical traces (sha256 `e470fd1c...`), and `scripts/compare_logit_traces.py --decided 2.0` reports 0 of 1024 disagreements and an identical perplexity. + +Confirmation on the change rebased onto main `ad844354` (2026-10-08), one guarded run per model: gemma-3-4b-it-4bit 2885 tok/s prefill, 62.14 tok/s decode, MLX peak 4.52 GB, active after prefill 2.56 GB, 340 bypasses and an empty cache; Meta-Llama-3.1-8B-Instruct-4bit 1119 tok/s prefill, 36.66 tok/s decode, MLX peak 6.92 GB, active 4.75 GB, and the same 320 misses, 0 hits and 2 cached entries. Both agree with the after runs above. Gemma's active memory at the end of the run reads 2.56 GB there against 2.73 GB on `87538835`; main moved between the two bases and that difference was not isolated. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index ec1dfa7a3..5750be44d 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -259,7 +259,12 @@ default applies (lablup/mlxcel#2152). An unset or empty variable takes the default silently. `MLX_ROCM_QMM_DEQUANT_M_THRESHOLD` (1 to 2147483647, default the built-in crossover) is the row count from which a quantized GEMM prefers dequantize + GEMM. `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE` (0 to 2147483647, default -8, `0` turns it off) is the number of dequantized weights kept for reuse. +8, `0` turns it off) is the number of dequantized weights kept for reuse; the +bf16 GEMMs sent above `MLX_ROCM_WMMA_QMM_MAX_M` never use that cache, since a +forward pass at those row counts cycles it without a hit (lablup/mlxcel#2151). +`MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1` prints the cache's hit, miss, insert, +eviction and bypass counts, its entry count and its bytes on stderr when the +process exits. `MLX_ROCM_GEMM_F32_SOLUTION_INDEX` and `MLX_ROCM_GEMM_BF16_SOLUTION_INDEX` (0 to 2147483647, default 0) pick a rocBLAS solution, and their `_BATCHED_` variants (same range) fall back to them. `MLX_ROCM_QMV_TILE_N` (1 to 32, diff --git a/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md b/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md index fd217f12e..d19b61431 100644 --- a/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md +++ b/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md @@ -38,7 +38,7 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code 25. **`MLX_ROCM_FFT_CACHE_SIZE` validated.** The fork's `LRUBytesKeyCache` constructor in `mlx/backend/rocm/lru_cache.h` read its capacity variable with an unchecked `std::stoul`, and its one env-driven user is the hipFFT plan cache in `fft.hip` (entry 18). Per value, confirmed with a host g++ 14 probe of the header: `0` was accepted, and the first `put()` then ran `while (size() >= capacity_)` on an empty list and called `back()` and `pop_back()` on it, which is undefined behavior (the probe threw `std::bad_alloc`; with `-D_GLIBCXX_ASSERTIONS` it aborts on `!this->empty()`); `abc`, an empty value and an overflowing value such as `99999999999999999999999` threw `std::invalid_argument` or `std::out_of_range` whose message is only `stoul`, out of the static's initializer in `FFT::eval_gpu`, so the static stayed uninitialized and every later FFT threw again; `-1` was taken silently as `SIZE_MAX` (a cache that never evicts), `8abc` silently as 8, and `0x10` as 0, with the same undefined behavior as `0` (on gfx1151 the first GPU rfft failed with `std::bad_alloc` for both). `capacity_from_env()` in the same header now parses it the way `gpu_watchdog_seconds()` in `device.cpp` parses `MLX_ROCM_GPU_WATCHDOG_SECS`: `std::strtol` base 10, the whole string must be the number, and the value must be 1 to `INT_MAX` (the bound CUDA's `int` read implies). Unset gives the default silently; anything else prints one stderr line, `[ROCm] ignoring invalid MLX_ROCM_FFT_CACHE_SIZE="" (expected a positive integer); using the default 128`, and uses the default, once, because the static then initializes. Both `LRUBytesKeyCache` and `LRUCache(size_t)` now throw `std::runtime_error("LRUCache requires capacity > 0.")` for a capacity of 0, as upstream CUDA's `LRUCache` does (`mlx/backend/cuda/lru_cache.h` at the pin). Upstream CUDA reads `MLX_CUDA_FFT_CACHE_SIZE` with `atoi`, so there `0` or junk throws that error at the first FFT and `-1` becomes `SIZE_MAX`; that was not copied. The comment on `fft_plan_cache()` in `fft.hip` states the rule; editing that file was also what made a warm build recompile it, because at the time each HIP object depended only on its own `.hip` source (fixed in item 26). `tests/rocm_fft_cache_env.rs` runs one child process per value (`0`, `-1`, `abc`, empty, `8abc`, `0x10`, an overflow past `long`, `2147483648`, `16`, `1`, `2147483647`, unset) and checks the warning count and an rfft against the CPU stream. Applies to the fork. 27. **CPU-stream BLAS runs single-threaded over fine-grained memory.** On an APU the fork's allocator (`mlx/backend/rocm/allocator.cpp`, `unified_malloc`) gives every array fine-grained device memory, CPU-stream arrays included, and `Buffer::raw_ptr()` hands the CPU that same pointer. MLX's CPU backend writes BLAS and LAPACK results straight into it (`cblas_sgemm` in `mlx/backend/cpu/gemms/cblas.cpp`, also `conv.cpp`, `masked_mm.cpp` and the LAPACK primitives), and multithreaded OpenBLAS returned wrong output columns there, differently from call to call. Measured on the gfx1151 host (Ryzen AI MAX+ 395, Debian OpenBLAS 0.3.29 pthread, 32 threads): a standalone C `cblas_sgemm` of `[1, 2880] x [2880, 2880]^T` was wrong in 37 of 50 calls with inputs and output from `hipExtMallocWithFlags(hipDeviceMallocFinegrained)`, in 10 of 300 with only the output there, and exact in every call with malloc or `hipHostMalloc` buffers or with `OPENBLAS_NUM_THREADS` at 1, 2 or 4. Through MLX the wrong columns were each the last column of one OpenBLAS thread's share of the output (shares of 93 columns: 650, 1022, 1859, ...), the column whose 64-byte line the next thread also writes. This was lablup/mlxcel#2072: `tests/rocm_mxfp4_quant.rs` failed `qmm 2880x2880 M=1` in 3 of 8 runs here (7 of 10 in the issue), and the fault was in its CPU reference, not the GPU kernel; with the tensors dumped from a failing run, the GPU result matched an independent dequantize-and-dot of the packed bytes and the CPU `matmul` did not. `unified_malloc` now calls `openblas_set_num_threads(1)` once, on its first fine-grained allocation, which precedes any CPU-stream BLAS call writing such a buffer; the symbol is declared weak, so a build against another BLAS links and changes nothing. The cost is CPU-stream BLAS speed: the CPU reads fine-grained memory slowly, and that gemv went from about 0.2 s to about 1.4 s per call. GPU work and quantized or half-precision CPU matmuls (MLX's own kernels, no BLAS) are unaffected, so model inference under `MLXCEL_DEVICE=cpu` barely touches BLAS. Keeping BLAS multithreaded by giving it a cacheable scratch output would need a change at every BLAS call site in MLX's CPU backend and was not taken. `tests/rocm_cpu_blas_finegrained.rs` repeats that gemv on the CPU stream against an f64 host reference; with this change reverted it failed in 5 of 5 runs, and with it `quantized_matmul_matches_dequantized_reference` passed 50 runs in a row. Applies to the fork. 28. **In-flight command batches bounded by what they allocate, and the cache limit enforced.** Measured on the gfx1151 host at pp512/tg128 (lablup/mlxcel#2062), the MLX peak (live buffers, `get_peak_memory`) was 20.60 GB for Meta-Llama-3.1-8B-Instruct-4bit, whose weights take 4.75 GB once loaded, and 23.56 GB for Qwen3-30B-A3B-4bit (17.17 GB), and the device-wide `mem_info_vram_used` rose by 21.19 GB and 24.09 GB. The excess was not buffer cache (0.4 to 1.7 GB at every phase boundary) but live transients inside the measured prefill: every buffer an operation allocates is kept by `CommandEncoder::add_temporary` until its batch's completion handler runs, the eager path (`use_hip_graphs()` is off) committed only every `MLX_MAX_OPS_PER_BUFFER` ops (2000) and never registered those commits with MLX's scheduler, so neither the scheduler's task cap nor `set_memory_limit` ever made the host wait, and the host could encode far ahead of the GPU, holding the transients of many batches at once. For the 8B, an f16 checkpoint, most of it is the f16 weight copy `QuantizedMatmul`'s dequantize-and-GEMM path allocates per matrix (with `MLX_ROCM_QMM_DEQUANT_GEMM=0` the peak was 6.78 GB, at a twentieth of the prefill speed); the MoE model is bf16, whose affine 4-bit matmuls took the fused WMMA kernel at the time and allocated no copy (item 29 moved its prefill projections to the dequantize path) (the variable left its peak at 23.56 GB), so its excess is other operations' outputs held the same way, which the bound below removes too. `MLX_MAX_OPS_PER_BUFFER=50` alone changed nothing, because more commits do not stop the host from running ahead. Now `gpu::eval` (`mlx/backend/rocm/eval.cpp`) adds what each primitive allocated during `eval_gpu` to its encoder's open batch, read from a thread-local byte counter in the allocator (`rocm::thread_allocated_bytes()`, so frees on the worker thread do not disturb it), and `CommandEncoder::maybe_commit` (`mlx/backend/rocm/device.cpp`) commits once the open batch has allocated a quarter of `MLX_ROCM_MAX_INFLIGHT_MB` (default 1024); after that commit and after the commit `gpu::finalize` makes at the end of every eval and `async_eval`, it records a HIP event, and blocks on the oldest such batch while the committed ones still exceed the budget. The wait is `hipEventSynchronize` (the device is in blocking-sync mode, so the thread sleeps), or a polled wait that gives up at `MLX_ROCM_GPU_WATCHDOG_SECS` when that is set; a faulted stream ends either with the fault, which is recorded as the stream's error. Nothing is tracked while a decode-step capture or any stream capture is in progress. The events are raw `hipEvent_t` owned by the encoder, not `HipEvent`: `HipEvent` returns its handle to a function-local static pool when destroyed, an encoder can be destroyed after that pool at process exit, and a first version built on it aborted `tests/rocm_mxfp4_quant.rs` at exit with `malloc_consolidate(): unaligned fastbin chunk detected`. `0` restores the old behavior. `tests/rocm_inflight_bound.rs` runs 40 f16 4-bit 8192x8192 matmuls over 512 rows in one evaluation (5 GiB of dequantized weight copies): the peak rose by 1.75 GiB with the default and by 6.25 GiB with `MLX_ROCM_MAX_INFLIGHT_MB=0`, which fails its 3 GiB bound. With the default the peaks were 6.14 GB and 18.57 to 18.58 GB over three harness runs each, decode throughput was unchanged (median 37.57 against 37.02 tok/s and 61.69 against 61.79 tok/s) and the 8B's prefill lost 1.8% (the MoE model's prefill varies more than that from run to run); budgets from 256 to 4096 MiB traded peak for prefill speed within a few percent (the table is in `docs/benchmark_results/rocm-memory-gfx1151-2026-09-30.md`). Separately, `RocmAllocator::set_cache_limit` stored `max_pool_size_` and nothing read it, so MLX's `set_cache_limit` (mlxcel's `MLXCEL_CACHE_LIMIT`) bounded nothing, while the exact-size cache (`min_utilization` 1.0) keeps every buffer size it has seen until `clear_cache()`. `malloc_async` now trims the cache on a cache miss once it exceeds `max_pool_size_`, the only place the footprint grows, so active plus cache stays within the live set plus the limit plus the one request; it trims to three quarters of the limit, so one round of blocking `hipFree` calls covers several misses of a workload whose shapes keep changing; `free()` and cache hits still never call `hipFree`, which the fork avoids for the training reasons in its comment there. The limit still defaults to the memory limit (76.8 GiB here), so the trim changes nothing until someone sets a limit; mlxcel sets 2 GiB on ROCm builds. `memory::tests::cache_limit_bounds_the_free_buffer_cache` in mlxcel-core frees 64 buffers of distinct sizes under a 1 MiB limit and checks the cache stays under 8 MiB; with the trim removed it held 49.8 MB and failed. The count is of allocations, not only of transients (a batch that grows the KV cache counts too), and a batch's buffers are dropped by the worker thread just after its event completes, so the bound is approximate. Applies to the fork. -29. **bf16 GEMMs of 128 rows or more leave the fused WMMA kernel on RDNA 3.5.** `QuantizedMatmul::eval_gpu` (`mlx/backend/rocm/quantized/qmm.hip`) sent every single bf16 affine GEMM with more than one row (4-, 6- or 8-bit, group 64, `N % 16 == 0`, `K % 64 == 0`) to `qmm_wmma_dense_kernel` on a device with native WMMA that is not a low-CU iGPU, however many rows it had. On gfx1151 that kernel lost to the dequantize + hipBLASLt path it bypasses on every bf16 shape measured from 128 rows up (1.0x to 3.1x at 128 rows, 1.6x to 2.4x at 256, 2.4x to 5.1x at 1024 and 2048; it won only on some shapes at 2 to 64 rows), so bf16 prefill ran at a fraction of what the same device did for f16 checkpoints (lablup/mlxcel#2081). It also made `quantized_matmul` and `dequantize` + `matmul` disagree for bf16 (499 of 1,048,576 outputs at `[1024, 2048] x [1024, 2048]`, a reduction-order difference), which is what mlxcel's dense prefill path checks. The dispatch decision is now one function, `select_qmm_route`, used by `eval_gpu`; it keeps the kernel below `MLX_ROCM_WMMA_QMM_MAX_M` rows (default 128 on the RDNA 3.5 tier, no ceiling on the unmeasured tiers) and sends larger GEMMs to the bf16 dequantize + hipBLASLt route, except where that route would be the fp8 path (RDNA 4), which keeps the kernel as before. `MLX_ROCM_WMMA_QMM=1` removes the ceiling, which is the old dispatch on any device that is not a low-CU iGPU (there it forces the kernel on, as it always did). `rocm.h` exports `quantized_matmul_runs_dequant_gemm` (stubbed in `no_rocm.cpp`), which answers from the same function whether a single GEMM of a given shape takes dequantize + hipBLASLt, for mlxcel to decide when its dense prefill path returns the same bytes. Same binary, `MLX_ROCM_WMMA_QMM=1` against the default, prefill tok/s over three ABBA runs: Gemma 3 4B 4-bit 1146 to 2289 at 512 tokens and 961 to 2815 at 2048, Qwen3-0.6B 4-bit 4581 to 7633 and 3059 to 4236, Qwen3-30B-A3B 4-bit (MoE, only the attention projections take this path) 311 to 326 and 283 to 297; decode unchanged, and MLX peak memory within 0.1 GB except Qwen3-0.6B at 512 tokens (1.04 to 1.46 GB), since the route allocates a bf16 weight copy per matrix that item 28 bounds. The route's dequantized-weight LRU (8 matrices or 256 MB, `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE`, `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`) gets no hits across a model's projections in prefill but still keeps its last entries alive afterwards, as it already did for f16 checkpoints; `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE=0` turns it off. `docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md` has the tables. Applies to the fork; the ceiling was measured on gfx1151 only. +29. **bf16 GEMMs of 128 rows or more leave the fused WMMA kernel on RDNA 3.5.** `QuantizedMatmul::eval_gpu` (`mlx/backend/rocm/quantized/qmm.hip`) sent every single bf16 affine GEMM with more than one row (4-, 6- or 8-bit, group 64, `N % 16 == 0`, `K % 64 == 0`) to `qmm_wmma_dense_kernel` on a device with native WMMA that is not a low-CU iGPU, however many rows it had. On gfx1151 that kernel lost to the dequantize + hipBLASLt path it bypasses on every bf16 shape measured from 128 rows up (1.0x to 3.1x at 128 rows, 1.6x to 2.4x at 256, 2.4x to 5.1x at 1024 and 2048; it won only on some shapes at 2 to 64 rows), so bf16 prefill ran at a fraction of what the same device did for f16 checkpoints (lablup/mlxcel#2081). It also made `quantized_matmul` and `dequantize` + `matmul` disagree for bf16 (499 of 1,048,576 outputs at `[1024, 2048] x [1024, 2048]`, a reduction-order difference), which is what mlxcel's dense prefill path checks. The dispatch decision is now one function, `select_qmm_route`, used by `eval_gpu`; it keeps the kernel below `MLX_ROCM_WMMA_QMM_MAX_M` rows (default 128 on the RDNA 3.5 tier, no ceiling on the unmeasured tiers) and sends larger GEMMs to the bf16 dequantize + hipBLASLt route, except where that route would be the fp8 path (RDNA 4), which keeps the kernel as before. `MLX_ROCM_WMMA_QMM=1` removes the ceiling, which is the old dispatch on any device that is not a low-CU iGPU (there it forces the kernel on, as it always did). `rocm.h` exports `quantized_matmul_runs_dequant_gemm` (stubbed in `no_rocm.cpp`), which answers from the same function whether a single GEMM of a given shape takes dequantize + hipBLASLt, for mlxcel to decide when its dense prefill path returns the same bytes. Same binary, `MLX_ROCM_WMMA_QMM=1` against the default, prefill tok/s over three ABBA runs: Gemma 3 4B 4-bit 1146 to 2289 at 512 tokens and 961 to 2815 at 2048, Qwen3-0.6B 4-bit 4581 to 7633 and 3059 to 4236, Qwen3-30B-A3B 4-bit (MoE, only the attention projections take this path) 311 to 326 and 283 to 297; decode unchanged, and MLX peak memory within 0.1 GB except Qwen3-0.6B at 512 tokens (1.04 to 1.46 GB), since the route allocates a bf16 weight copy per matrix that item 28 bounds. GEMMs sent to that route by the ceiling are `QmmRoute::DequantGemmAboveWmmaCeiling`, which runs the same dequantize and GEMM but dequantizes into a temporary freed after the GEMM instead of going through the dequantized-weight LRU (8 matrices or 256 MB, `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE`, `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`): a forward pass runs more distinct projections than it holds, so it never hit, and its last entries stayed alive through decode (lablup/mlxcel#2151). Gemma 3 4B 4-bit at pp2048/tg128: 0 hits in 340 lookups and 230 MiB still cached after prefill before; after, 340 bypasses, an empty cache, 0.24 GB less active memory after prefill (2.80 to 2.56 GB), MLX peak unchanged at 4.52 GB, prefill and decode within noise, and byte-identical logit traces. f16 and the other shapes on `DequantGemm` keep the cache unchanged. The cache is one struct with relaxed atomic hit, miss, insert, eviction and bypass counters; `rocm.h` exports `dequant_cache_stats()` (stubbed in `no_rocm.cpp`) returning them with the entry count and bytes, and `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1` prints them on stderr at exit. `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE=0` turns the cache off. `tests/rocm_qmm_dequant_cache.rs` checks both routes. `docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md` has the tables. Applies to the fork; the ceiling was measured on gfx1151 only. 30. **Custom kernels receive `_shape` and `_strides` by value.** `build_kernel` in the fork's `mlx/backend/rocm/custom_kernel.cpp` declared an input's shape and strides as `const int32_t* _shape` and `const int64_t* _strides`, while `CustomKernel::eval_gpu` passes them with `KernelArgs::append_ndim`, which pads the vector to `JIT_MAX_NDIM` (8) entries and hands `hipModuleLaunchKernel` a pointer to that storage, so the kernel received the values themselves and read the first eight bytes of the shape as an address. Any `fast::hip_kernel` source that mentioned `_shape` faulted the queue on its first launch: the HIP Gumbel-max sampler (lablup/mlxcel#2064), which reads `logits_shape[1]` as the CUDA and Metal ports do, died with `HSA_STATUS_ERROR_MEMORY_FAULT` at address `0x80000000000`, the page of `0x80000000800`, which is the shape `{2048, 2048}` read as a 64-bit pointer, on gfx1151. No HIP kernel in mlxcel used either argument before, which is why it went unseen. The generated header now defines `KernelShape` (8 `int32_t`) and `KernelStrides` (8 `int64_t`) with an `operator[]`, the parameters are declared with them by value, which is what upstream CUDA's `const __grid_constant__ Shape` / `Strides` does, and `elem_to_loc` gains an overload for them; a `static_assert` ties the 8 to `JIT_MAX_NDIM`. The launch also appends an input's shape, strides and ndim only when its ndim is above zero, the condition under which `build_kernel` declares them; before, a 0-d input whose name appeared as `_shape` in the source shifted every later argument (no mlxcel kernel has such an input). `sampling_fixed_key_tests` in mlxcel-core reads `logits_shape` through the Gumbel-max port and faults without this change; the `_strides` declaration has no user in mlxcel and was checked only by compiling it. Applies to the fork. 31. **mxfp4 in the expert-batched gather qmv.** `GatherQMM::eval_gpu` in `mlx/backend/rocm/quantized/qmm.hip` sent only affine weights to `gather_qmv_expert_batched_kernel` (item 9), so a sorted mxfp4 prefill took the per-row `gather_qmv_kernel` that item 10 routes non-affine modes to, rereading an expert's weights for every routed row: gpt-oss-20b spent 99.8% of a 512-token prefill there, about 940 ms per call (lablup/mlxcel#2106). The kernel now also takes mxfp4 (group size 32, one E8M0 scale byte per group, no biases) with bf16 activations, under the same gate (sorted, transposed, `M == 1`, `B >= 64`, `E <= 64`, `B / E >= 4`) and the same switch (`MLX_ROCM_GATHER_QMV_EXPERT_BATCHED=0` turns it off). Its `AFFINE=false` arm decodes each e2m1 nibble without a branch by placing its sign, exponent and mantissa bits in an fp16 (`fp4_e2m1_to_float_scaled`, the value times 2^-14) and scales each word's dot product back by 2^14 with the group's scale. That is one new instantiation, ``; f16 mxfp4 and mxfp8 keep the per-row kernel, and `qmm.hip` still compiles in about 38 s (37.9 and 37.8 s, against 38.1 and 37.9 s for `main`'s file, same flags). `tests/rocm_gather_qmm_expert_batched.rs` compares the sorted kernel with the unsorted per-row path and a dequantized f32 reference at gpt-oss-20b's expert shape (32 experts, top 4, `K = N = 2880`) and at a narrow shape with a partial last column block, in flat and broadcast index layouts: the two paths' relative errors agree to four digits (1.6e-3 to 1.7e-3), and with the sign bit dropped from the decode the first case fails at 1.38. On gfx1151, 512-token prefill of `gpt-oss-20b-MXFP4-Q4` went from 7.53 to 526.3 tok/s (medians of six alternated, guarded runs), with decode unchanged and no decided-position disagreement in `logit_trace` against the per-row path (`docs/benchmark_results/rocm-moe-prefill-mxfp4-gfx1151-2026-10-06.md`). Applies to the fork; kept in mlxcelverse under the 2026-10-06 fork policy, not proposed there. 33. **Custom-kernel JIT modules keyed by their generated source.** `CustomKernel::eval_gpu` in `mlx/backend/rocm/custom_kernel.cpp` passed `name_` to `rocm::get_jit_module`, which memoises by device index and that name (`jit_module.cpp`) and runs the source builder only on a miss. `hip_kernel` builds `name_` as `"custom_kernel_" + name + template_arguments_hash(template_args)`, but `build_kernel` also writes each input's and output's dtype into the source, and declares an input's `_shape`, `_strides` and `_ndim` parameters only when the input is not 0-d (item 30). A second call that differed only there reused the first call's module: an f16 or bf16 input was read through `const float*`, and a 0-d input after a 1-d one launched with one argument fewer than the module declares. Nothing threw. The module name is now `name_ + "_" + hex(std::hash{}(source_))`, copied from upstream CUDA's `mlx/backend/cuda/custom_kernel.cpp` at the pin (around line 313, added by upstream 9f5f7931, ml-explore/mlx#4273); the symbol looked up in the module keeps its plain name, so `rocprofv3` traces still show `custom_kernel_*`, and `get_jit_module` keeps its device-index prefix. mlxcel's own launchers were already shielded by naming input dtypes in `template_args` (`scripts/ci/check_kernel_dtype_keys.py`, kept as defense in depth), which cannot cover the 0-d case, an input whose dtype varies independently of the keyed one, or an output dtype that varies at one input dtype. `tests/rocm_custom_kernel_jit_key.rs` launches one kernel name with no template args (the bridge's `rocm_jit_key_probe`) at f32, f16 and bf16 inputs, at f32 then f16 outputs, and at a 1-d then a 0-d input in a child process; on gfx1151 both pass with this change, and with the name-only key restored the f16 call returns values unrelated to its input (`-0.0313`, `131336.0`, ... for `-1.25`, `3.0`, ...) and the 0-d child dies with SIGSEGV after its 1-d call. Kept in mlxcelverse under the 2026-10-06 fork policy (lablup/mlxcel#2149). diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/no_rocm.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/no_rocm.cpp index e870d9fcc..40a266304 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/no_rocm.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/no_rocm.cpp @@ -67,6 +67,10 @@ bool quantized_matmul_runs_dequant_gemm( return false; } +DequantCacheStats dequant_cache_stats() { + return DequantCacheStats{}; +} + } // namespace rocm namespace fast { diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip index d786b48b4..eb143d6db 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip @@ -9,6 +9,7 @@ #include "mlx/backend/rocm/kernel_utils.hpp" #include "mlx/backend/rocm/quantized/quantized.h" #include "mlx/backend/rocm/quantized/qmv_tiled_kernel.hip" +#include "mlx/backend/rocm/rocm.h" #include "mlx/primitives.h" #include @@ -35,6 +36,8 @@ #endif #include #include +#include +#include #include #include #include @@ -808,6 +811,13 @@ enum class QmmRoute { WmmaDense, // qmm_wmma_dense_kernel, fused dequant + rocWMMA DequantFp8Gemm, // dequant_fp8_gemm (weight and activation in e4m3) DequantGemm, // dequantize (cached) + dequant_rocblas_gemm in the x dtype + // The same GEMM for a bf16 shape the fused kernel would take below the row + // ceiling (wmma_qmm_max_m). The weight is dequantized into a temporary on + // every call and never enters the dequantized-weight cache: a forward pass + // runs more distinct projections at these row counts than the cache holds, + // so it never hit, and its last entries stayed alive through decode + // (lablup/mlxcel#2151). + DequantGemmAboveWmmaCeiling, Other, // qmv, batched qmv and the tiled kernels below }; @@ -905,7 +915,7 @@ inline QmmRoute select_qmm_route(const QmmRouteInputs& in, rocm::Device& d) { if (!hipblaslt_is_faster) { return QmmRoute::WmmaDense; } - return QmmRoute::DequantGemm; + return QmmRoute::DequantGemmAboveWmmaCeiling; } if (fp8()) { return QmmRoute::DequantFp8Gemm; @@ -978,6 +988,97 @@ inline size_t dequant_cache_max_bytes() { return max_bytes; } +// The dequantized-weight LRU of the DequantGemm route and its counters +// (lablup/mlxcel#2151). The counters are relaxed atomics, bumped once per +// GEMM on the route and never read on the hot path; rocm::dequant_cache_stats +// reads them with the entry count and bytes. MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1 +// prints them on stderr when the process exits. +struct DequantCacheCounters { + std::atomic hits{0}; + std::atomic misses{0}; + std::atomic inserts{0}; + std::atomic evictions{0}; + std::atomic bypasses{0}; +}; + +DequantCacheCounters g_dequant_cache_counters; + +inline void count_dequant_cache(std::atomic& counter) { + counter.fetch_add(1, std::memory_order_relaxed); +} + +struct DequantCache { + std::mutex mutex; + std::list lru; + size_t cached_bytes = 0; + std::unordered_map + entries; + + // Drops least recently used entries until both limits hold. Call with + // `mutex` held. + void evict_to(int capacity, size_t max_bytes) { + while (!lru.empty() && + (static_cast(entries.size()) > capacity || + cached_bytes > max_bytes)) { + auto evict_it = entries.find(lru.back()); + if (evict_it != entries.end()) { + cached_bytes -= evict_it->second.bytes; + entries.erase(evict_it); + count_dequant_cache(g_dequant_cache_counters.evictions); + } + lru.pop_back(); + } + } +}; + +void print_dequant_cache_stats(); + +// The process-wide cache. The exit report is registered after the cache is +// constructed, so it runs before the cache is destroyed. +DequantCache& dequant_cache() { + static DequantCache cache; + static const bool report_registered = [] { + const char* e = std::getenv("MLX_ROCM_QMM_DEQUANT_CACHE_STATS"); + if (e && e[0] == '1' && e[1] == '\0') { + std::atexit(print_dequant_cache_stats); + return true; + } + return false; + }(); + (void)report_registered; + return cache; +} + +rocm::DequantCacheStats read_dequant_cache_stats() { + rocm::DequantCacheStats stats{}; + const auto& c = g_dequant_cache_counters; + stats.hits = c.hits.load(std::memory_order_relaxed); + stats.misses = c.misses.load(std::memory_order_relaxed); + stats.inserts = c.inserts.load(std::memory_order_relaxed); + stats.evictions = c.evictions.load(std::memory_order_relaxed); + stats.bypasses = c.bypasses.load(std::memory_order_relaxed); + auto& cache = dequant_cache(); + std::lock_guard lock(cache.mutex); + stats.entries = cache.entries.size(); + stats.bytes = cache.cached_bytes; + return stats; +} + +void print_dequant_cache_stats() { + const auto s = read_dequant_cache_stats(); + std::fprintf( + stderr, + "[ROCm] qmm dequant cache: hits=%llu misses=%llu inserts=%llu " + "evictions=%llu bypassed=%llu entries=%zu bytes=%zu\n", + static_cast(s.hits), + static_cast(s.misses), + static_cast(s.inserts), + static_cast(s.evictions), + static_cast(s.bypasses), + s.entries, + s.bytes); +} + inline rocblas_operation to_rocblas_op(bool transpose) { return transpose ? rocblas_operation_transpose : rocblas_operation_none; } @@ -3411,7 +3512,13 @@ bool quantized_matmul_runs_dequant_gemm( // and fall back to different rocBLAS calls, so only the hipBLASLt case is // claimed. The claim assumes the hipBLASLt launch itself does not throw, // which would send each side to its own rocBLAS fallback. - return route == QmmRoute::DequantGemm && is_hipblaslt_available(); + return (route == QmmRoute::DequantGemm || + route == QmmRoute::DequantGemmAboveWmmaCeiling) && + is_hipblaslt_available(); +} + +DequantCacheStats dequant_cache_stats() { + return read_dequant_cache_stats(); } } // namespace rocm @@ -3524,7 +3631,9 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { // Dequant + rocBLAS GEMM path // Disable with MLX_ROCM_QMM_DEQUANT_GEMM=0 if needed - if (route == QmmRoute::DequantGemm || route == QmmRoute::DequantFp8Gemm) { + if (route == QmmRoute::DequantGemm || + route == QmmRoute::DequantGemmAboveWmmaCeiling || + route == QmmRoute::DequantFp8Gemm) { if (!((x_batch_count == 1) || (x_batch_count == batch_count))) { throw std::runtime_error( "Unsupported x batch shape for dequant GEMM fallback"); @@ -3567,15 +3676,16 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { bool cache_hit = false; int cache_cap = dequant_cache_capacity(); size_t cache_max_bytes = dequant_cache_max_bytes(); - if (cache_cap > 0 && cache_max_bytes > 0) { - static std::mutex cache_mutex; - static std::list lru; - static size_t cached_bytes = 0; - static std::unordered_map< - DequantCacheKey, - DequantCacheEntry, - DequantCacheKeyHasher> - cache; + // Above the WMMA ceiling the cache is skipped: the weight is dequantized + // into a temporary below, as with the cache turned off. + const bool use_cache = route == QmmRoute::DequantGemm && cache_cap > 0 && + cache_max_bytes > 0; + if (route == QmmRoute::DequantGemmAboveWmmaCeiling) { + dequant_cache(); // registers the exit report, if requested + count_dequant_cache(g_dequant_cache_counters.bypasses); + } + if (use_cache) { + auto& c = dequant_cache(); DequantCacheKey key{ reinterpret_cast(gpu_ptr(w)), @@ -3589,14 +3699,18 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { x.dtype()}; { - std::lock_guard lock(cache_mutex); - auto it = cache.find(key); - if (it != cache.end() && it->second.weight.shape() == w_dequant_shape) { - lru.splice(lru.begin(), lru, it->second.lru_it); + std::lock_guard lock(c.mutex); + auto it = c.entries.find(key); + if (it != c.entries.end() && + it->second.weight.shape() == w_dequant_shape) { + c.lru.splice(c.lru.begin(), c.lru, it->second.lru_it); w_dequant = it->second.weight; cache_hit = true; } } + count_dequant_cache( + cache_hit ? g_dequant_cache_counters.hits + : g_dequant_cache_counters.misses); if (!cache_hit) { w_dequant.set_data(mlx::core::rocm::malloc_async(w_dequant.nbytes(), enc)); @@ -3608,13 +3722,13 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { fp_dequantize(w, scales, w_dequant, group_size_, bits_, enc, s); } - std::lock_guard lock(cache_mutex); - auto it = cache.find(key); - if (it == cache.end()) { + std::lock_guard lock(c.mutex); + auto it = c.entries.find(key); + if (it == c.entries.end()) { size_t entry_bytes = w_dequant.nbytes(); if (entry_bytes <= cache_max_bytes) { - lru.push_front(key); - cache.emplace( + c.lru.push_front(key); + c.entries.emplace( key, DequantCacheEntry{ w_dequant, @@ -3622,47 +3736,30 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { scales, has_bias ? std::optional(*biases) : std::nullopt, entry_bytes, - lru.begin()}); - cached_bytes += entry_bytes; - - while (static_cast(cache.size()) > cache_cap || - cached_bytes > cache_max_bytes) { - auto evict = lru.back(); - auto evict_it = cache.find(evict); - if (evict_it != cache.end()) { - cached_bytes -= evict_it->second.bytes; - cache.erase(evict_it); - } - lru.pop_back(); - } + c.lru.begin()}); + c.cached_bytes += entry_bytes; + count_dequant_cache(g_dequant_cache_counters.inserts); + c.evict_to(cache_cap, cache_max_bytes); } } else { size_t entry_bytes = w_dequant.nbytes(); if (entry_bytes > cache_max_bytes) { - cached_bytes -= it->second.bytes; - lru.erase(it->second.lru_it); - cache.erase(it); + c.cached_bytes -= it->second.bytes; + c.lru.erase(it->second.lru_it); + c.entries.erase(it); + count_dequant_cache(g_dequant_cache_counters.evictions); } else { - cached_bytes -= it->second.bytes; + c.cached_bytes -= it->second.bytes; it->second.w_source = w; it->second.scales_source = scales; it->second.biases_source = has_bias ? std::optional(*biases) : std::nullopt; it->second.weight = w_dequant; it->second.bytes = entry_bytes; - cached_bytes += it->second.bytes; - lru.splice(lru.begin(), lru, it->second.lru_it); - - while (static_cast(cache.size()) > cache_cap || - cached_bytes > cache_max_bytes) { - auto evict = lru.back(); - auto evict_it = cache.find(evict); - if (evict_it != cache.end()) { - cached_bytes -= evict_it->second.bytes; - cache.erase(evict_it); - } - lru.pop_back(); - } + c.cached_bytes += it->second.bytes; + c.lru.splice(c.lru.begin(), c.lru, it->second.lru_it); + count_dequant_cache(g_dequant_cache_counters.inserts); + c.evict_to(cache_cap, cache_max_bytes); } } } diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/rocm.h b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/rocm.h index f28551395..93306ec40 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/rocm.h +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/rocm.h @@ -8,6 +8,7 @@ #include "mlx/utils.h" #include +#include #include #include @@ -79,4 +80,26 @@ MLX_API bool quantized_matmul_runs_dequant_gemm( int group_size, int bits); +// Counters of QuantizedMatmul's dequantized-weight cache, process-wide since +// start (lablup/mlxcel#2151). `hits` and `misses` count lookups on the cached +// dequantize + GEMM route, `inserts` the entries stored (a refresh of an entry +// under the same key counts), `evictions` the entries dropped for the +// MLX_ROCM_QMM_DEQUANT_CACHE_SIZE or MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES +// limit, and `bypasses` the bf16 GEMMs above the fused kernel's row ceiling, +// which dequantize into a temporary without touching the cache. `entries` and +// `bytes` are what the cache holds now. All zero without ROCm. +// MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1 prints the same numbers on stderr at +// process exit. +struct DequantCacheStats { + uint64_t hits; + uint64_t misses; + uint64_t inserts; + uint64_t evictions; + uint64_t bypasses; + size_t entries; + size_t bytes; +}; + +MLX_API DequantCacheStats dequant_cache_stats(); + } // namespace mlx::core::rocm diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 7ff6bc679..85a3bdd61 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -5187,6 +5187,23 @@ int32_t rocm_device_warp_size() { #endif } +rust::Vec rocm_dequant_cache_stats() { + rust::Vec out; +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND + const auto s = mlx::core::rocm::dequant_cache_stats(); + for (uint64_t v : {s.hits, s.misses, s.inserts, s.evictions, s.bypasses, + static_cast(s.entries), + static_cast(s.bytes)}) { + out.push_back(v); + } +#else + for (int i = 0; i < 7; ++i) { + out.push_back(0); + } +#endif + return out; +} + void set_rocm_port_warp_size_for_tests(int32_t warp_size) { mlxcel::set_rocm_port_warp_size_for_tests(static_cast(warp_size)); } diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 8f1122b3c..2740fdb78 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1419,6 +1419,12 @@ int32_t gpu_backend_kind(); // width `port_for` sees (0 restores the hardware value). bool rocm_port_allowed(bool any_wave_size, int32_t warp_size); int32_t rocm_device_warp_size(); + +// Counters of the ROCm QuantizedMatmul dequantized-weight cache +// (lablup/mlxcel#2151), as `mlx::core::rocm::dequant_cache_stats()` returns +// them: hits, misses, inserts, evictions, bypasses, entries, bytes, in that +// order. Seven zeros off ROCm. +rust::Vec rocm_dequant_cache_stats(); void set_rocm_port_warp_size_for_tests(int32_t warp_size); // Test-only (lablup/mlxcel#1804): a lazy array whose evaluation fails on the diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 51454685f..6c23d6d9f 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -2377,6 +2377,12 @@ mod ffi { /// `MLX_ROCM_FORCE_WARP_SIZE` launch-width override (issue #2147). fn rocm_device_warp_size() -> i32; + /// Raw counters of the ROCm QuantizedMatmul dequantized-weight + /// cache (issue #2151): hits, misses, inserts, evictions, bypasses, + /// entries, bytes. Seven zeros off ROCm. Prefer + /// [`crate::rocm_qmm_cache::dequant_cache_stats`]. + fn rocm_dequant_cache_stats() -> Vec; + /// Test-only seam (issue #2147): replaces the wavefront width the /// ROCm port tables are checked against for the rest of the process /// (0 restores the hardware value), so the wave64 refusal can be @@ -3946,6 +3952,10 @@ pub mod dtype; pub mod cuda_arch; pub mod rocm_arch; +// Counters of the ROCm QuantizedMatmul dequantized-weight cache (#2151), for +// tests and memory diagnostics. +pub mod rocm_qmm_cache; + // Deliberate ROCm GPU failures for tests (#1804): a typed wrapper over the // bridge's fault-probe fixture, so `tests/rocm_gpu_faults.rs` can provoke a // rejected launch and a queue fault without knowing the bridge's integers. diff --git a/src/lib/mlxcel-core/src/rocm_qmm_cache.rs b/src/lib/mlxcel-core/src/rocm_qmm_cache.rs new file mode 100644 index 000000000..20dff8cf9 --- /dev/null +++ b/src/lib/mlxcel-core/src/rocm_qmm_cache.rs @@ -0,0 +1,76 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Counters of the ROCm QuantizedMatmul dequantized-weight cache (issue +//! #2151). +//! +//! On ROCm, a quantized GEMM that takes the dequantize + GEMM route keeps the +//! dequantized weight in a process-wide LRU (8 matrices or 256 MB by default, +//! `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE` and `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`). +//! bf16 GEMMs at or above the fused WMMA kernel's row ceiling take the same +//! GEMM without the cache, and count as bypasses. These counters say whether +//! the cache earns the memory it holds. `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1` +//! prints the same numbers on stderr when the process exits, which is how a +//! benchmark run reads them. + +/// A snapshot of the cache's process-wide counters and its current size. +/// All zero on every backend other than ROCm. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct DequantCacheStats { + /// Lookups that found the dequantized weight. + pub hits: u64, + /// Lookups that did not, after which the weight was dequantized. + pub misses: u64, + /// Entries stored, including a refresh of an entry under the same key. + pub inserts: u64, + /// Entries dropped to keep the cache within its limits. + pub evictions: u64, + /// bf16 GEMMs above the WMMA row ceiling, dequantized without the cache. + pub bypasses: u64, + /// Entries the cache holds now. + pub entries: u64, + /// Bytes of dequantized weights the cache holds now. + pub bytes: u64, +} + +/// Reads the counters. Cheap: seven relaxed atomic loads and the cache's +/// size under its lock. +pub fn dequant_cache_stats() -> DequantCacheStats { + let raw = crate::ffi::rocm_dequant_cache_stats(); + let at = |i: usize| raw.get(i).copied().unwrap_or(0); + DequantCacheStats { + hits: at(0), + misses: at(1), + inserts: at(2), + evictions: at(3), + bypasses: at(4), + entries: at(5), + bytes: at(6), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hardware::{GpuBackendKind, gpu_backend_kind}; + + #[test] + fn stats_are_zero_off_rocm() { + if gpu_backend_kind() == GpuBackendKind::Rocm { + eprintln!("skipping: ROCm counts real GEMMs"); + return; + } + assert_eq!(dequant_cache_stats(), DequantCacheStats::default()); + } +} diff --git a/tests/rocm_qmm_dequant_cache.rs b/tests/rocm_qmm_dequant_cache.rs new file mode 100644 index 000000000..d797a24c3 --- /dev/null +++ b/tests/rocm_qmm_dequant_cache.rs @@ -0,0 +1,493 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! bf16 GEMMs routed above the fused WMMA kernel's row ceiling skip the +//! dequantized-weight cache (issue #2151). +//! +//! On RDNA 3.5, `select_qmm_route` in the ROCm overlay sends a bf16 affine +//! GEMM of `MLX_ROCM_WMMA_QMM_MAX_M` rows or more (128 by default) to +//! dequantize + hipBLASLt. That route used to store every dequantized weight +//! in the process-wide LRU (8 matrices or 256 MB), which a forward pass cycles +//! without a hit while its last entries stay alive through decode. The route +//! now dequantizes into a temporary and leaves the cache alone; the f16 and +//! non-WMMA shapes that reach the same GEMM keep the cache. +//! +//! The cache and its counters are process-wide, and the knobs are read once +//! per process, so every case runs in a fresh child process (the parent +//! re-invokes this binary on an ignored child test, one child at a time). The +//! child runs one quantized GEMM shape twice on the GPU, compares each result +//! with `dequantize` + `matmul` in f32 on the CPU stream, and prints the +//! counters `mlxcel_core::rocm_qmm_cache::dequant_cache_stats` read after each +//! pass. `MLX_ROCM_QMM_DEQUANT_M_THRESHOLD=1` makes the dequantize route +//! eligible at every row count, so only the ceiling and the dtypes decide the +//! route. The bf16 above-ceiling case also checks that the output bytes equal +//! `dequantize` + `matmul` on the GPU, the GEMM the route claims to run. +//! +//! With the route change reverted (the ceiling branch of `select_qmm_route` +//! returning `DequantGemm` again, measured on gfx1151) the two bf16 +//! above-ceiling cases fail on their first pass, which reads `misses=1 +//! inserts=1 bypasses=0 entries=1` instead of one bypass and an empty cache; +//! the forced-WMMA and f16 cases still pass. +//! +//! Skips on any other backend. Run on a ROCm host with: +//! +//! ```sh +//! cargo test --release --features rocm --test rocm_qmm_dequant_cache -- --test-threads=1 +//! ``` + +#![cfg(feature = "rocm")] + +use std::io::Read; +use std::process::{Command, Stdio}; +use std::thread; +use std::time::{Duration, Instant}; + +use mlxcel_core::hardware::{GpuBackendKind, gpu_backend_kind}; +use mlxcel_core::rocm_qmm_cache::{DequantCacheStats, dequant_cache_stats}; +use mlxcel_core::streams::DefaultDeviceGuard; +use mlxcel_core::{MlxArray, dtype}; + +const MAX_M: &str = "MLX_ROCM_WMMA_QMM_MAX_M"; +const WMMA: &str = "MLX_ROCM_WMMA_QMM"; +const THRESHOLD: &str = "MLX_ROCM_QMM_DEQUANT_M_THRESHOLD"; +const CACHE_SIZE: &str = "MLX_ROCM_QMM_DEQUANT_CACHE_SIZE"; +const CACHE_MAX_BYTES: &str = "MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES"; +const STATS: &str = "MLX_ROCM_QMM_DEQUANT_CACHE_STATS"; +/// Cleared in every child before a case applies its own, so the host +/// environment cannot leak in. +const ALL_VARS: &[&str] = &[MAX_M, WMMA, THRESHOLD, CACHE_SIZE, CACHE_MAX_BYTES, STATS]; + +const CHILD_ENV: &str = "MLXCEL_ROCM_QMM_DEQUANT_CACHE_CHILD"; +const CASE_ENV: &str = "MLXCEL_ROCM_QMM_DEQUANT_CACHE_CASE"; +const CHILD_TEST: &str = "child_dequant_cache_case"; +const CHILD_BUDGET: Duration = Duration::from_secs(180); +const DONE_MARKER: &str = "DEQUANT_CACHE done"; +const PASS_PREFIX: &str = "DEQUANT_CACHE pass="; + +const GROUP_SIZE: i32 = 64; +const BITS: i32 = 4; +/// Relative error of a half-precision GEMM against the f32 reference. +const TOLERANCE: f32 = 2e-2; + +fn on_rocm() -> bool { + gpu_backend_kind() == GpuBackendKind::Rocm +} + +/// What the counters must read after each of the two passes. +#[derive(Clone, Copy)] +struct Expect { + hits: u64, + misses: u64, + inserts: u64, + bypasses: u64, + entries: u64, +} + +struct Case { + name: &'static str, + env: &'static [(&'static str, &'static str)], + /// Activation and weight dtype. + dtype: i32, + rows: i32, + k: i32, + n: i32, + /// Whether the GPU output must equal `dequantize` + `matmul` bytewise. + dense_bytes: bool, + after: [Expect; 2], +} + +const NOTHING: Expect = Expect { + hits: 0, + misses: 0, + inserts: 0, + bypasses: 0, + entries: 0, +}; + +fn cases() -> Vec { + vec![ + // 256 rows is above the default ceiling of 128 on gfx1151: the GEMM + // is dequantized per call and the cache stays empty. + Case { + name: "bf16_above_default_ceiling", + env: &[(THRESHOLD, "1")], + dtype: dtype::BFLOAT16, + rows: 256, + k: 4096, + n: 4096, + dense_bytes: true, + after: [ + Expect { + bypasses: 1, + ..NOTHING + }, + Expect { + bypasses: 2, + ..NOTHING + }, + ], + }, + // A ceiling set through the environment also produces the bypass. + Case { + name: "bf16_above_env_ceiling", + env: &[(THRESHOLD, "1"), (MAX_M, "32")], + dtype: dtype::BFLOAT16, + rows: 64, + k: 1024, + n: 1024, + dense_bytes: true, + after: [ + Expect { + bypasses: 1, + ..NOTHING + }, + Expect { + bypasses: 2, + ..NOTHING + }, + ], + }, + // MLX_ROCM_WMMA_QMM=1 removes the ceiling: the fused kernel runs, and + // neither the cache nor the bypass is touched. + Case { + name: "bf16_forced_wmma", + env: &[(THRESHOLD, "1"), (WMMA, "1")], + dtype: dtype::BFLOAT16, + rows: 256, + k: 1024, + n: 1024, + dense_bytes: false, + after: [NOTHING, NOTHING], + }, + // f16 never takes the fused kernel, so it reaches the cached route: + // the first pass misses and inserts, the second hits. + Case { + name: "f16_cached", + env: &[(THRESHOLD, "1")], + dtype: dtype::FLOAT16, + rows: 128, + k: 1024, + n: 1024, + dense_bytes: false, + after: [ + Expect { + misses: 1, + inserts: 1, + entries: 1, + ..NOTHING + }, + Expect { + hits: 1, + misses: 1, + inserts: 1, + entries: 1, + ..NOTHING + }, + ], + }, + ] +} + +fn find_case(name: &str) -> Case { + cases() + .into_iter() + .find(|c| c.name == name) + .unwrap_or_else(|| panic!("unknown case {name}")) +} + +fn pseudo_random(count: usize, seed: u32) -> Vec { + let mut state = seed; + (0..count) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + ((state >> 8) as f32 / (1u32 << 24) as f32) * 2.0 - 1.0 + }) + .collect() +} + +fn max_abs(a: &MlxArray) -> f32 { + let f = mlxcel_core::astype(a, dtype::FLOAT32); + let m = mlxcel_core::max_all(&mlxcel_core::abs(&f)); + mlxcel_core::eval(&m); + mlxcel_core::item_f32(&m) +} + +fn format_stats(s: &DequantCacheStats) -> String { + format!( + "hits={} misses={} inserts={} evictions={} bypasses={} entries={} bytes={}", + s.hits, s.misses, s.inserts, s.evictions, s.bypasses, s.entries, s.bytes + ) +} + +/// Runs the case's GEMM twice on the GPU and checks results and counters. +fn run_case(case: &Case) { + let (rows, k, n) = (case.rows, case.k, case.n); + let w_f32 = mlxcel_core::from_slice_f32(&pseudo_random((n * k) as usize, 0x2151), &[n, k]); + let w = mlxcel_core::astype(&w_f32, case.dtype); + let q = mlxcel_core::quantize_weights_with_mode(&w, GROUP_SIZE, BITS, "affine"); + let packed = mlxcel_core::quantized_weights_w(&q); + let scales = mlxcel_core::quantized_weights_scales(&q); + let biases = mlxcel_core::quantized_weights_biases(&q); + let x_f32 = mlxcel_core::from_slice_f32( + &pseudo_random((rows * k) as usize, 0x2151 ^ rows as u32), + &[rows, k], + ); + let x = mlxcel_core::astype(&x_f32, case.dtype); + + let reference = { + let _cpu = DefaultDeviceGuard::cpu(); + // SAFETY: `biases` is a live array owned by this frame. + let w_deq = unsafe { + mlxcel_core::dequantize( + &packed, + &scales, + &*biases as *const MlxArray, + GROUP_SIZE, + BITS, + "affine", + ) + }; + let w_deq = mlxcel_core::astype(&w_deq, dtype::FLOAT32); + let x_ref = mlxcel_core::astype(&x, dtype::FLOAT32); + let out = mlxcel_core::matmul(&x_ref, &mlxcel_core::transpose(&w_deq)); + mlxcel_core::eval(&out); + out + }; + let scale = max_abs(&reference).max(1e-6); + + let dense = case.dense_bytes.then(|| { + let _gpu = DefaultDeviceGuard::gpu(); + // SAFETY: `biases` is a live array owned by this frame. + let w_deq = unsafe { + mlxcel_core::dequantize( + &packed, + &scales, + &*biases as *const MlxArray, + GROUP_SIZE, + BITS, + "affine", + ) + }; + let out = mlxcel_core::matmul(&x, &mlxcel_core::transpose(&w_deq)); + mlxcel_core::eval(&out); + mlxcel_core::array_to_raw_bytes(&out) + }); + // `dequantize` + `matmul` never goes through QuantizedMatmul, so the + // counters still read zero here. + assert_eq!( + dequant_cache_stats(), + DequantCacheStats::default(), + "{}: counters moved before the first quantized GEMM", + case.name + ); + + for (pass, expect) in case.after.iter().enumerate() { + let got = { + let _gpu = DefaultDeviceGuard::gpu(); + // SAFETY: every reference is a live array owned by this frame. + let out = unsafe { + mlxcel_core::quantized_matmul( + &x, + &packed, + &scales, + &*biases as *const MlxArray, + true, + GROUP_SIZE, + BITS, + "affine", + ) + }; + if let Err(err) = mlxcel_core::try_eval(&out) { + panic!("{} pass {pass} on the GPU failed: {err}", case.name); + } + out + }; + let stats = dequant_cache_stats(); + println!("{PASS_PREFIX}{pass} {}", format_stats(&stats)); + + if let Some(dense) = &dense { + assert!( + mlxcel_core::array_to_raw_bytes(&got) == *dense, + "{} pass {pass}: quantized_matmul differs from dequantize + matmul", + case.name + ); + } + let got = mlxcel_core::astype(&got, dtype::FLOAT32); + let rel = max_abs(&mlxcel_core::subtract(&got, &reference)) / scale; + // Written so a NaN fails too. + assert!( + rel < TOLERANCE, + "{} pass {pass}: GPU result is {rel:e} from the f32 reference", + case.name + ); + + let expected = ( + expect.hits, + expect.misses, + expect.inserts, + expect.bypasses, + expect.entries, + ); + let actual = ( + stats.hits, + stats.misses, + stats.inserts, + stats.bypasses, + stats.entries, + ); + assert_eq!( + actual, + expected, + "{} pass {pass}: (hits, misses, inserts, bypasses, entries) after the GEMM; {}", + case.name, + format_stats(&stats) + ); + assert_eq!(stats.evictions, 0, "{} pass {pass}: evictions", case.name); + if expect.entries == 0 { + assert_eq!(stats.bytes, 0, "{} pass {pass}: cached bytes", case.name); + } else { + let weight_bytes = (n as u64) * (k as u64) * 2; + assert_eq!( + stats.bytes, weight_bytes, + "{} pass {pass}: cached bytes", + case.name + ); + } + } +} + +struct ChildRun { + success: bool, + status: String, + stdout: String, + stderr: String, +} + +fn run_child(case: &Case) -> ChildRun { + let exe = std::env::current_exe().expect("path of this test binary"); + let mut cmd = Command::new(exe); + cmd.args([ + CHILD_TEST, + "--exact", + "--ignored", + "--nocapture", + "--test-threads=1", + ]) + .env(CHILD_ENV, "1") + .env(CASE_ENV, case.name) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + for var in ALL_VARS { + cmd.env_remove(var); + } + for (k, v) in case.env { + cmd.env(k, v); + } + let mut child = cmd.spawn().expect("spawn the child test process"); + + // Drain both pipes on their own threads so the child never blocks on a + // full pipe while the parent waits for it. + let mut stdout_pipe = child.stdout.take().expect("child stdout"); + let mut stderr_pipe = child.stderr.take().expect("child stderr"); + let stdout_reader = thread::spawn(move || { + let mut text = String::new(); + let _ = stdout_pipe.read_to_string(&mut text); + text + }); + let stderr_reader = thread::spawn(move || { + let mut text = String::new(); + let _ = stderr_pipe.read_to_string(&mut text); + text + }); + + let deadline = Instant::now() + CHILD_BUDGET; + let status = loop { + match child.try_wait() { + Ok(Some(status)) => break Some(status), + Ok(None) => {} + Err(err) => { + let _ = child.kill(); + let _ = child.wait(); + panic!("poll the child: {err}"); + } + } + if Instant::now() > deadline { + let _ = child.kill(); + let _ = child.wait(); + break None; + } + thread::sleep(Duration::from_millis(50)); + }; + let stdout = stdout_reader.join().expect("stdout reader thread"); + let stderr = stderr_reader.join().expect("stderr reader thread"); + let status = status.unwrap_or_else(|| { + panic!( + "the child for {} did not finish within {CHILD_BUDGET:?}\n--- child stdout ---\n{stdout}\n--- child stderr ---\n{stderr}", + case.name + ) + }); + ChildRun { + success: status.success(), + status: status.to_string(), + stdout, + stderr, + } +} + +fn check_child(case: &Case) -> Option { + let run = run_child(case); + if run.success && run.stdout.contains(DONE_MARKER) { + eprintln!("{}: ok", case.name); + return None; + } + Some(format!( + "{}: exited with {}\n--- child stdout ---\n{}\n--- child stderr ---\n{}", + case.name, run.status, run.stdout, run.stderr + )) +} + +/// One test, so the children run one after another and never share the GPU. +#[test] +fn ceiling_routed_gemms_skip_the_dequant_cache() { + if !on_rocm() { + eprintln!("skipping: not running on a ROCm device"); + return; + } + let cases = cases(); + let total = cases.len(); + let failures: Vec = cases.iter().filter_map(check_child).collect(); + assert!( + failures.is_empty(), + "{} of {total} cases failed:\n\n{}", + failures.len(), + failures.join("\n\n") + ); +} + +/// The body the parent runs in a fresh process per case. +#[test] +#[ignore = "spawned by ceiling_routed_gemms_skip_the_dequant_cache with one case per process"] +fn child_dequant_cache_case() { + if std::env::var_os(CHILD_ENV).is_none() { + eprintln!("skipping: run through ceiling_routed_gemms_skip_the_dequant_cache"); + return; + } + if !on_rocm() { + eprintln!("skipping: not running on a ROCm device"); + return; + } + let name = std::env::var(CASE_ENV).expect("the parent names the case"); + run_case(&find_case(&name)); + println!("{DONE_MARKER}"); +}