huggingface / huggingface/candle
Metal: K-quant mul_mv kernels store past ne01; Q2K over-dispatches 2x on top of it
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
The Metal K-quant `mul_mv` kernels store past `ne01`. There appear to be two
independent causes: a missing bound check that affects the whole family, and a
dispatch-table mismatch specific to Q2K that makes it much worse. There is a
self-contained failing test on a branch cut from current `main`.
I have deliberately **not** written a fix — the two obvious fixes are not
equivalent and interact, so the choice looks like a maintainer's call.
**Environment:** Apple M1 Max, macOS 26.6, `main` at
`6f74e7c390c717f8fd34f23ce02aceb058173370`, debug profile, Metal backend.
Reproduction branch:
[`oetiker/candle@repro/metal-q2k-mv-store-overrun`](https://github.com/oetiker/candle/tree/repro/metal-q2k-mv-store-overrun)
(commits `e4f69b54` and `60d1c2b5`; adds two tests to
`candle-metal-kernels/src/tests.rs`, changes no kernel and no dispatch code).
## There are two independent defects
### Defect 1 — missing store guard, family-wide
None of the K-quant `mul_mv` kernels bound-check their store:
```c
for (int row = 0; row < N_DST; ++row) {
all_sum = simd_sum(sumf[row]);
if (tiisg == 0) {
dst[r1*ne0 + im*ne0*ne1 + first_row + row] = all_sum; // quantized.metal:4667
}
}
```
The legacy-quant kernels in the same file do carry the check, at `:2374` and
`:2551`:
```c
if (tiisg == 0 && first_row + row < ne01) {
```
Consequence: any `ne01` that is not a multiple of the kernel's row stride
overruns by up to `stride − 1` rows. This affects q3_K, q4_K, q5_K and q6_K.
### Defect 2 — dispatch-table mismatch, Q2K only
`kernel_mul_mv_q2_K_f32_impl` advances **8** rows per threadgroup —
`first_row = (r0 * N_SIMDGROUP + sgitg) * N_DST` with `N_SIMDGROUP == 2` and
`N_DST == 4`, both defined once at `quantized.metal:2306-2307` with no `#undef`
anywhere in the file. But `call_quantized_matmul_mv_t` gives Q2K `align = 4`
(`candle-metal-kernels/src/kernels/quantized.rs:76-83`) and dispatches
`ceil(ne01 / align)` threadgroups (`:115`, `divide` = `div_ceil` at `:359`).
So Q2K launches **twice** the threadgroups its own row stride needs, and
overruns by ~`ne01` values at *every* size — including sizes that are exact
multiples of its 8-row stride, where defect 1 alone would cause no overrun at
all. Every other kernel's stride matches its `align`.
## Evidence
`cargo test -p candle-metal-kernels q2k_mul_mv_writes_past_ne01 -- --nocapture`,
verbatim:
```
running 1 test
thread 'tests::q2k_mul_mv_writes_past_ne01' (5324844) panicked at src/tests.rs:2576:5:
kernel_mul_mv_q2_K_f32 wrote 12 f32 values past ne01 = 12: indices [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
test tests::q2k_mul_mv_writes_past_ne01 ... FAILED
failures:
failures:
tests::q2k_mul_mv_writes_past_ne01
test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 61 filtered out; finished in 1.18s
```
`cargo test -p candle-metal-kernels kquant_mul_mv_store_overrun_survey -- --nocapture`,
verbatim. Numbers are f32 values written past `ne01`, into a 4× oversized
sentinel-filled destination:
```
--- K-quant mul_mv store overrun survey ---
Q2K ne01=8 overrun=8 in_range_ok=true
Q2K ne01=10 overrun=14 in_range_ok=true
Q2K ne01=12 overrun=12 in_range_ok=true
Q2K ne01=16 overrun=16 in_range_ok=true
Q2K ne01=17 overrun=23 in_range_ok=true
Q3K ne01=8 overrun=0 in_range_ok=true
Q3K ne01=10 overrun=2 in_range_ok=true
Q3K ne01=12 overrun=0 in_range_ok=true
Q3K ne01=16 overrun=0 in_range_ok=true
Q3K ne01=17 overrun=3 in_range_ok=true
Q4K ne01=8 overrun=0 in_range_ok=true
Q4K ne01=10 overrun=2 in_range_ok=true
Q4K ne01=12 overrun=0 in_range_ok=true
Q4K ne01=16 overrun=0 in_range_ok=true
Q4K ne01=17 overrun=3 in_range_ok=true
Q5K ne01=8 overrun=0 in_range_ok=true
Q5K ne01=10 overrun=2 in_range_ok=true
Q5K ne01=12 overrun=0 in_range_ok=true
Q5K ne01=16 overrun=0 in_range_ok=true
Q5K ne01=17 overrun=3 in_range_ok=true
Q6K ne01=8 overrun=0 in_range_ok=true
Q6K ne01=10 overrun=0 in_range_ok=true
Q6K ne01=12 overrun=0 in_range_ok=true
Q6K ne01=16 overrun=0 in_range_ok=true
Q6K ne01=17 overrun=1 in_range_ok=true
```
Reading: q3/q4/q5 (stride 4) overrun only at `ne01` ∈ {10, 17}, by `stride − 1`
or less — defect 1. q6_K (stride 2) only at odd `ne01`, by 1 — defect 1. Q2K
overruns at **all five** sizes, including 8 and 16 — defect 2. `in_range_ok`
confirms the computed rows themselves are correct throughout, so this is purely
a store-bounds problem, not an arithmetic one.
## How the tests work
Both call `call_quantized_matmul_mv_t` directly with an **all-zero** block of the
relevant quant type. `d == dmin == 0` makes every dequantized weight 0, so the
entire correct output is exactly `0.0` and all the information lives in the
canary — a destination buffer 4× the required size, pre-filled with a sentinel.
Each test asserts the in-range rows are `0.0` first, so a canary hit can never be
confused with a compute bug.
No quantizer, no GGUF, no model, no `candle-core`. Everything is inside
`candle-metal-kernels`.
## What is measured, what is inferred, what is only suspected
Keeping these apart deliberately, because only the first tier is demonstrated.
**MEASURED.** The kernels write past `ne01`. Deterministic, reproduced at five
sizes across five quant types, with the magnitudes in the table above. This is
what the two tests show and it is the whole of what they show.
**INFERRED, from verified source but not measured.**
`candle_core::quantized::metal::QMatMul::fwd_mv`
(`candle-core/src/quantized/metal.rs:313-334`) allocates the destination as
exactly `dst_shape.elem_count()` f32 — that is `m * n` — and then calls the
kernel once per batch row into slot `batch_id * n`. Given a measured Q2K overrun
of ~`n`, the final batch row's store reaches about `(m + 1) * n` in an `m * n`
buffer. For single-token decode (`m == 1`) the buffer is `n` f32 and the kernel
writes ~`2n`. That puts a write outside the allocation, into candle's Metal
buffer pool, where Metal does not bounds-check. We have **not** demonstrated
this; it is arithmetic on top of the measured overrun.
**SUSPECTED, undemonstrated.** candle creates its compute encoders with
`MTLDispatchType::Concurrent`
(`candle-metal-kernels/src/metal/command_buffer.rs:24` and `:35`), so the
per-batch-row dispatches inside one encoder carry no ordering guarantee relative
to each other. If so, one row's overrun can land after the next row's legitimate
store, making intermediate output slots racily wrong rather than merely
transiently clobbered. We have not built a test for this and are not claiming it.
*(An earlier version of this write-up asserted that serial dispatch ordering
repairs every slot but the last. That was wrong — it assumed the Metal default
rather than checking candle's encoder construction. Retracted.)*
## Two fixes, and why we did not pick one
1. **Guard the store** — add `&& first_row + row < ne01`, matching `:2374`.
Local to the kernels, changes no caller, fixes defect 1 everywhere.
2. **Correct the table** — set Q2K's `align` to 8. Fixes defect 2 at its root.
But `call_quantized_matmul_mv_t` is shared by every quantized mv caller in
`candle-core`, so this changes launch geometry for all of them.
These are not alternatives to each other, and the interaction matters:
**fixing the guard alone would MASK defect 2 rather than fix it.** Q2K would stop
writing out of bounds, but it would still dispatch twice the threadgroups it
needs, with half of them computing full dot products whose results are then
discarded by the new guard — a permanent ~2× waste on every Q2K matrix-vector
product, invisible because the numbers would now be correct.
Doing both is probably right. We have not assumed that.
## Reproducing
```
git clone https://github.com/oetiker/candle && cd candle
git checkout repro/metal-q2k-mv-store-overrun
cargo test -p candle-metal-kernels q2k_mul_mv_writes_past_ne01 -- --nocapture
cargo test -p candle-metal-kernels kquant_mul_mv_store_overrun_survey -- --nocapture
```
Requires Metal hardware. The first test fails; the second prints the table.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by checking the reproduction branch and running the two named tests in candle-metal-kernels/src/tests.rs on Metal hardware. Read the K-quant store sites in quantized.metal and the Q2K alignment and dispatch logic in candle-metal-kernels/src/kernels/quantized.rs. Done means the tests pass without canary writes, in-range results remain correct, and the agreed kernel and dispatch behavior is covered by regression tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100