82d06000b5
1. REQUIRED (backend_metal.mm): no Metal call runs under g_slab_mtx anymore. The set mutations (addAllocation/removeAllocation/commit) and the dirty flag moved to a dedicated g_resset_mtx; coli_metal_register/unregister do their g_slabs bookkeeping under g_slab_mtx exactly as stock, then call resset_add/resset_remove after dropping it (still before returning, so unregister's remove+commit still lands before the caller frees the host memory). The original shape -- commit under the slab lock the parallel OMP loader threads contend on -- was structurally identical to the mutex-over- live-Metal-call bug E4's audit round 2 fixed, the leading suspect for its replicated +12s expert-disk regression. The register->flush->resolve happens-before argument survives the split: resset_add completes inside coli_metal_register before it returns, the engine cannot dispatch an expert before its load returns, and add/flush are serialized by g_resset_mtx (comment at resset_add). The two mutexes are never held together, so no lock-order hazard exists. 2. REQUIRED (backend_metal.mm, backend_metal.h, glm.c): the resset_flush cost in moe_submit -- deliberately outside the g_t_setup window so the A/B harness counters keep their meaning -- was invisible. New g_t_resset_flush accumulator timed around the flush, exported via coli_metal_resset_stats(), printed by profile_print as a separate gate-on-only "METAL-RESSET: flush" line (mirrors E4's METAL-HEAP line; the METAL: line the harness parses keeps its exact format; stock output byte-identical -- the stats call returns 0 with the gate off and the whole block sits in the pre-existing #ifdef COLI_METAL arm). Register- side add/remove costs land in the existing t_ewait window; comment at resset_add says so instead of a second counter. 3. REQUIRED (SUMMARY.md; the moe_submit commit message was already reworded in place, pre-push): corrected two false attributions -- the hazard- tracking and thread-safety statements were credited to the SDK header, which is silent on both; the actual source is Apple's online MTLResidencySet class reference and residency-set adoption guide (fetched 2026-07-18). Note: the scaffolding commit's message was checked and contains no such claim, so it was left untouched. 4. DOC (SUMMARY.md): the pre-existing fslab OOM-unwind bug (glm.c ~1868-73, frees s->slab without unregister) has a strictly longer-lived exposure under E5 (permanent set member vs stock's transient per-CB declaration). Out of scope here; the upstream PR built from E5 must carry the one-line fix (reference: E4 branch commit 6753225).
148 lines
7.4 KiB
C
148 lines
7.4 KiB
C
#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);
|
|
/* E5 (COLI_METAL_RESSET=1): returns 1 when the queue-attached residency set is active and
|
|
* writes the cumulative seconds moe_submit spent committing pending set adds -- a cost that
|
|
* sits OUTSIDE the setup/gpu/scatter breakdown above. Returns 0 (and writes 0) when off. */
|
|
int coli_metal_resset_stats(double *flush_s);
|
|
|
|
/*
|
|
* 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
|