293 lines
11 KiB
Markdown
293 lines
11 KiB
Markdown
# metal
|
|
|
|
Apple Metal GPU compute backend. Accelerates tensor operations on macOS with Apple Silicon.
|
|
|
|
**Header:** `metal/fm_metal.h` (umbrella; per-family headers in `metal/kernels/`)
|
|
**Source:** `metal/fm_metal.m` (core) + one unit per kernel family in `metal/kernels/`
|
|
**Status:** Complete (17 tests passing)
|
|
|
|
---
|
|
|
|
## Overview
|
|
|
|
The metal module provides GPU-accelerated tensor operations via Apple's Metal framework:
|
|
- MPS GEMM/GEMV (MetalPerformanceShaders matrix multiply)
|
|
- Custom compute shaders (WKV-7, softmax, RMS norm, SiLU, elementwise ops)
|
|
- INT8 quantized matrix-vector product
|
|
- Command batching (accumulate ops into a single GPU submission)
|
|
- Unified memory buffers (zero-copy on Apple Silicon when page-aligned)
|
|
|
|
All operations use threshold-based dispatch: small ops stay on CPU, large batched ops go to GPU. On Apple Silicon's unified memory, single-token decode (M=1) is always faster on CPU due to ~100-200us GPU dispatch overhead.
|
|
|
|
---
|
|
|
|
## Build
|
|
|
|
```bash
|
|
# Build with Metal
|
|
make clean && make USE_BLAS=1 USE_METAL=1 rwkvx-generate
|
|
|
|
# Build training with Metal
|
|
make clean && make USE_BLAS=1 USE_METAL=1 rwkvx-train
|
|
|
|
# Run Metal tests
|
|
make USE_METAL=1 test-metal
|
|
|
|
# Q4_K GEMM correctness + timing harnesses
|
|
make test-gemm-q4-k # batched-decode GEMM (fm_sgemm_q4_k)
|
|
make test-gemm-q4-k-f16 # prefill GEMM (fm_gemm_q4_k_f16)
|
|
```
|
|
|
|
The `USE_METAL=1` flag links `-framework Metal -framework MetalPerformanceShaders -framework Foundation` and compiles `metal/fm_metal.m` and `metal/kernels/*.m` with Objective-C flags (no `-std=c11`, adds `-fobjc-arc`).
|
|
|
|
---
|
|
|
|
## Quick Start
|
|
|
|
```c
|
|
#include "metal/fm_metal.h"
|
|
|
|
// Initialize (auto-called by rx_model_load)
|
|
fm_init();
|
|
|
|
// Wrap existing data as Metal buffers (zero-copy if page-aligned)
|
|
float A[4] = {1,2,3,4};
|
|
float B[4] = {5,6,7,8};
|
|
float C[4] = {0};
|
|
|
|
fm_buffer *ba = fm_buffer_from_ptr(A, sizeof(A));
|
|
fm_buffer *bb = fm_buffer_from_ptr(B, sizeof(B));
|
|
fm_buffer *bc = fm_buffer_from_ptr(C, sizeof(C));
|
|
|
|
// GEMM: C = A @ B (2x2 matrices)
|
|
fm_sgemm(ba, bb, bc, 2, 2, 2, 1.0f, 0.0f);
|
|
|
|
// Read result back (zero-copy — same pointer if page-aligned)
|
|
float *result = (float*)fm_buffer_data(bc);
|
|
|
|
fm_buffer_free(ba);
|
|
fm_buffer_free(bb);
|
|
fm_buffer_free(bc);
|
|
fm_free();
|
|
```
|
|
|
|
---
|
|
|
|
## API Reference
|
|
|
|
### Lifecycle
|
|
|
|
```c
|
|
int fm_init(void); // Create device + command queue + pipeline cache
|
|
void fm_free(void); // Cleanup all Metal resources
|
|
bool fm_available(void); // Metal initialized and ready?
|
|
```
|
|
|
|
### Buffers
|
|
|
|
```c
|
|
// Wrap existing pointer (zero-copy if page-aligned, copies otherwise)
|
|
fm_buffer* fm_buffer_from_ptr(void *ptr, size_t bytes);
|
|
|
|
// Allocate new Metal buffer with shared memory
|
|
fm_buffer* fm_buffer_alloc(size_t bytes);
|
|
|
|
// Free buffer wrapper (does NOT free underlying data from fm_buffer_from_ptr)
|
|
void fm_buffer_free(fm_buffer *buf);
|
|
|
|
// Access buffer data from CPU
|
|
void* fm_buffer_data(fm_buffer *buf);
|
|
size_t fm_buffer_size(fm_buffer *buf);
|
|
```
|
|
|
|
### GEMM / GEMV (via MPS)
|
|
|
|
```c
|
|
// C = alpha * A @ B + beta * C — all row-major float32
|
|
// A: [M,K], B: [K,N], C: [M,N]
|
|
void fm_sgemm(fm_buffer *A, fm_buffer *B, fm_buffer *C,
|
|
int M, int N, int K, float alpha, float beta);
|
|
|
|
// y = alpha * A @ x + beta * y — A: [M,K] row-major, x: [K], y: [M]
|
|
void fm_sgemv(fm_buffer *A, fm_buffer *x, fm_buffer *y,
|
|
int M, int K, float alpha, float beta);
|
|
```
|
|
|
|
### Custom Compute Kernels
|
|
|
|
```c
|
|
// WKV-7 single timestep — state: [n_head, head_size, head_size], modified in-place
|
|
void fm_wkv7_one(fm_buffer *r, fm_buffer *w, fm_buffer *k, fm_buffer *v,
|
|
fm_buffer *neg_kk, fm_buffer *kk_a,
|
|
fm_buffer *state, fm_buffer *y_out,
|
|
int n_head, int head_size);
|
|
|
|
// WKV-7 sequence — inputs: [T, n_head * head_size]
|
|
void fm_wkv7_seq(fm_buffer *r, fm_buffer *w, fm_buffer *k, fm_buffer *v,
|
|
fm_buffer *neg_kk, fm_buffer *kk_a,
|
|
fm_buffer *state, fm_buffer *y_out,
|
|
int n_head, int head_size, int t_len);
|
|
|
|
// Softmax over last dimension — x: [rows, cols], out: [rows, cols]
|
|
void fm_softmax(fm_buffer *x, fm_buffer *out, int rows, int cols);
|
|
|
|
// RMS normalization — x: [n], weight: [n] (may be NULL), out: [n]
|
|
void fm_rms_norm(fm_buffer *x, fm_buffer *weight, fm_buffer *out,
|
|
int n, float eps);
|
|
|
|
// SiLU activation: out = x * sigmoid(x)
|
|
void fm_silu(fm_buffer *x, fm_buffer *out, int n);
|
|
|
|
// Elementwise
|
|
void fm_add(fm_buffer *a, fm_buffer *b, fm_buffer *out, int n);
|
|
void fm_mul(fm_buffer *a, fm_buffer *b, fm_buffer *out, int n);
|
|
|
|
// INT8 quantized matvec — x_q: [cols] int8, w: [rows,cols] int8
|
|
void fm_matvec_q8(fm_buffer *x_q, fm_buffer *w, fm_buffer *w_scales,
|
|
fm_buffer *out, float x_scale, int rows, int cols);
|
|
|
|
// INT8 quantized GEMM (batched prefill). X: [M,K], W_q8: [N,K], scales: [N]
|
|
void fm_sgemm_q8(fm_buffer *X, fm_buffer *W_q8, fm_buffer *scales,
|
|
fm_buffer *Out, int M, int N, int K);
|
|
```
|
|
|
|
### Q4_K quantized GEMM (weights stay compressed on the GPU)
|
|
|
|
These kernels operate directly on raw Q4_K super-block bytes (144 bytes per 256
|
|
weights, `row_stride_bytes = (cols/256)*144`). The weight never leaves its
|
|
quantized form in memory, so only ~0.5625 bytes per weight are read.
|
|
|
|
```c
|
|
// Q4_K matvec (M=1 decode). Fused dequant + dot, one threadgroup per row.
|
|
void fm_matvec_q4_k(fm_buffer *w, fm_buffer *x, fm_buffer *out,
|
|
int rows, int cols, int row_stride_bytes);
|
|
|
|
// Q4_K x B GEMM (batched decode): out[B, rows] = x[B, cols] @ dequant(w)^T.
|
|
// Tiled simdgroup_matrix GEMM (F32 accumulation), B padded to a multiple of 8
|
|
// internally, tested up to B = 32.
|
|
void fm_sgemm_q4_k(fm_buffer *w, fm_buffer *x, fm_buffer *out,
|
|
int rows, int cols, int row_stride_bytes, int B);
|
|
|
|
// Q4_K GEMM for PREFILL (large M): out[M, N] = x[M, K] @ dequant(w)[N, K]^T.
|
|
// Dequantizes the weight to F16 and runs an F16 MPS matmul.
|
|
void fm_gemm_q4_k_f16(fm_buffer *w, fm_buffer *x_f32, fm_buffer *out_f32,
|
|
int M, int N, int K, int row_stride_bytes);
|
|
```
|
|
|
|
**Decode GEMM (`fm_sgemm_q4_k`).** One threadgroup owns a 64-row output tile
|
|
(256 threads, 8 SIMD groups). Per 64-column slice it stages the batch
|
|
activations into threadgroup memory once (coalesced) and cooperatively
|
|
dequantizes the Q4_K weights into a transposed tile, both reused across every
|
|
batch and row sub-tile, then runs the inner product on `simdgroup_float8x8`
|
|
with F32 accumulation. This amortizes the Q4_K dequant across the batch (the GPU
|
|
analog of a CPU AMX panel). Correct to err/RMS ~6e-6 against a CPU dequant-dot
|
|
reference. About 8x faster than a naive one-row-per-SIMD-group version, and
|
|
every tested shape clears the CPU AMX per-weight budget in isolation. This is a
|
|
per-weight microbenchmark result; it is the seam for a concurrent CPU+GPU decode
|
|
batch, not on its own a full decode-step win.
|
|
|
|
**Prefill GEMM (`fm_gemm_q4_k_f16`).** For large M (hundreds to thousands of
|
|
tokens) the weight is dequantized to F16 and the matmul runs as an F16 MPS GEMM
|
|
(F16 has ~2x the GPU throughput of F32). F16 is safe here because prefill
|
|
activations are post-RMSNorm and bounded, unlike the outlier-prone vision
|
|
residuals that force F32 elsewhere. Correct to err/RMS ~0.004 against a CPU F32
|
|
dequant-dot reference. At M=1369 it runs ~4 ms/call versus ~22 ms/weight for the
|
|
CPU AMX prefill path (~5.5x per weight); at prefill sizes the per-call
|
|
dispatch/sync overhead is negligible.
|
|
|
|
### Command Batching
|
|
|
|
```c
|
|
int fm_batch_begin(void); // Start accumulating GPU commands
|
|
void fm_batch_end(void); // Commit + wait once
|
|
bool fm_batch_active(void); // Check if batching
|
|
```
|
|
|
|
---
|
|
|
|
## Dispatch Thresholds
|
|
|
|
Operations below these thresholds stay on CPU:
|
|
|
|
| Threshold | Value | What it controls |
|
|
|-----------|-------|------------------|
|
|
| `FM_GEMM_THRESHOLD` | 128M ops (M*K*N) | `tm_matmul()` — needs M >= ~20 for 2560-dim |
|
|
| `FM_GEMV_THRESHOLD` | 16M ops (M*K) | `fm_sgemv()` |
|
|
| `FM_ELEM_THRESHOLD` | 1M elements | `tm_softmax()`, `tm_rms_norm()`, `tm_silu()` |
|
|
| `FM_WKV7_THRESHOLD` | disabled | `rx_wkv7_one()` — always M=1, always CPU |
|
|
| `FM_Q8_THRESHOLD` | disabled | INT8 matvec — always M=1, always CPU |
|
|
|
|
### Why single-token decode stays on CPU
|
|
|
|
On Apple Silicon, CPU and GPU share the same memory bus. Single-vector operations (M=1) are memory-bandwidth bound — both CPU and GPU hit the same bandwidth ceiling. But GPU adds ~100-200us dispatch overhead per kernel. With ~160 kernel calls per token, that's 16-32ms of pure overhead, turning a 50ms/token baseline into 66-82ms/token.
|
|
|
|
Metal only wins when batch size (M) is large enough to amortize dispatch overhead — typically M >= 16-20 for GEMM on 2560-dim matrices.
|
|
|
|
---
|
|
|
|
## Integration Points
|
|
|
|
Metal dispatch is integrated at these call sites in the inference engine:
|
|
|
|
| File | Function | Kernel |
|
|
|------|----------|--------|
|
|
| `tensor/tensor.c` | `tm_matmul()` | `fm_sgemm` |
|
|
| `tensor/tensor.c` | `tm_softmax()` | `fm_softmax` |
|
|
| `tensor/tensor.c` | `tm_rms_norm()` | `fm_rms_norm` |
|
|
| `tensor/tensor.c` | `tm_silu()` | `fm_silu` |
|
|
| `rwkvx/rwkvx.c` | `rx_wkv7_one()` | `fm_wkv7_one` |
|
|
| `rwkvx/rwkvx.c` | `fast_matvec_q8_preq()` | `fm_matvec_q8` |
|
|
| `rwkvx/rwkvx.c` | `rx_model_load()` | `fm_init()` (auto-init) |
|
|
|
|
### C API embedding
|
|
|
|
```bash
|
|
# macOS with Metal
|
|
cc -DUSE_BLAS -DACCELERATE_NEW_LAPACK -DUSE_METAL -DUSE_THREADS \
|
|
-O3 -march=native -ffast-math \
|
|
your_app.c rwkvx/rwkvx.c tensor/tensor.c safetensors/safetensors.c \
|
|
hybridmem/hybridmem.c hdc/hdc.c threadpool/threadpool.c \
|
|
metal/fm_metal.m metal/kernels/*.m \
|
|
-framework Accelerate -framework Metal \
|
|
-framework MetalPerformanceShaders -framework Foundation \
|
|
-lm -lpthread
|
|
```
|
|
|
|
Note: `fm_metal.m` and `metal/kernels/*.m` are Objective-C — must not be compiled with `-std=c11`.
|
|
|
|
---
|
|
|
|
## Metal Shaders
|
|
|
|
Shaders are embedded as string fragments, one per kernel family, in the matching
|
|
`metal/kernels/fm_kernel_<family>.m`. `fm_kernel_library_source()` concatenates them
|
|
into a single source compiled once at `fm_init()`, so nothing has to ship or locate a
|
|
`.metallib`. Pipeline states are cached by kernel name plus function-constant values.
|
|
|
|
The fragments share one MSL translation unit, so a `#define` in one family is visible
|
|
to every family after it — prefix macros with the family name and `#undef` them at the
|
|
end of the fragment.
|
|
|
|
| Kernel | Thread layout | Description |
|
|
|--------|--------------|-------------|
|
|
| `wkv7_one` | 1 threadgroup/head, head_size threads | Single-timestep WKV-7 |
|
|
| `wkv7_seq` | 1 threadgroup/head, head_size threads | Sequence WKV-7 |
|
|
| `softmax_row` | 1 threadgroup/row, 256 threads | SIMD-reduced row softmax |
|
|
| `rms_norm` | 1 threadgroup, 256 threads | SIMD-reduced RMS norm |
|
|
| `silu` | 1 thread/element | SiLU activation |
|
|
| `elementwise_add` | 1 thread/element | Vector addition |
|
|
| `elementwise_mul` | 1 thread/element | Vector multiplication |
|
|
| `matvec_q8` | 1 threadgroup/row, 256 threads | INT8 quantized matvec |
|
|
| `gemm_q8` | tiled | INT8 quantized GEMM (batched prefill) |
|
|
| `matvec_q4_k` | 1 threadgroup/row, 1 SIMD group | Q4_K fused dequant+dot matvec (decode) |
|
|
| `gemm_q4_k` | 1 threadgroup/64-row tile, 256 threads | Q4_K x B tiled simdgroup_matrix GEMM (batched decode) |
|
|
| `dequant_q4_k_f16` | 1 thread/super-block | Q4_K weight to F16 (prefill GEMM, feeds MPS F16 matmul) |
|
|
| `cvt_f32_f16` / `cvt_f16_f32` | 1 thread/element | F32/F16 conversion (prefill GEMM staging) |
|
|
|
|
---
|
|
|
|
## See Also
|
|
|
|
- [tensor](tensor.md) - Tensor operations (Metal dispatch integrated here)
|
|
- [quantization](quantization.md) - INT8/INT4 quantization
|