Skip to content

metal: tensor-op GEMMs on the M5's neural accelerators (Metal 4 matmul2d) - #796

Merged
midagedev merged 8 commits into
mainfrom
mac-nax
Oct 24, 2026
Merged

midagedev merged 8 commits into
mainfrom
mac-nax

Conversation

@midagedev

@midagedev midagedev commented Oct 24, 2026 •

Copy link
Copy Markdown
Owner

The prompt pass's GEMMs on the M5's neural accelerators (Metal 4 tensor ops). Part of #593.

On the contributor's machines (#593) bloomery grows 1.25–1.29× from the M3 Ultra to the M5 Ultra, llama.cpp 2.0–2.1× and mlx-lm 2.2–3.1×. Our GEMMs multiply with simdgroup_matrix only, so they do not reach the M5's matrix units. MLX reaches them through its NAX kernels, llama.cpp through its tensor mul_mm.

  • kernels/gemm_tensor.metal: the dense quantized GEMM (Q4_K, Q6_K, A4; f16 and f32 operands; store, residual, segmented) and the routed grouped GEMM (pair and plain) on mpp::tensor_ops::matmul2d, under the same host names, buffers and grid as today's. Design: MLX's NAX quantized matmul (weights dequantized to threadgroup memory, a 64×64 tile over 2×2 SIMD groups, the weight operand transposed). It stages the activations too (the f16 cast raises the activation fault), keeps relaxed precision off, and runs routed ops 16 places high. Written anew; no MLX code copied (THIRD_PARTY_NOTICES names the design).
  • The shared GEMM ABI (geometry, argument blocks, pair_swiglu) moves to kernels/gemm_shared.metal. Today's kernels are unchanged.
  • MatrixUnits (library.rs) is the one rule: Metal 4 and an Apple GPU generation ≥ 17 (≥ 18 for the phone class), MLX's is_nax_available. Such a device compiles the tensor source at language 4.0 beside the 3.2 library, and the GEMM pipelines come from it. Every other device keeps today's kernels. No flag.
  • metal_gemm_time times the tensor arms beside today's (dense shapes and --routed), and its header prints the rule's verdict. generate_metal prints gemm_path=.

Bits: the tensor op's sum order is the hardware's, so the new kernels are held to a band for any order of the K products (accumulation term K instead of K/8 + 8; derivation in clauses_gemm.rs), against the same f64 oracle. Every tensor clause runs the kernels forced, on any Metal 4 device. The M1 Pro and the M4 Pro have no matrix units but run the tensor ops correctly, so the clauses run here.

Gates (M1 Pro, fixture), each rc 0 at this branch's tree on top of #792 (before that PR's review fix, which touches only the delta ring copy): cargo fmt --all --check, check-unsafe, check-comments, just check-recipes, just mac-metal-lint, just lab-metal, just lab-metal-e2e. New clauses and the mutant each catches:

