diff --git a/.gitignore b/.gitignore index 09f3aaa..b208456 100644 --- a/.gitignore +++ b/.gitignore @@ -44,3 +44,8 @@ stats*.txt /quant_test.py /profile_run.py /sweep.py +c/backend_metal.o +c/backend_metal_test +c/tests/test_tier +c/mio_env/ +c/bench/ diff --git a/README.md b/README.md index 4042d38..90e0e9e 100644 --- a/README.md +++ b/README.md @@ -238,6 +238,31 @@ threaded HTTP is continuous batching. RAM admission accounts for every configure Use `COLI_KV_SLOTS=N` as the environment equivalent. Start with a small value: at the default 4096-token context, every slot costs hundreds of MB. +### Experimental Metal backend (Apple Silicon) + +On Apple Silicon the decode profile is matmul-bound, and unified memory removes the +PCIe copy tax that keeps CUDA's streaming experts on the CPU — so colibrì has an +opt-in Metal backend that runs the **routed-expert SwiGLU (batched, zero-copy from +the RAM slabs)**, the **fused decode attention** (full MLA layer in one command +buffer, S≤4), and **prefill's large GEMMs** on the GPU. Token-exact vs the CPU path. + +```bash +cd c +make glm METAL=1 # macOS only; no Xcode needed (shader compiles at runtime) +make metal-test # standalone kernel/attention correctness vs CPU reference +COLI_METAL=1 COLI_MODEL=/path/glm52_i4 ./coli chat --ram 96 +``` + +Measured on an M4 Max (128 GB, warm cache, MTP on): CPU 0.30 → Metal **0.42 tok/s (~1.4×)** +(best config adds `DIRECT=1`; ~3× vs this machine's first cold run). +Key design points: Metal's ~5 ms submit latency makes per-matmul dispatch a loss — +everything is batched into few command buffers per layer, and the resident experts' +GPU work is submitted *before* the missed experts' disk reads so I/O and compute +overlap. `COLI_METAL_GEMM_MIN` tunes the prefill GEMM row threshold (default 16). +Streaming, cache, MTP, DSA and the persistence formats are unchanged; every GPU +path falls back to the CPU per-block on any fault. Numerics are dequant→f32-MAC +(same as the CUDA tier); greedy outputs are byte-identical to the CPU engine. + ### Experimental resident CUDA backend colibrì includes an opt-in CUDA backend for model-resident tensors. Streaming diff --git a/c/Makefile b/c/Makefile index 965b506..3833e22 100644 --- a/c/Makefile +++ b/c/Makefile @@ -80,18 +80,39 @@ LDFLAGS += -L$(CUDA_HOME)/lib64 -Wl,-rpath,$(CUDA_HOME)/lib64 -lcudart -lstdc++ CUDA_OBJ = backend_cuda.o endif +# METAL=1 adds an opt-in Apple-GPU backend (macOS only). The shader is compiled at +# runtime, so no Xcode / offline metal compiler is required. Default build unchanged. +METAL ?= 0 +METAL_OBJ = +METALXX = clang++ -x objective-c++ -fobjc-arc -O3 +ifeq ($(METAL),1) +ifneq ($(UNAME_S),Darwin) +$(error METAL=1 is supported only on macOS) +endif +CFLAGS += -DCOLI_METAL +LDFLAGS += -framework Metal -framework Foundation -lc++ +METAL_OBJ = backend_metal.o +endif + all: glm$(EXE) # phony 'glm' → 'glm.exe' on Windows (so 'make glm' and 'coli build' work on every platform) glm: glm$(EXE) -glm$(EXE): glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h $(CUDA_OBJ) - $(CC) $(CFLAGS) glm.c $(CUDA_OBJ) -o glm$(EXE) $(LDFLAGS) +glm$(EXE): glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h $(CUDA_OBJ) $(METAL_OBJ) + $(CC) $(CFLAGS) glm.c $(CUDA_OBJ) $(METAL_OBJ) -o glm$(EXE) $(LDFLAGS) backend_cuda.o: backend_cuda.cu backend_cuda.h @command -v $(NVCC) >/dev/null 2>&1 || { echo "nvcc not found: set CUDA_HOME or NVCC" >&2; exit 1; } $(NVCC) $(NVCCFLAGS) -c backend_cuda.cu -o $@ +backend_metal.o: backend_metal.mm backend_metal.h + $(METALXX) -c backend_metal.mm -o $@ + +metal-test: tests/test_backend_metal.mm backend_metal.mm backend_metal.h + $(METALXX) tests/test_backend_metal.mm backend_metal.mm -framework Metal -framework Foundation -o backend_metal_test + ./backend_metal_test + cuda-test: backend_cuda.cu backend_cuda.h tests/test_backend_cuda.cu @command -v $(NVCC) >/dev/null 2>&1 || { echo "nvcc not found: set CUDA_HOME or NVCC" >&2; exit 1; } $(NVCC) $(NVCCFLAGS) backend_cuda.cu tests/test_backend_cuda.cu -o backend_cuda_test$(EXE) @@ -137,7 +158,7 @@ check: $(MAKE) test clean: - rm -f olmoe$(EXE) glm$(EXE) iobench$(EXE) backend_cuda.o backend_cuda_test$(EXE) $(TEST_BINS) + rm -f olmoe$(EXE) glm$(EXE) iobench$(EXE) backend_cuda.o backend_cuda_test$(EXE) backend_metal.o backend_metal_test $(TEST_BINS) rm -rf tests/__pycache__ .PHONY: all glm cuda-test portable test-c test-python test check clean diff --git a/c/backend_metal.h b/c/backend_metal.h new file mode 100644 index 0000000..4383b3b --- /dev/null +++ b/c/backend_metal.h @@ -0,0 +1,143 @@ +#ifndef COLIBRI_BACKEND_METAL_H +#define COLIBRI_BACKEND_METAL_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* + * Apple-GPU (Metal) backend for colibrì. Apple Silicon has one GPU and unified + * memory, so there is no device list and no host<->device copy: resident weights + * are read zero-copy from the RAM they already occupy. The shader is compiled at + * runtime (newLibraryWithSource:), so no Xcode / offline metal compiler is needed. + */ + +/* Opaque, persistent GPU handle for one resident quantized tensor. */ +typedef struct ColiMetalTensor ColiMetalTensor; + +/* Returns 1 if a Metal device is available and pipelines compiled, else 0. */ +int coli_metal_init(void); +void coli_metal_shutdown(void); +int coli_metal_available(void); +/* Bytes of unified memory in use by wrapped tensors, and their count. */ +void coli_metal_stats(size_t *tensor_count, size_t *tensor_bytes); +int coli_metal_mem_info(size_t *used_bytes, size_t *total_bytes); + +/* + * y[S,O] = (x[S,I] @ W[O,I]^T) * scale[o]. + * fmt matches QT in glm.c: 0=f32, 1=int8, 2=int4(packed), 3=int2(packed). + * The first successful call wraps W and its row scales in GPU-visible buffers; + * later calls reuse them (weights are assumed stable at the same address). + * Returns 1 on success, 0 if Metal is unavailable or fmt is invalid. + */ +int coli_metal_matmul(ColiMetalTensor **tensor, + float *y, const float *x, + const void *weights, const float *scales, + int fmt, int S, int I, int O); + +void coli_metal_tensor_free(ColiMetalTensor *tensor); +size_t coli_metal_tensor_bytes(const ColiMetalTensor *tensor); + +/* + * Register a page-aligned host allocation (expert slab / scale slab) so the batched + * MoE path can read it zero-copy: the backend wraps it once in an MTLBuffer + * (newBufferWithBytesNoCopy) and resolves any pointer inside [base,base+len) to a GPU + * address. Call after (re)allocating a slab; call unregister before freeing it. + * base must be aligned to 16384 (Apple page) and len a multiple of it. + */ +void coli_metal_spin_start(void); /* COLI_METAL_SPIN=1 keep-alive experiment */ +void coli_metal_spin_stop(void); +void coli_metal_register(void *base, size_t len); +void coli_metal_unregister(void *base); + +/* + * Fused decode (S=1) attention for one layer, entirely on the GPU in one command buffer: + * q_a -> rmsnorm -> q_b -> RoPE ; kv_a -> latent rmsnorm@pos + krot RoPE@pos (cache write) ; + * MLA absorption core ; o_proj. Weights (q_a/q_b/kv_a/kv_b/o) and the Lc/Rc caches must be + * registered (page-aligned) for zero-copy resolve. GLM-5.2 dims compiled in. Handles st0==0 + * full-range only. Returns 1 on success, 0 to signal CPU fallback. + */ +/* + * Full decode layer in ONE command buffer: in_ln -> attention -> residual -> post_ln -> + * shared expert -> router+top-K (exact phase-A semantics). x updated in place; nrm_out + * is the expert input; sh_out the shared-expert output; idx/w/keff the routing. + * Returns 0 -> caller runs the whole layer on the CPU path. + */ +int coli_metal_layer_decode(float *x, + const float *in_ln, const float *post_ln, + const void *qa_w, const float *qa_s, int qa_fmt, const float *qa_ln, + const void *qb_w, const float *qb_s, int qb_fmt, + const void *kva_w, const float *kva_s, int kva_fmt, const float *kva_ln, + const void *kvb_w, const float *kvb_s, int kvb_fmt, + const void *o_w, const float *o_s, int o_fmt, + const void *shg_w, const float *shg_s, int shg_fmt, + const void *shu_w, const float *shu_s, int shu_fmt, + const void *shd_w, const float *shd_s, int shd_fmt, + const float *router_w, const float *router_bias, + int E, int K, int Ksel, float topp, int normk, float rscale, + float *Lc, float *Rc, int S, int pos_base, int st0, + float eps, float theta, float ascale, + float *inrm_out, float *nrm_out, float *sh_out, int *idx_out, float *w_out, int *keff_out); + +int coli_metal_gemm(float *y, const float *x, const void *weights, const float *scales, + int fmt, int S, int I, int O); /* large-batch sync GEMM; 0 -> CPU */ +void coli_metal_attn_counts(uint64_t *ok, double *wall, double *kernel); +void coli_metal_attn_lat(double *ksched, double *gsched); +int coli_metal_attn_decode(const float *x, + const void *qa_w, const float *qa_s, int qa_fmt, const float *qa_ln, + const void *qb_w, const float *qb_s, int qb_fmt, + const void *kva_w, const float *kva_s, int kva_fmt, const float *kva_ln, + const void *kvb_w, const float *kvb_s, int kvb_fmt, + const void *o_w, const float *o_s, int o_fmt, + float *Lc, float *Rc, int S, int pos_base, int st0, float eps, float theta, float ascale, float *out); + +/* Diagnostics: GPU blocks executed, CPU-fallback blocks, experts run on GPU. */ +void coli_metal_moe_counts(uint64_t *ok, uint64_t *fb, uint64_t *experts); +void coli_metal_moe_times(double *setup, double *gpu, double *scatter); +double coli_metal_moe_kernel_time(void); + +/* + * Batched routed-expert SwiGLU for one MoE block, in ONE command buffer. + * For each expert e in [0,nb): computes hh_e[nr_e, D] = down( silu(gate(xg_e)) * up(xg_e) ) + * and scatter-adds rw * hh_e into out. All experts share the command buffer so the + * ~150us Metal launch latency is paid once per block, not per matmul. + * + * D = hidden size, Iinter = moe intermediate size + * g/u/d[e] = pointers to expert e's gate/up/down quantized weights (in RAM slabs) + * gs/us/ds[e] = pointers to expert e's per-row scales + * fmt = quant format (shared across experts) + * xg = packed activations [total_rows, D]; xoff[e] = row offset of expert e + * nr[e] = rows for expert e; rows[]/rw[] map packed rows back to out positions + * out = [S, D] accumulate target + * Returns 1 on success, 0 to signal the caller to fall back to the CPU path. + */ +int coli_metal_moe_block(int nb, int D, int Iinter, int fmt, + const void *const *g, const void *const *u, const void *const *d, + const float *const *gs, const float *const *us, const float *const *ds, + const float *xg, const int *xoff, const int *nr, + const int *rows, const float *rw, + float *out, int S); + +/* + * Async two-phase variant: begin encodes+commits the block (own scratch, no wait) and + * returns a handle, so the CPU can load missed experts from disk WHILE the GPU computes + * the resident ones; end waits, checks for GPU faults, scatter-adds into out, and frees + * the handle. begin returns NULL (nothing submitted) on unresolved slab / bad fmt / R==0; + * end returns 0 on GPU fault (caller redoes those experts on CPU). + */ +typedef struct ColiMetalMoeHandle ColiMetalMoeHandle; +ColiMetalMoeHandle* coli_metal_moe_block_begin(int nb, int D, int Iinter, int fmt, + const void *const *g, const void *const *u, const void *const *d, + const float *const *gs, const float *const *us, const float *const *ds, + const float *xg, const int *xoff, const int *nr, + const int *rows, const float *rw); +int coli_metal_moe_block_end(ColiMetalMoeHandle *h, float *out); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/c/backend_metal.mm b/c/backend_metal.mm new file mode 100644 index 0000000..51ed8a0 --- /dev/null +++ b/c/backend_metal.mm @@ -0,0 +1,786 @@ +// Apple-GPU (Metal) backend for colibrì. Runtime-compiled shader (no Xcode needed), +// zero-copy over unified memory. See backend_metal.h and docs/plans/2026-07-10-*. +#import +#import +#include "backend_metal.h" +#include +#include +#include + +// ---- shader: general quantized GEMV, one threadgroup per output element (o,si) ---- +// y[si,o] = (sum_i dequant(W[o,i]) * x[si,i]) * scale[o]. fmt: 0=f32 1=i8 2=i4 3=i2. +static const char *SHADER = R"METAL( +#include +using namespace metal; + +kernel void mm_gemv(device const uchar* w [[buffer(0)]], // raw weight bytes + device const float* scale [[buffer(1)]], // [O] + device const float* x [[buffer(2)]], // [S,I] + device float* y [[buffer(3)]], // [S,O] + constant int& S [[buffer(4)]], constant int& I [[buffer(5)]], + constant int& O [[buffer(6)]], constant int& fmt [[buffer(7)]], + constant int& NT [[buffer(8)]], + uint tg [[threadgroup_position_in_grid]], + uint slane [[thread_index_in_simdgroup]], + uint sgid [[simdgroup_index_in_threadgroup]]) { + // one SIMDGROUP per output element, 4 per threadgroup, 8-value loads (see moe_gemv) + long row = (long)tg*4 + sgid; if (row >= NT) return; + int o = row % O, si = row / O; + device const float* xr = x + (long)si * I; + device const float4* x4 = (device const float4*)xr; + int I8 = (I & 7) ? 0 : (I/8); + float acc = 0.0f; + if (fmt == 1) { // int8 + device const char* wr = (device const char*)(w) + (long)o * I; + device const char4* w4 = (device const char4*)wr; + for (int c = slane; c < I8; c += 32) acc += dot(float4(w4[2*c]),x4[2*c]) + dot(float4(w4[2*c+1]),x4[2*c+1]); + for (int i = I8*8 + slane; i < I; i += 32) acc += float(wr[i]) * xr[i]; + } else if (fmt == 2) { // int4 packed, rb=(I+1)/2 + int rb = (I+1)/2; + device const uchar* wr = w + (long)o * rb; + device const uchar4* w4 = (device const uchar4*)wr; + for (int c = slane; c < I8; c += 32) { uchar4 b = w4[c]; + float4 w0 = float4(float(int(b.x&0xF)-8), float(int(b.x>>4)-8), float(int(b.y&0xF)-8), float(int(b.y>>4)-8)); + float4 w1 = float4(float(int(b.z&0xF)-8), float(int(b.z>>4)-8), float(int(b.w&0xF)-8), float(int(b.w>>4)-8)); + acc += dot(w0,x4[2*c]) + dot(w1,x4[2*c+1]); + } + for (int i = I8*8 + slane; i < I; i += 32) { + uchar b = wr[i>>1]; int v = (i&1) ? (b>>4) : (b&0xF); acc += float(v-8) * xr[i]; + } + } else if (fmt == 3) { // int2 packed, rb=(I+3)/4 + int rb = (I+3)/4; + device const uchar* wr = w + (long)o * rb; + for (int i = slane; i < I; i += 32) { + uchar b = wr[i>>2]; int v = (b >> (2*(i&3))) & 0x3; acc += float(v-2) * xr[i]; + } + } else { // f32 + device const float* wr = (device const float*)(w) + (long)o * I; + device const float4* w4 = (device const float4*)wr; + for (int c = slane; c < I8; c += 32) acc += dot(w4[2*c],x4[2*c]) + dot(w4[2*c+1],x4[2*c+1]); + for (int i = I8*8 + slane; i < I; i += 32) acc += wr[i] * xr[i]; + } + acc = simd_sum(acc); + if (slane == 0) y[row] = acc * scale[o]; +} + +// Batched bindless expert GEMV: each row gr belongs to expert erow[gr], whose weight and +// scale live at gpuAddresses waddr[e]/saddr[e] (zero-copy in the RAM slab). fmt 1=i8, 2=i4. +// One SIMDGROUP per output row, 4 rows/threadgroup, 8-value loads: measured 1.5-2.1x over +// one-threadgroup-per-row with uchar2 loads (358-389 GB/s on engine-like block shapes). +kernel void moe_gemv(device const ulong* waddr [[buffer(0)]], device const ulong* saddr [[buffer(1)]], + device const int* erow [[buffer(2)]], device const float* xin [[buffer(3)]], + device float* yout [[buffer(4)]], + constant int& O [[buffer(5)]], constant int& K [[buffer(6)]], + constant int& Kin [[buffer(7)]], constant int& fmt [[buffer(8)]], + constant int& NT [[buffer(9)]], + uint tg [[threadgroup_position_in_grid]], + uint slane [[thread_index_in_simdgroup]], + uint sgid [[simdgroup_index_in_threadgroup]]) { + long row = (long)tg*4 + sgid; if (row >= NT) return; + int gr = row / O, o = row % O; int e = erow[gr]; int K8 = (K & 7) ? 0 : (K/8); + device const float* xr = xin + (long)gr * Kin; + device const float* sc = (device const float*)(saddr[e]); + device const float4* x4 = (device const float4*)xr; + float acc = 0.0f; + if (fmt == 2) { int rb=(K+1)/2; device const uchar* w=(device const uchar*)(waddr[e])+(long)o*rb; + device const uchar4* w4=(device const uchar4*)w; + for(int c=slane;c>4)-8),float(int(b.y&0xF)-8),float(int(b.y>>4)-8)); + float4 w1=float4(float(int(b.z&0xF)-8),float(int(b.z>>4)-8),float(int(b.w&0xF)-8),float(int(b.w>>4)-8)); + acc+=dot(w0,x4[2*c])+dot(w1,x4[2*c+1]); } + for(int i=K8*8+slane;i>1]; int v=(i&1)?(b>>4):(b&0xF); acc+=float(v-8)*xr[i]; } + } else { device const char* w=(device const char*)(waddr[e])+(long)o*K; + device const char4* w4=(device const char4*)w; + for(int c=slane;c0;k>>=1){ if(lid>1]; int val=(i&1)?(b>>4):(b&0xF); return float(val-8)*sc[row]; } +kernel void a_qabs(device const uchar* kvb [[buffer(0)]], device const float* sc [[buffer(1)]], + device const float* q [[buffer(2)]], device float* qabs [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + int s=gid/(A_H*A_KVL), r=gid%(A_H*A_KVL), h=r/A_KVL, i=r%A_KVL; int rbase=h*A_ROWSH; + device const float* qp=q+(long)s*A_QHH+(long)h*A_QH; + float a=0; for(int d=0;d PB+s){ sc[o]=-1e30f; return; } // causal mask + device const float* qa=qabs+(long)(s*A_H+h)*A_KVL; device const float* Lt=Lc+(long)t*A_KVL; + device const float* qr=q+(long)s*A_QHH+(long)h*A_QH+A_NOPE; device const float* Rt=Rc+(long)t*A_ROPE; + float a=0; for(int i=0;i0;k>>=1){ if(lid0;k>>=1){ if(lid sig=1/(1+exp(-logit)). One simdgroup/row. +kernel void r_router(device const float* rw [[buffer(0)]], device const float* x [[buffer(1)]], + device float* sig [[buffer(2)]], constant int& E [[buffer(3)]], + constant int& D [[buffer(4)]], constant int& NT [[buffer(5)]], + uint tg [[threadgroup_position_in_grid]], + uint slane [[thread_index_in_simdgroup]], uint sgid [[simdgroup_index_in_threadgroup]]) { + long row=(long)tg*4+sgid; if(row>=NT) return; + int e=row%E, s=row/E; + device const float4* w4=(device const float4*)(rw+(long)e*D); + device const float4* x4=(device const float4*)(x+(long)s*D); + float acc=0; int D4=D/4; + for(int c=slane;cbv){bv=ch;best=e;} } + id_[kk]=best; ww[kk]=sg[best]; + } + int Ke=Ksel; + if(topp>0.0f && topp<1.0f){ + for(int a=1;a=0 && ww[b]=topp*tot){ Ke=kk+1; break; } } + } + keff[s]=Ke; + if(normk){ float sm=0; for(int kk=0;kk w; // weights (wrapped, zero-copy when page-aligned) + id s; // scales + int fmt, I, O; size_t wbytes; +}; + +static id g_dev; +static id g_queue; +static id g_gemv, g_moe_gemv, g_moe_silu; +static id g_a_rms, g_a_rope, g_a_copy, g_a_qabs, g_a_score, g_a_smax, g_a_clat, g_a_ctx; +static id g_a_add, g_r_router, g_r_top8; +static size_t g_tensor_count, g_tensor_bytes; +static uint64_t g_moe_ok, g_moe_fb, g_moe_experts; // GPU blocks / CPU-fallback blocks / experts on GPU +static double g_t_setup, g_t_gpu, g_t_scatter, g_t_kernel; // per-block time breakdown (seconds) +static const int TG = 128; +static MTLResourceOptions g_res_opts = MTLResourceStorageModeShared; // COLI_METAL_UNTRACKED=1 adds HazardTrackingModeUntracked +#include +static double mnow(){ static mach_timebase_info_data_t tb; if(tb.denom==0) mach_timebase_info(&tb); + return (double)mach_absolute_time()*tb.numer/tb.denom/1e9; } + +extern "C" void coli_metal_moe_counts(uint64_t *ok, uint64_t *fb, uint64_t *ex) { + if(ok)*ok=g_moe_ok; if(fb)*fb=g_moe_fb; if(ex)*ex=g_moe_experts; +} +extern "C" void coli_metal_moe_times(double *setup, double *gpu, double *scatter) { + if(setup)*setup=g_t_setup; if(gpu)*gpu=g_t_gpu; if(scatter)*scatter=g_t_scatter; +} +extern "C" double coli_metal_moe_kernel_time(void){ return g_t_kernel; } +static uint64_t g_attn_ok; static double g_attn_wall, g_attn_kernel, g_attn_sched, g_attn_ksched; +extern "C" void coli_metal_attn_counts(uint64_t *ok, double *wall, double *kernel){ + if(ok)*ok=g_attn_ok; if(wall)*wall=g_attn_wall; if(kernel)*kernel=g_attn_kernel; } +extern "C" void coli_metal_attn_lat(double *ksched, double *gsched){ + if(ksched)*ksched=g_attn_ksched; if(gsched)*gsched=g_attn_sched; } + +// Registry of page-aligned host slabs wrapped zero-copy for the batched MoE path. +struct Slab { void *base; size_t len; id buf; }; +static std::vector g_slabs; +static std::mutex g_slab_mtx; // expert_load registers slabs from parallel OpenMP threads +// Persistent scratch buffers (grow-only) for the MoE pipeline. +static id g_gg, g_uu, g_hh, g_xg; static size_t g_gg_cap, g_uu_cap, g_hh_cap, g_xg_cap; +static id ensure(id b, size_t *cap, size_t need) { + if (b && *cap >= need) return b; + *cap = need; return [g_dev newBufferWithLength:need options:g_res_opts]; +} + +static size_t fmt_bytes(int fmt, int I, int O) { + if (fmt == 1) return (size_t)O * I; + if (fmt == 2) return (size_t)O * ((I+1)/2); + if (fmt == 3) return (size_t)O * ((I+3)/4); + return (size_t)O * I * sizeof(float); +} + +// Wrap host memory zero-copy if page-aligned, else copy into a shared buffer. +static id wrap(const void *p, size_t n) { + size_t pg = 16384; // Apple Silicon page + if (((uintptr_t)p % pg) == 0 && (n % pg) == 0) + return [g_dev newBufferWithBytesNoCopy:(void*)p length:n options:MTLResourceStorageModeShared deallocator:nil]; + return [g_dev newBufferWithBytes:p length:n options:MTLResourceStorageModeShared]; +} + +extern "C" int coli_metal_init(void) { + if (g_dev) return 1; + if (getenv("COLI_METAL_UNTRACKED") && atoi(getenv("COLI_METAL_UNTRACKED"))) + g_res_opts = MTLResourceStorageModeShared | MTLResourceHazardTrackingModeUntracked; + @autoreleasepool { + g_dev = MTLCreateSystemDefaultDevice(); + if (!g_dev) return 0; + g_queue = [g_dev newCommandQueue]; + NSError *err = nil; + id lib = [g_dev newLibraryWithSource:[NSString stringWithUTF8String:SHADER] + options:nil error:&err]; + if (!lib) { fprintf(stderr, "[metal] shader compile failed: %s\n", + err ? [[err localizedDescription] UTF8String] : "?"); g_dev = nil; return 0; } + g_gemv = [g_dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"mm_gemv"] error:&err]; + g_moe_gemv = [g_dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"moe_gemv"] error:&err]; + g_moe_silu = [g_dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"moe_silu"] error:&err]; + auto P=[&](const char*n){ return [g_dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@(n)] error:&err]; }; + g_a_rms=P("a_rmsnorm"); g_a_rope=P("a_rope"); g_a_copy=P("a_copy"); + g_a_qabs=P("a_qabs"); g_a_score=P("a_score"); g_a_smax=P("a_smax"); g_a_clat=P("a_clat"); g_a_ctx=P("a_ctx"); + g_a_add=P("a_add"); g_r_router=P("r_router"); g_r_top8=P("r_top8"); + if(!g_a_add||!g_r_router||!g_r_top8){ fprintf(stderr,"[metal] tail pipelines failed\n"); g_dev=nil; return 0; } + if (!g_gemv || !g_moe_gemv || !g_moe_silu || !g_a_rms || !g_a_rope || !g_a_copy || + !g_a_qabs || !g_a_score || !g_a_smax || !g_a_clat || !g_a_ctx) { + fprintf(stderr, "[metal] pipeline failed\n"); g_dev = nil; return 0; } + } + return 1; +} + +extern "C" void coli_metal_register(void *base, size_t len) { + if (!g_dev || !base) return; + id b = [g_dev newBufferWithBytesNoCopy:base length:len + options:g_res_opts deallocator:nil]; + if (!b) return; + std::lock_guard lk(g_slab_mtx); // called from parallel expert_load threads + for (auto &s : g_slabs) if (s.base == base) { s.len = len; s.buf = b; return; } + g_slabs.push_back({base, len, b}); +} +extern "C" void coli_metal_unregister(void *base) { + std::lock_guard lk(g_slab_mtx); + for (size_t i=0;i resolve(const void *p, uint64_t *addr) { + std::lock_guard lk(g_slab_mtx); + uintptr_t u=(uintptr_t)p; + for (auto &s : g_slabs) { uintptr_t b=(uintptr_t)s.base; + if (u>=b && u +#include +static std::atomic g_spin_run{false}; +static std::thread g_spin_thr; +extern "C" void coli_metal_spin_start(void) { + if (!g_dev || g_spin_run.exchange(true)) return; + g_spin_thr = std::thread([]{ + id q = [g_dev newCommandQueue]; // own queue: never blocks real work + id b = [g_dev newBufferWithLength:4096 options:MTLResourceStorageModeShared]; + while (g_spin_run.load()) { + @autoreleasepool { + id cb=[q commandBuffer]; + id e=[cb computeCommandEncoder]; + [e setComputePipelineState:g_moe_silu]; + [e setBuffer:b offset:0 atIndex:0]; [e setBuffer:b offset:0 atIndex:1]; + [e dispatchThreads:MTLSizeMake(1024,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + [e endEncoding]; [cb commit]; [cb waitUntilCompleted]; + } + } + }); + g_spin_thr.detach(); // never joinable at exit (joinable global -> std::terminate) +} +extern "C" void coli_metal_spin_stop(void) { g_spin_run.store(false); } + +extern "C" void coli_metal_shutdown(void) { coli_metal_spin_stop(); g_gemv=nil; g_queue=nil; g_dev=nil; g_tensor_count=g_tensor_bytes=0; } +extern "C" int coli_metal_available(void) { return g_dev != nil; } +extern "C" void coli_metal_stats(size_t *c, size_t *b) { if(c)*c=g_tensor_count; if(b)*b=g_tensor_bytes; } +extern "C" int coli_metal_mem_info(size_t *used, size_t *total) { + if (!g_dev) return 0; + if (used) *used = (size_t)[g_dev currentAllocatedSize]; + if (total) *total = (size_t)[g_dev recommendedMaxWorkingSetSize]; + return 1; +} + +extern "C" int coli_metal_matmul(ColiMetalTensor **tp, float *y, const float *x, + const void *weights, const float *scales, + int fmt, int S, int I, int O) { + if (!g_dev || fmt < 0 || fmt > 3) return 0; + @autoreleasepool { + ColiMetalTensor *t = *tp; + if (!t) { + t = new ColiMetalTensor(); + t->fmt = fmt; t->I = I; t->O = O; t->wbytes = fmt_bytes(fmt, I, O); + t->w = wrap(weights, t->wbytes); + t->s = wrap(scales, (size_t)O * sizeof(float)); + *tp = t; + g_tensor_count++; g_tensor_bytes += t->wbytes; + } + id bx = [g_dev newBufferWithBytes:x length:(size_t)S*I*sizeof(float) options:MTLResourceStorageModeShared]; + id by = [g_dev newBufferWithLength:(size_t)S*O*sizeof(float) options:MTLResourceStorageModeShared]; + id cb = [g_queue commandBuffer]; + id e = [cb computeCommandEncoder]; + [e setComputePipelineState:g_gemv]; + [e setBuffer:t->w offset:0 atIndex:0]; [e setBuffer:t->s offset:0 atIndex:1]; + [e setBuffer:bx offset:0 atIndex:2]; [e setBuffer:by offset:0 atIndex:3]; + int NT=S*O; + [e setBytes:&S length:4 atIndex:4]; [e setBytes:&I length:4 atIndex:5]; + [e setBytes:&O length:4 atIndex:6]; [e setBytes:&fmt length:4 atIndex:7]; + [e setBytes:&NT length:4 atIndex:8]; + [e dispatchThreadgroups:MTLSizeMake(((size_t)NT+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + [e endEncoding]; [cb commit]; [cb waitUntilCompleted]; + memcpy(y, [by contents], (size_t)S*O*sizeof(float)); + } + return 1; +} + +// ---- fused decode attention scratch (GLM-5.2 dims) ---- +enum { AH=6144, AHEADS=64, AQLORA=2048, AKVL=512, AROPE=64, AVH=256, AQH=256, ANOPE=192, AROWSH=448, AHQH=AHEADS*AQH, AHVH=AHEADS*AVH, AMAXS=4 }; +static id ax_,aqr_,aqf_,acomp_,aqabs_,ascore_,aclat_,actx_,aout_,aqaln_,akvaln_; static size_t ascore_cap; +static id axr_,anrm_,ash1_,ash2_,ashout_,asig_,aidx_,aw_,akeff_; // full-layer tail +static void attn_scratch_init(){ + if(ax_) return; + auto L=[&](size_t n){ return [g_dev newBufferWithLength:n*AMAXS options:g_res_opts]; }; + ax_=L(AH*4); aqr_=L(AQLORA*4); aqf_=L(AHQH*4); acomp_=L((AKVL+AROPE)*4); + aqabs_=L((size_t)AHEADS*AKVL*4); aclat_=L((size_t)AHEADS*AKVL*4); actx_=L(AHVH*4); aout_=L(AH*4); + aqaln_=L(AQLORA*4/AMAXS); akvaln_=L(AKVL*4/AMAXS); // norm weights are per-tensor, not per-row + axr_=L(AH*4); anrm_=L(AH*4); ash1_=L(2048*4); ash2_=L(2048*4); ashout_=L(AH*4); + asig_=L(256*4); aidx_=L(8*4); aw_=L(8*4); akeff_=L(4); +} +// y[S,O] = quantized-weight(w) applied to xin[S,I]. Weights are registered (page-aligned, +// zero-copy) at model load; resolve to (buffer,offset). Returns false to fall back to CPU. +static bool bind_gemv(id e, const void* w, const float* s, int fmt, int I, int O, + id xin, id yout, int S){ + uint64_t wa=0,sa=0; id wb=resolve(w,&wa); id sb=resolve(s,&sa); + if(!wb||!sb) return false; + size_t woff=wa-(uint64_t)[wb gpuAddress], soff=sa-(uint64_t)[sb gpuAddress]; + [e useResource:wb usage:MTLResourceUsageRead]; [e useResource:sb usage:MTLResourceUsageRead]; + [e setComputePipelineState:g_gemv]; + [e setBuffer:wb offset:woff atIndex:0]; [e setBuffer:sb offset:soff atIndex:1]; + [e setBuffer:xin offset:0 atIndex:2]; [e setBuffer:yout offset:0 atIndex:3]; + int NT=S*O; + [e setBytes:&S length:4 atIndex:4]; [e setBytes:&I length:4 atIndex:5]; [e setBytes:&O length:4 atIndex:6]; [e setBytes:&fmt length:4 atIndex:7]; + [e setBytes:&NT length:4 atIndex:8]; + [e dispatchThreadgroups:MTLSizeMake(((size_t)NT+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; +} + +// Weight-pointer bundle for one layer's attention (+optional layer tail). All pointers +// must be inside registered allocations. +typedef struct { + const void *qa_w; const float *qa_s; int qa_fmt; const float *qa_ln; + const void *qb_w; const float *qb_s; int qb_fmt; + const void *kva_w; const float *kva_s; int kva_fmt; const float *kva_ln; + const void *kvb_w; const float *kvb_s; int kvb_fmt; + const void *o_w; const float *o_s; int o_fmt; +} AttnW; + +// Encode the fused attention chain into encoder e. Input: ax_ holds the NORMED x [S,AH]. +// Output: aout_ holds attention output [S,AH]. Returns false on unresolved weights. +static bool encode_attention(id e, const AttnW *W, + id Lb, size_t loff, id Rb, size_t roff, + id kvbW, size_t kvbwoff, id kvbS, size_t kvbsoff, + int S, int pos_base, float eps, float theta, float ascale) { + int T=pos_base+S; + memcpy([aqaln_ contents],W->qa_ln,AQLORA*4); memcpy([akvaln_ contents],W->kva_ln,AKVL*4); + size_t Loff=loff+(size_t)pos_base*AKVL*4, Roff=roff+(size_t)pos_base*AROPE*4; + auto BAR=[&]{ [e memoryBarrierWithScope:MTLBarrierScopeBuffers]; }; + auto rms=[&](id b,size_t off,id w,int n,int nrows){ [e setComputePipelineState:g_a_rms]; + [e setBuffer:b offset:off atIndex:0]; [e setBuffer:w offset:0 atIndex:1]; [e setBytes:&n length:4 atIndex:2]; [e setBytes:&eps length:4 atIndex:3]; + [e dispatchThreadgroups:MTLSizeMake(nrows,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; }; + auto rope=[&](id b,size_t off,int base,int rs,int hs,int nh){ [e setComputePipelineState:g_a_rope]; [e setBuffer:b offset:off atIndex:0]; + [e setBytes:&base length:4 atIndex:1]; [e setBytes:&rs length:4 atIndex:2]; [e setBytes:&hs length:4 atIndex:3]; [e setBytes:&nh length:4 atIndex:4]; [e setBytes:&pos_base length:4 atIndex:5]; [e setBytes:&theta length:4 atIndex:6]; + [e dispatchThreads:MTLSizeMake((size_t)S*nh*(AROPE/2),1,1) threadsPerThreadgroup:MTLSizeMake(64,1,1)]; }; + auto cpy=[&](int off,id dst,size_t doff,int n){ int ss=AKVL+AROPE; [e setComputePipelineState:g_a_copy]; + [e setBuffer:acomp_ offset:0 atIndex:0]; [e setBytes:&off length:4 atIndex:1]; [e setBytes:&ss length:4 atIndex:2]; + [e setBuffer:dst offset:doff atIndex:3]; [e setBytes:&n length:4 atIndex:4]; [e setBytes:&n length:4 atIndex:5]; + [e dispatchThreads:MTLSizeMake((size_t)S*n,1,1) threadsPerThreadgroup:MTLSizeMake(64,1,1)]; }; + bind_gemv(e,W->qa_w,W->qa_s,W->qa_fmt,AH,AQLORA,ax_,aqr_,S); + bind_gemv(e,W->kva_w,W->kva_s,W->kva_fmt,AH,AKVL+AROPE,ax_,acomp_,S); BAR(); + rms(aqr_,0,aqaln_,AQLORA,S); cpy(0,Lb,Loff,AKVL); cpy(AKVL,Rb,Roff,AROPE); BAR(); + bind_gemv(e,W->qb_w,W->qb_s,W->qb_fmt,AQLORA,AHQH,aqr_,aqf_,S); rms(Lb,Loff,akvaln_,AKVL,S); rope(Rb,Roff,0,AROPE,0,1); BAR(); + rope(aqf_,0,ANOPE,AHQH,AQH,AHEADS); BAR(); + [e setComputePipelineState:g_a_qabs]; [e setBuffer:kvbW offset:kvbwoff atIndex:0]; [e setBuffer:kvbS offset:kvbsoff atIndex:1]; [e setBuffer:aqf_ offset:0 atIndex:2]; [e setBuffer:aqabs_ offset:0 atIndex:3]; + [e dispatchThreads:MTLSizeMake((size_t)S*AHEADS*AKVL,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); + [e setComputePipelineState:g_a_score]; [e setBuffer:aqabs_ offset:0 atIndex:0]; [e setBuffer:Lb offset:loff atIndex:1]; [e setBuffer:Rb offset:roff atIndex:2]; [e setBuffer:aqf_ offset:0 atIndex:3]; [e setBuffer:ascore_ offset:0 atIndex:4]; + [e setBytes:&T length:4 atIndex:5]; [e setBytes:&ascale length:4 atIndex:6]; [e setBytes:&pos_base length:4 atIndex:7]; + [e dispatchThreads:MTLSizeMake((size_t)S*AHEADS*T,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); + [e setComputePipelineState:g_a_smax]; [e setBuffer:ascore_ offset:0 atIndex:0]; [e setBytes:&T length:4 atIndex:1]; + [e dispatchThreadgroups:MTLSizeMake((size_t)S*AHEADS,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); + [e setComputePipelineState:g_a_clat]; [e setBuffer:ascore_ offset:0 atIndex:0]; [e setBuffer:Lb offset:loff atIndex:1]; [e setBuffer:aclat_ offset:0 atIndex:2]; [e setBytes:&T length:4 atIndex:3]; + [e dispatchThreads:MTLSizeMake((size_t)S*AHEADS*AKVL,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); + [e setComputePipelineState:g_a_ctx]; [e setBuffer:kvbW offset:kvbwoff atIndex:0]; [e setBuffer:kvbS offset:kvbsoff atIndex:1]; [e setBuffer:aclat_ offset:0 atIndex:2]; [e setBuffer:actx_ offset:0 atIndex:3]; + [e dispatchThreads:MTLSizeMake((size_t)S*AHEADS*AVH,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); + bind_gemv(e,W->o_w,W->o_s,W->o_fmt,AHVH,AH,actx_,aout_,S); + return true; +} +// Resolve Lc/Rc + kv_b (+pre-check the projection weights). Returns false -> CPU fallback. +static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, + id *Lb, size_t *loff, id *Rb, size_t *roff, + id *kvbW, size_t *kvbwoff, id *kvbS, size_t *kvbsoff) { + uint64_t la=0,ra=0,kva=0,ksa=0; + *Lb=resolve(Lc,&la); *Rb=resolve(Rc,&ra); *kvbW=resolve(W->kvb_w,&kva); *kvbS=resolve(W->kvb_s,&ksa); + if(!*Lb||!*Rb||!*kvbW||!*kvbS) return false; + uint64_t d; const void* ws[]={W->qa_w,W->qa_s,W->qb_w,W->qb_s,W->kva_w,W->kva_s,W->o_w,W->o_s}; + for(auto p:ws) if(!resolve(p,&d)) return false; + *loff=la-(uint64_t)[*Lb gpuAddress]; *roff=ra-(uint64_t)[*Rb gpuAddress]; + *kvbwoff=kva-(uint64_t)[*kvbW gpuAddress]; *kvbsoff=ksa-(uint64_t)[*kvbS gpuAddress]; + return true; +} + +extern "C" int coli_metal_attn_decode(const float* x, + const void* qa_w,const float* qa_s,int qa_fmt,const float* qa_ln, + const void* qb_w,const float* qb_s,int qb_fmt, + const void* kva_w,const float* kva_s,int kva_fmt,const float* kva_ln, + const void* kvb_w,const float* kvb_s,int kvb_fmt, + const void* o_w,const float* o_s,int o_fmt, + float* Lc,float* Rc,int S,int pos_base,int st0,float eps,float theta,float ascale,float* out){ + if(!g_dev) return 0; + if(st0!=0 || S<1 || S>AMAXS) return 0; // partial-KV / S>4 -> CPU + int T=pos_base+S; + @autoreleasepool { + attn_scratch_init(); + AttnW W={qa_w,qa_s,qa_fmt,qa_ln,qb_w,qb_s,qb_fmt,kva_w,kva_s,kva_fmt,kva_ln,kvb_w,kvb_s,kvb_fmt,o_w,o_s,o_fmt}; + id Lb,Rb,kvbW,kvbS; size_t loff,roff,kvbwoff,kvbsoff; + if(!resolve_attn(&W,Lc,Rc,&Lb,&loff,&Rb,&roff,&kvbW,&kvbwoff,&kvbS,&kvbsoff)) return 0; + ascore_=ensure(ascore_,&ascore_cap,(size_t)S*AHEADS*T*4); + memcpy([ax_ contents],x,(size_t)S*AH*4); + id cb=[g_queue commandBuffer]; id e=[cb computeCommandEncoder]; + [e useResource:Lb usage:MTLResourceUsageRead|MTLResourceUsageWrite]; [e useResource:Rb usage:MTLResourceUsageRead|MTLResourceUsageWrite]; + [e useResource:kvbW usage:MTLResourceUsageRead]; [e useResource:kvbS usage:MTLResourceUsageRead]; + if(!encode_attention(e,&W,Lb,loff,Rb,roff,kvbW,kvbwoff,kvbS,kvbsoff,S,pos_base,eps,theta,ascale)) return 0; + double tc=mnow(); + [e endEncoding]; [cb commit]; [cb waitUntilCompleted]; + if(cb.status==MTLCommandBufferStatusError){ fprintf(stderr,"[metal] attn cmdbuf error: %s\n", cb.error?[[cb.error localizedDescription]UTF8String]:"?"); return 0; } + g_attn_ok++; g_attn_wall += mnow()-tc; g_attn_kernel += [cb GPUEndTime]-[cb GPUStartTime]; + g_attn_sched += [cb GPUStartTime]-[cb kernelStartTime]; g_attn_ksched += [cb kernelStartTime]-tc; + memcpy(out,[aout_ contents],(size_t)S*AH*4); + } + return 1; +} + +// Full decode layer on the GPU in ONE command buffer: +// in_ln rmsnorm -> fused attention -> residual add (x updated) -> post_ln rmsnorm -> +// shared expert (gate/up/silu/down) -> router (f32 matvec+sigmoid) -> exact top-K select. +// CPU keeps: expert resolve/disk loads + expert CBs + scatter (unchanged). Outputs: +// x (updated in place), nrm=post_ln(x) (expert input), sh_out (shared-expert output), +// idx/w/keff (routing). Returns 0 -> CPU fallback (whole layer falls back). +extern "C" int coli_metal_layer_decode(float *x, + const float *in_ln, const float *post_ln, + const void* qa_w,const float* qa_s,int qa_fmt,const float* qa_ln, + const void* qb_w,const float* qb_s,int qb_fmt, + const void* kva_w,const float* kva_s,int kva_fmt,const float* kva_ln, + const void* kvb_w,const float* kvb_s,int kvb_fmt, + const void* o_w,const float* o_s,int o_fmt, + const void* shg_w,const float* shg_s,int shg_fmt, + const void* shu_w,const float* shu_s,int shu_fmt, + const void* shd_w,const float* shd_s,int shd_fmt, + const float *router_w, const float *router_bias, + int E, int K, int Ksel, float topp, int normk, float rscale, + float *Lc, float *Rc, int S, int pos_base, int st0, + float eps, float theta, float ascale, + float *inrm_out, float *nrm_out, float *sh_out, int *idx_out, float *w_out, int *keff_out) { + if(!g_dev) return 0; + if(st0!=0 || S<1 || S>AMAXS || E!=256 || K!=8) return 0; + int T=pos_base+S; const int SI=2048; + @autoreleasepool { + attn_scratch_init(); + AttnW W={qa_w,qa_s,qa_fmt,qa_ln,qb_w,qb_s,qb_fmt,kva_w,kva_s,kva_fmt,kva_ln,kvb_w,kvb_s,kvb_fmt,o_w,o_s,o_fmt}; + id Lb,Rb,kvbW,kvbS; size_t loff,roff,kvbwoff,kvbsoff; + if(!resolve_attn(&W,Lc,Rc,&Lb,&loff,&Rb,&roff,&kvbW,&kvbwoff,&kvbS,&kvbsoff)) return 0; + uint64_t ina=0,pna=0,rwa=0,rba=0,d; + id inB=resolve(in_ln,&ina), pnB=resolve(post_ln,&pna); + id rwB=resolve(router_w,&rwa), rbB=resolve(router_bias,&rba); + if(!inB||!pnB||!rwB||!rbB) return 0; + { const void* ws[]={shg_w,shg_s,shu_w,shu_s,shd_w,shd_s}; + for(auto p:ws) if(!resolve(p,&d)) return 0; } + size_t inoff=ina-(uint64_t)[inB gpuAddress], pnoff=pna-(uint64_t)[pnB gpuAddress]; + size_t rwoff=rwa-(uint64_t)[rwB gpuAddress], rboff=rba-(uint64_t)[rbB gpuAddress]; + ascore_=ensure(ascore_,&ascore_cap,(size_t)S*AHEADS*T*4); + memcpy([axr_ contents],x,(size_t)S*AH*4); + + id cb=[g_queue commandBuffer]; id e=[cb computeCommandEncoder]; + [e useResource:Lb usage:MTLResourceUsageRead|MTLResourceUsageWrite]; [e useResource:Rb usage:MTLResourceUsageRead|MTLResourceUsageWrite]; + [e useResource:kvbW usage:MTLResourceUsageRead]; [e useResource:kvbS usage:MTLResourceUsageRead]; + [e useResource:inB usage:MTLResourceUsageRead]; [e useResource:pnB usage:MTLResourceUsageRead]; + [e useResource:rwB usage:MTLResourceUsageRead]; [e useResource:rbB usage:MTLResourceUsageRead]; + auto BAR=[&]{ [e memoryBarrierWithScope:MTLBarrierScopeBuffers]; }; + auto rmsw=[&](id b,id wb,size_t woff,int n,int nrows){ [e setComputePipelineState:g_a_rms]; + [e setBuffer:b offset:0 atIndex:0]; [e setBuffer:wb offset:woff atIndex:1]; [e setBytes:&n length:4 atIndex:2]; [e setBytes:&eps length:4 atIndex:3]; + [e dispatchThreadgroups:MTLSizeMake(nrows,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; }; + auto copyrow=[&](id src,id dst,int n){ int off=0,ss=n; [e setComputePipelineState:g_a_copy]; + [e setBuffer:src offset:0 atIndex:0]; [e setBytes:&off length:4 atIndex:1]; [e setBytes:&ss length:4 atIndex:2]; + [e setBuffer:dst offset:0 atIndex:3]; [e setBytes:&n length:4 atIndex:4]; [e setBytes:&n length:4 atIndex:5]; + [e dispatchThreads:MTLSizeMake((size_t)S*n,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; }; + // 1) in_ln: ax_ = rmsnorm(x) + copyrow(axr_,ax_,AH); BAR(); rmsw(ax_,inB,inoff,AH,S); BAR(); + // 2) attention (ax_ -> aout_) + if(!encode_attention(e,&W,Lb,loff,Rb,roff,kvbW,kvbwoff,kvbS,kvbsoff,S,pos_base,eps,theta,ascale)) return 0; + BAR(); + // 3) residual: axr_ += aout_ ; then nrm = post_ln(x_new) + [e setComputePipelineState:g_a_add]; [e setBuffer:axr_ offset:0 atIndex:0]; [e setBuffer:aout_ offset:0 atIndex:1]; + [e dispatchThreads:MTLSizeMake((size_t)S*AH,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); + copyrow(axr_,anrm_,AH); BAR(); rmsw(anrm_,pnB,pnoff,AH,S); BAR(); + // 4) shared expert gate/up + router (all read anrm_, independent) + bind_gemv(e,shg_w,shg_s,shg_fmt,AH,SI,anrm_,ash1_,S); + bind_gemv(e,shu_w,shu_s,shu_fmt,AH,SI,anrm_,ash2_,S); + { int NT=S*E, D=AH; [e setComputePipelineState:g_r_router]; + [e setBuffer:rwB offset:rwoff atIndex:0]; [e setBuffer:anrm_ offset:0 atIndex:1]; [e setBuffer:asig_ offset:0 atIndex:2]; + [e setBytes:&E length:4 atIndex:3]; [e setBytes:&D length:4 atIndex:4]; [e setBytes:&NT length:4 atIndex:5]; + [e dispatchThreadgroups:MTLSizeMake(((size_t)NT+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; } + BAR(); + // 5) silu(gate)*up + exact top-K select + [e setComputePipelineState:g_moe_silu]; [e setBuffer:ash1_ offset:0 atIndex:0]; [e setBuffer:ash2_ offset:0 atIndex:1]; + [e dispatchThreads:MTLSizeMake((size_t)S*SI,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + { [e setComputePipelineState:g_r_top8]; + [e setBuffer:asig_ offset:0 atIndex:0]; [e setBuffer:rbB offset:rboff atIndex:1]; + [e setBuffer:aidx_ offset:0 atIndex:2]; [e setBuffer:aw_ offset:0 atIndex:3]; [e setBuffer:akeff_ offset:0 atIndex:4]; + [e setBytes:&E length:4 atIndex:5]; [e setBytes:&K length:4 atIndex:6]; [e setBytes:&Ksel length:4 atIndex:7]; + [e setBytes:&topp length:4 atIndex:8]; [e setBytes:&normk length:4 atIndex:9]; [e setBytes:&rscale length:4 atIndex:10]; + [e dispatchThreads:MTLSizeMake(S,1,1) threadsPerThreadgroup:MTLSizeMake(S,1,1)]; } + BAR(); + // 6) shared down + bind_gemv(e,shd_w,shd_s,shd_fmt,SI,AH,ash1_,ashout_,S); + double tc=mnow(); + [e endEncoding]; [cb commit]; [cb waitUntilCompleted]; + if(cb.status==MTLCommandBufferStatusError){ fprintf(stderr,"[metal] layer cmdbuf error: %s\n", cb.error?[[cb.error localizedDescription]UTF8String]:"?"); return 0; } + g_attn_ok++; g_attn_wall += mnow()-tc; g_attn_kernel += [cb GPUEndTime]-[cb GPUStartTime]; + g_attn_sched += [cb GPUStartTime]-[cb kernelStartTime]; g_attn_ksched += [cb kernelStartTime]-tc; + memcpy(x,[axr_ contents],(size_t)S*AH*4); + memcpy(inrm_out,[ax_ contents],(size_t)S*AH*4); + memcpy(nrm_out,[anrm_ contents],(size_t)S*AH*4); + memcpy(sh_out,[ashout_ contents],(size_t)S*AH*4); + memcpy(idx_out,[aidx_ contents],(size_t)S*K*4); + memcpy(w_out,[aw_ contents],(size_t)S*K*4); + memcpy(keff_out,[akeff_ contents],(size_t)S*4); + } + return 1; +} + +// Sync GEMM for large row-batches (prefill): y[S,O] = x[S,I] @ W^T * scale. Weights must be +// registered (zero-copy); x/y go through grow-only shared scratch. Returns 0 -> CPU fallback. +static id g_gx, g_gy; static size_t g_gx_cap, g_gy_cap; +extern "C" int coli_metal_gemm(float *y, const float *x, const void *wp, const float *sp, + int fmt, int S, int I, int O) { + if (!g_dev || (fmt!=1 && fmt!=2)) return 0; + @autoreleasepool { + uint64_t wa=0,sa=0; id wb=resolve(wp,&wa), sb=resolve(sp,&sa); + if(!wb||!sb) return 0; + size_t woff=wa-(uint64_t)[wb gpuAddress], soff=sa-(uint64_t)[sb gpuAddress]; + g_gx=ensure(g_gx,&g_gx_cap,(size_t)S*I*4); g_gy=ensure(g_gy,&g_gy_cap,(size_t)S*O*4); + memcpy([g_gx contents],x,(size_t)S*I*4); + id cb=[g_queue commandBuffer]; id e=[cb computeCommandEncoder]; + [e useResource:wb usage:MTLResourceUsageRead]; [e useResource:sb usage:MTLResourceUsageRead]; + [e setComputePipelineState:g_gemv]; + [e setBuffer:wb offset:woff atIndex:0]; [e setBuffer:sb offset:soff atIndex:1]; + [e setBuffer:g_gx offset:0 atIndex:2]; [e setBuffer:g_gy offset:0 atIndex:3]; + int NT=S*O; + [e setBytes:&S length:4 atIndex:4]; [e setBytes:&I length:4 atIndex:5]; + [e setBytes:&O length:4 atIndex:6]; [e setBytes:&fmt length:4 atIndex:7]; + [e setBytes:&NT length:4 atIndex:8]; + [e dispatchThreadgroups:MTLSizeMake(((size_t)NT+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + [e endEncoding]; [cb commit]; [cb waitUntilCompleted]; + if(cb.status==MTLCommandBufferStatusError){ fprintf(stderr,"[metal] gemm cmdbuf error (S=%d O=%d)\n",S,O); return 0; } + memcpy(y,[g_gy contents],(size_t)S*O*4); + } + return 1; +} + +extern "C" void coli_metal_tensor_free(ColiMetalTensor *t) { + if (!t) return; + g_tensor_count--; g_tensor_bytes -= t->wbytes; + t->w = nil; t->s = nil; delete t; +} +extern "C" size_t coli_metal_tensor_bytes(const ColiMetalTensor *t) { return t ? t->wbytes : 0; } + +// Batched routed-expert SwiGLU for one block in ONE command buffer. Returns 0 (CPU fallback) +// if Metal is off or any expert pointer is not in a registered slab. +// Encode + commit a MoE block (no wait). Writes hh[R,D] into hh_buf. Returns nil on +// unresolved slab / bad fmt (caller falls back to CPU). +static id moe_submit(int nb, int D, int Iinter, int fmt, + const void *const *g, const void *const *u, const void *const *d, + const float *const *gs, const float *const *us, const float *const *ds, + const float *xg, const int *xoff, const int *nr, int R, + id xg_buf, id gg_buf, id uu_buf, id hh_buf) { + if (!g_dev || (fmt != 1 && fmt != 2)) return nil; + double ts_start = mnow(); + std::vector ag(nb),au(nb),ad(nb),sgv(nb),suv(nb),sdv(nb); + std::vector> use; use.reserve(nb*2); + auto add_use=[&](id b){ for(auto&x:use) if(x==b) return; use.push_back(b); }; + for (int e=0;e b; + if(!(b=resolve(g[e],&ag[e]))) {g_moe_fb++; return nil;} add_use(b); + if(!(b=resolve(u[e],&au[e]))) {g_moe_fb++; return nil;} add_use(b); + if(!(b=resolve(d[e],&ad[e]))) {g_moe_fb++; return nil;} add_use(b); + if(!(b=resolve(gs[e],&sgv[e]))) {g_moe_fb++; return nil;} add_use(b); + if(!(b=resolve(us[e],&suv[e]))) {g_moe_fb++; return nil;} add_use(b); + if(!(b=resolve(ds[e],&sdv[e]))) {g_moe_fb++; return nil;} add_use(b); + } + std::vector erow(R); for(int e=0;e bag=shb(ag.data(),nb*8), bau=shb(au.data(),nb*8), bad=shb(ad.data(),nb*8); + id bsg=shb(sgv.data(),nb*8), bsu=shb(suv.data(),nb*8), bsd=shb(sdv.data(),nb*8); + id berow=shb(erow.data(),R*4); + memcpy([xg_buf contents], xg, (size_t)R*D*4); + + id cb=[g_queue commandBuffer]; id e=[cb computeCommandEncoder]; + for(auto&b:use) [e useResource:b usage:MTLResourceUsageRead]; + auto gemv=[&](id wa,id sa,id xin,id y,int O,int K,int Kin){ + int NT=R*O; + [e setComputePipelineState:g_moe_gemv]; + [e setBuffer:wa offset:0 atIndex:0];[e setBuffer:sa offset:0 atIndex:1];[e setBuffer:berow offset:0 atIndex:2]; + [e setBuffer:xin offset:0 atIndex:3];[e setBuffer:y offset:0 atIndex:4]; + [e setBytes:&O length:4 atIndex:5];[e setBytes:&K length:4 atIndex:6];[e setBytes:&Kin length:4 atIndex:7];[e setBytes:&fmt length:4 atIndex:8]; + [e setBytes:&NT length:4 atIndex:9]; + [e dispatchThreadgroups:MTLSizeMake(((size_t)NT+3)/4,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; }; + gemv(bag,bsg,xg_buf,gg_buf,Iinter,D,D); // gate + gemv(bau,bsu,xg_buf,uu_buf,Iinter,D,D); // up + [e memoryBarrierWithScope:MTLBarrierScopeBuffers]; + [e setComputePipelineState:g_moe_silu]; + [e setBuffer:gg_buf offset:0 atIndex:0];[e setBuffer:uu_buf offset:0 atIndex:1]; + [e dispatchThreads:MTLSizeMake((size_t)R*Iinter,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + [e memoryBarrierWithScope:MTLBarrierScopeBuffers]; + gemv(bad,bsd,gg_buf,hh_buf,D,Iinter,Iinter); // down + g_t_setup += mnow() - ts_start; + [e endEncoding];[cb commit]; + return cb; +} + +// Wait + error-check + scatter-add hh into out. Returns 0 on GPU fault. +static int moe_finish(id cb, id hh_buf, int nb, int R, int D, + const int *rows, const float *rw, float *out) { + double t0 = mnow(); + [cb waitUntilCompleted]; + double ts_gpu = mnow(); g_t_gpu += ts_gpu - t0; + g_t_kernel += [cb GPUEndTime] - [cb GPUStartTime]; + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "[metal] moe_block cmdbuf error (nb=%d R=%d): %s\n", nb, R, + cb.error ? [[cb.error localizedDescription] UTF8String] : "?"); + g_moe_fb++; return 0; + } + const float *hh=(const float*)[hh_buf contents]; + for(int gr=0;gr cb = moe_submit(nb,D,Iinter,fmt,g,u,d,gs,us,ds,xg,xoff,nr,R,g_xg,g_gg,g_uu,g_hh); + if (!cb) return 0; + return moe_finish(cb,g_hh,nb,R,D,rows,rw,out); + } +} + +// Async two-phase API: begin submits the block (own scratch, no wait) so the CPU can +// overlap disk loads with GPU compute; end waits + scatters. Handle owns everything. +struct ColiMetalMoeHandle { + id cb; id hh; + std::vector rows; std::vector rwv; + int nb, R, D; +}; +extern "C" ColiMetalMoeHandle* coli_metal_moe_block_begin(int nb, int D, int Iinter, int fmt, + const void *const *g, const void *const *u, const void *const *d, + const float *const *gs, const float *const *us, const float *const *ds, + const float *xg, const int *xoff, const int *nr, + const int *rows, const float *rw) { + @autoreleasepool { + int R = 0; for (int e=0;e bxg=[g_dev newBufferWithLength:(size_t)R*D*4 options:g_res_opts]; + id bgg=[g_dev newBufferWithLength:(size_t)R*Iinter*4 options:g_res_opts]; + id buu=[g_dev newBufferWithLength:(size_t)R*Iinter*4 options:g_res_opts]; + id bhh=[g_dev newBufferWithLength:(size_t)R*D*4 options:g_res_opts]; + id cb = moe_submit(nb,D,Iinter,fmt,g,u,d,gs,us,ds,xg,xoff,nr,R,bxg,bgg,buu,bhh); + if (!cb) return nullptr; + ColiMetalMoeHandle *h = new ColiMetalMoeHandle(); + h->cb=cb; h->hh=bhh; h->rows.assign(rows,rows+R); h->rwv.assign(rw,rw+R); + h->nb=nb; h->R=R; h->D=D; + return h; + } +} +extern "C" int coli_metal_moe_block_end(ColiMetalMoeHandle *h, float *out) { + if (!h) return 0; + int ok; + @autoreleasepool { ok = moe_finish(h->cb,h->hh,h->nb,h->R,h->D,h->rows.data(),h->rwv.data(),out); } + h->cb=nil; h->hh=nil; delete h; + return ok; +} diff --git a/c/glm.c b/c/glm.c index 24412f9..981e6d7 100644 --- a/c/glm.c +++ b/c/glm.c @@ -30,6 +30,7 @@ #if defined(__APPLE__) || defined(__linux__) #include #include /* mlock: inchioda le pagine in RAM / wire pages into RAM */ +#include /* fstat per mmap degli shard (COLI_MMAP) */ #endif #include "st.h" #include "tok.h" @@ -44,6 +45,15 @@ static inline int omp_get_thread_num(void){ return 0; } #ifdef COLI_CUDA #include "backend_cuda.h" #endif +#ifdef COLI_METAL +#include "backend_metal.h" +#include +static int g_metal_enabled; +static int g_metal_gemm_min=16; /* COLI_METAL_GEMM_MIN: min rows to send a matmul_qt GEMM to GPU */ +/* routing precalcolata dalla GPU (layer CB): moe() la usa e salta la FASE A */ +static const int *g_pre_idx; static const float *g_pre_w; static const int *g_pre_keff; +static const float *g_pre_sh; /* output dello shared expert gia' calcolato su GPU */ +#endif #ifdef __AVX2__ #include static inline float hsum256(__m256 v){ /* somma orizzontale di 8 float */ @@ -534,6 +544,15 @@ static void quant_scratch(size_t xn, size_t sn, int8_t **xq, float **sx){ } static void matmul_qt(float *y, const float *x, QT *w, int S){ +#ifdef COLI_METAL + /* Large row-batches (prefill: kv_b reconstruction, o_proj, dense MLP, step_all logits) + * amortize Metal's ~5ms submit latency; small-S decode matmuls stay on CPU (NEON wins). + * Weights must be registered (all dense QT allocs are, via qalloc). */ + if(g_metal_enabled && S>=g_metal_gemm_min && (w->fmt==1||w->fmt==2) && !omp_in_parallel()){ + const void *wp = w->fmt==1 ? (const void*)w->q8 : (const void*)w->q4; + if(coli_metal_gemm(y,x,wp,w->s,w->fmt,S,w->I,w->O)) return; + } +#endif #ifdef COLI_CUDA /* The CUDA backend owns persistent copies only for model-resident tensors. * Streaming expert slots are reused for different IDs and must never enter @@ -658,6 +677,17 @@ static int64_t la_hit[3], la_tot[3]; static int la_pred[2][130][16]; static signed char la_val[2][130]; static int g_pilot=0; /* PILOT=1: prefetch pilotato dal router (vedi pilot_prefetch) */ static int g_pilot_k=8; /* PILOT_K=k: prefetcha solo le prime k predizioni per posizione */ +/* Aligned allocator for dense QT weights/scales: under METAL, page-align + register so the + * GPU reads them zero-copy (no upload duplicate). Plain malloc otherwise. */ +static void *qalloc(size_t n){ +#ifdef COLI_METAL + if(g_metal_enabled){ void *p; size_t r=(n+16383)&~(size_t)16383; + if(posix_memalign(&p,16384,r)){fprintf(stderr,"OOM qalloc\n");exit(1);} + coli_metal_register(p,r); return p; } +#endif + return malloc(n); +} +static float *qsalloc(int O){ return (float*)qalloc((size_t)O*sizeof(float)); } static int g_pilot_real=0;/* PILOT_REAL=1: il pilota fa LOAD VERI cross-layer dentro ecache[L+1] * (non il semplice WILLNEED). Implica PILOT=1. Default OFF: hint-only. */ /* Handshake main<->pilota per il load-vero cross-layer. Invariante di sicurezza in DUE parti: @@ -680,9 +710,9 @@ static _Atomic long g_pilot_drops=0; /* predizioni scartate perche' il main static void qt_alloc(QT *t, int O, int I, int bits){ t->O=O; t->I=I; t->qf=NULL; t->q8=NULL; t->q4=NULL; t->s=NULL; if(bits>=16){ t->fmt=0; t->qf=falloc((int64_t)O*I); } - else if(bits>=5 || g_nopack){ t->fmt=1; t->q8=malloc((int64_t)O*I); t->s=falloc(O); } - else if(bits>=3){ t->fmt=2; t->q4=malloc((int64_t)O*((I+1)/2)); t->s=falloc(O); } - else { t->fmt=3; t->q4=malloc((int64_t)O*((I+3)/4)); t->s=falloc(O); } + else if(bits>=5 || g_nopack){ t->fmt=1; t->q8=qalloc((int64_t)O*I); t->s=qsalloc(O); } + else if(bits>=3){ t->fmt=2; t->q4=qalloc((int64_t)O*((I+1)/2)); t->s=qsalloc(O); } + else { t->fmt=3; t->q4=qalloc((int64_t)O*((I+3)/4)); t->s=qsalloc(O); } } static void qt_fill(QT *t, const float *w, int bits){ if(t->fmt==0) memcpy(t->qf, w, (int64_t)t->O*t->I*sizeof(float)); @@ -789,8 +819,8 @@ static void qt_from_disk(Model *m, const char *name, int O, int I, int bits, int if(st_has(&m->S,sn)){ int64_t nb=st_nbytes(&m->S,name); int fmt = (nb==(int64_t)O*I)?1 : (nb==(int64_t)O*((I+1)/2))?2 : 3; /* int8 / int4 / int2 dai byte */ - if(fmt==1){ if(t->fmt!=1||!t->q8){ t->fmt=1; t->O=O; t->I=I; t->q8=malloc(nb); t->s=falloc(O); } st_read_raw(&m->S,name,t->q8,drop); } - else { if(t->fmt!=fmt||!t->q4){ t->fmt=fmt; t->O=O; t->I=I; t->q4=malloc(nb); t->s=falloc(O); } st_read_raw(&m->S,name,t->q4,drop); } + if(fmt==1){ if(t->fmt!=1||!t->q8){ t->fmt=1; t->O=O; t->I=I; t->q8=qalloc(nb); t->s=qsalloc(O); } st_read_raw(&m->S,name,t->q8,drop); } + else { if(t->fmt!=fmt||!t->q4){ t->fmt=fmt; t->O=O; t->I=I; t->q4=qalloc(nb); t->s=qsalloc(O); } st_read_raw(&m->S,name,t->q4,drop); } st_read_f32(&m->S,sn,t->s,drop); } else { if(!t->qf && !t->q8 && !t->q4) qt_alloc(t,O,I,bits); @@ -811,7 +841,8 @@ static QT qt_load(Model *m, const char *name, int O, int I, int bits){ } static float *ld(Model *m, const char *name){ /* tensore 1D f32 residente (norme/bias) */ int64_t n=st_numel(&m->S,name); if(n<0){fprintf(stderr,"missing %s\n",name);exit(1);} - float *p=falloc(n); st_read_f32(&m->S,name,p,0); return p; + float *p=(float*)qalloc((size_t)n*sizeof(float)); /* registrato per la GPU sotto METAL */ + st_read_f32(&m->S,name,p,0); return p; } static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits){ @@ -969,6 +1000,31 @@ static void embed_row(Model *m, int tok, float *x){ for(int i=0;i>2]; int sh=(i&3)*2; x[i]=(float)((int)((byte>>sh)&3)-2)*s; } } +/* COLI_MMAP=1: gli expert diventano VISTE dentro mmap dei file safetensors (niente pread, + * niente slab, niente copia: la page cache del kernel E' la cache). Le mappe sono + * registrate con Metal (newBufferWithBytesNoCopy su pagine file-backed, come llama.cpp), + * quindi la GPU legge gli stessi byte. Fallback allo slab path su disallineamento. */ +static int g_mmap=0; +static struct { int fd; void *base; size_t len; } g_maps[512]; static int g_nmaps; +static pthread_mutex_t g_map_mtx = PTHREAD_MUTEX_INITIALIZER; /* expert_load e' OMP-parallel */ +static void *map_of_fd(int fd){ + pthread_mutex_lock(&g_map_mtx); + for(int i=0;i UNA pread coalescente da ~19 MB dentro `slab` (+ le scale in fslab); i QT sono * viste dentro lo slab (zero copie). Fallback per modelli non quantizzati (oracolo tiny). @@ -1003,16 +1059,73 @@ static int expert_load(Model *m, int layer, int eid, ESlot *s, int fatal){ snprintf(qn,sizeof(qn),"%s.qs",nm[k]); tq[k]=st_find(&m->S,qn); if(!tw[k]||!tq[k]){ fprintf(stderr,"missing %s\n",nm[k]); if(fatal) exit(1); return -1; } } + if(g_mmap){ + void *bw[3],*bq[3]; int okm=1; + for(int k=0;k<3;k++){ + bw[k]=map_of_fd(tw[k]->fd); bq[k]=map_of_fd(tq[k]->fd); + if(!bw[k]||!bq[k]||((tw[k]->off)&3)||((tq[k]->off)&3)) okm=0; + } + if(okm){ + QT *qt[3]={&s->g,&s->u,&s->d}; int OO[3]={I,I,D}, II[3]={D,D,I}; + for(int k=0;k<3;k++){ + int64_t nb=tw[k]->nbytes; + int fmt=(nb==(int64_t)OO[k]*II[k])?1:(nb==(int64_t)OO[k]*((II[k]+1)/2))?2:3; + qt[k]->fmt=fmt; qt[k]->O=OO[k]; qt[k]->I=II[k]; qt[k]->qf=NULL; + qt[k]->q8=(int8_t*)((char*)bw[k]+tw[k]->off); qt[k]->q4=(uint8_t*)((char*)bw[k]+tw[k]->off); + qt[k]->s=(float*)((char*)bq[k]+tq[k]->off); + } + /* CPU pre-touch: fault the pages in HERE (cheap, parallel, overlapped with the + * resident-experts GPU submit) so the GPU never demand-faults file-backed pages + * (measured catastrophic). madvise starts async readahead, the touch guarantees + * residency. This is pread's I/O without the copy and without the slab. */ + for(int k=0;k<3;k++){ + char *p=(char*)bw[k]+tw[k]->off; size_t n=(size_t)tw[k]->nbytes; + madvise((void*)((uintptr_t)p & ~16383UL), n+16384, MADV_WILLNEED); + volatile char acc=0; + for(size_t i=0;ioff; size_t nq=(size_t)tq[k]->nbytes; + for(size_t i=0;ieid=eid; return 0; + } + } int64_t wtot=tw[0]->nbytes+tw[1]->nbytes+tw[2]->nbytes; int64_t ftot=(tq[0]->nbytes+tq[1]->nbytes+tq[2]->nbytes)/4; /* rialloca se lo slot (riusato tra layer) e' troppo piccolo per QUESTO expert: * pread oltre la mappatura = short-read o CORRUZIONE silenziosa dei vicini */ if(!s->slab || wtot+8192 > s->slab_cap){ +#ifdef COLI_METAL + /* page-align + zero-copy wrap: the GPU reads this slab in place (unified memory) */ + if(s->slab && g_metal_enabled) coli_metal_unregister(s->slab); + compat_aligned_free(s->slab); + size_t need=((size_t)wtot+8192+16383)&~(size_t)16383; + if(posix_memalign((void**)&s->slab,16384,need)){fprintf(stderr,"OOM slab\n"); if(fatal) exit(1); s->slab=NULL; s->slab_cap=0; return -1;} + s->slab_cap=need; + if(g_metal_enabled) coli_metal_register(s->slab,need); +#else compat_aligned_free(s->slab); if(posix_memalign((void**)&s->slab,4096,wtot+8192)){fprintf(stderr,"OOM slab\n"); if(fatal) exit(1); s->slab=NULL; s->slab_cap=0; return -1;} s->slab_cap=wtot+8192; +#endif } if(!s->fslab || ftot > s->fslab_cap){ +#ifdef COLI_METAL + /* page-align + register: the GPU reads the scales in place (unified memory). + * Honours `fatal` exactly like the CPU arm below — a speculative pilot load + * that hits OOM must unwind into a clean hidden slot, never exit(). */ + if(s->fslab && g_metal_enabled) coli_metal_unregister(s->fslab); + free(s->fslab); + size_t fb=(((size_t)ftot*sizeof(float))+16383)&~(size_t)16383; + if(ftot<0 || (uint64_t)ftot > SIZE_MAX/sizeof(float) || + posix_memalign((void**)&s->fslab,16384,fb)){ + fprintf(stderr,"OOM fslab\n"); if(fatal) exit(1); + compat_aligned_free(s->slab); s->slab=NULL; s->slab_cap=0; /* clean, hidden slot (eid stays -1) */ + s->fslab=NULL; s->fslab_cap=0; return -1; + } + s->fslab_cap=ftot; + if(g_metal_enabled) coli_metal_register(s->fslab,fb); +#else free(s->fslab); if(fatal){ s->fslab=falloc(ftot); } /* main path: byte-identical exit-on-OOM */ else { /* speculative pilot: checked alloc, never exit() */ @@ -1020,11 +1133,12 @@ static int expert_load(Model *m, int layer, int eid, ESlot *s, int fatal){ if(ftot<0 || (uint64_t)ftot > SIZE_MAX/sizeof(float) || !(s->fslab=malloc((size_t)ftot*sizeof(float)))){ fprintf(stderr,"OOM fslab\n"); - free(s->slab); s->slab=NULL; s->slab_cap=0; /* leave a clean, hidden slot (eid stays -1) */ + compat_aligned_free(s->slab); s->slab=NULL; s->slab_cap=0; /* leave a clean, hidden slot (eid stays -1) */ s->fslab=NULL; s->fslab_cap=0; return -1; } } s->fslab_cap=ftot; +#endif } int ord[3]={0,1,2}; /* ordina per offset nel file */ for(int a=0;a<3;a++) for(int bb=a+1;bb<3;bb++) if(tw[ord[bb]]->offoff){ int t=ord[a]; ord[a]=ord[bb]; ord[bb]=t; } @@ -1214,6 +1328,33 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba Cfg *c=&m->c; int H=c->n_heads, D=c->hidden, qh=c->qk_head, vh=c->v_head; int kvb_dim=H*(c->qk_nope+vh), Tk=pos_base+S; double ta0=now_s(); +#ifdef COLI_METAL + /* Fused decode attention on GPU: whole layer in one command buffer (keeps the GPU hot). + * S<=4 absorption path with st0==0, DSA selection inactive, and GLM-5.2 int4 dims. */ + if(g_metal_enabled && S<=4 && (g_absorb==1||(g_absorb<0&&S<=4)) && m->kv_start[layer]==0 + && D==6144 && H==64 && c->q_lora==2048 && c->kv_lora==512 && c->qk_nope==192 + && c->qk_rope==64 && vh==256 && l->kv_b.fmt==2){ + int sel_active = m->has_dsa && layern_layers && c->idx_type[layer] && (pos_base+S) > c->index_topk; + if(!sel_active){ + if(m->has_dsa && layern_layers && c->idx_type[layer]){ /* index keys for future selection */ + for(int s=0;sIc[layer]+(int64_t)pos*c->index_hd; + matmul_qt(kd, x+(int64_t)s*D, &m->ix_wk[layer], 1); + layernorm(kd, m->ix_knw[layer], m->ix_knb[layer], c->index_hd, 1e-6f); + rope_interleave(kd, pos, c); } + } + #define WP_(q) ((q).fmt==1?(const void*)(q).q8:(const void*)(q).q4) + int ok = coli_metal_attn_decode(x, + WP_(l->q_a), l->q_a.s, l->q_a.fmt, l->q_a_ln, + WP_(l->q_b), l->q_b.s, l->q_b.fmt, + WP_(l->kv_a), l->kv_a.s, l->kv_a.fmt, l->kv_a_ln, + WP_(l->kv_b), l->kv_b.s, l->kv_b.fmt, + WP_(l->o), l->o.s, l->o.fmt, + m->Lc[layer], m->Rc[layer], S, pos_base, m->kv_start[layer], c->eps, c->theta, c->attn_scale, out); + #undef WP_ + if(ok){ m->t_attn += now_s()-ta0; return; } + } + } +#endif float *ctx=falloc((int64_t)S*H*vh); float *Q=falloc((int64_t)S*H*qh); /* query (roped) dei token nuovi */ float *QR=falloc((int64_t)S*c->q_lora), *comp=falloc(c->kv_lora+c->qk_rope); @@ -1396,6 +1537,21 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){ /* ---- FASE A: routing di tutte le S posizioni ---- */ int *idxs=malloc((size_t)S*K*sizeof(int)); float *ws=malloc((size_t)S*K*sizeof(float)); int *keff=malloc(S*sizeof(int)); +#ifdef COLI_METAL + if(g_pre_idx){ /* routing gia' calcolata dal layer CB (GPU) */ + memcpy(idxs,g_pre_idx,(size_t)S*K*sizeof(int)); + memcpy(ws,g_pre_w,(size_t)S*K*sizeof(float)); + memcpy(keff,g_pre_keff,(size_t)S*sizeof(int)); + for(int s=0;sereq+=keff[s]; + for(int kk=0;kkeusage[layer][idxs[(int64_t)s*K+kk]]++; + if(m->eheat[layer][idxs[(int64_t)s*K+kk]]eheat[layer][idxs[(int64_t)s*K+kk]]++; + } + for(int d=0;drouter, 1, D, E); @@ -1447,6 +1603,7 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){ /* ---- FASE C/D: risolvi (pin/cache/disco) e calcola, a blocchi di 64 unici ---- */ float *xg=falloc((int64_t)S*D), *gg=falloc((int64_t)S*I), *uu=falloc((int64_t)S*I), *hh=falloc((int64_t)S*D); int *rows=malloc(S*sizeof(int)); float *rw=malloc(S*sizeof(float)); + int shared_on_gpu=0; (void)shared_on_gpu; /* set by the Metal path when Phase E was fused */ for(int base=0;basehits++; Sl[z].used=(uint64_t)__atomic_add_fetch(&m->eclock,1,__ATOMIC_RELAXED); use[j]=&Sl[z]; break; } } if(!use[j]){ qof[j]=nmiss; use[j]=&m->ws[nmiss]; missk[nmiss++]=j; m->miss++; } } + int metal_done=0; +#ifdef COLI_METAL + /* GPU/disk OVERLAP: submit the RESIDENT experts (pin/LRU hits, + shared expert on + * the first block) to the GPU BEFORE loading the missed experts from disk, so the + * preads run while the GPU computes; the missed subset follows in a second submit. + * Per-subset CPU fallback on unresolved slab / bad fmt / GPU fault. */ + int is_miss[64]={0}; ColiMetalMoeHandle *mh=NULL; + int cpu_res=1, cpu_miss=1, mh_shared=0, nbb=0, Rtot=0, mfmt=-1, sh_in=0; + const void *MG[65],*MU[65],*MD[65]; const float *MGS[65],*MUS[65],*MDS[65]; + int xoffb[65],nrb[65]; + float *mxg=NULL; int *mrows=NULL; float *mrw=NULL; + /* subset builder: experts with is_miss==WANTMISS (+ shared expert when TRY_SH) */ + #define MB_BUILD(WANTMISS, TRY_SH) do{ \ + nbb=0; Rtot=0; mfmt=-1; sh_in=0; \ + for(int j=0;jg.fmt; \ + MG[nbb]=e->g.fmt==1?(const void*)e->g.q8:(const void*)e->g.q4; \ + MU[nbb]=e->u.fmt==1?(const void*)e->u.q8:(const void*)e->u.q4; \ + MD[nbb]=e->d.fmt==1?(const void*)e->d.q8:(const void*)e->d.q4; \ + MGS[nbb]=e->g.s; MUS[nbb]=e->u.s; MDS[nbb]=e->d.s; \ + xoffb[nbb]=Rtot; nrb[nbb]=cnt; Rtot+=cnt; nbb++; \ + } \ + if(TRY_SH){ int shf = mfmt<0 ? l->sh_gate.fmt : mfmt; \ + if(c->n_shared==1 && sI==I && l->sh_gate.fmt==shf && l->sh_up.fmt==shf && l->sh_down.fmt==shf){ \ + if(mfmt<0) mfmt=shf; \ + MG[nbb]=shf==1?(const void*)l->sh_gate.q8:(const void*)l->sh_gate.q4; \ + MU[nbb]=shf==1?(const void*)l->sh_up.q8 :(const void*)l->sh_up.q4; \ + MD[nbb]=shf==1?(const void*)l->sh_down.q8:(const void*)l->sh_down.q4; \ + MGS[nbb]=l->sh_gate.s; MUS[nbb]=l->sh_up.s; MDS[nbb]=l->sh_down.s; \ + xoffb[nbb]=Rtot; nrb[nbb]=S; Rtot+=S; nbb++; sh_in=1; } } \ + int p=0; \ + for(int j=0;j0){ + double t0=now_s(); + mh=coli_metal_moe_block_begin(nbb,D,I,mfmt,MG,MU,MD,MGS,MUS,MDS,mxg,xoffb,nrb,mrows,mrw); + m->t_emm += now_s()-t0; + if(mh){ cpu_res=0; mh_shared=sh_in; } + } else cpu_res=0; + } +#endif + /* Expert loads run HERE, after the resident-experts GPU submit above: under METAL the + * preads overlap the GPU compute (that submit is async). With METAL off the submit block + * is a no-op / compiled out, so this sits exactly where dev put it and CPU behaviour is + * unchanged. */ if(nmiss){ if(g_pipe){ /* PIPE: launch loads async, matmul overlaps them */ if(!g_pp.started) pipe_init(m); @@ -1481,11 +1699,47 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){ if(!found) expert_prefetch(m,layer,eid); } } +#ifdef COLI_METAL + if(g_metal_enabled){ + /* PIPE drain. Two reasons this barrier is mandatory here, and not optional: + * 1) MB_BUILD(1) hands the missed experts' slabs straight to the GPU — a slot still + * being pread by an I/O worker would be matmul-ed half-loaded. + * 2) PIPE's only drain barrier is the per-expert pipe_wait() in the CPU matmul loop + * below, which metal_done SKIPS ENTIRELY. Without this, a still-writing worker + * would race the end-of-block LRU swap that recycles ws[]. + * pipe_wait() is an idempotent spin on ready[q], so the per-expert waits below stay + * correct (and free) when a subset falls back to the CPU. */ + if(g_pipe && nmiss){ double tw=now_s(); + for(int q=0;qt_edisk += now_s()-tw; } + MB_BUILD(1, 0); /* missed experts, now loaded */ + if(nbb>0){ + double t0=now_s(); + if(coli_metal_moe_block(nbb,D,I,mfmt,MG,MU,MD,MGS,MUS,MDS,mxg,xoffb,nrb,mrows,mrw,out,S)) cpu_miss=0; + m->t_emm += now_s()-t0; + } else cpu_miss=0; + if(mh){ double t0=now_s(); + if(coli_metal_moe_block_end(mh,out)){ if(mh_shared) shared_on_gpu=1; } + else cpu_res=1; + m->t_emm += now_s()-t0; mh=NULL; } + metal_done = (!cpu_res && !cpu_miss); + free(mxg); free(mrows); free(mrw); + } + #undef MB_BUILD +#endif + if(!metal_done) for(int j=0;j=1 routing invariant. */ + * its ws[] slab, so correctness does not depend on the nr>=1 routing invariant. + * Stays ABOVE the METAL skip: a subset that fell back to the CPU still needs its + * slot drained here, and under METAL the block-level drain above already ran (this + * spin is then a no-op). */ if(g_pipe && qof[j]>=0){ double tw=now_s(); pipe_wait(qof[j]); m->t_edisk += now_s()-tw; } +#ifdef COLI_METAL + /* skip the subsets already computed on GPU */ + if(g_metal_enabled && ((is_miss[j] && !cpu_miss) || (!is_miss[j] && !cpu_res))) continue; +#endif int nr=0; /* righe (posizioni) che usano questo expert */ for(int s=0;sws[q]; m->ws[q]=tmp; dst->used=(uint64_t)__atomic_add_fetch(&m->eclock,1,__ATOMIC_RELAXED); } } } - /* ---- FASE E: shared expert, un matmul a S righe ---- */ + /* ---- FASE E: shared expert, un matmul a S righe (skipped se fuso nel blocco GPU) ---- */ float *sg=falloc((int64_t)S*sI), *su=falloc((int64_t)S*sI); - matmul_qt(sg, x, &l->sh_gate, S); - matmul_qt(su, x, &l->sh_up, S); - for(int64_t z=0;z<(int64_t)S*sI;z++) sg[z]=siluf(sg[z])*su[z]; - matmul_qt(hh, sg, &l->sh_down, S); - for(int64_t z=0;z<(int64_t)S*D;z++) out[z]+=hh[z]; +#ifdef COLI_METAL + if(g_pre_sh){ for(int64_t z=0;z<(int64_t)S*D;z++) out[z]+=g_pre_sh[z]; shared_on_gpu=1; } +#endif + if(!shared_on_gpu){ + matmul_qt(sg, x, &l->sh_gate, S); + matmul_qt(su, x, &l->sh_up, S); + for(int64_t z=0;z<(int64_t)S*sI;z++) sg[z]=siluf(sg[z])*su[z]; + matmul_qt(hh, sg, &l->sh_down, S); + for(int64_t z=0;z<(int64_t)S*D;z++) out[z]+=hh[z]; + } free(logit); free(choice); free(idxs); free(ws); free(keff); free(uniq); free(xg); free(gg); free(uu); free(hh); free(rows); free(rw); free(sg); free(su); } @@ -1660,6 +1919,60 @@ static void layer_forward(Model *m, Layer *l, int li, float *x, int S, int pos_b if(g_spec && g_prefetch && l->sparse && m->enr[li]>0) for(int z=0;zenr[li];z++) expert_prefetch(m,li,m->eroute[li][z]); if(g_looka && S==1 && lin_layers && l->sparse) la_predict(m,li,x,0); +#ifdef COLI_METAL + /* FULL-LAYER CB: in_ln + attention + residuo + post_ln + shared expert + router/top-K + * in un solo submit GPU; la CPU legge il routing e fa solo resolve/disk/expert-CB. + * Fallback: qualsiasi condizione mancante -> percorso CPU intero qui sotto. */ + if(g_metal_enabled && S<=4 && lin_layers && l->sparse + && (g_absorb==1||(g_absorb<0&&S<=4)) && m->kv_start[li]==0 + && D==6144 && c->n_heads==64 && c->q_lora==2048 && c->kv_lora==512 + && c->qk_nope==192 && c->qk_rope==64 && c->v_head==256 && l->kv_b.fmt==2 + && c->n_experts==256 && c->topk==8 && c->n_shared==1 && c->moe_inter==2048){ + int sel_active = m->has_dsa && c->idx_type[li] && (pos_base+S) > c->index_topk; + if(!sel_active){ + static float *linrm,*lnrm,*lsh,*lw; static int *lidx,*lkeff; + if(!linrm){ linrm=falloc(4*(int64_t)D); lnrm=falloc(4*(int64_t)D); lsh=falloc(4*(int64_t)D); + lidx=malloc(4*8*sizeof(int)); lw=malloc(4*8*sizeof(float)); lkeff=malloc(4*sizeof(int)); } + int Ksel = g_topk>0 ? (g_topk<8?g_topk:8) : 8; + float tp = (g_topp>0 && g_topp<1.f) ? g_topp : 0.f; + double ta0=now_s(); + #define WP_(q) ((q).fmt==1?(const void*)(q).q8:(const void*)(q).q4) + int ok = coli_metal_layer_decode(x, l->in_ln, l->post_ln, + WP_(l->q_a), l->q_a.s, l->q_a.fmt, l->q_a_ln, + WP_(l->q_b), l->q_b.s, l->q_b.fmt, + WP_(l->kv_a), l->kv_a.s, l->kv_a.fmt, l->kv_a_ln, + WP_(l->kv_b), l->kv_b.s, l->kv_b.fmt, + WP_(l->o), l->o.s, l->o.fmt, + WP_(l->sh_gate), l->sh_gate.s, l->sh_gate.fmt, + WP_(l->sh_up), l->sh_up.s, l->sh_up.fmt, + WP_(l->sh_down), l->sh_down.s, l->sh_down.fmt, + l->router, l->router_bias, + c->n_experts, c->topk, Ksel, tp, c->norm_topk, c->routed_scale, + m->Lc[li], m->Rc[li], S, pos_base, m->kv_start[li], + c->eps, c->theta, c->attn_scale, + linrm, lnrm, lsh, lidx, lw, lkeff); + #undef WP_ + if(ok){ + m->t_attn += now_s()-ta0; + if(m->has_dsa && c->idx_type[li]){ /* index key per selezioni future */ + for(int s=0;sIc[li]+(int64_t)pos*c->index_hd; + matmul_qt(kd, linrm+(int64_t)s*D, &m->ix_wk[li], 1); + layernorm(kd, m->ix_knw[li], m->ix_knb[li], c->index_hd, 1e-6f); + rope_interleave(kd, pos, c); + } + } + if(g_pilot && S<=8 && li+1n_layers && m->L[li+1].sparse) pilot_prefetch(m,li+1,x,S); + if(g_looka && S==1 && li+1n_layers && m->L[li+1].sparse) la_predict(m,li+1,x,1); + g_pre_idx=lidx; g_pre_w=lw; g_pre_keff=lkeff; g_pre_sh=lsh; + moe(m,l,li,lnrm,S,tmp); + g_pre_idx=NULL; g_pre_w=NULL; g_pre_keff=NULL; g_pre_sh=NULL; + for(int64_t j=0;j<(int64_t)S*D;j++) x[j]+=tmp[j]; + return; + } + } + } +#endif for(int s=0;sin_ln, D, c->eps); attention(m,l,li,nrm,S,pos_base,tmp); for(int64_t j=0;j<(int64_t)S*D;j++) x[j]+=tmp[j]; @@ -1690,7 +2003,11 @@ static void layers_forward(Model *m, float *x, int S, int pos_base){ static void kv_alloc(Model *m, int max_t){ Cfg *c=&m->c; KVState *k=m->kv; - if(k->Lc){ for(int i=0;in_layers+1;i++){ free(k->Lc[i]); free(k->Rc[i]); } free(k->Lc); free(k->Rc); } + if(k->Lc){ for(int i=0;in_layers+1;i++){ +#ifdef COLI_METAL + if(g_metal_enabled){ coli_metal_unregister(k->Lc[i]); coli_metal_unregister(k->Rc[i]); } +#endif + free(k->Lc[i]); free(k->Rc[i]); } free(k->Lc); free(k->Rc); } if(k->Ic){ for(int i=0;in_layers;i++) free(k->Ic[i]); free(k->Ic); k->Ic=NULL; } if(m->has_dsa){ k->Ic=calloc(c->n_layers,sizeof(float*)); @@ -1700,7 +2017,20 @@ static void kv_alloc(Model *m, int max_t){ int NR=c->n_layers+1; /* riga extra: KV del layer MTP */ k->Lc=calloc(NR,sizeof(float*)); k->Rc=calloc(NR,sizeof(float*)); for(int i=0;iLc[i]=falloc((int64_t)max_t*c->kv_lora); - k->Rc[i]=falloc((int64_t)max_t*c->qk_rope); } + k->Rc[i]=falloc((int64_t)max_t*c->qk_rope); +#ifdef COLI_METAL + /* page-align + register Lc/Rc for zero-copy GPU attention. falloc isn't 16K-aligned, + * so re-allocate aligned and register the exact byte length. */ + if(g_metal_enabled){ + size_t lb=(((size_t)max_t*c->kv_lora*sizeof(float))+16383)&~(size_t)16383; + size_t rb=(((size_t)max_t*c->qk_rope*sizeof(float))+16383)&~(size_t)16383; + free(k->Lc[i]); free(k->Rc[i]); void *lp,*rp; + if(posix_memalign(&lp,16384,lb)||posix_memalign(&rp,16384,rb)){fprintf(stderr,"OOM kv\n");exit(1);} + k->Lc[i]=lp; k->Rc[i]=rp; + coli_metal_register(k->Lc[i],lb); coli_metal_register(k->Rc[i],rb); + } +#endif + } m->Lc=k->Lc; m->Rc=k->Rc; m->Ic=k->Ic; m->max_t=k->max_t; m->kv_start=k->kv_start; } @@ -2094,6 +2424,15 @@ static void profile_print(Model *m, double elapsed){ printf("PROFILE: expert-disk %.3fs | expert-matmul %.3fs | attention %.3fs " "(including kvb %.3fs) | lm_head %.3fs | other %.3fs\n", m->t_edisk,m->t_emm,m->t_attn,m->t_kvb,m->t_head,elapsed-accounted); +#ifdef COLI_METAL + if(g_metal_enabled){ uint64_t ok=0,fb=0,ex=0; double su=0,gp=0,sc=0; + coli_metal_moe_counts(&ok,&fb,&ex); coli_metal_moe_times(&su,&gp,&sc); + { uint64_t aok=0; double aw=0,ak=0; coli_metal_attn_counts(&aok,&aw,&ak); + if(aok){ double ks=0,gs=0; coli_metal_attn_lat(&ks,&gs); + printf("METAL-ATTN: layer GPU %llu | gpu-wall %.2fs (kernel %.2fs | cpu-sched %.2fs gpu-sched %.2fs)\n",(unsigned long long)aok,aw,ak,ks,gs); } } + printf("METAL: blocchi GPU %llu | fallback CPU %llu | expert su GPU %llu | setup %.2fs gpu-wall %.2fs (kernel %.2fs) scatter %.2fs\n", + (unsigned long long)ok,(unsigned long long)fb,(unsigned long long)ex,su,gp,coli_metal_moe_kernel_time(),sc); } +#endif } /* Fixed-token decode benchmark: prefill all but the prompt's last token, then @@ -2815,6 +3154,8 @@ int main(int argc, char **argv){ g_nopack = getenv("NOPACK")?1:0; g_drop = getenv("DROP")?1:0; g_prefetch = getenv("PREFETCH")?atoi(getenv("PREFETCH")):0; + g_mmap = getenv("COLI_MMAP")?atoi(getenv("COLI_MMAP")):0; + if(g_mmap) fprintf(stderr,"[MMAP] expert = viste zero-copy nei file (page cache = cache)\n"); g_topk = getenv("TOPK")?atoi(getenv("TOPK")):0; g_topp = getenv("TOPP")?atof(getenv("TOPP")):0; g_mlock = getenv("MLOCK")?atoi(getenv("MLOCK")):-1; /* -1 auto (ON macOS), 0 off, 1 force / auto (ON macOS), 0 off, 1 force */ @@ -2874,6 +3215,20 @@ int main(int argc, char **argv){ fprintf(stderr,"CUDA was requested, but this binary is CPU-only; rebuild with: make CUDA=1\n"); return 2; } +#endif +#ifdef COLI_METAL + if(getenv("COLI_METAL") && atoi(getenv("COLI_METAL"))){ + g_metal_enabled = coli_metal_init(); + if(!g_metal_enabled){ fprintf(stderr,"[METAL] backend requested but not available\n"); return 2; } + fprintf(stderr,"[METAL] mode: batched routed experts on GPU (unified-memory zero-copy)\n"); + if(getenv("COLI_METAL_SPIN") && atoi(getenv("COLI_METAL_SPIN"))){ coli_metal_spin_start(); fprintf(stderr,"[METAL] keep-alive spinner ON\n"); } + if(getenv("COLI_METAL_GEMM_MIN")) g_metal_gemm_min=atoi(getenv("COLI_METAL_GEMM_MIN")); + } +#else + if(getenv("COLI_METAL") && atoi(getenv("COLI_METAL"))){ + fprintf(stderr,"METAL was requested, but this binary has no Metal backend; rebuild with: make METAL=1\n"); + return 2; + } #endif printf("== GLM C engine (glm_moe_dsa), cache=%d experts/layer | experts@%d-bit dense@%d-bit | idot: " IDOT_KERNEL " ==\n", cap, ebits, dbits); g_mem_avail_boot = mem_available_gb(); diff --git a/c/tests/test_backend_metal.mm b/c/tests/test_backend_metal.mm new file mode 100644 index 0000000..ffdd39b --- /dev/null +++ b/c/tests/test_backend_metal.mm @@ -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 +#include +#include +#include +#include + +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>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 W(wn); std::vector 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 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& 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 xoff(nb),nr(nrv); for(int e=0;e slab(nb), fslab(nb); + std::vector g(nb),u(nb),d(nb); std::vector 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 xg((size_t)R*D); for(auto&v:xg) v=((rand()%2000)-1000)/1000.f; + std::vector rows(R); std::vector rw(R); + for(int gr=0;gr position 0 + int S=1; + // CPU reference + std::vector refout((size_t)S*D,0.f), gg(I),uu(I),hh(D); + for(int e=0;e 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