Metal backend (Apple Silicon): batched experts + fused attention on GPU, unified-memory zero-copy, gated behind COLI_METAL — 2.06 tok/s M5 Max (#72, #87, #103)
* docs: Metal expert-matmul backend design (Apple Silicon) Empirically-validated design for a batched MoE expert-matmul Metal backend. Microbenchmarks (scratchpad) establish: runtime-compiled Metal needs no Xcode; V3 (float4 + threadgroup reduction) kernel is correct and fast; synchronous per-matmul dispatch loses to CPU due to ~150us Metal launch latency, so the win is batched full-layer dispatch (854us/layer, 707 GFLOP/s) reading expert slabs zero-copy from unified memory. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: backend infrastructure + kernel-correctness test (M1) Add backend_metal.{h,mm} — an opt-in Apple-GPU backend built with METAL=1 on macOS. Runtime-compiled shader (no Xcode needed), zero-copy over unified memory. Implements coli_metal_matmul (general quantized GEMV, f32/int8/int4/int2) via a threadgroup-reduction + float4 kernel; batched moe_block is stubbed (returns 0 -> CPU fallback) for M2. tests/test_backend_metal.mm validates all formats and edge shapes (odd S, non-mult-4 dims) against a CPU reference (nerr ~2e-6). Makefile gains a METAL=1 Darwin branch and a metal-test target. Default build unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: batched moe_block + zero-copy slab registry (M2 backend) Implement coli_metal_moe_block: gate/up/silu/down for a whole expert block in ONE command buffer, with GPU memory barriers between stages and BINDLESS gpuAddress pointers so each expert is read zero-copy from its own RAM slab (exceeds Metal's ~31 buffer-binding limit). coli_metal_register/unregister wrap page-aligned slabs via newBufferWithBytesNoCopy and resolve interior pointers to GPU addresses. Per-row ragged expert routing supported; CPU does the final weighted scatter-add. test_backend_metal validates decode + ragged blocks vs a CPU reference (nerr ~2e-6). Still gated off in glm.c until the moe() wiring lands. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: wire batched moe_block into glm.c, token-exact (M2 integration) moe() now dispatches each routed-expert block through the GPU in one command buffer when COLI_METAL=1, reading expert weights zero-copy from page-aligned RAM slabs (registered in expert_load). Falls back to CPU per-block on any unresolved slab or GPU fault. Default build byte-identical (all #ifdef COLI_METAL). Fixes a heap-corruption crash: expert_load registers slabs from parallel OpenMP threads, so the slab registry is now mutex-guarded (buffer creation stays outside the lock). Added command-buffer error checking (fall back to CPU on GPU fault) and a COLI_METAL_DEBUG one-shot trace. Validated token-exact vs the CPU path (greedy): identical 12-token output; expert-matmul time 29.9s -> 21.1s with pinned experts still on CPU. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: instrument moe_block (GPU/CPU split, wall-vs-kernel time) Add diagnostics printed on the PROFILO line under COLI_METAL: GPU vs CPU-fallback block counts, experts-on-GPU, and a per-block time split (setup / gpu-wall / kernel / scatter). Reveals that with a warm cache all experts run on the GPU (0 fallback) and expert-matmul drops ~1.3x vs CPU, but ~62% of GPU wall-time is idle/scheduling latency (3.1s kernel of 8.3s wall over 396 sporadic submits) — the GPU powers down between blocks because attention runs on the CPU per layer. Points the next optimization at keeping the GPU hot (offload attention). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: Metal backend measured results + next levers Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: Phase 2 fused decode attention plan + absorption-core validated Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: fused decode attention on GPU, token-exact (Phase 2) coli_metal_attn_decode runs a full S=1 decode attention layer in ONE command buffer: q_a -> rmsnorm -> q_b -> RoPE ; kv_a -> latent rmsnorm@pos + krot RoPE@pos (cache write) ; MLA absorption core (qabs/score/softmax/clat/ctx) ; o_proj. The absorption-core kernels were validated in isolation (nerr ~1e-6) before wiring. Projection matmuls reuse the mm_gemv kernel; attention weights are uploaded+cached (serial path, no lock); Lc/Rc caches are page-aligned + registered in kv_alloc for zero-copy GPU read/write. GLM-5.2 dims compiled in; falls back to CPU for S>1 (prefill/MTP verify), st0!=0, active DSA selection (context>topk), or mismatched dims. DSA index-key write stays on CPU so future selection still works. Validated token-exact vs CPU (identical greedy output); attention time 16.5s -> 10.5s (~1.57x), end-to-end 0.20 -> 0.28 tok/s. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: Phase 2 fused attention complete + known limits Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: attention coverage/latency instrumentation + honest results Add per-layer fused-attention counters (METAL-ATTN line): GPU layer count, gpu-wall and true kernel time. Measurement (DRAFT=0, all-S=1 decode) shows the fused attention triggers on all decode layers but is submit-latency-bound: gpu-wall 3.70s vs kernel 0.63s (83% idle latency over 546 sporadic command buffers). Attention time is neutral vs CPU; the earlier MTP-on "16.5->10.5" was run-to-run variance. Design doc corrected with the honest result: both offloads are gated by Metal's ~5ms cold-GPU submit latency; reducing submit count (fuse attention+experts per layer) is the real lever. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: fused attention handles S<=4 (covers MTP verify forwards) Extend coli_metal_attn_decode from S=1 to S<=4: the core kernels (qabs/score/ smax/clat/ctx) gain a query-row dimension with per-row causal masking (query s attends keys [0, pos_base+s]); rmsnorm/rope/copy became row-aware; projections run S rows via mm_gemv. This covers the default MTP config (draft=3 -> S=4 verify forwards), which previously fell back to CPU attention entirely. Token-exact vs CPU (identical greedy output, MTP on). Perf is inconclusive at short context: still submit-latency-bound (attn gpu-wall 5.5s vs kernel 0.9s) and the measurement is dominated by disk-streaming variance (+/-15s between runs). Next: measure with a fully-warm cache to isolate compute, then reduce submit count. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: clean warm A/B shows real ~1.4x (experts+S<=4 attention), token-exact Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: interleave attention q/kv paths, 7->4 barriers (iter 2) The q-path (q_a->rmsnorm->q_b->rope) and kv-path (kv_a->copy->rmsnorm+rope) are independent until the absorption core, but were serialized by memory barriers. Interleave them into 4 barrier-separated stages so the GPU overlaps independent dispatches. Token-exact; attention gpu-wall 3.04s -> 2.73s (~10%). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: zero-copy attention weights + fuse shared expert into GPU block (iter 2) Dense QT weights/scales now allocate page-aligned + registered (qalloc) under METAL, so the fused attention reads q_a/q_b/kv_a/kv_b/o zero-copy instead of uploading ~6 GB of duplicates (RSS -3 GB, upload copies gone). bind_gemv resolves registered pointers (buffer,offset) with a pre-check guard. Phase E's shared expert (identical shapes to a routed expert: gate/up [I,D], down [D,I], same int4 container) is appended to the first Metal moe_block as an extra expert with rw=1.0 over all S rows — removes 3 CPU matmuls per layer and fills the same GPU submit. CPU Phase E still runs on any fallback. Zero-copy validated token-exact: 35.1s -> 29.7s (0.34 tok/s) warm. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: iteration 2 findings Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: iter 2 final ~1.56x + iter 3 plan (disk/GPU overlap) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: overlap disk loads with GPU compute inside the layer (iter 3) Split each MoE block into two GPU submits: the RESIDENT experts (pin/LRU hits, plus the fused shared expert) are encoded and committed BEFORE the missed experts' OMP pread loop, so the GPU computes while the disk reads; the missed subset follows in a second (sync) submit once loaded. New two-phase backend API (coli_metal_moe_block_begin/end) with handle-owned scratch so the async submit cannot collide with the sync path's static buffers; moe_submit/moe_finish are shared by both. Per-subset CPU fallback preserved (resident and missed fall back independently on unresolved slab or GPU fault). Token-exact. Warm 96GB: expert-matmul 8.96 -> 4.92s (resident compute now hidden inside the disk window; expert idle latency ~5.7s -> ~0.9s), total 28.97s (0.35 tok/s) vs CPU 50.2s = ~1.73x. Note: 'make glm METAL=1' after a default build does NOT rebuild (target looks up-to-date) — touch glm.c or clean when switching build flavors. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: iter 3 disk/GPU overlap results (~1.73x) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: keep-alive spinner experiment (env-gated) + latency decomposition COLI_METAL_SPIN=1 keeps trivial GPU work in flight on a separate queue to probe whether inter-submit idle is clock ramp-down; thread is detached (a joinable global thread std::terminate'd the process at exit). First contended A/B was inconclusive but showed the spinner does NOT collapse attention wall per-call (~16ms both ways), so ramp-down is not the whole story. METAL-ATTN now decomposes latency: cpu-sched (commit->kernelStart) vs gpu-sched (kernelStart->GPUStart) vs kernel execution, to pinpoint where the ~13ms/call goes. Default behavior unchanged (spinner off). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: standalone regression tests for fused decode attention run_attn builds full-size fake GLM-5.2 attention weights (int4, page-aligned, registered), replicates glm.c's absorb-branch math exactly on the CPU (q_a -> rmsnorm -> q_b -> rope; kv_a -> latent rmsnorm + krot rope -> cache; per-head qabs/score/softmax/clat/ctx; o_proj), and checks coli_metal_attn_decode against it at S=1/3/4 and pos_base 0/12/37 — including the Lc/Rc cache write-back, which end-to-end runs cannot isolate. All pass (nerr ~5e-6, cache ~1.4e-5). The whole Metal path (gemv, moe_block, fused attention) is now testable without the model. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: route large row-batch matmul_qt GEMMs to the GPU (prefill) matmul_qt now dispatches to a new coli_metal_gemm when S >= COLI_METAL_GEMM_MIN (default 16), the weight is int8/int4 and registered (all dense QT allocs are, via qalloc), and we're not inside an OpenMP region (mirrors the CUDA guard). Decode-sized matmuls stay on the CPU where NEON wins vs submit latency; prefill's big GEMMs (kv_b reconstruction at S=Tk, o_proj, dense MLP, step_all's S x vocab logits) amortize it — microbench showed ~6x over the CPU idot at S=16. Standalone test: registered int4 GEMM S=64 vs cpu_ref (nerr 2.9e-6). Machine busy again; end-to-end token-exactness + threshold sweep pending idle. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * README: document the experimental Metal backend (Apple Silicon) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: 1.5-2.1x faster moe_gemv (simdgroup-per-row + 8-value loads) Replace one-threadgroup-per-output-row (128 threads reducing via threadgroup memory) with one SIMDGROUP per output row, 4 rows per threadgroup, and uchar4 loads (8 nibbles / 8 int8 per lane-iteration). Removes the threadgroup barrier + shared-memory reduction entirely (simd_sum only) and doubles load width. Engine-like block-shape microbench (pure GPU time): S=4 block 2548->1739us, S=1 block 934->437us, big block 4582->3414us — 358-389 GB/s vs 182-264. Row-bound guard added (NT) since the grid rounds up to 4 rows/TG. All backend tests pass (moe_block nerr 2.4e-6, attention unchanged). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: mm_gemv simdgroup-per-row + 8-value loads (attention projections, prefill GEMM) Apply the moe_gemv V2 transformation to the general quantized GEMV: one simdgroup per output element (4/threadgroup), 8-value loads for i8/i4/f32, no threadgroup reduction. Same measured 1.5-2.1x class of win; serves the fused-attention projections (q_a/q_b/kv_a/o), coli_metal_gemm (prefill), and coli_metal_matmul. All three dispatch sites updated (NT row-bound guard, grid ceil(NT/4) x 128). Full test suite green, incl. non-mult-of-8 tail paths (2050x6146) and all fmts. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: experimental COLI_MMAP=1 — experts as zero-copy views into mmap'd files Lazily mmap each safetensors file (PROT_READ, MAP_SHARED, mutex-guarded — expert loads are OMP-parallel), register the mapping with Metal, and make expert_load a pointer assignment into the map: no pread, no slab, no copy; the OS page cache is the cache. Alignment guards fall back to the slab path. Default OFF. First validation (machine at load 66 + 46GB swap): token-exact, RSS 58 -> 10.5 GB as designed, but GPU wall exploded (~130 MB/s effective) — the GPU demand-faults file-backed pages, catastrophic when memory pressure evicts them. Needs an idle-machine A/B to judge fairly (llama.cpp's identical technique relies on pages staying resident); possible fixes if slow even idle: CPU pre-touch of missed experts' pages before the GPU submit, or madvise/mlock windows. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: CPU pre-touch for COLI_MMAP expert pages (fix GPU demand-faulting) In mmap mode, fault the missed expert's pages in on the CPU inside expert_load (madvise WILLNEED for async readahead + a page-stride touch): this is pread's I/O without the copy and without the slab, it runs inside the existing OMP loop that overlaps with the resident-experts GPU submit (iter 3), and it guarantees the GPU only ever reads resident pages — GPU demand-faulting of file-backed pages measured catastrophic (~130 MB/s). Read-only addition: outputs unchanged from the validated mmap run; perf pending the idle-machine A/B. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: idle-machine suite results (~1.33x same-session; mmap negative result) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: COLI_METAL_UNTRACKED experiment (negative result, default off) Env-gated MTLResourceHazardTrackingModeUntracked on registered wraps + scratch to test whether cross-CB hazard tracking causes the ~10ms/CB gpu-sched delay. Idle A/B: no effect (gpu-sched 3.9 vs 3.4s, noise), token-exact. Together with the spinner negative, this pins the attention CB delay as inherent scheduler/wake overhead on an empty pipeline — removable only by eliminating the CB boundary, which CPU-side routing at ~58% hit-rate forces. Metal side is at its floor: kernel 3.5s+0.8s (near BW ceiling), sched ~3.2s, disk ~15s dominant (10 tok). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: loop conclusion — best config DIRECT=1+COLI_METAL=1, 0.42 tok/s (~1.4x) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: refactor attention into encode_attention()+resolve_attn() (layer-CB prep) Behavior-preserving: attn_decode is now a thin wrapper; all attention tests byte-identical. Prepares embedding the chain in a full-layer command buffer. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * metal: full decode layer in ONE command buffer (token-exact) coli_metal_layer_decode runs the whole layer prelude on the GPU in a single submit: in_ln rmsnorm -> fused attention -> residual add -> post_ln rmsnorm -> shared expert (gate/up/silu/down) -> router (f32 simdgroup matvec + sigmoid) -> exact phase-A top-K selection (greedy argmax over sigmoid+bias with CPU tie order, --topp truncation, norm_topk, routed_scale) in a serial-per-row kernel. The CPU's per-layer work shrinks to: read 8 expert IDs, resolve/load, expert CBs (disk/GPU overlap unchanged), scatter. moe() consumes the precomputed routing (g_pre_*: skips phase A, keeps eusage/eheat/ereq counters for the learning cache) and adds the GPU shared-expert output instead of computing phase E. ld() tensors (norms/router/bias) now allocate registered so the GPU reads them zero-copy. DSA index keys still computed on CPU from the in_ln-normed x (new inrm output). Every missing condition falls back to the full CPU layer. Validated token-exact vs CPU (identical greedy output, MTP on). Profile: "altro" 3.8s -> 0.53s (12 tok); 0.42 tok/s despite disk-variance headwind. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs: Phase 3 full-layer CB results — 0.43 tok/s record, token-exact Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * gitignore: Metal build artifacts, venv, bench datasets Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * remove internal design docs before PR Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Code Arranger
@@ -0,0 +1,221 @@
|
||||
// Kernel-correctness test for the Metal backend: coli_metal_matmul vs CPU reference
|
||||
// (dequant->f32 MAC * per-row scale) for f32/int8/int4/int2 across real GLM shapes.
|
||||
#include "../backend_metal.h"
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
|
||||
enum { F32=0, I8=1, I4=2, I2=3 };
|
||||
|
||||
static void cpu_ref(int fmt, const void *W, const float *s, const float *x,
|
||||
float *y, int S, int I, int O) {
|
||||
const int8_t *q8 = (const int8_t*)W; const uint8_t *q4 = (const uint8_t*)W;
|
||||
const float *qf = (const float*)W;
|
||||
int rb4=(I+1)/2, rb2=(I+3)/4;
|
||||
for (int o=0;o<O;o++) for (int si=0;si<S;si++){
|
||||
const float *xr = x + (size_t)si*I; float acc=0;
|
||||
for (int i=0;i<I;i++){
|
||||
float w;
|
||||
if (fmt==I8) w=(float)q8[(size_t)o*I+i];
|
||||
else if (fmt==I4){ uint8_t b=q4[(size_t)o*rb4+(i>>1)]; int v=(i&1)?(b>>4):(b&0xF); w=(float)(v-8); }
|
||||
else if (fmt==I2){ uint8_t b=q4[(size_t)o*rb2+(i>>2)]; int v=(b>>(2*(i&3)))&0x3; w=(float)(v-2); }
|
||||
else w=qf[(size_t)o*I+i];
|
||||
acc += w*xr[i];
|
||||
}
|
||||
y[(size_t)si*O+o]=acc*s[o];
|
||||
}
|
||||
}
|
||||
|
||||
static int run(int fmt, int O, int I, int S, const char *name) {
|
||||
int rb4=(I+1)/2, rb2=(I+3)/4;
|
||||
size_t wn = (fmt==I8)?(size_t)O*I : (fmt==I4)?(size_t)O*rb4 : (fmt==I2)?(size_t)O*rb2 : (size_t)O*I*sizeof(float);
|
||||
std::vector<uint8_t> W(wn); std::vector<float> Wf;
|
||||
srand(99);
|
||||
if (fmt==F32){ Wf.resize((size_t)O*I); for(auto&v:Wf) v=((rand()%2000)-1000)/1000.f; }
|
||||
else for(auto&b:W) b=(uint8_t)((fmt==I8)?((rand()%255)-127):(rand()&0xFF));
|
||||
const void *Wp = (fmt==F32)?(const void*)Wf.data():(const void*)W.data();
|
||||
std::vector<float> s(O), x((size_t)S*I), yr((size_t)S*O), yg((size_t)S*O);
|
||||
for(auto&v:s) v=(fmt==F32)?1.0f:(0.01f+(rand()%100)/10000.f);
|
||||
for(auto&v:x) v=((rand()%2000)-1000)/1000.f;
|
||||
cpu_ref(fmt, Wp, s.data(), x.data(), yr.data(), S, I, O);
|
||||
ColiMetalTensor *t=nullptr;
|
||||
if (!coli_metal_matmul(&t, yg.data(), x.data(), Wp, s.data(), fmt, S, I, O)) {
|
||||
printf(" %-22s FAIL (matmul returned 0)\n", name); return 1; }
|
||||
double maxabs=0, ymax=0;
|
||||
for(size_t i=0;i<(size_t)S*O;i++){ maxabs=fmax(maxabs,fabs(yg[i]-yr[i])); ymax=fmax(ymax,fabs(yr[i])); }
|
||||
double nerr=maxabs/(ymax+1e-9);
|
||||
int ok = nerr < 1e-4;
|
||||
printf(" %-22s nerr=%.2e %s\n", name, nerr, ok?"ok":"*** MISMATCH");
|
||||
coli_metal_tensor_free(t);
|
||||
return ok?0:1;
|
||||
}
|
||||
|
||||
static float deq4(const uint8_t* w,int i){ uint8_t b=w[i>>1]; int v=(i&1)?(b>>4):(b&0xF); return (float)(v-8); }
|
||||
static size_t roundpg(size_t n){ size_t p=16384; return ((n+p-1)/p)*p; }
|
||||
|
||||
// Validate coli_metal_moe_block against a CPU reference (gate/up/silu/down + weighted scatter-add).
|
||||
static int run_moe(const std::vector<int>& nrv, const char* name) {
|
||||
const int D=6144, I=2048, fmt=2; int rbG=(D+1)/2, rbD=(I+1)/2, nb=(int)nrv.size();
|
||||
int R=0; std::vector<int> xoff(nb),nr(nrv); for(int e=0;e<nb;e++){ xoff[e]=R; R+=nrv[e]; }
|
||||
srand(2024+nb);
|
||||
// per-expert page-aligned slab [Wg|Wu|Wd] and fslab [Sg|Su|Sd]; register both.
|
||||
std::vector<void*> slab(nb), fslab(nb);
|
||||
std::vector<const void*> g(nb),u(nb),d(nb); std::vector<const float*> gs(nb),us(nb),ds(nb);
|
||||
size_t wlen=roundpg((size_t)I*rbG*2 + (size_t)D*rbD), flen=roundpg(((size_t)I*2+D)*sizeof(float));
|
||||
for(int e=0;e<nb;e++){
|
||||
posix_memalign(&slab[e],16384,wlen); posix_memalign(&fslab[e],16384,flen);
|
||||
uint8_t* sp=(uint8_t*)slab[e]; for(size_t i=0;i<(size_t)I*rbG*2+(size_t)D*rbD;i++) sp[i]=(uint8_t)(rand()&0xFF);
|
||||
float* fp=(float*)fslab[e]; for(size_t i=0;i<(size_t)I*2+D;i++) fp[i]=0.01f+(rand()%50)/50000.f;
|
||||
g[e]=sp; u[e]=sp+(size_t)I*rbG; d[e]=sp+(size_t)I*rbG*2;
|
||||
gs[e]=fp; us[e]=fp+I; ds[e]=fp+2*I;
|
||||
coli_metal_register(slab[e],wlen); coli_metal_register(fslab[e],flen);
|
||||
}
|
||||
std::vector<float> xg((size_t)R*D); for(auto&v:xg) v=((rand()%2000)-1000)/1000.f;
|
||||
std::vector<int> rows(R); std::vector<float> rw(R);
|
||||
for(int gr=0;gr<R;gr++){ rows[gr]=0; rw[gr]=0.1f+(rand()%100)/100.f; } // decode: all -> position 0
|
||||
int S=1;
|
||||
// CPU reference
|
||||
std::vector<float> refout((size_t)S*D,0.f), gg(I),uu(I),hh(D);
|
||||
for(int e=0;e<nb;e++) for(int r=0;r<nr[e];r++){ int gr=xoff[e]+r; const float* xr=&xg[(size_t)gr*D];
|
||||
const uint8_t* wg=(const uint8_t*)g[e]; const uint8_t* wu=(const uint8_t*)u[e]; const uint8_t* wd=(const uint8_t*)d[e];
|
||||
for(int o=0;o<I;o++){ float a=0; for(int k=0;k<D;k++) a+=deq4(wg+(size_t)o*rbG,k)*xr[k]; gg[o]=a*gs[e][o]; }
|
||||
for(int o=0;o<I;o++){ float a=0; for(int k=0;k<D;k++) a+=deq4(wu+(size_t)o*rbG,k)*xr[k]; uu[o]=a*us[e][o]; }
|
||||
for(int o=0;o<I;o++){ float v=gg[o]; gg[o]=(v/(1.f+expf(-v)))*uu[o]; }
|
||||
for(int o=0;o<D;o++){ float a=0; for(int k=0;k<I;k++) a+=deq4(wd+(size_t)o*rbD,k)*gg[k]; hh[o]=a*ds[e][o]; }
|
||||
float* os=&refout[(size_t)rows[gr]*D]; for(int o=0;o<D;o++) os[o]+=rw[gr]*hh[o];
|
||||
}
|
||||
std::vector<float> gout((size_t)S*D,0.f);
|
||||
int ok = coli_metal_moe_block(nb,D,I,fmt,g.data(),u.data(),d.data(),gs.data(),us.data(),ds.data(),
|
||||
xg.data(),xoff.data(),nr.data(),rows.data(),rw.data(),gout.data(),S);
|
||||
double maxabs=0,ymax=0; for(size_t i=0;i<gout.size();i++){ maxabs=fmax(maxabs,fabs(gout[i]-refout[i])); ymax=fmax(ymax,fabs(refout[i])); }
|
||||
double nerr=maxabs/(ymax+1e-9); int pass = ok && nerr<1e-4;
|
||||
printf(" %-22s R=%d nerr=%.2e %s\n", name, R, nerr, pass?"ok":"*** MISMATCH");
|
||||
for(int e=0;e<nb;e++){ coli_metal_unregister(slab[e]); coli_metal_unregister(fslab[e]); free(slab[e]); free(fslab[e]); }
|
||||
return pass?0:1;
|
||||
}
|
||||
|
||||
// ---- fused decode attention vs a CPU reference replicating glm.c's exact math ----
|
||||
// GLM-5.2 dims (hardcoded in the backend): hidden=6144 H=64 q_lora=2048 kv_lora=512
|
||||
// nope=192 rope=64 vh=256; theta=10000 ascale=1/16 eps=1e-5.
|
||||
enum { TH=6144, THH=64, TQL=2048, TKVL=512, TNOPE=192, TROPE=64, TVH=256, TQH=256, TROWSH=448 };
|
||||
static void t_rms(float*o,const float*x,const float*w,int n,float eps){ double ms=0; for(int i=0;i<n;i++) ms+=(double)x[i]*x[i];
|
||||
float r=1.f/sqrtf((float)(ms/n)+eps); for(int i=0;i<n;i++) o[i]=x[i]*r*w[i]; }
|
||||
static void t_rope(float*v,int pos,float th){ int hl=TROPE/2; float in[TROPE]; memcpy(in,v,sizeof(in));
|
||||
for(int j=0;j<hl;j++){ float inv=powf(th,-2.f*j/TROPE), a=in[2*j], b=in[2*j+1], cs=cosf(pos*inv), sn=sinf(pos*inv);
|
||||
v[j]=a*cs-b*sn; v[hl+j]=b*cs+a*sn; } }
|
||||
static void t_gemv4(float*y,const float*x,const uint8_t*w,const float*sc,int O,int I){ int rb=(I+1)/2;
|
||||
for(int o=0;o<O;o++){ const uint8_t*r=w+(size_t)o*rb; float a=0;
|
||||
for(int i=0;i<I;i++){ uint8_t b=r[i>>1]; int v=(i&1)?(b>>4):(b&0xF); a+=(float)(v-8)*x[i]; } y[o]=a*sc[o]; } }
|
||||
struct TW { uint8_t*w; float*s; size_t wb, sb; };
|
||||
static TW t_mkw(int O,int I){ TW t; int rb=(I+1)/2;
|
||||
t.wb=((size_t)O*rb+16383)&~(size_t)16383; t.sb=((size_t)O*4+16383)&~(size_t)16383;
|
||||
posix_memalign((void**)&t.w,16384,t.wb); posix_memalign((void**)&t.s,16384,t.sb);
|
||||
for(size_t i=0;i<(size_t)O*rb;i++) t.w[i]=(uint8_t)(rand()&0xFF);
|
||||
for(int i=0;i<O;i++) t.s[i]=0.01f+(rand()%40)/40000.f;
|
||||
coli_metal_register(t.w,t.wb); coli_metal_register(t.s,t.sb); return t; }
|
||||
static int run_attn(int S, int pos_base, const char* name){
|
||||
const float eps=1e-5f, theta=10000.f, ascale=1.f/16.f;
|
||||
srand(4242+S+pos_base);
|
||||
TW qa=t_mkw(TQL,TH), qb=t_mkw(THH*TQH,TQL), kva=t_mkw(TKVL+TROPE,TH), kvb=t_mkw(THH*TROWSH,TKVL), o=t_mkw(TH,THH*TVH);
|
||||
std::vector<float> qaln(TQL), kvaln(TKVL);
|
||||
for(auto&v:qaln) v=0.5f+(rand()%1000)/1000.f; for(auto&v:kvaln) v=0.5f+(rand()%1000)/1000.f;
|
||||
int T=pos_base+S; size_t lcb=(((size_t)T*TKVL*4)+16383)&~(size_t)16383, rcb=(((size_t)T*TROPE*4)+16383)&~(size_t)16383;
|
||||
float *Lc,*Rc; posix_memalign((void**)&Lc,16384,lcb); posix_memalign((void**)&Rc,16384,rcb);
|
||||
coli_metal_register(Lc,lcb); coli_metal_register(Rc,rcb);
|
||||
// pre-existing cache history [0,pos_base): random normed latents + roped krot
|
||||
for(int t=0;t<pos_base;t++){ for(int i=0;i<TKVL;i++) Lc[(size_t)t*TKVL+i]=((rand()%2000)-1000)/1500.f;
|
||||
for(int i=0;i<TROPE;i++) Rc[(size_t)t*TROPE+i]=((rand()%2000)-1000)/1500.f; }
|
||||
std::vector<float> x((size_t)S*TH); for(auto&v:x) v=((rand()%2000)-1000)/1000.f;
|
||||
std::vector<float> Lr((size_t)T*TKVL), Rr((size_t)T*TROPE); // reference cache copies
|
||||
memcpy(Lr.data(),Lc,(size_t)pos_base*TKVL*4); memcpy(Rr.data(),Rc,(size_t)pos_base*TROPE*4);
|
||||
// CPU reference: mirrors glm.c attention() absorb branch (per new token, then per head)
|
||||
std::vector<float> Q((size_t)S*THH*TQH), ref((size_t)S*TH);
|
||||
for(int s=0;s<S;s++){ int pos=pos_base+s;
|
||||
std::vector<float> qr(TQL), comp(TKVL+TROPE);
|
||||
t_gemv4(qr.data(),&x[(size_t)s*TH],qa.w,qa.s,TQL,TH); t_rms(qr.data(),qr.data(),qaln.data(),TQL,eps);
|
||||
t_gemv4(&Q[(size_t)s*THH*TQH],qr.data(),qb.w,qb.s,THH*TQH,TQL);
|
||||
for(int h=0;h<THH;h++) t_rope(&Q[(size_t)s*THH*TQH+(size_t)h*TQH+TNOPE],pos,theta);
|
||||
t_gemv4(comp.data(),&x[(size_t)s*TH],kva.w,kva.s,TKVL+TROPE,TH);
|
||||
t_rms(&Lr[(size_t)pos*TKVL],comp.data(),kvaln.data(),TKVL,eps);
|
||||
memcpy(&Rr[(size_t)pos*TROPE],&comp[TKVL],TROPE*4); t_rope(&Rr[(size_t)pos*TROPE],pos,theta);
|
||||
}
|
||||
int rb=(TKVL+1)/2;
|
||||
for(int s=0;s<S;s++){ int pos=pos_base+s; std::vector<float> ctx((size_t)THH*TVH);
|
||||
for(int h=0;h<THH;h++){ int rbase=h*TROWSH;
|
||||
const float* qp=&Q[(size_t)s*THH*TQH+(size_t)h*TQH]; const float* qro=qp+TNOPE;
|
||||
std::vector<float> qabs(TKVL,0);
|
||||
for(int d=0;d<TNOPE;d++){ const uint8_t*r=kvb.w+(size_t)(rbase+d)*rb; float sc=kvb.s[rbase+d];
|
||||
for(int i=0;i<TKVL;i++){ uint8_t b=r[i>>1]; int v=(i&1)?(b>>4):(b&0xF); qabs[i]+=qp[d]*(float)(v-8)*sc; } }
|
||||
std::vector<float> a(pos+1);
|
||||
for(int t=0;t<=pos;t++){ const float*Lt=&Lr[(size_t)t*TKVL]; const float*Rt=&Rr[(size_t)t*TROPE];
|
||||
float v=0; for(int i=0;i<TKVL;i++) v+=qabs[i]*Lt[i]; for(int d=0;d<TROPE;d++) v+=qro[d]*Rt[d]; a[t]=v*ascale; }
|
||||
float mx=-1e30f; for(float v:a) mx=fmaxf(mx,v); float sum=0; for(float&v:a){ v=expf(v-mx); sum+=v; } for(float&v:a) v/=sum;
|
||||
std::vector<float> cl(TKVL,0);
|
||||
for(int t=0;t<=pos;t++){ const float*Lt=&Lr[(size_t)t*TKVL]; for(int i=0;i<TKVL;i++) cl[i]+=a[t]*Lt[i]; }
|
||||
for(int j=0;j<TVH;j++){ const uint8_t*r=kvb.w+(size_t)(rbase+TNOPE+j)*rb; float sc=kvb.s[rbase+TNOPE+j];
|
||||
float v=0; for(int i=0;i<TKVL;i++){ uint8_t b=r[i>>1]; int vv=(i&1)?(b>>4):(b&0xF); v+=cl[i]*(float)(vv-8)*sc; }
|
||||
ctx[(size_t)h*TVH+j]=v; } }
|
||||
t_gemv4(&ref[(size_t)s*TH],ctx.data(),o.w,o.s,TH,THH*TVH);
|
||||
}
|
||||
std::vector<float> got((size_t)S*TH);
|
||||
int ok=coli_metal_attn_decode(x.data(), qa.w,qa.s,2,qaln.data(), qb.w,qb.s,2,
|
||||
kva.w,kva.s,2,kvaln.data(), kvb.w,kvb.s,2, o.w,o.s,2,
|
||||
Lc,Rc,S,pos_base,0,eps,theta,ascale,got.data());
|
||||
double ma=0,ym=0; for(size_t i=0;i<ref.size();i++){ ma=fmax(ma,fabs(got[i]-ref[i])); ym=fmax(ym,fabs(ref[i])); }
|
||||
// also verify the cache write-back (Lc/Rc for the new positions)
|
||||
double mc=0; for(int s=0;s<S;s++){ int pos=pos_base+s;
|
||||
for(int i=0;i<TKVL;i++) mc=fmax(mc,fabs(Lc[(size_t)pos*TKVL+i]-Lr[(size_t)pos*TKVL+i]));
|
||||
for(int i=0;i<TROPE;i++) mc=fmax(mc,fabs(Rc[(size_t)pos*TROPE+i]-Rr[(size_t)pos*TROPE+i])); }
|
||||
double nerr=ma/(ym+1e-9);
|
||||
int pass = ok && nerr<2e-4 && mc<1e-4;
|
||||
printf(" %-24s nerr=%.2e cache=%.2e %s\n", name, nerr, mc, pass?"ok":"*** MISMATCH");
|
||||
auto freew=[&](TW&t){ coli_metal_unregister(t.w); coli_metal_unregister(t.s); free(t.w); free(t.s); };
|
||||
freew(qa); freew(qb); freew(kva); freew(kvb); freew(o);
|
||||
coli_metal_unregister(Lc); coli_metal_unregister(Rc); free(Lc); free(Rc);
|
||||
return pass?0:1;
|
||||
}
|
||||
|
||||
int main(void) {
|
||||
if (!coli_metal_init()) { printf("Metal unavailable (skipping)\n"); return 0; }
|
||||
printf("Metal backend kernel tests:\n");
|
||||
int fail=0;
|
||||
fail |= run(I8, 2048,6144,1, "int8 gate/up S=1");
|
||||
fail |= run(I4, 2048,6144,1, "int4 gate/up S=1");
|
||||
fail |= run(I4, 6144,2048,1, "int4 down S=1");
|
||||
fail |= run(I2, 2048,6144,1, "int2 gate/up S=1");
|
||||
fail |= run(F32,1024,6144,1, "f32 S=1");
|
||||
fail |= run(I8, 2048,6144,4, "int8 gate/up S=4");
|
||||
fail |= run(I4, 2048,6144,7, "int4 gate/up S=7 (odd)");
|
||||
fail |= run(I4, 2050,6146,3, "int4 non-mult-4 dims");
|
||||
printf("Metal batched moe_block tests:\n");
|
||||
fail |= run_moe({1,1,1,1,1,1,1,1}, "moe decode nb=8");
|
||||
fail |= run_moe({3,1,4,2,1,5}, "moe ragged nb=6");
|
||||
printf("Metal large-batch gemm test:\n");
|
||||
{ // registered int4 weights, S=64: coli_metal_gemm vs cpu_ref
|
||||
srand(77); int O=2048,I=6144,S=64,rb=(I+1)/2;
|
||||
size_t wb=(((size_t)O*rb)+16383)&~(size_t)16383, sb2=(((size_t)O*4)+16383)&~(size_t)16383;
|
||||
uint8_t*W; float*Sc; posix_memalign((void**)&W,16384,wb); posix_memalign((void**)&Sc,16384,sb2);
|
||||
for(size_t i=0;i<(size_t)O*rb;i++) W[i]=(uint8_t)(rand()&0xFF);
|
||||
for(int i=0;i<O;i++) Sc[i]=0.01f+(rand()%50)/50000.f;
|
||||
coli_metal_register(W,wb); coli_metal_register(Sc,sb2);
|
||||
std::vector<float> x((size_t)S*I), yr((size_t)S*O), yg((size_t)S*O);
|
||||
for(auto&v:x) v=((rand()%2000)-1000)/1000.f;
|
||||
cpu_ref(I4,W,Sc,x.data(),yr.data(),S,I,O);
|
||||
int ok=coli_metal_gemm(yg.data(),x.data(),W,Sc,2,S,I,O);
|
||||
double ma=0,ym=0; for(size_t i=0;i<yr.size();i++){ ma=fmax(ma,fabs(yg[i]-yr[i])); ym=fmax(ym,fabs(yr[i])); }
|
||||
int pass = ok && ma/(ym+1e-9)<1e-4;
|
||||
printf(" gemm S=64 int4 nerr=%.2e %s\n", ma/(ym+1e-9), pass?"ok":"*** MISMATCH");
|
||||
fail |= !pass;
|
||||
coli_metal_unregister(W); coli_metal_unregister(Sc); free(W); free(Sc);
|
||||
}
|
||||
printf("Metal fused attention tests:\n");
|
||||
fail |= run_attn(1, 0, "attn S=1 pos=0");
|
||||
fail |= run_attn(1, 37, "attn S=1 pos=37");
|
||||
fail |= run_attn(4, 12, "attn S=4 pos=12 (MTP)");
|
||||
fail |= run_attn(3, 0, "attn S=3 pos=0");
|
||||
printf(fail? "metal backend tests: FAILED\n" : "metal backend tests: ok\n");
|
||||
coli_metal_shutdown();
|
||||
return fail;
|
||||
}
|
||||
Reference in New Issue
Block a user