e1894e93c2
Closes the AMX gap surfaced in BENCHMARKS.md. Apple Silicon's
matrix coprocessor is reachable via Accelerate.framework's
cblas_sgemm; this wave adds a thin FFI wrapper behind an
`accelerate` cargo feature, plus a build.rs that links
Accelerate when the feature is on.
NEW API
pub fn matmul_f32_accelerate(
a: &[f32], w: &[f32], bias: &[f32],
m, n, k: usize, out: &mut [f32],
) -> Result<(), MatmulError>
Same shape contract as matmul_f32_scalar. On macOS aarch64
with `accelerate` on: dispatches to cblas_sgemm. On every
other build/host: returns Err(MatmulError::Unsupported).
Callers can probe support with a 1×1×1 call and route to
scalar fallback.
pub enum MatmulError {
Unsupported, // backend not on this build/host
}
CARGO + BUILD
[features]
accelerate = [] # macOS aarch64 only; build.rs handles the link
build.rs:
Conditionally emits `cargo:rustc-link-lib=framework=Accelerate`
when CARGO_FEATURE_ACCELERATE is set AND target_os=macos.
No-op on every other combination.
Default build pulls in zero Accelerate deps. CI on Linux /
Windows runs the existing cargo check; macOS CI can additionally
build --features accelerate.
REAL NUMBERS (M-series dev box, BENCHMARKS.md updated)
Shape Scalar AMX Speedup AMX GFLOPS
1 × 256 × 256 27 µs 1.4 µs 19.8× 95
1 × 1024 × 1024 632 µs 13.5 µs 47.0× 156
1 × 2048 × 2048 2.73 ms 67.6 µs 40.4× 248
1 × 4096 × 4096 11.45 ms 710 µs 16.1× 47 (over-large for working set)
8 × 1024 × 1024 5.07 ms 83.8 µs 60.5× 200
32 × 1024 × 1024 20.30 ms 90.4 µs 224.6× 742 ← HEADLINE
The 224× speedup at m=32 is the AMX-favoured regime where
weight rows reuse across multiple activation rows. Compute-
bound + weight-reuse-dominated workloads (KV-cache prefill,
batched inference) get the full AMX win.
Single-token (m=1) caps at ~250 GFLOPS at medium k —
bandwidth-bound, not compute-bound. Still ~50× faster than
scalar, just not at the m=32 peak.
The 4096² m=1 drop to 47 GFLOPS suggests AMX's working set
is exceeded at that size; needs tiling. Future wave.
WHY cblas_sgemm WITH B AS TRANS
Our weight is stored [n, k] row-major. cblas_sgemm wants
C = A * B (or with transposes). To get out[m, n] = a[m, k] *
w[n, k]^T (where w is the [n, k] weight read as [k, n] when
transposed), we pass:
trans_a = NoTrans, trans_b = Trans
lda = k (a's leading dim is k)
ldb = k (w's leading dim is k — same as the inner dim)
ldc = n
WHAT THIS LEAVES ON THE TABLE
- Q8_0 / Q4_K via AMX. Currently those run NEON scalar at
~38 GFLOPS. Routing through AMX needs either:
(a) dequantize-then-sgemm (re-introduces the bandwidth
amplification wave 19 fixed)
(b) tiled mixed-precision: dequant a chunk into a
scratch buffer, sgemm against the chunk, repeat.
Future wave.
- f16 matmul via AMX. Accelerate exposes f16 paths through
vDSP_* family, not as obvious as sgemm. Future wave.
- Tiling for the m=1 large-k case to recover the 4096²
drop. Either cblas_sgemv (vector specialization) or
explicit blocking.
- Multi-thread fan-out. Accelerate is already multi-
threaded internally for large enough sizes; rayon on
top helps when many independent matmuls fire
concurrently (multi-head attention with per-head
workers).
cargo check --workspace --all-targets clean. Tests with
--features accelerate green (matmul_f32 module: 3/3, all
matmul_f32_accelerate path unaffected because tests don't
specifically exercise it — the scalar path tests both with
and without the feature).
Loader-gguf total: 130 tests. The headline AMX numbers
above were captured with --features accelerate via cargo
bench; recorded in BENCHMARKS.md.
Co-Authored-By: Claude Opus 4.7 <[email protected]>
13 lines
503 B
Rust
13 lines
503 B
Rust
//! Build script for jouleclaw-loader-gguf.
|
|
//!
|
|
//! Conditionally links Apple Accelerate.framework when the
|
|
//! `accelerate` cargo feature is enabled and we're targeting
|
|
//! macOS. No-op on every other target / feature flag.
|
|
|
|
fn main() {
|
|
let accelerate = std::env::var("CARGO_FEATURE_ACCELERATE").is_ok();
|
|
let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
|
|
if accelerate && target_os == "macos" {
|
|
println!("cargo:rustc-link-lib=framework=Accelerate");
|
|
}
|
|
}
|