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
@@ -30,6 +30,7 @@
|
||||
#if defined(__APPLE__) || defined(__linux__)
|
||||
#include <sys/resource.h>
|
||||
#include <sys/mman.h> /* mlock: inchioda le pagine in RAM / wire pages into RAM */
|
||||
#include <sys/stat.h> /* 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 <omp.h>
|
||||
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 <immintrin.h>
|
||||
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<D;i++){ uint8_t byte=q[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<g_nmaps;i++) if(g_maps[i].fd==fd){ void *b=g_maps[i].base; pthread_mutex_unlock(&g_map_mtx); return b; }
|
||||
void *base=NULL; struct stat st;
|
||||
if(g_nmaps<512 && fstat(fd,&st)==0){
|
||||
size_t len=((size_t)st.st_size+16383)&~(size_t)16383;
|
||||
void *p=mmap(NULL,len,PROT_READ,MAP_SHARED,fd,0);
|
||||
if(p!=MAP_FAILED){
|
||||
base=p; g_maps[g_nmaps].fd=fd; g_maps[g_nmaps].base=p; g_maps[g_nmaps].len=len; g_nmaps++;
|
||||
#ifdef COLI_METAL
|
||||
if(g_metal_enabled) coli_metal_register(p,len);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
pthread_mutex_unlock(&g_map_mtx);
|
||||
return base;
|
||||
}
|
||||
|
||||
/* carica un expert nello slot. Container pre-quantizzato: le 3 matrici sono contigue nel
|
||||
* file -> 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;i<n;i+=4096) acc+=p[i];
|
||||
acc+=p[n-1]; (void)acc;
|
||||
char *q=(char*)bq[k]+tq[k]->off; size_t nq=(size_t)tq[k]->nbytes;
|
||||
for(size_t i=0;i<nq;i+=4096) acc+=q[i];
|
||||
}
|
||||
s->eid=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]]->off<tw[ord[a]]->off){ 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 && layer<c->n_layers && c->idx_type[layer] && (pos_base+S) > c->index_topk;
|
||||
if(!sel_active){
|
||||
if(m->has_dsa && layer<c->n_layers && c->idx_type[layer]){ /* index keys for future selection */
|
||||
for(int s=0;s<S;s++){ int pos=pos_base+s; float *kd=m->Ic[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;s<S;s++){
|
||||
m->ereq+=keff[s];
|
||||
for(int kk=0;kk<keff[s];kk++){
|
||||
m->eusage[layer][idxs[(int64_t)s*K+kk]]++;
|
||||
if(m->eheat[layer][idxs[(int64_t)s*K+kk]]<UINT32_MAX) m->eheat[layer][idxs[(int64_t)s*K+kk]]++;
|
||||
}
|
||||
for(int d=0;d<D;d++) out[(int64_t)s*D+d]=0;
|
||||
}
|
||||
} else
|
||||
#endif
|
||||
for(int s=0;s<S;s++){
|
||||
const float *xs=x+(int64_t)s*D;
|
||||
matmul(logit, xs, l->router, 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;base<nu;base+=64){
|
||||
int nb = nu-base<64 ? nu-base : 64;
|
||||
ESlot *use[64]; int missk[64]; int qof[64]; int nmiss=0;
|
||||
@@ -1457,6 +1614,67 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
|
||||
for(int z=0;z<nn;z++) if(Sl[z].eid==eid){ m->hits++; 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;j<nb;j++){ if(is_miss[j]!=(WANTMISS)) continue; \
|
||||
int eid=uniq[base+j]; ESlot *e=use[j]; int cnt=0; \
|
||||
for(int s=0;s<S;s++) for(int kk=0;kk<keff[s];kk++) \
|
||||
if(idxs[(int64_t)s*K+kk]==eid){ cnt++; break; } \
|
||||
if(!cnt) continue; \
|
||||
if(mfmt<0) mfmt=e->g.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;j<nb;j++){ if(is_miss[j]!=(WANTMISS)) continue; int eid=uniq[base+j]; \
|
||||
for(int s=0;s<S;s++) for(int kk=0;kk<keff[s];kk++) \
|
||||
if(idxs[(int64_t)s*K+kk]==eid){ \
|
||||
memcpy(mxg+(int64_t)p*D, x+(int64_t)s*D, D*sizeof(float)); \
|
||||
mrows[p]=s; mrw[p]=ws[(int64_t)s*K+kk]; p++; break; } } \
|
||||
if(sh_in) for(int s=0;s<S;s++){ \
|
||||
memcpy(mxg+(int64_t)p*D, x+(int64_t)s*D, D*sizeof(float)); \
|
||||
mrows[p]=s; mrw[p]=1.0f; p++; } \
|
||||
}while(0)
|
||||
if(g_metal_enabled){
|
||||
for(int q=0;q<nmiss;q++) is_miss[missk[q]]=1;
|
||||
mxg=falloc((int64_t)(nb+1)*S*D);
|
||||
mrows=malloc((size_t)(nb+1)*S*sizeof(int)); mrw=malloc((size_t)(nb+1)*S*sizeof(float));
|
||||
MB_BUILD(0, base==0 && !g_pre_sh);
|
||||
if(nbb>0){
|
||||
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;q<nmiss;q++) pipe_wait(q);
|
||||
m->t_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<nb;j++){ int eid=uniq[base+j]; ESlot *e=use[j];
|
||||
/* Drain this miss's async load BEFORE the nr==0 early-exit below: every
|
||||
* dispatched slot must be waited before the end-of-block LRU swap can reuse
|
||||
* its ws[] slab, so correctness does not depend on the nr>=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;s<S;s++) for(int kk=0;kk<keff[s];kk++)
|
||||
if(idxs[(int64_t)s*K+kk]==eid){ rows[nr]=s; rw[nr]=ws[(int64_t)s*K+kk]; nr++; break; }
|
||||
@@ -1515,13 +1769,18 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
|
||||
ESlot tmp=*dst; *dst=m->ws[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;z<m->enr[li];z++) expert_prefetch(m,li,m->eroute[li][z]);
|
||||
if(g_looka && S==1 && li<c->n_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 && li<c->n_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;s<S;s++){ int pos=pos_base+s;
|
||||
float *kd=m->Ic[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+1<c->n_layers && m->L[li+1].sparse) pilot_prefetch(m,li+1,x,S);
|
||||
if(g_looka && S==1 && li+1<c->n_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;s<S;s++) rmsnorm(nrm+(int64_t)s*D, x+(int64_t)s*D, l->in_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;i<c->n_layers+1;i++){ free(k->Lc[i]); free(k->Rc[i]); } free(k->Lc); free(k->Rc); }
|
||||
if(k->Lc){ for(int i=0;i<c->n_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;i<c->n_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;i<NR;i++){ k->Lc[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();
|
||||
|
||||
Reference in New Issue
Block a user