clause mutant red on it
tensor_gemm_q4k tensor_k_step_skipped
tensor_gemm_q6k tensor_rows_offset
tensor_gemm_a4 tensor_resid_stores
tensor_gemm_seg tensor_seg_out
tensor_gemm_range tensor_range_no_raise
tensor_grouped_q4k tensor_pair_swapped, tensor_routed_short_block
tensor_grouped_q6k tensor_routed_k_step_skipped
tensor_grouped_a4 tensor_routed_block_rows, tensor_tile_no_raise
q_tensor (e2e) tensor_pass_pair_swapped
gate_metal_qwen35moe_kernels --mutants-all: PASS (124 mutants, each red on its clause)
clause q_tensor: PASS (tensor-op GEMMs: 4 teacher-forced tokens after 1,024 ids: worst layer output 1.497e-1 of the stepped run's (band 0.26), 8 router sites with other picks (cap 12), 2 of 5 picks equal)
gate_metal_qwen35moe_e2e --mutants-all: PASS (12 mutants, each red on its clause)

q_tensor prints the same worst and flip count as q_free (the simdgroup path). The pass does run the tensor kernels: the dense and routed pipelines both come from the forced library, and tensor_pass_pair_swapped turns q_tensor red. Why the printed values match was not measured.

M4 Pro (no matrix units; development-loop numbers, not public), metal_gemm_time, one buffer at a time:

today's kernels tensor-op
dense f16 Q4_K 8192×2048, 2048 rows 10.050 ms 10.045 ms
dense f32 Q4_K 8192×2048, 2048 rows 11.088 ms 11.374 ms
routed layer, U = 2048 18.39 ms 20.34 ms
routed layer, U = 512 5.93 ms 7.53 ms

The M4 Pro never selects them. The M5 gain is not measured yet. Derived for the M5 Ultra: the GEMMs 485 → ~170 ms of a 2048-row pass (bounded then by the staging per K step), the pass ~572 → ~255 ms. It merges before an M5 run (the maintainer's call): an M5 selects these kernels by default, and the contributor's M5 run is the first measurement of their speed and bits on matrix units.

The code and gates were written with Claude Code rounds and reviewed by the maintainer.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: MUST_FIX: The change duplicates the activation fault policy and the mutant-program contract.

Must fix

1. Reuse the staged activation conversion (kind B, R14)

crates/metal/kernels/gemm_tensor.metal:147 repeats the half-range check, fault raise and NaN substitution in crates/metal/kernels/gemm.metal:219 and crates/metal/kernels/gemm_routed.metal:147.
These are three owners of the same conversion policy. Sharing the existing helper now removes two copies before that policy changes.

Move tensor_stage_x unchanged into gemm_shared.metal. Replace the dense stage_b lambda's inner g loop with:

tensor_stage_x<T>(xr, b_dst[b], faults, layer);

Replace the routed stage lambda's inner g loop with:

tensor_stage_x<T>(xr, bd, faults, args.layer);

In routed_range_off, replace both edit strings with the shared helper's check:

"const bool4 bad = !(fabs(v) <= 65504.f);",
"const bool4 bad = bool4(false);",

2. Share the mutant-program support (kind B, R14)

crates/metal/src/bin/gate_metal_qwen35moe_e2e/mutants.rs:180 and crates/metal/src/bin/gate_metal_qwen35moe_kernels/mutants.rs:1210 implement the same program selection, occurrence validation and replacement. The new Program type and table walk are also copied.
One shared implementation removes five repeated definitions and keeps future validation changes in one place.

Keep each file's imports and mutation tables. Replace its Program, Mutant, all, find and source_of definitions with:

include!("../shared/mutants.rs");

Put their common definitions in src/bin/shared/mutants.rs:

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Program { Main, Tensor }

pub struct Mutant {
    pub name: &'static str,
    pub clause: &'static str,
    pub why: &'static str,
    pub edits: &'static [(&'static str, &'static str, usize)],
}

pub fn all() -> impl Iterator<Item = (&'static Mutant, Program)> {
    MUTANTS.iter().map(|m| (m, Program::Main))
        .chain(TENSOR_MUTANTS.iter().map(|m| (m, Program::Tensor)))
}

pub fn find(name: &str) -> Option<(&'static Mutant, Program)> {
    all().find(|(m, _)| m.name == name)
}

pub fn source_of(mutant: &Mutant, program: Program) -> Result<String, String> {
    let mut text = match program {
        Program::Main => SOURCE,
        Program::Tensor => TENSOR_SOURCE,
    }.to_string();
    for (find, replace, count) in mutant.edits {
        let found = text.matches(find).count();
        if found != *count {
            return Err(format!(
                "mutant {}: {find:?} occurs {found} times, not {count}",
                mutant.name
            ));
        }
        text = text.replace(find, replace);
    }
    Ok(text)
}

Non-blocking

⚠ Supply M5 validation: crates/metal/kernels/gemm_tensor.metal:398 rounds routed runs to 16 places; the old kernel uses 8-place fragments. For runs of 1–8 places, this adds 8 × 64 × K products per tile, multiplied by short tiles per pass. The geometry and float-sum-order classes require the owning gates and a timed A/B. Attach an M5 run with both commits, the model, commands and output to establish the gain before enabling this path there.

Checked

Duplicates: the three activation conversions and both sets of five mutant-support definitions were compared. Both proposed extractions reduce repeated code and definitions.
Performance: the tensor kernel's 16-row blocks and the simdgroup kernel's 8-row switch confirm the short-run product count.
Scope: full review; neither revert nor mirror exclusions apply. No ledger lessons are in force.
Escape catalog: no entry applies to these findings.

Notes

Runtime gate results are author-reported. Metal runtime gates were not executed during validation.

🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the last review:

Verdict: CLEAN: The new commits resolve both duplication findings; no must-fix item remains.

Checked

Duplicates: 6 moved items checked; each now has one owner. The dryer pair, LibraryPath::name and GemmPath::name, renders different enums and is unchanged this round.
Complexity: the functions flagged by crapper are unchanged this round. The extracted loop preserves its branches.
Silent failure: the shared conversion retains the half-range check, fault raise and NaN staging. The shared mutant support rejects incorrect occurrence counts.
Tests: no test cases or assertions added or changed. Existing range mutants still target the shared check and fault raise; they were not executed here.
Gates: no bands or pins changed. The three affected mutation strings each match once in their assembled program.
Adversarial checks: the half-range comparison still rejects NaN and overflow and accepts the finite-half boundary. Both programs declare the macro and fault helper before staging; both shared Rust includes resolve.
Escape catalog: no C1–C4 pattern matches the incremental diff.

Notes

Metal compilation and runtime gates were not run here. The PR's gate results are author-reported.

🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)

github-actions[bot]
github-actions Bot previously approved these changes Oct 24, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: CLEAN: carried from e33d9cbc, not re-reviewed.

The approval of e33d9cbc stands for 4264bf24. This pull request's 7 commits are the same patches in the same order: git patch-id --verbatim (the changed lines, the lines around them and whitespace) and the subject are equal for every commit. That review ran under the reviewer instructions, tools and slag ledger that review this head, and it approved.

The base moved and also changed THIRD_PARTY_NOTICES.md, which this pull request changes too. Each of its hunks still applies with the same context lines.

A carry covers each patch's own lines and the lines around them. What the base changed elsewhere is covered by CI on this head and by the landing gates.

🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)

