danielkrause/colibri-strix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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)

Browse Source 
* 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
2026-07-12 23:57:10 -07:00
committed by GitHub
parent 342ceaf368
commit 3716e4006a
7 changed files with 1575 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+5
View File
@@ -44,3 +44,8 @@ stats*.txt
/quant_test.py /quant_test.py
/profile_run.py /profile_run.py
/sweep.py /sweep.py
c/backend_metal.o
c/backend_metal_test
c/tests/test_tier
c/mio_env/
c/bench/
README.md
+25
View File
@@ -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 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. 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 ### Experimental resident CUDA backend
colibrì includes an opt-in CUDA backend for model-resident tensors. Streaming colibrì includes an opt-in CUDA backend for model-resident tensors. Streaming
c/Makefile
+24 -3
View File
@@ -80,18 +80,39 @@ LDFLAGS += -L$(CUDA_HOME)/lib64 -Wl,-rpath,$(CUDA_HOME)/lib64 -lcudart -lstdc++
CUDA_OBJ = backend_cuda.o CUDA_OBJ = backend_cuda.o
endif 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) all: glm$(EXE)
# phony 'glm' → 'glm.exe' on Windows (so 'make glm' and 'coli build' work on every platform) # phony 'glm' → 'glm.exe' on Windows (so 'make glm' and 'coli build' work on every platform)
glm: glm$(EXE) glm: glm$(EXE)
glm$(EXE): glm.c st.h json.h tok.h tok_unicode.h compat.h grammar.h $(CUDA_OBJ) 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) -o glm$(EXE) $(LDFLAGS) $(CC) $(CFLAGS) glm.c $(CUDA_OBJ) $(METAL_OBJ) -o glm$(EXE) $(LDFLAGS)
backend_cuda.o: backend_cuda.cu backend_cuda.h 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; } @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 $@ $(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 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; } @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) $(NVCC) $(NVCCFLAGS) backend_cuda.cu tests/test_backend_cuda.cu -o backend_cuda_test$(EXE)
@@ -137,7 +158,7 @@ check:
$(MAKE) test $(MAKE) test
clean: 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__ rm -rf tests/__pycache__
.PHONY: all glm cuda-test portable test-c test-python test check clean .PHONY: all glm cuda-test portable test-c test-python test check clean
c/backend_metal.h
+143
View File
@@ -0,0 +1,143 @@
#ifndef COLIBRI_BACKEND_METAL_H
#define COLIBRI_BACKEND_METAL_H
#include <stddef.h>
#include <stdint.h>
#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
c/backend_metal.mm
+786
View File
@@ -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 <Metal/Metal.h>
#import <Foundation/Foundation.h>
#include "backend_metal.h"
#include <cstring>
#include <vector>
#include <mutex>
// ---- 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 <metal_stdlib>
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<K8;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=K8*8+slane;i<K;i+=32){ uchar b=w[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;c<K8;c+=32) acc+=dot(float4(w4[2*c]),x4[2*c])+dot(float4(w4[2*c+1]),x4[2*c+1]);
for(int i=K8*8+slane;i<K;i+=32) acc+=float(w[i])*xr[i];
}
acc=simd_sum(acc);
if(slane==0) yout[row]=acc*sc[o];
}
kernel void moe_silu(device float* g [[buffer(0)]], device const float* u [[buffer(1)]],
uint i [[thread_position_in_grid]]) { float v=g[i]; g[i]=(v/(1.0f+exp(-v)))*u[i]; }
// ===== Fused decode attention (GLM-5.2 dims, S=1) =====
constant int A_HID=6144, A_H=64, A_QLORA=2048, A_KVL=512, A_NOPE=192, A_ROPE=64, A_VH=256;
constant int A_QH=256 /*nope+rope*/, A_ROWSH=448 /*nope+vh*/;
// per-row in-place RMSNorm: row = threadgroup index, x[row*n + i]. grid = nrows threadgroups.
kernel void a_rmsnorm(device float* x [[buffer(0)]], device const float* w [[buffer(1)]],
constant int& n [[buffer(2)]], constant float& eps [[buffer(3)]],
uint row [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint tgsz [[threads_per_threadgroup]]) {
device float* xr=x+(long)row*n; threadgroup float red[256];
float s=0; for(int i=lid;i<n;i+=tgsz) s+=xr[i]*xr[i];
red[lid]=s; threadgroup_barrier(mem_flags::mem_threadgroup);
for(uint k=tgsz/2;k>0;k>>=1){ if(lid<k) red[lid]+=red[lid+k]; threadgroup_barrier(mem_flags::mem_threadgroup); }
float r=rsqrt(red[0]/n+eps); threadgroup_barrier(mem_flags::mem_threadgroup);
for(int i=lid;i<n;i+=tgsz) xr[i]=xr[i]*r*w[i];
}
// interleaved partial RoPE. vv = v + base + s*rowstride + h*headstride, pos = PB+s. grid = S*nheads*(ROPE/2).
kernel void a_rope(device float* v [[buffer(0)]], constant int& base [[buffer(1)]],
constant int& rowstride [[buffer(2)]], constant int& headstride [[buffer(3)]],
constant int& nheads [[buffer(4)]], constant int& PB [[buffer(5)]],
constant float& theta [[buffer(6)]], uint gid [[thread_position_in_grid]]) {
int hlf=A_ROPE/2; int idx=gid/hlf, j=gid%hlf; int s=idx/nheads, h=idx%nheads; int pos=PB+s;
device float* vv=v+(long)base+(long)s*rowstride+(long)h*headstride;
float inv=pow(theta, -2.0f*j/A_ROPE); float ang=pos*inv, cs=cos(ang), sn=sin(ang);
float a=vv[2*j], b=vv[2*j+1]; vv[j]=a*cs-b*sn; vv[hlf+j]=b*cs+a*sn;
}
// per-row copy: dst[s*dststride + i] = src[s*srcstride + off + i]. grid = S*n.
kernel void a_copy(device const float* src [[buffer(0)]], constant int& off [[buffer(1)]], constant int& srcstride [[buffer(2)]],
device float* dst [[buffer(3)]], constant int& dststride [[buffer(4)]], constant int& n [[buffer(5)]],
uint gid [[thread_position_in_grid]]) { int s=gid/n, i=gid%n; dst[(long)s*dststride+i]=src[(long)s*srcstride+off+i]; }
// ---- absorption core (S query rows, per-row causal). q:[S,H*QH]; qabs/clat:[S*H,KVL];
// sc:[S*H,T]; ctx:[S*H,VH]. Query row s (abs pos PB+s) attends keys [0, PB+s]. ----
constant int A_QHH=A_H*A_QH;
inline float a_deqrow(device const uchar* base, int row, int i, device const float* sc){
device const uchar* w=base+(long)row*((A_KVL+1)/2); uchar b=w[i>>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<A_NOPE;d++) a+=qp[d]*a_deqrow(kvb,rbase+d,i,sc); qabs[(long)(s*A_H+h)*A_KVL+i]=a;
}
kernel void a_score(device const float* qabs [[buffer(0)]], device const float* Lc [[buffer(1)]],
device const float* Rc [[buffer(2)]], device const float* q [[buffer(3)]],
device float* sc [[buffer(4)]], constant int& T [[buffer(5)]], constant float& ascale [[buffer(6)]],
constant int& PB [[buffer(7)]], uint gid [[thread_position_in_grid]]) {
int s=gid/(A_H*T), r=gid%(A_H*T), h=r/T, t=r%T; long o=(long)(s*A_H+h)*T+t;
if(t > 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;i<A_KVL;i++) a+=qa[i]*Lt[i]; for(int d=0;d<A_ROPE;d++) a+=qr[d]*Rt[d]; sc[o]=a*ascale;
}
kernel void a_smax(device float* sc [[buffer(0)]], constant int& T [[buffer(1)]],
uint sh [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint tgsz [[threads_per_threadgroup]]) {
device float* s=sc+(long)sh*T; threadgroup float red[256];
float m=-1e30f; for(int t=lid;t<T;t+=tgsz) m=max(m,s[t]); red[lid]=m; threadgroup_barrier(mem_flags::mem_threadgroup);
for(uint k=tgsz/2;k>0;k>>=1){ if(lid<k) red[lid]=max(red[lid],red[lid+k]); threadgroup_barrier(mem_flags::mem_threadgroup);}
float mx=red[0]; threadgroup_barrier(mem_flags::mem_threadgroup);
float sum=0; for(int t=lid;t<T;t+=tgsz){ float e=exp(s[t]-mx); s[t]=e; sum+=e; } red[lid]=sum; threadgroup_barrier(mem_flags::mem_threadgroup);
for(uint k=tgsz/2;k>0;k>>=1){ if(lid<k) red[lid]+=red[lid+k]; threadgroup_barrier(mem_flags::mem_threadgroup);}
float tot=red[0]; threadgroup_barrier(mem_flags::mem_threadgroup); for(int t=lid;t<T;t+=tgsz) s[t]/=tot;
}
kernel void a_clat(device const float* sc [[buffer(0)]], device const float* Lc [[buffer(1)]],
device float* clat [[buffer(2)]], constant int& T [[buffer(3)]], uint gid [[thread_position_in_grid]]) {
int sh=gid/A_KVL, i=gid%A_KVL; device const float* s=sc+(long)sh*T; float a=0;
for(int t=0;t<T;t++) a+=s[t]*Lc[(long)t*A_KVL+i]; clat[(long)sh*A_KVL+i]=a;
}
kernel void a_ctx(device const uchar* kvb [[buffer(0)]], device const float* sc [[buffer(1)]],
device const float* clat [[buffer(2)]], device float* ctx [[buffer(3)]], uint gid [[thread_position_in_grid]]) {
int sh=gid/A_VH, j=gid%A_VH, h=sh%A_H; int row=h*A_ROWSH+A_NOPE+j; device const float* cl=clat+(long)sh*A_KVL;
float a=0; for(int i=0;i<A_KVL;i++) a+=cl[i]*a_deqrow(kvb,row,i,sc); ctx[(long)sh*A_VH+j]=a;
}
// ===== full-layer tail kernels =====
// y[i] += a[i] (residual add), grid = n
kernel void a_add(device float* y [[buffer(0)]], device const float* a [[buffer(1)]],
uint i [[thread_position_in_grid]]) { y[i] += a[i]; }
// router: logit[s][e] = x[s].w_e (f32 rows [E,D]) -> 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;c<D4;c+=32) acc+=dot(w4[c],x4[c]);
acc=simd_sum(acc);
if(slane==0) sig[row]=1.0f/(1.0f+exp(-acc));
}
// exact replica of glm.c phase-A selection per row s (serial, deterministic ties):
// choice=sig+bias; greedy top-Ksel by choice; w=sig[best]; optional topp truncation
// (insertion-sort desc + cumulative); optional norm_topk; * routed_scale.
kernel void r_top8(device const float* sig [[buffer(0)]], device const float* bias [[buffer(1)]],
device int* idx [[buffer(2)]], device float* w [[buffer(3)]],
device int* keff [[buffer(4)]], constant int& E [[buffer(5)]],
constant int& K [[buffer(6)]], constant int& Ksel [[buffer(7)]],
constant float& topp [[buffer(8)]], constant int& normk [[buffer(9)]],
constant float& rscale [[buffer(10)]],
uint s [[thread_position_in_grid]]) {
device const float* sg=sig+(long)s*E;
device int* id_=idx+(long)s*K; device float* ww=w+(long)s*K;
for(int kk=0;kk<Ksel;kk++){ int best=-1; float bv=-1e30f;
for(int e=0;e<E;e++){ bool tk=false; for(int j=0;j<kk;j++) if(id_[j]==e){tk=true;break;}
float ch=sg[e]+bias[e];
if(!tk && ch>bv){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<Ksel;a++){ int ii=id_[a]; float wv=ww[a]; int b=a-1;
while(b>=0 && ww[b]<wv){ ww[b+1]=ww[b]; id_[b+1]=id_[b]; b--; } ww[b+1]=wv; id_[b+1]=ii; }
float tot=1e-20f; for(int kk=0;kk<Ksel;kk++) tot+=ww[kk];
float cum=0; for(int kk=0;kk<Ksel;kk++){ cum+=ww[kk]; if(cum>=topp*tot){ Ke=kk+1; break; } }
}
keff[s]=Ke;
if(normk){ float sm=0; for(int kk=0;kk<Ke;kk++) sm+=ww[kk]; sm+=1e-20f; for(int kk=0;kk<Ke;kk++) ww[kk]/=sm; }
for(int kk=0;kk<Ke;kk++) ww[kk]*=rscale;
}
)METAL";
struct ColiMetalTensor {
id<MTLBuffer> w; // weights (wrapped, zero-copy when page-aligned)
id<MTLBuffer> s; // scales
int fmt, I, O; size_t wbytes;
};
static id<MTLDevice> g_dev;
static id<MTLCommandQueue> g_queue;
static id<MTLComputePipelineState> g_gemv, g_moe_gemv, g_moe_silu;
static id<MTLComputePipelineState> 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<MTLComputePipelineState> 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 <mach/mach_time.h>
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<MTLBuffer> buf; };
static std::vector<Slab> 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<MTLBuffer> 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<MTLBuffer> ensure(id<MTLBuffer> 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<MTLBuffer> 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<MTLLibrary> 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<MTLBuffer> b = [g_dev newBufferWithBytesNoCopy:base length:len
options:g_res_opts deallocator:nil];
if (!b) return;
std::lock_guard<std::mutex> 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<std::mutex> lk(g_slab_mtx);
for (size_t i=0;i<g_slabs.size();i++) if (g_slabs[i].base==base) { g_slabs[i].buf=nil; g_slabs.erase(g_slabs.begin()+i); return; }
}
// Resolve a host pointer inside a registered slab to (buffer, gpuAddress). Returns nil if unknown.
static id<MTLBuffer> resolve(const void *p, uint64_t *addr) {
std::lock_guard<std::mutex> 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<b+s.len) { *addr = (uint64_t)[s.buf gpuAddress] + (u-b); return s.buf; } }
return nil;
}
// Keep-alive spinner (COLI_METAL_SPIN=1): keeps trivial GPU work in flight so the GPU
// doesn't ramp its clock down between the engine's short per-layer bursts. Experiment to
// quantify how much of the observed submit latency is clock ramp-down.
#include <thread>
#include <atomic>
static std::atomic<bool> 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<MTLCommandQueue> q = [g_dev newCommandQueue]; // own queue: never blocks real work
id<MTLBuffer> b = [g_dev newBufferWithLength:4096 options:MTLResourceStorageModeShared];
while (g_spin_run.load()) {
@autoreleasepool {
id<MTLCommandBuffer> cb=[q commandBuffer];
id<MTLComputeCommandEncoder> 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<MTLBuffer> bx = [g_dev newBufferWithBytes:x length:(size_t)S*I*sizeof(float) options:MTLResourceStorageModeShared];
id<MTLBuffer> by = [g_dev newBufferWithLength:(size_t)S*O*sizeof(float) options:MTLResourceStorageModeShared];
id<MTLCommandBuffer> cb = [g_queue commandBuffer];
id<MTLComputeCommandEncoder> 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<MTLBuffer> ax_,aqr_,aqf_,acomp_,aqabs_,ascore_,aclat_,actx_,aout_,aqaln_,akvaln_; static size_t ascore_cap;
static id<MTLBuffer> 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<MTLComputeCommandEncoder> e, const void* w, const float* s, int fmt, int I, int O,
id<MTLBuffer> xin, id<MTLBuffer> yout, int S){
uint64_t wa=0,sa=0; id<MTLBuffer> wb=resolve(w,&wa); id<MTLBuffer> 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<MTLComputeCommandEncoder> e, const AttnW *W,
id<MTLBuffer> Lb, size_t loff, id<MTLBuffer> Rb, size_t roff,
id<MTLBuffer> kvbW, size_t kvbwoff, id<MTLBuffer> 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<MTLBuffer> b,size_t off,id<MTLBuffer> 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<MTLBuffer> 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<MTLBuffer> 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<MTLBuffer> *Lb, size_t *loff, id<MTLBuffer> *Rb, size_t *roff,
id<MTLBuffer> *kvbW, size_t *kvbwoff, id<MTLBuffer> *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<MTLBuffer> 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<MTLCommandBuffer> cb=[g_queue commandBuffer]; id<MTLComputeCommandEncoder> 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<MTLBuffer> 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<MTLBuffer> inB=resolve(in_ln,&ina), pnB=resolve(post_ln,&pna);
id<MTLBuffer> 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<MTLCommandBuffer> cb=[g_queue commandBuffer]; id<MTLComputeCommandEncoder> 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<MTLBuffer> b,id<MTLBuffer> 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<MTLBuffer> src,id<MTLBuffer> 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<MTLBuffer> 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<MTLBuffer> 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<MTLCommandBuffer> cb=[g_queue commandBuffer]; id<MTLComputeCommandEncoder> 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<MTLCommandBuffer> 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<MTLBuffer> xg_buf, id<MTLBuffer> gg_buf, id<MTLBuffer> uu_buf, id<MTLBuffer> hh_buf) {
if (!g_dev || (fmt != 1 && fmt != 2)) return nil;
double ts_start = mnow();
std::vector<uint64_t> ag(nb),au(nb),ad(nb),sgv(nb),suv(nb),sdv(nb);
std::vector<id<MTLBuffer>> use; use.reserve(nb*2);
auto add_use=[&](id<MTLBuffer> b){ for(auto&x:use) if(x==b) return; use.push_back(b); };
for (int e=0;e<nb;e++) {
id<MTLBuffer> 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<int> erow(R); for(int e=0;e<nb;e++) for(int r=0;r<nr[e];r++) erow[xoff[e]+r]=e;
auto shb=[&](const void*p,size_t n){ return [g_dev newBufferWithBytes:p length:n options:MTLResourceStorageModeShared]; };
id<MTLBuffer> bag=shb(ag.data(),nb*8), bau=shb(au.data(),nb*8), bad=shb(ad.data(),nb*8);
id<MTLBuffer> bsg=shb(sgv.data(),nb*8), bsu=shb(suv.data(),nb*8), bsd=shb(sdv.data(),nb*8);
id<MTLBuffer> berow=shb(erow.data(),R*4);
memcpy([xg_buf contents], xg, (size_t)R*D*4);
id<MTLCommandBuffer> cb=[g_queue commandBuffer]; id<MTLComputeCommandEncoder> e=[cb computeCommandEncoder];
for(auto&b:use) [e useResource:b usage:MTLResourceUsageRead];
auto gemv=[&](id<MTLBuffer> wa,id<MTLBuffer> sa,id<MTLBuffer> xin,id<MTLBuffer> 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<MTLCommandBuffer> cb, id<MTLBuffer> 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<R;gr++){ float *os=out+(size_t)rows[gr]*D, w=rw[gr]; const float *hr=hh+(size_t)gr*D;
for(int dd=0;dd<D;dd++) os[dd]+=w*hr[dd]; }
g_t_scatter += mnow() - ts_gpu;
g_moe_ok++; g_moe_experts += nb;
return 1;
}
extern "C" 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) {
(void)S;
@autoreleasepool {
int R = 0; for (int e=0;e<nb;e++) R += nr[e];
if (R == 0) return 1;
g_xg = ensure(g_xg,&g_xg_cap,(size_t)R*D*4);
g_gg = ensure(g_gg,&g_gg_cap,(size_t)R*Iinter*4);
g_uu = ensure(g_uu,&g_uu_cap,(size_t)R*Iinter*4);
g_hh = ensure(g_hh,&g_hh_cap,(size_t)R*D*4);
id<MTLCommandBuffer> 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<MTLCommandBuffer> cb; id<MTLBuffer> hh;
std::vector<int> rows; std::vector<float> 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<nb;e++) R += nr[e];
if (R == 0 || !g_dev) return nullptr;
id<MTLBuffer> bxg=[g_dev newBufferWithLength:(size_t)R*D*4 options:g_res_opts];
id<MTLBuffer> bgg=[g_dev newBufferWithLength:(size_t)R*Iinter*4 options:g_res_opts];
id<MTLBuffer> buu=[g_dev newBufferWithLength:(size_t)R*Iinter*4 options:g_res_opts];
id<MTLBuffer> bhh=[g_dev newBufferWithLength:(size_t)R*D*4 options:g_res_opts];
id<MTLCommandBuffer> 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;
}
c/glm.c
+371 -16
View File
@@ -30,6 +30,7 @@
#if defined(__APPLE__) || defined(__linux__) #if defined(__APPLE__) || defined(__linux__)
#include <sys/resource.h> #include <sys/resource.h>
#include <sys/mman.h> /* mlock: inchioda le pagine in RAM / wire pages into RAM */ #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 #endif
#include "st.h" #include "st.h"
#include "tok.h" #include "tok.h"
@@ -44,6 +45,15 @@ static inline int omp_get_thread_num(void){ return 0; }
#ifdef COLI_CUDA #ifdef COLI_CUDA
#include "backend_cuda.h" #include "backend_cuda.h"
#endif #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__ #ifdef __AVX2__
#include <immintrin.h> #include <immintrin.h>
static inline float hsum256(__m256 v){ /* somma orizzontale di 8 float */ 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){ 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 #ifdef COLI_CUDA
/* The CUDA backend owns persistent copies only for model-resident tensors. /* The CUDA backend owns persistent copies only for model-resident tensors.
* Streaming expert slots are reused for different IDs and must never enter * 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 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=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 */ 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] 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. */ * (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: /* 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){ 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; 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); } 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>=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=malloc((int64_t)O*((I+1)/2)); t->s=falloc(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=malloc((int64_t)O*((I+3)/4)); t->s=falloc(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){ 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)); 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)){ if(st_has(&m->S,sn)){
int64_t nb=st_nbytes(&m->S,name); 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 */ 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); } 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=malloc(nb); t->s=falloc(O); } st_read_raw(&m->S,name,t->q4,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); st_read_f32(&m->S,sn,t->s,drop);
} else { } else {
if(!t->qf && !t->q8 && !t->q4) qt_alloc(t,O,I,bits); 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) */ 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);} 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){ 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; } 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 /* 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 * 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). * 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); 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(!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 wtot=tw[0]->nbytes+tw[1]->nbytes+tw[2]->nbytes;
int64_t ftot=(tq[0]->nbytes+tq[1]->nbytes+tq[2]->nbytes)/4; 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: /* rialloca se lo slot (riusato tra layer) e' troppo piccolo per QUESTO expert:
* pread oltre la mappatura = short-read o CORRUZIONE silenziosa dei vicini */ * pread oltre la mappatura = short-read o CORRUZIONE silenziosa dei vicini */
if(!s->slab || wtot+8192 > s->slab_cap){ 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); 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;} 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; s->slab_cap=wtot+8192;
#endif
} }
if(!s->fslab || ftot > s->fslab_cap){ 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); free(s->fslab);
if(fatal){ s->fslab=falloc(ftot); } /* main path: byte-identical exit-on-OOM */ if(fatal){ s->fslab=falloc(ftot); } /* main path: byte-identical exit-on-OOM */
else { /* speculative pilot: checked alloc, never exit() */ 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) || if(ftot<0 || (uint64_t)ftot > SIZE_MAX/sizeof(float) ||
!(s->fslab=malloc((size_t)ftot*sizeof(float)))){ !(s->fslab=malloc((size_t)ftot*sizeof(float)))){
fprintf(stderr,"OOM fslab\n"); 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=NULL; s->fslab_cap=0; return -1;
} }
} }
s->fslab_cap=ftot; s->fslab_cap=ftot;
#endif
} }
int ord[3]={0,1,2}; /* ordina per offset nel file */ 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; } 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; 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; int kvb_dim=H*(c->qk_nope+vh), Tk=pos_base+S;
double ta0=now_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 *ctx=falloc((int64_t)S*H*vh);
float *Q=falloc((int64_t)S*H*qh); /* query (roped) dei token nuovi */ 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); 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 ---- */ /* ---- 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 *idxs=malloc((size_t)S*K*sizeof(int)); float *ws=malloc((size_t)S*K*sizeof(float));
int *keff=malloc(S*sizeof(int)); 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++){ for(int s=0;s<S;s++){
const float *xs=x+(int64_t)s*D; const float *xs=x+(int64_t)s*D;
matmul(logit, xs, l->router, 1, D, E); 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 ---- */ /* ---- 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); 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 *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){ for(int base=0;base<nu;base+=64){
int nb = nu-base<64 ? nu-base : 64; int nb = nu-base<64 ? nu-base : 64;
ESlot *use[64]; int missk[64]; int qof[64]; int nmiss=0; 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; } } 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++; } 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(nmiss){
if(g_pipe){ /* PIPE: launch loads async, matmul overlaps them */ if(g_pipe){ /* PIPE: launch loads async, matmul overlaps them */
if(!g_pp.started) pipe_init(m); 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); 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]; 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 /* 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 * 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; } 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 */ int nr=0; /* righe (posizioni) che usano questo expert */
for(int s=0;s<S;s++) for(int kk=0;kk<keff[s];kk++) 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; } 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); } 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); float *sg=falloc((int64_t)S*sI), *su=falloc((int64_t)S*sI);
matmul_qt(sg, x, &l->sh_gate, S); #ifdef COLI_METAL
matmul_qt(su, x, &l->sh_up, S); 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; }
for(int64_t z=0;z<(int64_t)S*sI;z++) sg[z]=siluf(sg[z])*su[z]; #endif
matmul_qt(hh, sg, &l->sh_down, S); if(!shared_on_gpu){
for(int64_t z=0;z<(int64_t)S*D;z++) out[z]+=hh[z]; 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(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); 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) 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]); 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); 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); 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); attention(m,l,li,nrm,S,pos_base,tmp);
for(int64_t j=0;j<(int64_t)S*D;j++) x[j]+=tmp[j]; 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){ static void kv_alloc(Model *m, int max_t){
Cfg *c=&m->c; Cfg *c=&m->c;
KVState *k=m->kv; 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(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){ if(m->has_dsa){
k->Ic=calloc(c->n_layers,sizeof(float*)); 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 */ int NR=c->n_layers+1; /* riga extra: KV del layer MTP */
k->Lc=calloc(NR,sizeof(float*)); k->Rc=calloc(NR,sizeof(float*)); 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); 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; 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 " printf("PROFILE: expert-disk %.3fs | expert-matmul %.3fs | attention %.3fs "
"(including kvb %.3fs) | lm_head %.3fs | other %.3fs\n", "(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); 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 /* 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_nopack = getenv("NOPACK")?1:0;
g_drop = getenv("DROP")?1:0; g_drop = getenv("DROP")?1:0;
g_prefetch = getenv("PREFETCH")?atoi(getenv("PREFETCH")):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_topk = getenv("TOPK")?atoi(getenv("TOPK")):0;
g_topp = getenv("TOPP")?atof(getenv("TOPP")):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 */ 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"); fprintf(stderr,"CUDA was requested, but this binary is CPU-only; rebuild with: make CUDA=1\n");
return 2; 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 #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); 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(); g_mem_avail_boot = mem_available_gb();
c/tests/test_backend_metal.mm
+221
View File
@@ -0,0 +1,221 @@
// Kernel-correctness test for the Metal backend: coli_metal_matmul vs CPU reference
// (dequant->f32 MAC * per-row scale) for f32/int8/int4/int2 across real GLM shapes.
#include "../backend_metal.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <vector>
enum { F32=0, I8=1, I4=2, I2=3 };
static void cpu_ref(int fmt, const void *W, const float *s, const float *x,
float *y, int S, int I, int O) {
const int8_t *q8 = (const int8_t*)W; const uint8_t *q4 = (const uint8_t*)W;
const float *qf = (const float*)W;
int rb4=(I+1)/2, rb2=(I+3)/4;
for (int o=0;o<O;o++) for (int si=0;si<S;si++){
const float *xr = x + (size_t)si*I; float acc=0;
for (int i=0;i<I;i++){
float w;
if (fmt==I8) w=(float)q8[(size_t)o*I+i];
else if (fmt==I4){ uint8_t b=q4[(size_t)o*rb4+(i>>1)]; int v=(i&1)?(b>>4):(b&0xF); w=(float)(v-8); }
else if (fmt==I2){ uint8_t b=q4[(size_t)o*rb2+(i>>2)]; int v=(b>>(2*(i&3)))&0x3; w=(float)(v-2); }
else w=qf[(size_t)o*I+i];
acc += w*xr[i];
}
y[(size_t)si*O+o]=acc*s[o];
}
}
static int run(int fmt, int O, int I, int S, const char *name) {
int rb4=(I+1)/2, rb2=(I+3)/4;
size_t wn = (fmt==I8)?(size_t)O*I : (fmt==I4)?(size_t)O*rb4 : (fmt==I2)?(size_t)O*rb2 : (size_t)O*I*sizeof(float);
std::vector<uint8_t> W(wn); std::vector<float> Wf;
srand(99);
if (fmt==F32){ Wf.resize((size_t)O*I); for(auto&v:Wf) v=((rand()%2000)-1000)/1000.f; }
else for(auto&b:W) b=(uint8_t)((fmt==I8)?((rand()%255)-127):(rand()&0xFF));
const void *Wp = (fmt==F32)?(const void*)Wf.data():(const void*)W.data();
std::vector<float> s(O), x((size_t)S*I), yr((size_t)S*O), yg((size_t)S*O);
for(auto&v:s) v=(fmt==F32)?1.0f:(0.01f+(rand()%100)/10000.f);
for(auto&v:x) v=((rand()%2000)-1000)/1000.f;
cpu_ref(fmt, Wp, s.data(), x.data(), yr.data(), S, I, O);
ColiMetalTensor *t=nullptr;
if (!coli_metal_matmul(&t, yg.data(), x.data(), Wp, s.data(), fmt, S, I, O)) {
printf(" %-22s FAIL (matmul returned 0)\n", name); return 1; }
double maxabs=0, ymax=0;
for(size_t i=0;i<(size_t)S*O;i++){ maxabs=fmax(maxabs,fabs(yg[i]-yr[i])); ymax=fmax(ymax,fabs(yr[i])); }
double nerr=maxabs/(ymax+1e-9);
int ok = nerr < 1e-4;
printf(" %-22s nerr=%.2e %s\n", name, nerr, ok?"ok":"*** MISMATCH");
coli_metal_tensor_free(t);
return ok?0:1;
}
static float deq4(const uint8_t* w,int i){ uint8_t b=w[i>>1]; int v=(i&1)?(b>>4):(b&0xF); return (float)(v-8); }
static size_t roundpg(size_t n){ size_t p=16384; return ((n+p-1)/p)*p; }
// Validate coli_metal_moe_block against a CPU reference (gate/up/silu/down + weighted scatter-add).
static int run_moe(const std::vector<int>& nrv, const char* name) {
const int D=6144, I=2048, fmt=2; int rbG=(D+1)/2, rbD=(I+1)/2, nb=(int)nrv.size();
int R=0; std::vector<int> xoff(nb),nr(nrv); for(int e=0;e<nb;e++){ xoff[e]=R; R+=nrv[e]; }
srand(2024+nb);
// per-expert page-aligned slab [Wg|Wu|Wd] and fslab [Sg|Su|Sd]; register both.
std::vector<void*> slab(nb), fslab(nb);
std::vector<const void*> g(nb),u(nb),d(nb); std::vector<const float*> gs(nb),us(nb),ds(nb);
size_t wlen=roundpg((size_t)I*rbG*2 + (size_t)D*rbD), flen=roundpg(((size_t)I*2+D)*sizeof(float));
for(int e=0;e<nb;e++){
posix_memalign(&slab[e],16384,wlen); posix_memalign(&fslab[e],16384,flen);
uint8_t* sp=(uint8_t*)slab[e]; for(size_t i=0;i<(size_t)I*rbG*2+(size_t)D*rbD;i++) sp[i]=(uint8_t)(rand()&0xFF);
float* fp=(float*)fslab[e]; for(size_t i=0;i<(size_t)I*2+D;i++) fp[i]=0.01f+(rand()%50)/50000.f;
g[e]=sp; u[e]=sp+(size_t)I*rbG; d[e]=sp+(size_t)I*rbG*2;
gs[e]=fp; us[e]=fp+I; ds[e]=fp+2*I;
coli_metal_register(slab[e],wlen); coli_metal_register(fslab[e],flen);
}
std::vector<float> xg((size_t)R*D); for(auto&v:xg) v=((rand()%2000)-1000)/1000.f;
std::vector<int> rows(R); std::vector<float> rw(R);
for(int gr=0;gr<R;gr++){ rows[gr]=0; rw[gr]=0.1f+(rand()%100)/100.f; } // decode: all -> position 0
int S=1;
// CPU reference
std::vector<float> refout((size_t)S*D,0.f), gg(I),uu(I),hh(D);
for(int e=0;e<nb;e++) for(int r=0;r<nr[e];r++){ int gr=xoff[e]+r; const float* xr=&xg[(size_t)gr*D];
const uint8_t* wg=(const uint8_t*)g[e]; const uint8_t* wu=(const uint8_t*)u[e]; const uint8_t* wd=(const uint8_t*)d[e];
for(int o=0;o<I;o++){ float a=0; for(int k=0;k<D;k++) a+=deq4(wg+(size_t)o*rbG,k)*xr[k]; gg[o]=a*gs[e][o]; }
for(int o=0;o<I;o++){ float a=0; for(int k=0;k<D;k++) a+=deq4(wu+(size_t)o*rbG,k)*xr[k]; uu[o]=a*us[e][o]; }
for(int o=0;o<I;o++){ float v=gg[o]; gg[o]=(v/(1.f+expf(-v)))*uu[o]; }
for(int o=0;o<D;o++){ float a=0; for(int k=0;k<I;k++) a+=deq4(wd+(size_t)o*rbD,k)*gg[k]; hh[o]=a*ds[e][o]; }
float* os=&refout[(size_t)rows[gr]*D]; for(int o=0;o<D;o++) os[o]+=rw[gr]*hh[o];
}
std::vector<float> gout((size_t)S*D,0.f);
int ok = coli_metal_moe_block(nb,D,I,fmt,g.data(),u.data(),d.data(),gs.data(),us.data(),ds.data(),
xg.data(),xoff.data(),nr.data(),rows.data(),rw.data(),gout.data(),S);
double maxabs=0,ymax=0; for(size_t i=0;i<gout.size();i++){ maxabs=fmax(maxabs,fabs(gout[i]-refout[i])); ymax=fmax(ymax,fabs(refout[i])); }
double nerr=maxabs/(ymax+1e-9); int pass = ok && nerr<1e-4;
printf(" %-22s R=%d nerr=%.2e %s\n", name, R, nerr, pass?"ok":"*** MISMATCH");
for(int e=0;e<nb;e++){ coli_metal_unregister(slab[e]); coli_metal_unregister(fslab[e]); free(slab[e]); free(fslab[e]); }
return pass?0:1;
}
// ---- fused decode attention vs a CPU reference replicating glm.c's exact math ----
// GLM-5.2 dims (hardcoded in the backend): hidden=6144 H=64 q_lora=2048 kv_lora=512
// nope=192 rope=64 vh=256; theta=10000 ascale=1/16 eps=1e-5.
enum { TH=6144, THH=64, TQL=2048, TKVL=512, TNOPE=192, TROPE=64, TVH=256, TQH=256, TROWSH=448 };
static void t_rms(float*o,const float*x,const float*w,int n,float eps){ double ms=0; for(int i=0;i<n;i++) ms+=(double)x[i]*x[i];
float r=1.f/sqrtf((float)(ms/n)+eps); for(int i=0;i<n;i++) o[i]=x[i]*r*w[i]; }
static void t_rope(float*v,int pos,float th){ int hl=TROPE/2; float in[TROPE]; memcpy(in,v,sizeof(in));
for(int j=0;j<hl;j++){ float inv=powf(th,-2.f*j/TROPE), a=in[2*j], b=in[2*j+1], cs=cosf(pos*inv), sn=sinf(pos*inv);
v[j]=a*cs-b*sn; v[hl+j]=b*cs+a*sn; } }
static void t_gemv4(float*y,const float*x,const uint8_t*w,const float*sc,int O,int I){ int rb=(I+1)/2;
for(int o=0;o<O;o++){ const uint8_t*r=w+(size_t)o*rb; float a=0;
for(int i=0;i<I;i++){ uint8_t b=r[i>>1]; int v=(i&1)?(b>>4):(b&0xF); a+=(float)(v-8)*x[i]; } y[o]=a*sc[o]; } }
struct TW { uint8_t*w; float*s; size_t wb, sb; };
static TW t_mkw(int O,int I){ TW t; int rb=(I+1)/2;
t.wb=((size_t)O*rb+16383)&~(size_t)16383; t.sb=((size_t)O*4+16383)&~(size_t)16383;
posix_memalign((void**)&t.w,16384,t.wb); posix_memalign((void**)&t.s,16384,t.sb);
for(size_t i=0;i<(size_t)O*rb;i++) t.w[i]=(uint8_t)(rand()&0xFF);
for(int i=0;i<O;i++) t.s[i]=0.01f+(rand()%40)/40000.f;
coli_metal_register(t.w,t.wb); coli_metal_register(t.s,t.sb); return t; }
static int run_attn(int S, int pos_base, const char* name){
const float eps=1e-5f, theta=10000.f, ascale=1.f/16.f;
srand(4242+S+pos_base);
TW qa=t_mkw(TQL,TH), qb=t_mkw(THH*TQH,TQL), kva=t_mkw(TKVL+TROPE,TH), kvb=t_mkw(THH*TROWSH,TKVL), o=t_mkw(TH,THH*TVH);
std::vector<float> qaln(TQL), kvaln(TKVL);
for(auto&v:qaln) v=0.5f+(rand()%1000)/1000.f; for(auto&v:kvaln) v=0.5f+(rand()%1000)/1000.f;
int T=pos_base+S; size_t lcb=(((size_t)T*TKVL*4)+16383)&~(size_t)16383, rcb=(((size_t)T*TROPE*4)+16383)&~(size_t)16383;
float *Lc,*Rc; posix_memalign((void**)&Lc,16384,lcb); posix_memalign((void**)&Rc,16384,rcb);
coli_metal_register(Lc,lcb); coli_metal_register(Rc,rcb);
// pre-existing cache history [0,pos_base): random normed latents + roped krot
for(int t=0;t<pos_base;t++){ for(int i=0;i<TKVL;i++) Lc[(size_t)t*TKVL+i]=((rand()%2000)-1000)/1500.f;
for(int i=0;i<TROPE;i++) Rc[(size_t)t*TROPE+i]=((rand()%2000)-1000)/1500.f; }
std::vector<float> x((size_t)S*TH); for(auto&v:x) v=((rand()%2000)-1000)/1000.f;
std::vector<float> Lr((size_t)T*TKVL), Rr((size_t)T*TROPE); // reference cache copies
memcpy(Lr.data(),Lc,(size_t)pos_base*TKVL*4); memcpy(Rr.data(),Rc,(size_t)pos_base*TROPE*4);
// CPU reference: mirrors glm.c attention() absorb branch (per new token, then per head)
std::vector<float> Q((size_t)S*THH*TQH), ref((size_t)S*TH);
for(int s=0;s<S;s++){ int pos=pos_base+s;
std::vector<float> qr(TQL), comp(TKVL+TROPE);
t_gemv4(qr.data(),&x[(size_t)s*TH],qa.w,qa.s,TQL,TH); t_rms(qr.data(),qr.data(),qaln.data(),TQL,eps);
t_gemv4(&Q[(size_t)s*THH*TQH],qr.data(),qb.w,qb.s,THH*TQH,TQL);
for(int h=0;h<THH;h++) t_rope(&Q[(size_t)s*THH*TQH+(size_t)h*TQH+TNOPE],pos,theta);
t_gemv4(comp.data(),&x[(size_t)s*TH],kva.w,kva.s,TKVL+TROPE,TH);
t_rms(&Lr[(size_t)pos*TKVL],comp.data(),kvaln.data(),TKVL,eps);
memcpy(&Rr[(size_t)pos*TROPE],&comp[TKVL],TROPE*4); t_rope(&Rr[(size_t)pos*TROPE],pos,theta);
}
int rb=(TKVL+1)/2;
for(int s=0;s<S;s++){ int pos=pos_base+s; std::vector<float> ctx((size_t)THH*TVH);
for(int h=0;h<THH;h++){ int rbase=h*TROWSH;
const float* qp=&Q[(size_t)s*THH*TQH+(size_t)h*TQH]; const float* qro=qp+TNOPE;
std::vector<float> qabs(TKVL,0);
for(int d=0;d<TNOPE;d++){ const uint8_t*r=kvb.w+(size_t)(rbase+d)*rb; float sc=kvb.s[rbase+d];
for(int i=0;i<TKVL;i++){ uint8_t b=r[i>>1]; int v=(i&1)?(b>>4):(b&0xF); qabs[i]+=qp[d]*(float)(v-8)*sc; } }
std::vector<float> a(pos+1);
for(int t=0;t<=pos;t++){ const float*Lt=&Lr[(size_t)t*TKVL]; const float*Rt=&Rr[(size_t)t*TROPE];
float v=0; for(int i=0;i<TKVL;i++) v+=qabs[i]*Lt[i]; for(int d=0;d<TROPE;d++) v+=qro[d]*Rt[d]; a[t]=v*ascale; }
float mx=-1e30f; for(float v:a) mx=fmaxf(mx,v); float sum=0; for(float&v:a){ v=expf(v-mx); sum+=v; } for(float&v:a) v/=sum;
std::vector<float> cl(TKVL,0);
for(int t=0;t<=pos;t++){ const float*Lt=&Lr[(size_t)t*TKVL]; for(int i=0;i<TKVL;i++) cl[i]+=a[t]*Lt[i]; }
for(int j=0;j<TVH;j++){ const uint8_t*r=kvb.w+(size_t)(rbase+TNOPE+j)*rb; float sc=kvb.s[rbase+TNOPE+j];
float v=0; for(int i=0;i<TKVL;i++){ uint8_t b=r[i>>1]; int vv=(i&1)?(b>>4):(b&0xF); v+=cl[i]*(float)(vv-8)*sc; }
ctx[(size_t)h*TVH+j]=v; } }
t_gemv4(&ref[(size_t)s*TH],ctx.data(),o.w,o.s,TH,THH*TVH);
}
std::vector<float> got((size_t)S*TH);
int ok=coli_metal_attn_decode(x.data(), qa.w,qa.s,2,qaln.data(), qb.w,qb.s,2,
kva.w,kva.s,2,kvaln.data(), kvb.w,kvb.s,2, o.w,o.s,2,
Lc,Rc,S,pos_base,0,eps,theta,ascale,got.data());
double ma=0,ym=0; for(size_t i=0;i<ref.size();i++){ ma=fmax(ma,fabs(got[i]-ref[i])); ym=fmax(ym,fabs(ref[i])); }
// also verify the cache write-back (Lc/Rc for the new positions)
double mc=0; for(int s=0;s<S;s++){ int pos=pos_base+s;
for(int i=0;i<TKVL;i++) mc=fmax(mc,fabs(Lc[(size_t)pos*TKVL+i]-Lr[(size_t)pos*TKVL+i]));
for(int i=0;i<TROPE;i++) mc=fmax(mc,fabs(Rc[(size_t)pos*TROPE+i]-Rr[(size_t)pos*TROPE+i])); }
double nerr=ma/(ym+1e-9);
int pass = ok && nerr<2e-4 && mc<1e-4;
printf(" %-24s nerr=%.2e cache=%.2e %s\n", name, nerr, mc, pass?"ok":"*** MISMATCH");
auto freew=[&](TW&t){ coli_metal_unregister(t.w); coli_metal_unregister(t.s); free(t.w); free(t.s); };
freew(qa); freew(qb); freew(kva); freew(kvb); freew(o);
coli_metal_unregister(Lc); coli_metal_unregister(Rc); free(Lc); free(Rc);
return pass?0:1;
}
int main(void) {
if (!coli_metal_init()) { printf("Metal unavailable (skipping)\n"); return 0; }
printf("Metal backend kernel tests:\n");
int fail=0;
fail |= run(I8, 2048,6144,1, "int8 gate/up S=1");
fail |= run(I4, 2048,6144,1, "int4 gate/up S=1");
fail |= run(I4, 6144,2048,1, "int4 down S=1");
fail |= run(I2, 2048,6144,1, "int2 gate/up S=1");
fail |= run(F32,1024,6144,1, "f32 S=1");
fail |= run(I8, 2048,6144,4, "int8 gate/up S=4");
fail |= run(I4, 2048,6144,7, "int4 gate/up S=7 (odd)");
fail |= run(I4, 2050,6146,3, "int4 non-mult-4 dims");
printf("Metal batched moe_block tests:\n");
fail |= run_moe({1,1,1,1,1,1,1,1}, "moe decode nb=8");
fail |= run_moe({3,1,4,2,1,5}, "moe ragged nb=6");
printf("Metal large-batch gemm test:\n");
{ // registered int4 weights, S=64: coli_metal_gemm vs cpu_ref
srand(77); int O=2048,I=6144,S=64,rb=(I+1)/2;
size_t wb=(((size_t)O*rb)+16383)&~(size_t)16383, sb2=(((size_t)O*4)+16383)&~(size_t)16383;
uint8_t*W; float*Sc; posix_memalign((void**)&W,16384,wb); posix_memalign((void**)&Sc,16384,sb2);
for(size_t i=0;i<(size_t)O*rb;i++) W[i]=(uint8_t)(rand()&0xFF);
for(int i=0;i<O;i++) Sc[i]=0.01f+(rand()%50)/50000.f;
coli_metal_register(W,wb); coli_metal_register(Sc,sb2);
std::vector<float> x((size_t)S*I), yr((size_t)S*O), yg((size_t)S*O);
for(auto&v:x) v=((rand()%2000)-1000)/1000.f;
cpu_ref(I4,W,Sc,x.data(),yr.data(),S,I,O);
int ok=coli_metal_gemm(yg.data(),x.data(),W,Sc,2,S,I,O);
double ma=0,ym=0; for(size_t i=0;i<yr.size();i++){ ma=fmax(ma,fabs(yg[i]-yr[i])); ym=fmax(ym,fabs(yr[i])); }
int pass = ok && ma/(ym+1e-9)<1e-4;
printf(" gemm S=64 int4 nerr=%.2e %s\n", ma/(ym+1e-9), pass?"ok":"*** MISMATCH");
fail |= !pass;
coli_metal_unregister(W); coli_metal_unregister(Sc); free(W); free(Sc);
}
printf("Metal fused attention tests:\n");
fail |= run_attn(1, 0, "attn S=1 pos=0");
fail |= run_attn(1, 37, "attn S=1 pos=37");
fail |= run_attn(4, 12, "attn S=4 pos=12 (MTP)");
fail |= run_attn(3, 0, "attn S=3 pos=0");
printf(fail? "metal backend tests: FAILED\n" : "metal backend tests: ok\n");
coli_metal_shutdown();
return fail;
}