metal: PTQ1_0 exact-float trit extraction (+30-38% decode) and the missing per-expert mat-vec - #157
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
The oversized explanatory comment conflicts with the repository's explicit concise-comment convention.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Optimizes Metal PTQ1_0 decoding and restores per-expert matrix-vector support.
Changes:
- Replaces integer trit extraction with exact floating-point extraction.
- Adds the missing PTQ1_0
MUL_MAT_IDkernel instantiation.
File summaries
| File | Description |
|---|---|
ggml/src/ggml-metal/kernels/mul_mv.metal |
Optimizes PTQ1_0 decoding and enables per-expert dispatch. |
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| // Dot against coefficients already staged in registers by the caller and reused across | ||
| // all nr0 rows. A packed byte is a base-3 fraction of 256: with u = b/256, trit n is | ||
| // t_n = g_{n+1} - 3*g_n, g_k = floor(3^k * u) | ||
| // and 3^k*b <= 61965 is exact in fp32, so the floors are exact and this matches the | ||
| // integer recurrence bit for bit over all 256 bytes and all five positions. Summing | ||
| // t_n*y_n over a byte's trits then collapses to | ||
| // sum_{k=1..4} g_k*(y_{k-1} - 3*y_k) + g_5*y_4 | ||
| // whose coefficients depend only on the activations, so the caller stages those in | ||
| // place of the raw y. The inner loop is one floor and one fma per trit and never | ||
| // leaves the float pipe, which matters because this ISA cannot co-issue integer and | ||
| // floating-point work; the recurrence spent four integer ops and a convert per trit. | ||
| // sumy is subtracted once for the -1 offset, exactly as before. |
…urrence
The mat-vec pulled each trit out of its byte with the base-3 remainder
recurrence, w = v*3, t = w >> 8, v = w & 0xFF: four integer ops and a convert
per trit, on an ISA that cannot co-issue integer and floating-point work. The
kernel was already known to be instruction-bound rather than bandwidth-bound.
A packed byte is a base-3 fraction of 256. With u = b/256, trit n is
g_{n+1} - 3*g_n for g_k = floor(3^k*u), and every 3^k*b is below 2^24, so the
floors are exact and this matches the recurrence bit for bit over all 256
bytes and all five positions. Summing t_n*y_n over a byte's trits collapses to
sum_k g_k*(y_{k-1} - 3*y_k) + g_5*y_4, whose coefficients depend only on the
activations, so the caller stages those in place of the raw y and reuses them
across every row. The inner loop is one floor and one fma per trit with no
integer arithmetic. The qh trit uses the two-floor single-trit form, replacing
a lane-divergent loop that stepped the recurrence up to three times.
Verified against the recurrence for all bytes before the kernel was touched.
test-backend-ops MUL_MAT: 45 PTQ1_0 cases pass, zero failures. Output is not
bit-identical: the coefficient staging reassociates the fp32 sum.
Measured on an M5 Pro, llama-bench r=3, two interleaved passes, dense models:
small decode 221 -> 287 tok/s (+29.5%)
large decode 59 -> 82 tok/s (+37.8%)
Prefill unchanged; the mat-mul path does not use this routine.
…eady claims supports_op accepts PTQ1_0 for MUL_MAT_ID and the dispatcher carries an nsg/nr0 entry for it, but the kernel_mul_mv_id_ptq1_0_f32 pipeline was never instantiated, so the first per-expert PTQ1_0 mat-vec on Metal fails the library lookup and the process segfaults. test-backend-ops hit this on its first PTQ1_0 MUL_MAT_ID case and died there, which truncated the Metal MUL_MAT_ID run to 24 cases while still printing a clean tail; nothing after that case, for any type, was being exercised. Add the instantiation next to the other low-bit per-expert kernels. It reuses kernel_mul_mv_ptq1_0_f32_impl, so it inherits the exact-float extraction from the previous commit. The full Metal MUL_MAT_ID run now completes: 1022 cases, 75 of them PTQ1_0, zero failures.
e297bd2 to
11bf2a9
Compare
|
Rebased onto Measured on an M5 Max against a clean
PTQ1_0 decode is now at or above the PQ2_0 pack of the same size on the 2B, 4B and 9B. Correctness: Merging now so it makes the next release. Left for follow-up PRs, not blocking:
|
Two changes to the Metal PTQ1_0 mat-vec: a decode speedup, and a missing kernel whose absence segfaults any PTQ1_0 mixture-of-experts model on Metal today.
Decode: extract trits in exact float
The mat-vec pulled each trit out of its byte with the base-3 remainder recurrence,
w = v*3; t = w >> 8; v = w & 0xFF. That is four integer ops and a convert per trit, on an ISA that cannot co-issue integer and floating-point work, and the kernel's own comment records that it is instruction-bound rather than bandwidth-bound. So the integer work was the ceiling.A packed byte is a base-3 fraction of 256. With
u = b/256, tritnisg_{n+1} - 3*g_nforg_k = floor(3^k * u), and every3^k * bis below 2^24, so the floors are exact: this matches the recurrence bit for bit over all 256 bytes and all five positions. Summingt_n * y_nover a byte then collapses tosum_k g_k*(y_{k-1} - 3*y_k) + g_5*y_4, whose coefficients depend only on the activations, so the caller stages those in place of the rawyand reuses them across every row exactly as it did before. The inner loop is one floor and one fma per trit with no integer arithmetic. Theqhtrit uses the two-floor single-trit form, which replaces a lane-divergent loop that stepped the recurrence up to three times.Measured on an M5 Pro,
llama-bench -p 512 -n 128 -r 3, two interleaved passes, two dense models roughly five times apart in size:For scale, on the smaller model the same-binary TQ2_0 decode is 290, so PTQ1_0 now decodes at 99 percent of TQ2_0 at 1.75 versus 2.06 bits per weight. Prefill is untouched because the mat-mul path does not use this routine.
Output is not bit-identical to the previous kernel and should not be expected to be: staging coefficients reassociates the fp32 sum. The per-element values are exact; only summation order moves.
Correctness bug: the per-expert kernel did not exist
supports_opaccepts PTQ1_0 forMUL_MAT_IDand the dispatcher has annsg/nr0entry for it, butkernel_mul_mv_id_ptq1_0_f32was never instantiated. The first per-expert PTQ1_0 mat-vec on Metal fails the library lookup (kernel not found in any metal library) and the process segfaults. On currentprism-v7that is any PTQ1_0 MoE model.It also hid itself in the test suite.
test-backend-ops -o MUL_MAT_IDon Metal dies on its first PTQ1_0 case, which happens to come right after thepq2_0case in the type list, so the run printed 24 clean cases and stopped. Nothing after that point, for any type, was being exercised, and the truncated output looks like a complete pass unless you count the cases or check the exit code. With the instantiation added the same run completes: 1022MUL_MAT_IDcases, 75 of them PTQ1_0, zero failures, 3/3 backends.The instantiation reuses the existing impl, so it inherits the extraction change above.
Verification
The identity was checked against the recurrence for all 256 bytes and all positions before the kernel was touched, and the collapse checked exact against the raw trit sum.
test-backend-opson Metal:MUL_MAT45/45,GET_ROWS4/4,MUL_MAT_ID75/75 PTQ1_0 cases, zero failures. The A/B used two build directories with the embedded shader confirmed distinct bystringsimmediately before each measurement, since the shader is baked intolibggml-metal.dyliband a stale or rebuilt arm silently invalidates the comparison.Heads-up for the reviewer
There is uncommitted PTQ1_0 Metal work in another session's worktree on
feat/tq1_0-g128(a qh-only simdgroup broadcast, per its own notes). It touches the sameptq1_0_dot_regand will need reconciling with this if both land.