Fix Metal small-batch matmul parity for GLM verification - #1078
Conversation
📝 WalkthroughWalkthroughChangesVerification parity and backend execution
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant FullModelSampling
participant VerifyTargetPlan
participant ParityChecks
participant StageSession
CLI->>FullModelSampling: verification widths and continuation steps
FullModelSampling->>VerifyTargetPlan: choose target plan
VerifyTargetPlan-->>FullModelSampling: width-specific verify tokens
FullModelSampling->>ParityChecks: run requested parity checks
ParityChecks->>StageSession: batched verification and serial continuation
StageSession-->>ParityChecks: predictions, state, token, and signal results
ParityChecks-->>FullModelSampling: detailed parity results
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review. |
i386
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
crates/skippy-bench/src/verify_window_local.rs (2)
287-287: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename
max_widthto reflect that it receives a token count.Callers pass
target_token_count(args)(max width + continuation steps), and the loop then producescount + 1targets. Themax_widthname makes the+ 1slack hard to reason about.Also applies to: 296-297
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-bench/src/verify_window_local.rs` at line 287, Rename the max_width parameter and all related references in the affected function to indicate it represents a target token count, not a width. Update callers such as target_token_count(args) and the loop’s count-plus-one logic consistently, preserving the existing behavior.
469-504: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueContinuation parity is decided by an internal
bail!, makingcontinuation_matchedalwaystrue.Both the serial and batched calls already assert each step equals the plan, so the
batched_continuation == serial_continuationcomparison at Line 399 can never be false. Consider returning the decoded tokens without the per-step assertion and letting the caller's comparison (plus a plan comparison) be the single gate, so the reported flag carries real information.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-bench/src/verify_window_local.rs` around lines 469 - 504, The per-step bail in decode_expected_continuation prevents continuation parity mismatches from reaching the caller, making continuation_matched uninformative. Remove the predicted-versus-expected assertion from decode_expected_continuation while preserving token decoding and collection, and let the caller compare serial and batched continuations against each other and the target plan as the sole validation gate.third_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch (1)
76-92: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNested ternaries for
nsg/r0ptgare dense but internally consistent; consider extracting for readability.The logic checks out (F32's
nr0=2fixed row-count matches the 256B smem sizing in device.cpp; Q4_K/Q6_K's2*nsgmatches per-simdgroup row assignment), but the multi-way nested ternary spanning type checks is hard to scan at a glance. A small named helper (e.g.compute_batch_invariant_row_params(type, ne00)) returning{nsg, r0ptg}would make future changes to this dispatch logic safer to review.♻️ Illustrative refactor direction
- const int nsg = op->src[0]->type == GGML_TYPE_F32 ? - std::min(4, (ne00 + 127) / 128) : 2; + const int nsg = compute_mul_mv_ext_nsg(op->src[0]->type, ne00); @@ - const int16_t r0ptg = batch_invariant_ext_type ? - (op->src[0]->type == GGML_TYPE_F32 ? 2 : - op->src[0]->type == GGML_TYPE_Q5_0 ? N_R0_Q5_0*nsg : 2*nsg) : - nypsg*nsg; + const int16_t r0ptg = compute_mul_mv_ext_r0ptg( + op->src[0]->type, batch_invariant_ext_type, nsg, nypsg);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@third_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch` around lines 76 - 92, Extract the batch-invariant row-parameter selection currently embedded in the nsg and r0ptg ternaries within ggml_metal_op_mul_mat into a small named helper, such as compute_batch_invariant_row_params, returning both values. Preserve the existing type- and ne00-dependent behavior, and keep the non-batch-invariant nypsg*nsg path unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/skippy-bench/src/verify_window_local.rs`:
- Around line 341-454: Split run_parity_check into semantically named helpers
that keep each function below Clippy’s line-count threshold. Introduce a shared
ParitySide collection helper for serial and batched execution, state export,
token signals, and signal-window formatting, reusing one layer_end conversion;
extract the bail! message construction into format_parity_failure. Preserve the
existing comparisons, diagnostics, and VerifyParityCheck output.
- Around line 403-410: Update the parity result logic around
signal_window_matched so it participates in the matched conjunction and
signal-window divergence cannot pass silently. Also ensure the returned report
preserves the actual aggregate match status, returning matched: false for failed
checks instead of always reporting true after an early bail; retain the existing
failure diagnostics and successful result behavior.
- Around line 1088-1091: Update the throughput reporting around the sample-width
selection and fields such as batched_width2, split_inprocess_width2, and
verified_tokens_per_sec so calculations use the actual selected verify width
from verify_widths rather than assuming 2. Ensure labels and JSON field names
accurately represent configurable widths, preserving or intentionally updating
the report contract consistently.
- Around line 456-467: Update serial_decode_expected to fail immediately when
decode_step_frame_sampled_mtp returns a negative predicted value, matching
choose_target_plan’s behavior; do not skip the sentinel or continue collecting
later predictions, and return an appropriate error using the existing
error-context conventions.
---
Nitpick comments:
In `@crates/skippy-bench/src/verify_window_local.rs`:
- Line 287: Rename the max_width parameter and all related references in the
affected function to indicate it represents a target token count, not a width.
Update callers such as target_token_count(args) and the loop’s count-plus-one
logic consistently, preserving the existing behavior.
- Around line 469-504: The per-step bail in decode_expected_continuation
prevents continuation parity mismatches from reaching the caller, making
continuation_matched uninformative. Remove the predicted-versus-expected
assertion from decode_expected_continuation while preserving token decoding and
collection, and let the caller compare serial and batched continuations against
each other and the target plan as the sole validation gate.
In
`@third_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch`:
- Around line 76-92: Extract the batch-invariant row-parameter selection
currently embedded in the nsg and r0ptg ternaries within ggml_metal_op_mul_mat
into a small named helper, such as compute_batch_invariant_row_params, returning
both values. Preserve the existing type- and ne00-dependent behavior, and keep
the non-batch-invariant nypsg*nsg path unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 865562ed-baf7-4823-ac1c-0604d58f97f8
📒 Files selected for processing (3)
crates/skippy-bench/src/cli.rscrates/skippy-bench/src/verify_window_local.rsthird_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch
|
Reviewed the Metal patch against the pinned llama.cpp source rather than taking the dispatch geometry on faith. VerifiedThe kernel geometry checks out against the pin:
Issues1. Hardcoded throughput denominator is now wrong ( 2. 3. 4. 5. Signal-window equality via 6. Silent semantic change. Risks worth settling before mergeBlast radius. This isn't verifier-scoped — it changes Metal dispatch policy for every model this fork runs. Batch sizes 9–16 on F32/Q5_0/Q4_K/Q6_K previously reached Unenforced width ceiling. Parity holds only for Silent drift against upstream. The patch re-derives Minor
Description mismatchThe body claims the PR "remove[s] the temporary serial verifier fallback and Metal kernel-disable guards", but the diff against TestsCLI parsing coverage is good, and |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
third_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch (1)
221-315: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
r1ptgtemplate parameter is vestigial in all four batch-invariant kernels.Each new kernel (
(void) r1ptg;at lines 233, 250, 266, 283) ignores the template parameter and delegates to the same canonical_implcall regardless of its value, sokernel_mul_mv_ext_*_r1_2through_r1_5compile to four identical bodies per type. This is harmless (host names are 1:1 replacements of what was removed, and the canonical_implalready indexes rows viatgpig.y), but it's now purely a naming/lookup artifact rather than a real specialization axis — worth a comment or eventual consolidation to avoid confusing future maintainers into thinkingr1ptgstill affects behavior here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@third_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch` around lines 221 - 315, The r1ptg template parameter is unused in all batch-invariant kernels and only exists to generate host-name variants. Add a concise comment near the four kernel templates or their instantiations documenting that r1ptg is intentionally retained solely for naming/lookup compatibility and does not affect implementation behavior; do not alter the canonical _impl calls or generated host names.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In
`@third_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch`:
- Around line 221-315: The r1ptg template parameter is unused in all
batch-invariant kernels and only exists to generate host-name variants. Add a
concise comment near the four kernel templates or their instantiations
documenting that r1ptg is intentionally retained solely for naming/lookup
compatibility and does not affect implementation behavior; do not alter the
canonical _impl calls or generated host names.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: f65b281c-42e8-4c48-9a9d-d652e7808d21
📒 Files selected for processing (2)
crates/skippy-bench/src/verify_window_local.rsthird_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch
🚧 Files skipped from review as they are similar to previous changes (1)
- crates/skippy-bench/src/verify_window_local.rs
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/skippy-bench/src/cli.rs`:
- Around line 247-252: Enforce the documented maximum value of 16 at the CLI
boundary for both VerifyWindowLocalArgs::verify_widths elements and
sample_width. Add clap validation or equivalent argument constraints so invalid
values are rejected before program entry, while preserving the existing defaults
and valid-value behavior.
In
`@third_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch`:
- Around line 17-26: Update ggml_metal_library_get_pipeline_mul_mv_ext to add a
Q5_0-specific smem override for the batch-invariant pipeline, matching the
existing Q5_0 kernel path with 32*N_R0_Q5_0*sizeof(float). Ensure the override
applies when the source types select Q5_0 and preserves the existing F32
override.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ce7e6c5a-bc01-495b-9d97-a6f09f4efd22
📒 Files selected for processing (3)
crates/skippy-bench/src/cli.rscrates/skippy-bench/src/verify_window_local.rsthird_party/llama.cpp/patches/0043-ggml-metal-make-small-batch-matmul-batch-invariant.patch
🚧 Files skipped from review as they are similar to previous changes (1)
- crates/skippy-bench/src/verify_window_local.rs
|
Thanks for the detailed review. I've addressed the findings in
The earlier Linux CUDA failure was runner infrastructure ( |
* origin/main: Fix Metal small-batch matmul parity for GLM verification (#1078) Handle K-only transposed KV page import and export (#1084) Refresh llama.cpp upstream patch queue (#1085) chore: improve embedded native-runtime compatibility guidance (#1043) fix(console-ui): chat transcript snapping during live status updates (#1083) ci: bump Linux CUDA slim container to gha-convention base runner image fix: record activation cache prefix identities (#1041) fix: read-only model download caches (#1042) ci: disable sccache for Windows ROCm native runtime build (#1087) ci: fix v0.74 release GPU builds (sccache disk-only + force_hosted_runners) (#1086) Make release sccache failures non-fatal (#1079) Keep client-only nodes out of model election (#1074) # Conflicts: # crates/mesh-llm-host-runtime/src/runtime/auto_join.rs # crates/mesh-llm-host-runtime/src/runtime/tests/auto_join.rs # third_party/llama.cpp/patches/0004-Add-lanes-external-media-and-chat-grammar-support.patch
3 participants
Fixes #1076.
Summary
This PR fixes GLM-DSA batched
verify_tokensparity at the Metal matrix-multiplication boundary while preserving genuine one-call multirow execution.mul_mv_extkernels that use canonical one-row arithmetic and reduction order;--sample-widthinskippy-bench verify-window-local, require that timed widths are parity-checked, and reject widths above the proven ceiling during CLI parsing;No environment-variable guard or serial verifier fallback is required by this fix.
Root cause
Metal changed matrix multiplication implementation as verification width grew:
mul_mv;mul_mv_extwith a different reduction shape;mul_mm;Those paths were within normal backend numerical tolerances, but were not batch-invariant. The drift changed canonical exported state and token signals, and could later change greedy continuation. Direct sparse prefill and FlashAttention were not the root cause.
The repaired kernels keep one batched graph and parallel GPU execution while making every row follow the canonical one-row computation.
Scope
The batch-invariant path is bounded to widths 16 or lower:
Existing dispatch remains unchanged outside those source types and widths.
verify-window-localrejects larger widths until equivalent exactness is implemented and proven.The parity harness deliberately builds an always-accepted greedy target stream so it measures kernel/state parity rather than proposal quality. Native-MTP proposal quality and acceptance remain the responsibility of the dedicated MTP benchmarks.
GenerationSignalWindowis now compared as a typed value. It remains advisory because one batched call and N serial calls necessarily have different call-history window shapes; canonical model state andTokenSignalremain exactness gates.Physical validation
Validation used Apple M3 Ultra Metal, the public GLM-4.7 Flash MTP Q4_K_M GGUF, greedy sampling, one session, F16 K/V, and FlashAttention auto.
Exactness
Widths 1, 2, 4, 9, and 16 each passed:
This is 5/5 parity cells and 1,280 checked continuation decodes. Report SHA-256:
bfd74dae283b798cacb961d38ddae8dd187feffc7c1c986620059f2209fa67c4.The earlier regression matrix also covered 506-token and 20,006-token prompts at widths 1/2/4/9, with 8/8 parity cells and 2,048 continuation decodes.
Before/after performance
Each row below uses 64 measured iterations after 16 warmups. The baseline omits only patch 0043; Rust harness, model, prompt, cache types, FlashAttention mode, and hardware are identical.
The baseline failed exported-state and token-signal parity at both widths. The patched path passed both, so correctness did not require a performance regression.
Gates
cargo fmt --all -- --check;cargo test --locked -p skippy-bench(73passed);cargo clippy --locked -p skippy-bench --all-targets -- -D warnings;274/274MUL_MATcases passed.Summary by CodeRabbit