midagedev and others added 7 commits October 25, 2026 07:18
…p ones; the GEMM ABI moves to gemm_shared.metal

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ
…a GPU with matrix units

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ
…, a band for any sum order

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ
… the tensor-op kernels

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ
… --routed

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ
…, the one checked cast in gemm_shared.metal

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ
…source_of) shared by both gates

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ
github-actions[bot]
github-actions Bot previously approved these changes Oct 24, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the last review: there were no must-fix or non-blocking items. The rebase resolution retains the base's load-timing fields alongside the GEMM-path documentation.

Verdict: CLEAN: the interdiff preserves the timing output and the previously reviewed GEMM selection.

Checked

Duplicates: 0 new items in the interdiff. Both dryer candidates (LibraryPath::name and GemmPath::name) were read; they map different facts and share only structure.
Complexity: crapper flags device dispatch in Library::compile, tensor handling in both gate mains, and tensor arms in metal_gemm_time::{mac::one,routed::run}. None changed in this interdiff.
Silent failure: 0 new sites. Session-open errors still propagate before the summary is printed.
Tests: no tests changed in the interdiff, so no new mutation targets. A read-only comparison confirms the executable body equals base plus the previously reviewed GEMM-path print.
Gates: no pins, bands or cases moved in the interdiff.
Adversarial checks: timing-field loss is blocked by the preserved output statements. A mismatched GEMM report is blocked because generate_metal.rs:768 passes the reported library to the session. Zero decode passes retain the base's explicit handling.
Escape catalog: no C1–C4 patterns matched.

Notes

Metal runtime gates were not run here; they require Apple Silicon and macOS 26.2 or later. The PR states that M5 timing remains unmeasured.

🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)

…te's mutants.rs, not a module-level include!

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q5zTV27LbaVZFTjnkdmTjJ

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the last review: there were no must-fix or non-blocking items.

Verdict: CLEAN: the incremental refactor preserves mutant enumeration, lookup and error propagation.

Checked

Duplicates: 4 added adapters checked by Grep; table traversal and source editing retain their shared owner. The dryer pair maps distinct library-path and GEMM-path facts.
Complexity: crapper flags dispatch in Library::compile, tensor handling in both gate mains, and tensor arms in metal_gemm_time::{mac::one,routed::run}. These functions are unchanged in this increment.
Silent failure: 0 new sites. Unknown names and edit-count mismatches still produce named failures.
Tests: no assertions, clauses or mutant edits changed. Read-only checks confirm 124 kernel mutants and 12 e2e mutants, unique names and preserved program tags. No new assertion mutation target.
Gates: no pins, bands or cases changed.
Adversarial checks: both module paths resolve to the shared file; tensor entries remain chained after main entries; unknown-name lookup still reaches exit code 64.
Performance: gate tooling only; no inference hot path changed.
Escape catalog: no C1–C4 patterns matched.

Notes

Metal runtime gates were not run here; they require Apple Silicon and macOS 26.2 or later.

🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)

@midagedev
midagedev merged commit be9bd11 into main Oct 24, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant