diff --git a/c/Makefile b/c/Makefile index 97bf41ea9..f608325a6 100644 --- a/c/Makefile +++ b/c/Makefile @@ -961,10 +961,15 @@ tests/test_pipe_block$(EXE): tests/test_pipe_block.c colibri.c st.h uring.h json tests/test_pilot_ring$(EXE): tests/test_pilot_ring.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) +tests/test_moe_gs_guard$(EXE): tests/test_moe_gs_guard.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h + $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) tests/test_omp_tune$(EXE): tests/test_omp_tune.c omp_tune.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) +tests/test_kvb_notice$(EXE): tests/test_kvb_notice.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h + $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) + test-c: $(TEST_BINS) $(PYTHON) tools/run_tests.py $(TEST_BINS) diff --git a/c/backend_metal.h b/c/backend_metal.h index 2c09ba769..552adf83a 100644 --- a/c/backend_metal.h +++ b/c/backend_metal.h @@ -80,7 +80,7 @@ int coli_metal_layer_decode(float *x, const void *qa_w, const float *qa_s, int qa_fmt, int qa_gs, const float *qa_ln, const void *qb_w, const float *qb_s, int qb_fmt, int qb_gs, const void *kva_w, const float *kva_s, int kva_fmt, int kva_gs, const float *kva_ln, - const void *kvb_w, const float *kvb_s, int kvb_fmt, + const void *kvb_w, const float *kvb_s, int kvb_fmt, int kvb_gs, const void *o_w, const float *o_s, int o_fmt, int o_gs, const void *shg_w, const float *shg_s, int shg_fmt, int shg_gs, const void *shu_w, const float *shu_s, int shu_fmt, int shu_gs, @@ -116,7 +116,7 @@ int coli_metal_attn_decode(const float *x, const void *qa_w, const float *qa_s, int qa_fmt, int qa_gs, const float *qa_ln, const void *qb_w, const float *qb_s, int qb_fmt, int qb_gs, const void *kva_w, const float *kva_s, int kva_fmt, int kva_gs, const float *kva_ln, - const void *kvb_w, const float *kvb_s, int kvb_fmt, + const void *kvb_w, const float *kvb_s, int kvb_fmt, int kvb_gs, const void *o_w, const float *o_s, int o_fmt, int o_gs, float *Lc, float *Rc, int S, int pos_base, int st0, float eps, float theta, float ascale, float *out); @@ -137,18 +137,16 @@ int coli_metal_resset_stats(double *flush_s); * * 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). NOTE: fmt=4 (grouped int4) is - * NOT yet supported here -- gates to {1,2} and returns 0 (CPU fallback) - * for fmt=4 experts, same as before this stage. Grouped-int4 gained GPU - * support in mm_gemv (coli_metal_matmul/coli_metal_gemm/bind_gemv) only; - * extending the batched routed-expert path is future work (see PR_BODY.md). + * gs/us/ds[e] = pointers to expert e's per-row (fmt=1/2) or per-group (fmt=4) scales + * fmt = quant format (shared across experts): 1=int8, 2=int4 per-row, 4=int4 + * grouped. qgs is the fmt=4 group size (ignored, pass 0, for fmt!=4). + * qgs = fmt=4 group size shared across experts in this block (0 for fmt!=4) * 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, +int coli_metal_moe_block(int nb, int D, int Iinter, int fmt, int qgs, 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, @@ -163,7 +161,7 @@ int coli_metal_moe_block(int nb, int D, int Iinter, int fmt, * 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, +ColiMetalMoeHandle* coli_metal_moe_block_begin(int nb, int D, int Iinter, int fmt, int qgs, 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, diff --git a/c/backend_metal.mm b/c/backend_metal.mm index 545e03507..4a9685b9d 100644 --- a/c/backend_metal.mm +++ b/c/backend_metal.mm @@ -154,7 +154,10 @@ kernel void mm_gemv(device const uchar* w [[buffer(0)]], // raw weight by } // 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. +// scale live at gpuAddresses waddr[e]/saddr[e] (zero-copy in the RAM slab). fmt 1=i8, 2=i4 +// per-row, 4=i4 grouped (scale layout [O][ng], ng=ceil(K/qgs) -- same convention as mm_gemv +// fmt=4 above, but folded into the vectorized uchar4/float4 dot-product loop this kernel's +// fmt=2 branch already uses, since moe_gemv has no scalar strided branch to reuse). // 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)]], @@ -162,7 +165,7 @@ kernel void moe_gemv(device const ulong* waddr [[buffer(0)]], device const ulong 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)]], + constant int& NT [[buffer(9)]], constant int& qgs [[buffer(10)]], uint tg [[threadgroup_position_in_grid]], uint slane [[thread_index_in_simdgroup]], uint sgid [[simdgroup_index_in_threadgroup]]) { @@ -210,13 +213,25 @@ kernel void moe_gemv(device const ulong* waddr [[buffer(0)]], device const ulong as_type(uint(b.z)<<16), as_type(uint(b.w)<<16)); acc+=dot(w0,x4[2*c])+dot(w1,x4[2*c+1]); } for(int i=K8*8+slane;i(uint(w[i])<<16)*xr[i]; + } else if (fmt == 4) { // grouped int4: per-expert scale [O][ng] + int rb=(K+1)/2, ng=(K+qgs-1)/qgs; device const uchar* w=(device const uchar*)(waddr[e])+(long)o*rb; + device const float* sr=sc+(long)o*ng; // grouped scales for this output row + device const uchar4* w4=(device const uchar4*)w; + for(int c=slane;c>4)-8),float(int(b.y&0xF)-8),float(int(b.y>>4)-8)); + float4 w1=float4(float(int(b.z&0xF)-8),float(int(b.z>>4)-8),float(int(b.w&0xF)-8),float(int(b.w>>4)-8)); + int g0=(8*c+0)/qgs,g1=(8*c+1)/qgs,g2=(8*c+2)/qgs,g3=(8*c+3)/qgs; + int g4=(8*c+4)/qgs,g5=(8*c+5)/qgs,g6=(8*c+6)/qgs,g7=(8*c+7)/qgs; + acc+=dot(w0*float4(sr[g0],sr[g1],sr[g2],sr[g3]),x4[2*c]) + +dot(w1*float4(sr[g4],sr[g5],sr[g6],sr[g7]),x4[2*c+1]); } + for(int i=K8*8+slane;i>1]; int v=(i&1)?(b>>4):(b&0xF); acc+=float(v-8)*xr[i]*sr[i/qgs]; } } 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>1]; int val=(i&1)?(b>>4):(b&0xF); return float(val-8)*sc[row]; } +// kv_b inline dequant of column i of output row `row`. fmt=2 -> one scale per row; +// fmt=4 -> grouped int4, one scale per gs-wide group along the A_KVL input dim +// (scale layout [O][ng], ng=ceil(A_KVL/gs)), matching QT fmt=4 / mm_gemv above. +inline float a_deqrow(device const uchar* base, int row, int i, device const float* sc, int fmt, int gs){ + device const uchar* w=base+(long)row*((A_KVL+1)/2); uchar b=w[i>>1]; int val=(i&1)?(b>>4):(b&0xF); + float s = (fmt==4) ? sc[(long)row*((A_KVL+gs-1)/gs) + i/gs] : sc[row]; + return float(val-8)*s; } 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)]], + constant int& fmt [[buffer(4)]], constant int& gs [[buffer(5)]], 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 e, const void* w, const float // Weight-pointer bundle for one layer's attention (+optional layer tail). All pointers // must be inside registered allocations. *_gs: fmt=4 group size for the corresponding -// weight (0 if that weight isn't grouped). kv_b has none: it never flows through bind_gemv -// -- a_qabs/a_ctx dequantize it inline with a PER-ROW-only helper (a_deqrow), so a fmt=4 -// kv_b is out of scope for this stage (see PR_BODY.md UNCERTAINTIES). +// weight (0 if that weight isn't grouped). kv_b never flows through bind_gemv -- a_qabs/ +// a_ctx dequantize it inline via a_deqrow, which is fmt/gs-aware (fmt=2 per-row, fmt=4 +// grouped along A_KVL); kvb_gs is that group size (0 for fmt=2). typedef struct { const void *qa_w; const float *qa_s; int qa_fmt; int qa_gs; const float *qa_ln; const void *qb_w; const float *qb_s; int qb_fmt; int qb_gs; const void *kva_w; const float *kva_s; int kva_fmt; int kva_gs; const float *kva_ln; - const void *kvb_w; const float *kvb_s; int kvb_fmt; + const void *kvb_w; const float *kvb_s; int kvb_fmt; int kvb_gs; const void *o_w; const float *o_s; int o_fmt; int o_gs; } AttnW; @@ -890,6 +913,7 @@ static bool encode_attn_projections(id e, const AttnW bind_gemv(e,W->qb_w,W->qb_s,W->qb_fmt,W->qb_gs,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 setBytes:&W->kvb_fmt length:4 atIndex:4]; [e setBytes:&W->kvb_gs length:4 atIndex:5]; [e dispatchThreads:MTLSizeMake((size_t)S*AHEADS*AKVL,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); return true; } @@ -902,7 +926,8 @@ static bool encode_attn_projections(id e, const AttnW static bool encode_attn_core_chunk(id e, id Lb, size_t loff, id Rb, size_t roff, id kvbW, size_t kvbwoff, id kvbS, size_t kvbsoff, - int r0, int ch, int T, int pos_base, float ascale) { + int r0, int ch, int T, int pos_base, float ascale, + int kvb_fmt, int kvb_gs) { size_t qabs_off=(size_t)r0*AHEADS*AKVL*4, qf_off=(size_t)r0*AHQH*4, ctx_off=(size_t)r0*AHVH*4; int PB=pos_base; auto BAR=[&]{ [e memoryBarrierWithScope:MTLBarrierScopeBuffers]; }; @@ -915,6 +940,7 @@ static bool encode_attn_core_chunk(id e, [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)ch*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:ctx_off atIndex:3]; + [e setBytes:&kvb_fmt length:4 atIndex:4]; [e setBytes:&kvb_gs length:4 atIndex:5]; [e dispatchThreads:MTLSizeMake((size_t)ch*AHEADS*AVH,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); return true; } @@ -925,7 +951,7 @@ static bool encode_attention(id e, const AttnW *W, id kvbW, size_t kvbwoff, id kvbS, size_t kvbsoff, int S, int T, int pos_base, float eps, float theta, float ascale) { if(!encode_attn_projections(e,W,Lb,loff,Rb,roff,kvbW,kvbwoff,kvbS,kvbsoff,S,pos_base,eps,theta)) return false; - if(!encode_attn_core_chunk(e,Lb,loff,Rb,roff,kvbW,kvbwoff,kvbS,kvbsoff,0,S,T,pos_base,ascale)) return false; + if(!encode_attn_core_chunk(e,Lb,loff,Rb,roff,kvbW,kvbwoff,kvbS,kvbsoff,0,S,T,pos_base,ascale,W->kvb_fmt,W->kvb_gs)) return false; bind_gemv(e,W->o_w,W->o_s,W->o_fmt,W->o_gs,AHVH,AH,actx_,aout_,S); return true; } @@ -947,7 +973,7 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, const void* qa_w,const float* qa_s,int qa_fmt,int qa_gs,const float* qa_ln, const void* qb_w,const float* qb_s,int qb_fmt,int qb_gs, const void* kva_w,const float* kva_s,int kva_fmt,int kva_gs,const float* kva_ln, - const void* kvb_w,const float* kvb_s,int kvb_fmt, + const void* kvb_w,const float* kvb_s,int kvb_fmt,int kvb_gs, const void* o_w,const float* o_s,int o_fmt,int o_gs, float* Lc,float* Rc,int S,int pos_base,int st0,float eps,float theta,float ascale,float* out){ if(!g_dev) return 0; @@ -955,7 +981,7 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, int T=pos_base+S; @autoreleasepool { attn_scratch_init(); - AttnW W={qa_w,qa_s,qa_fmt,qa_gs,qa_ln,qb_w,qb_s,qb_fmt,qb_gs,kva_w,kva_s,kva_fmt,kva_gs,kva_ln,kvb_w,kvb_s,kvb_fmt,o_w,o_s,o_fmt,o_gs}; + AttnW W={qa_w,qa_s,qa_fmt,qa_gs,qa_ln,qb_w,qb_s,qb_fmt,qb_gs,kva_w,kva_s,kva_fmt,kva_gs,kva_ln,kvb_w,kvb_s,kvb_fmt,kvb_gs,o_w,o_s,o_fmt,o_gs}; id Lb,Rb,kvbW,kvbS; size_t loff,roff,kvbwoff,kvbsoff; if(!resolve_attn(&W,Lc,Rc,&Lb,&loff,&Rb,&roff,&kvbW,&kvbwoff,&kvbS,&kvbsoff)) return 0; @@ -994,7 +1020,7 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, const void* qa_w,const float* qa_s,int qa_fmt,int qa_gs,const float* qa_ln, const void* qb_w,const float* qb_s,int qb_fmt,int qb_gs, const void* kva_w,const float* kva_s,int kva_fmt,int kva_gs,const float* kva_ln, - const void* kvb_w,const float* kvb_s,int kvb_fmt, + const void* kvb_w,const float* kvb_s,int kvb_fmt,int kvb_gs, const void* o_w,const float* o_s,int o_fmt,int o_gs, const void* shg_w,const float* shg_s,int shg_fmt,int shg_gs, const void* shu_w,const float* shu_s,int shu_fmt,int shu_gs, @@ -1009,7 +1035,7 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, int T=pos_base+S; const int SI=2048; @autoreleasepool { attn_scratch_init(); - AttnW W={qa_w,qa_s,qa_fmt,qa_gs,qa_ln,qb_w,qb_s,qb_fmt,qb_gs,kva_w,kva_s,kva_fmt,kva_gs,kva_ln,kvb_w,kvb_s,kvb_fmt,o_w,o_s,o_fmt,o_gs}; + AttnW W={qa_w,qa_s,qa_fmt,qa_gs,qa_ln,qb_w,qb_s,qb_fmt,qb_gs,kva_w,kva_s,kva_fmt,kva_gs,kva_ln,kvb_w,kvb_s,kvb_fmt,kvb_gs,o_w,o_s,o_fmt,o_gs}; id Lb,Rb,kvbW,kvbS; size_t loff,roff,kvbwoff,kvbsoff; if(!resolve_attn(&W,Lc,Rc,&Lb,&loff,&Rb,&roff,&kvbW,&kvbwoff,&kvbS,&kvbsoff)) return 0; uint64_t ina=0,pna=0,rwa=0,rba=0,d; @@ -1195,12 +1221,12 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, // if Metal is off or any expert pointer is not in a registered slab. // Encode + commit a MoE block (no wait). Writes hh[R,D] into hh_buf. Returns nil on // unresolved slab / bad fmt (caller falls back to CPU). -static id moe_submit(int nb, int D, int Iinter, int fmt, +static id moe_submit(int nb, int D, int Iinter, int fmt, int qgs, const void *const *g, const void *const *u, const void *const *d, const float *const *gs, const float *const *us, const float *const *ds, const float *xg, const int *xoff, const int *nr, int R, id xg_buf, id gg_buf, id uu_buf, id hh_buf) { - if (!g_dev || (fmt != 1 && fmt != 2 && fmt != 5 && fmt != 6)) return nil; + if (!g_dev || (fmt != 1 && fmt != 2 && fmt != 4 && fmt != 5 && fmt != 6)) return nil; if (fmt == 6) { /* e8 kernel assumes clean block tiling, and every FWHT tile of the * down input (CPU tiling rule, e8_rot_rows) must fit threadgroup mem */ if ((D & 255) || (Iinter & 31)) return nil; @@ -1211,6 +1237,18 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, off += n; } } + // COLI_METAL_MOE_EXACT=1: route fmt=4 routed experts to the CPU reference path (bit-exact, + // matches matmul_i4_grouped) instead of the fast batched float4 kernel. Opt-in; default is + // the fast GPU path. Returning nil makes moe() fall back to CPU for fmt=4 experts, exactly + // like an unresolved slab -- fmt=1/2 stay on the GPU, attention/dense are untouched. + // (PR #587 gate-2: token-exact mode for trajectories whose gap dips under the drift tail.) + { static int g_moe_exact = -1; + if (g_moe_exact < 0) { const char *e = getenv("COLI_METAL_MOE_EXACT"); g_moe_exact = (e && e[0] && e[0] != '0'); } + if (g_moe_exact) return nil; } /* exact mode is path-scoped, not fmt-scoped: the resident + * tier (fmt=1 on mixed containers) carries the same + * accumulation-order drift, so ALL routed experts fall to + * CPU under the flag (measured: 4/5 prompt flips -> 0/5, + * real g64 744B container, #587) */ if (g_resset_enabled) { // E5: commit any pending slab adds before we may skip useResource: double t0 = mnow(); resset_flush(); g_t_resset_flush += mnow() - t0; // METAL-RESSET line } @@ -1253,7 +1291,7 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, [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 setBytes:&NT length:4 atIndex:9];[e setBytes:&qgs length:4 atIndex:10]; [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 @@ -1305,7 +1343,7 @@ static int moe_finish(id cb, id hh_buf, int nb, int return 1; } -extern "C" int coli_metal_moe_block(int nb, int D, int Iinter, int fmt, +extern "C" int coli_metal_moe_block(int nb, int D, int Iinter, int fmt, int qgs, 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, @@ -1318,7 +1356,7 @@ static int moe_finish(id cb, id hh_buf, int nb, int 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 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); + id cb = moe_submit(nb,D,Iinter,fmt,qgs,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); } @@ -1331,7 +1369,7 @@ static int moe_finish(id cb, id hh_buf, int nb, int std::vector rows; std::vector rwv; int nb, R, D; }; -extern "C" ColiMetalMoeHandle* coli_metal_moe_block_begin(int nb, int D, int Iinter, int fmt, +extern "C" ColiMetalMoeHandle* coli_metal_moe_block_begin(int nb, int D, int Iinter, int fmt, int qgs, 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, @@ -1343,7 +1381,7 @@ static int moe_finish(id cb, id hh_buf, int nb, int id bgg=[g_dev newBufferWithLength:(size_t)R*Iinter*4 options:g_res_opts]; id buu=[g_dev newBufferWithLength:(size_t)R*Iinter*4 options:g_res_opts]; id bhh=[g_dev newBufferWithLength:(size_t)R*D*4 options:g_res_opts]; - id cb = moe_submit(nb,D,Iinter,fmt,g,u,d,gs,us,ds,xg,xoff,nr,R,bxg,bgg,buu,bhh); + id cb = moe_submit(nb,D,Iinter,fmt,qgs,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); diff --git a/c/colibri.c b/c/colibri.c index 2723e6f9b..7b0af931d 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -83,6 +83,12 @@ static inline void omp_set_num_threads(int n){ (void)n; } #ifdef COLI_VULKAN #include "backend_vulkan.h" #endif +/* Declared unconditionally (not just under COLI_METAL): on a non-Metal build it just sits + * at 0 forever, which is the correct value there (no Metal backend => never enabled). Kept + * outside the #ifdef so portable code — e.g. kvb_fmt_gate_notice below — can read "is Metal + * active" without needing COLI_METAL (and the backend_metal.h / Metal framework link it + * would drag in) on every platform's test build. */ +static int g_metal_enabled; #ifdef COLI_METAL #include "backend_metal.h" /* No here: the guarded include above already provides it under _OPENMP @@ -93,6 +99,7 @@ static inline void omp_set_num_threads(int n){ (void)n; } * single-threaded"). */ 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 */ +static int g_moe_exact=0; /* output dello shared expert gia' calcolato su GPU (solo Metal layer-CB) */ static const float *g_pre_sh; #endif @@ -881,7 +888,7 @@ static void matmul_qt_ex(float *y, const float *x, QT *w, int S, int allow_idot) * fmt 1/2/4 in this build, so it fails CLOSED to the CPU branch below (matmul_fp8) * rather than being silently misread. coli_metal_gemm() also carries its own * internal fmt!=1&&fmt!=2&&fmt!=4 guard, so this is belt-and-braces. */ - if(g_metal_enabled && S>=g_metal_gemm_min && !spec_pinned() && (w->fmt==1||w->fmt==2||w->fmt==4) && !omp_in_parallel()){ + if(g_metal_enabled && S>=g_metal_gemm_min && !spec_pinned() && (w->fmt==1||w->fmt==2||(w->fmt==4&&!g_moe_exact)) && !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,w->gs)) return; } @@ -1769,6 +1776,29 @@ static void layer_cuda_shard_kvb(Layer *l,int H,int Q,int V){ } #endif +/* KV_B FMT-GATE NOTICE (#kvb): kernel contract — the fused Metal attention kernels + * (attention_rows' coli_metal_attn_decode and layer_forward_rows' coli_metal_layer_decode, + * both gated on `l->kv_b.fmt==2||l->kv_b.fmt==4`) only run against kv_b_proj stored INT4, + * either per-row (fmt=2) or grouped (fmt=4, since #587's kv_b grouped-int4 addition); any + * OTHER format silently falls through to the CPU absorb path for decode attention on that + * layer. v1-class mixed-precision containers can mint kv_b_proj at a format outside that + * pair, so this is a real trap, not a hypothetical: print it once (called from model_init, + * after all layer tensors are resolved) so it isn't silent. NOTICE ONLY — no gate/behavior + * change here. */ +static void kvb_fmt_gate_notice(Model *m){ + Cfg *c=&m->c; + if(!g_metal_enabled) return; + int bad=0, first_fmt=-1; + for(int i=0;in_layers;i++) if(m->L[i].kv_b.fmt!=2 && m->L[i].kv_b.fmt!=4){ + bad++; if(first_fmt<0) first_fmt=m->L[i].kv_b.fmt; } + if(bad) fprintf(stderr, + "[METAL] kv_b_proj is fmt=%d (not int4) on %d/%d layer%s: the fused Metal " + "attention path requires kv_b in int4, per-row (fmt=2) or grouped (fmt=4), so " + "those layers run decode attention on the CPU absorb path instead. Requantize " + "kv_b to int4 (per-row --kvb-bits 4, or grouped) to re-enable the fused path.\n", + first_fmt, bad, c->n_layers, bad==1?"":"s"); +} + static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits){ memset(m,0,sizeof(*m)); m->ebits=ebits; m->dbits=dbits; load_cfg(&m->c,snap); @@ -1844,6 +1874,7 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits } #undef P } + kvb_fmt_gate_notice(m); /* once per load, after all layer kv_b resolved */ /* testa MTP (layer n_layers): presente solo se convertita con --mtp */ { /* MTP attiva SOLO se il set e' COMPLETO (i tensori vivono su 3 shard: durante la @@ -3271,7 +3302,7 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p * flow through the shared per-fmt shader. */ if(g_metal_enabled && !kvs && g_absorb!=0 && (S<=4 || g_metal_prefill) && 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 + && c->qk_rope==64 && vh==256 && (l->kv_b.fmt==2||(l->kv_b.fmt==4&&!g_moe_exact)) && metal_fused_fmt_ok(l->q_a.fmt) && metal_fused_fmt_ok(l->q_b.fmt) && metal_fused_fmt_ok(l->kv_a.fmt) && metal_fused_fmt_ok(l->o.fmt)){ int sel_active = m->has_dsa && layern_layers && c->idx_type[layer] && (pos_base+S) > c->index_topk; @@ -3287,7 +3318,7 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p WP_(l->q_a), l->q_a.s, l->q_a.fmt, l->q_a.gs, l->q_a_ln, WP_(l->q_b), l->q_b.s, l->q_b.fmt, l->q_b.gs, WP_(l->kv_a), l->kv_a.s, l->kv_a.fmt, l->kv_a.gs, l->kv_a_ln, - WP_(l->kv_b), l->kv_b.s, l->kv_b.fmt, + WP_(l->kv_b), l->kv_b.s, l->kv_b.fmt, l->kv_b.gs, WP_(l->o), l->o.s, l->o.fmt, l->o.gs, m->Lc[layer], m->Rc[layer], S, pos_base, m->kv_start[layer], c->eps, c->theta, c->attn_scale, out); #undef WP_ @@ -3823,6 +3854,18 @@ static void metal_stage_rot_e8(float *mxg, const int *mrows, int p, int D){ } } #endif +/* moe()'s MB_BUILD subset builder (review F2): coli_metal_moe_block(_begin) take ONE + * fmt/qgs for an entire GPU batch of experts (+ optionally the fused shared expert), so + * a candidate member at (fmt,gs) is only safe to fold into a batch already established + * at (est_fmt,est_gs) if it can't disagree about layout: fmt=4 (grouped int4) is the + * only format with a group size to disagree ON, so this is a no-op (always compatible) + * for est_fmt!=4 (nothing established yet, or the batch isn't grouped) or fmt!=4 (the + * candidate itself isn't grouped -- its own fmt mismatch is a separate, pre-existing gap + * this guard does not cover, see MB_BUILD's comment). Pulled out of MB_BUILD as its own + * function so it's independently testable (see tests/test_moe_gs_guard.c). */ +static int mb_gs_compat(int est_fmt, int est_gs, int fmt, int gs){ + return !(est_fmt==4 && fmt==4 && gs!=est_gs); +} static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int with_shared){ if(g_pilot_real){ /* barriera cross-layer: prendi possesso di QUESTO layer e aspetta @@ -4253,19 +4296,37 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int * batch's rows all share s, so this is one FWHT per token in practice. The down * input is rotated on-GPU by moe_fwht inside moe_submit. */ 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; + int cpu_res=1, cpu_miss=1, mh_shared=0, nbb=0, Rtot=0, mfmt=-1, mgs=0, mgs_ok=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) */ + /* subset builder: experts with is_miss==WANTMISS (+ shared expert when TRY_SH). + * mgs_ok (review F2): moe_submit takes ONE fmt/qgs for the WHOLE batch, so a + * fmt=4 subset needs every member's group size to agree with the first expert's + * (mgs) -- fmt equality alone (already checked below for the shared expert) isn't + * enough once fmt=4 has a group size to disagree on. Routed experts get the same + * guard against each other (first-expert-gs consistency): a v1-class mixed- + * precision container could in principle mint routed experts at different gs + * within one layer (qt_resolve_fmt derives fmt/gs per-tensor from the file, with + * no uniformity enforced across experts at load time), even though a normal + * single-pass conversion never would. mgs_ok=0 does NOT drop the mismatched + * expert's rows: it leaves nbb's bookkeeping untouched and instead suppresses the + * GPU submit call at the two call sites below, which leaves cpu_res/cpu_miss at + * their initial 1 -- the same state a genuine GPU submission failure leaves them + * in, so the existing CPU fallback loop (metal_done false) redoes this whole + * nb-expert block correctly regardless of the mismatch. (Not checked, and out of + * scope for this guard: per-expert fmt agreement across g/u/d, and gs agreement + * across sh_gate/sh_up/sh_down or across a single expert's own g/u/d -- see + * WORKER_REPORT UNCERTAINTIES.) */ #define MB_BUILD(WANTMISS, TRY_SH) do{ \ - nbb=0; Rtot=0; mfmt=-1; sh_in=0; \ + nbb=0; Rtot=0; mfmt=-1; mgs=0; mgs_ok=1; sh_in=0; \ for(int j=0;jg.fmt; \ + if(mfmt<0){ mfmt=e->g.fmt; mgs=e->g.gs; } \ + else if(!mb_gs_compat(mfmt,mgs,e->g.fmt,e->g.gs)) mgs_ok=0; \ 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; \ @@ -4273,8 +4334,9 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int 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; \ + if(c->n_shared==1 && sI==I && l->sh_gate.fmt==shf && l->sh_up.fmt==shf && l->sh_down.fmt==shf \ + && mb_gs_compat(mfmt,mgs,shf,l->sh_gate.gs)){ \ + if(mfmt<0){ mfmt=shf; mgs=l->sh_gate.gs; } \ 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; \ @@ -4295,13 +4357,15 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int mxg=falloc((int64_t)(nb+1)*S*D); mrows=xalloc((size_t)(nb+1)*S*sizeof(int),"moe mrows"); mrw=xalloc((size_t)(nb+1)*S*sizeof(float),"moe mrw"); MB_BUILD(0, base==0 && !g_pre_sh); - if(nbb>0){ + if(nbb>0 && mgs_ok){ double t0=now_s(); if(mfmt==6) metal_stage_rot_e8(mxg,mrows,Rtot,D); - mh=coli_metal_moe_block_begin(nbb,D,I,mfmt,MG,MU,MD,MGS,MUS,MDS,mxg,xoffb,nrb,mrows,mrw); + mh=coli_metal_moe_block_begin(nbb,D,I,mfmt,mgs,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; + } else if(!nbb) cpu_res=0; /* nbb==0: nothing in this subset. nbb>0 && !mgs_ok + * (F2): gs-heterogeneous fmt=4 subset -- leave + * cpu_res=1 so the CPU loop below redoes it. */ } #endif /* Expert loads run HERE, after the resident-experts GPU submit above: under METAL the @@ -4358,12 +4422,12 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int for(int q=0;qt_ewait += now_s()-tw; } MB_BUILD(1, 0); /* missed experts, now loaded */ - if(nbb>0){ + if(nbb>0 && mgs_ok){ double t0=now_s(); if(mfmt==6) metal_stage_rot_e8(mxg,mrows,Rtot,D); - 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; + if(coli_metal_moe_block(nbb,D,I,mfmt,mgs,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; + } else if(!nbb) cpu_miss=0; /* see the resident-subset call site above */ 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; @@ -5549,7 +5613,7 @@ static void layer_forward_rows(Model *m, Layer *l, int li, float *x, int S, int if(g_metal_enabled && !kvs && S<=4 && lin_layers && l->sparse && (g_absorb==1||(g_absorb<0&&S<=4)) && m->kv_start[li]==0 && D==6144 && c->n_heads==64 && c->q_lora==2048 && c->kv_lora==512 - && c->qk_nope==192 && c->qk_rope==64 && c->v_head==256 && l->kv_b.fmt==2 + && c->qk_nope==192 && c->qk_rope==64 && c->v_head==256 && (l->kv_b.fmt==2||(l->kv_b.fmt==4&&!g_moe_exact)) && c->n_experts==256 && c->topk==8 && c->n_shared==1 && c->moe_inter==2048 && metal_fused_fmt_ok(l->q_a.fmt) && metal_fused_fmt_ok(l->q_b.fmt) && metal_fused_fmt_ok(l->kv_a.fmt) && metal_fused_fmt_ok(l->o.fmt) @@ -5568,7 +5632,7 @@ static void layer_forward_rows(Model *m, Layer *l, int li, float *x, int S, int WP_(l->q_a), l->q_a.s, l->q_a.fmt, l->q_a.gs, l->q_a_ln, WP_(l->q_b), l->q_b.s, l->q_b.fmt, l->q_b.gs, WP_(l->kv_a), l->kv_a.s, l->kv_a.fmt, l->kv_a.gs, l->kv_a_ln, - WP_(l->kv_b), l->kv_b.s, l->kv_b.fmt, + WP_(l->kv_b), l->kv_b.s, l->kv_b.fmt, l->kv_b.gs, WP_(l->o), l->o.s, l->o.fmt, l->o.gs, WP_(l->sh_gate), l->sh_gate.s, l->sh_gate.fmt, l->sh_gate.gs, WP_(l->sh_up), l->sh_up.s, l->sh_up.fmt, l->sh_up.gs, @@ -9324,6 +9388,7 @@ int main(int argc, char **argv){ 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")); + { const char *e=getenv("COLI_METAL_MOE_EXACT"); g_moe_exact=(e&&e[0]&&e[0]!='0'); } } #else if(getenv("COLI_METAL") && atoi(getenv("COLI_METAL"))){ diff --git a/c/sample.h b/c/sample.h index c1187cf36..332e13e77 100644 --- a/c/sample.h +++ b/c/sample.h @@ -5,6 +5,7 @@ #include #include +#include #include #include "tok.h" @@ -114,7 +115,22 @@ static int pick_tok(const float *lo, int V, int ban){ for (int k = 0; k < 5; k++) fprintf(stderr," %d:%.6f", id[k], v[k]); fprintf(stderr,"\n"); } - if (g_temp <= 0) return argmax_v(lo, V); + if (g_temp <= 0) { + int a = argmax_v(lo, V); + /* COLI_LOGIT_GAP=1: probe-only dump of the top-2 logits at each greedy + * pick, to tell a near-tie (float accumulation order) apart from a real + * divergence. Read-only: the token returned is unchanged. */ + static int gap_dbg = -1; static int gap_pos = 0; + if (gap_dbg < 0) { const char *e = getenv("COLI_LOGIT_GAP"); gap_dbg = (e && *e && *e != '0'); } + if (gap_dbg) { + int b = -1; float bv = -INFINITY; + for (int i = 0; i < V; i++) { float x = lo[i]; if (i != a && x == x && x > bv) { bv = x; b = i; } } + double t1 = (double)lo[a], t2 = (b >= 0) ? (double)bv : 0.0; + fprintf(stderr, "LOGITGAP pos=%d top1=%d:%.6f top2=%d:%.6f gap=%.9f\n", + gap_pos++, a, t1, b, t2, (b >= 0) ? t1 - t2 : 0.0); + } + return a; + } dist_build(lo, V); return dist_sample(V, ban); } diff --git a/c/tests/test_backend_metal.mm b/c/tests/test_backend_metal.mm index 2c1b4f105..6b66cd71b 100644 --- a/c/tests/test_backend_metal.mm +++ b/c/tests/test_backend_metal.mm @@ -269,7 +269,7 @@ static int run_fp8_moe_gate(const char *name) { const float *gs[1] = {(const float*)bad}, *us[1] = {(const float*)bad}, *ds[1] = {(const float*)bad}; float xg[8]={0}, out[8]={0}, rw[1]={1.0f}; int xoff[1]={0}, nr[1]={1}, rows[1]={0}; - int rc = coli_metal_moe_block(1, 8, 8, FP8, g, u, d, gs, us, ds, xg, xoff, nr, rows, rw, out, 1); + int rc = coli_metal_moe_block(1, 8, 8, FP8, 0, g, u, d, gs, us, ds, xg, xoff, nr, rows, rw, out, 1); int ok = (rc == 0); printf(" %-42s rc=%d (expect 0/CPU-fallback) %s\n", name, rc, ok?"ok":"*** MISMATCH (should have refused)"); return ok?0:1; @@ -279,42 +279,50 @@ static int run_fp8_moe_gate(const char *name) { static size_t roundpg(size_t n){ size_t p=16384; return ((n+p-1)/p)*p; } // Validate coli_metal_moe_block against a CPU reference (gate/up/silu/down + weighted scatter-add). -static int run_moe(const std::vector& nrv, const char* name) { - const int D=6144, I=2048, fmt=2; int rbG=(D+1)/2, rbD=(I+1)/2, nb=(int)nrv.size(); +// qgs==0 -> fmt=2 (per-row scale). qgs>0 -> fmt=4 grouped int4: per-expert scale slab is +// [O][ng] (ng=ceil(K/qgs)) for each of gate/up (K=D) and down (K=Iinter). +static int run_moe(const std::vector& nrv, int qgs, const char* name) { + const int D=6144, I=2048; int fmt = qgs>0 ? 4 : 2; + int rbG=(D+1)/2, rbD=(I+1)/2, nb=(int)nrv.size(); + int ngG = qgs>0 ? (D+qgs-1)/qgs : 1, ngD = qgs>0 ? (I+qgs-1)/qgs : 1; // scales/row for gate-up / down int R=0; std::vector xoff(nb),nr(nrv); for(int e=0;e0 ? s[(size_t)o*ngG + k/qgs] : s[o]; }; + auto scaD=[&](const float* s,int o,int k){ return qgs>0 ? s[(size_t)o*ngD + k/qgs] : s[o]; }; // per-expert page-aligned slab [Wg|Wu|Wd] and fslab [Sg|Su|Sd]; register both. std::vector slab(nb), fslab(nb); std::vector g(nb),u(nb),d(nb); std::vector gs(nb),us(nb),ds(nb); - size_t wlen=roundpg((size_t)I*rbG*2 + (size_t)D*rbD), flen=roundpg(((size_t)I*2+D)*sizeof(float)); + size_t nsc=(size_t)I*ngG*2 + (size_t)D*ngD; // gate + up + down scale counts + size_t wlen=roundpg((size_t)I*rbG*2 + (size_t)D*rbD), flen=roundpg(nsc*sizeof(float)); for(int e=0;e xg((size_t)R*D); for(auto&v:xg) v=((rand()%2000)-1000)/1000.f; std::vector rows(R); std::vector rw(R); for(int gr=0;gr position 0 int S=1; - // CPU reference + // CPU reference (grouped scale folded per-term; for fmt=2 that reduces to a*s[o]) std::vector refout((size_t)S*D,0.f), gg(I),uu(I),hh(D); for(int e=0;e gout((size_t)S*D,0.f); - int ok = coli_metal_moe_block(nb,D,I,fmt,g.data(),u.data(),d.data(),gs.data(),us.data(),ds.data(), + int ok = coli_metal_moe_block(nb,D,I,fmt,qgs,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& nrv, const char* name) { std::vector xg_gpu(xg); for(int gr=0;gr 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(), + int ok = coli_metal_moe_block(nb,D,I,fmt,0,g.data(),u.data(),d.data(),gs.data(),us.data(),ds.data(), xg_gpu.data(),xoff.data(),nr.data(),rows.data(),rw.data(),gout.data(),S); double maxabs=0,ymax=0; for(size_t i=0;i>1]; int v=(i&1)?(b>>4):(b&0xF); a+=(float)(v-8)*x[i]*scl[i/gs]; } y[o]=a; } } -static int run_attn(int S, int pos_base, const char* name){ +// kvb_gs==0 -> kv_b as fmt=2 (per-row scale); kvb_gs>0 -> kv_b as fmt=4 grouped int4 +// (exercises a_deqrow's grouped-scale path in a_qabs/a_ctx, #587's kv_b addition). +static int run_attn(int S, int pos_base, int kvb_gs, 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); + int kvb_fmt = kvb_gs>0 ? 4 : 2, kvng = kvb_gs>0 ? (TKVL+kvb_gs-1)/kvb_gs : 1; + TW qa=t_mkw(TQL,TH), qb=t_mkw(THH*TQH,TQL), kva=t_mkw(TKVL+TROPE,TH); + TW kvb = kvb_gs>0 ? t_mkw_g(THH*TROWSH,TKVL,kvb_gs) : t_mkw(THH*TROWSH,TKVL); + TW o=t_mkw(TH,THH*TVH); + // per-column kv_b scale: grouped (fmt=4) picks scale[row*ng + i/gs], else per-row. + auto kvb_sc=[&](int row,int i)->float{ return kvb_gs>0 ? kvb.s[(size_t)row*kvng + i/kvb_gs] : kvb.s[row]; }; std::vector 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; @@ -439,22 +454,22 @@ static int run_attn(int S, int pos_base, const char* name){ for(int h=0;h qabs(TKVL,0); - for(int d=0;d>1]; int v=(i&1)?(b>>4):(b&0xF); qabs[i]+=qp[d]*(float)(v-8)*sc; } } + for(int d=0;d>1]; int v=(i&1)?(b>>4):(b&0xF); qabs[i]+=qp[d]*(float)(v-8)*kvb_sc(rbase+d,i); } } std::vector 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 cl(TKVL,0); for(int t=0;t<=pos;t++){ const float*Lt=&Lr[(size_t)t*TKVL]; for(int i=0;i>1]; int vv=(i&1)?(b>>4):(b&0xF); v+=cl[i]*(float)(vv-8)*sc; } + for(int j=0;j>1]; int vv=(i&1)?(b>>4):(b&0xF); v+=cl[i]*(float)(vv-8)*kvb_sc(rbase+TNOPE+j,i); } 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 got((size_t)S*TH); int ok=coli_metal_attn_decode(x.data(), qa.w,qa.s,2,0,qaln.data(), qb.w,qb.s,2,0, - kva.w,kva.s,2,0,kvaln.data(), kvb.w,kvb.s,2, o.w,o.s,2,0, + kva.w,kva.s,2,0,kvaln.data(), kvb.w,kvb.s,kvb_fmt,kvb_gs, o.w,o.s,2,0, Lc,Rc,S,pos_base,0,eps,theta,ascale,got.data()); double ma=0,ym=0; for(size_t i=0;i0 cases below, not here. // // Two blind spots closed here (review round 1 -- see PR_BODY.md sec 10): // (a) pos_base must be >0 for any S=1 case. At pos_base=0, T=pos_base+S=1: softmax @@ -616,7 +632,7 @@ static int run_attn_grouped(int S, int pos_base, int gs, const char* name){ } std::vector got((size_t)S*TH); int ok=coli_metal_attn_decode(x.data(), qa.w,qa.s,4,gs,qaln.data(), qb.w,qb.s,2,0, // <- qa fmt=4/gs - kva.w,kva.s,2,0,kvaln.data(), kvb.w,kvb.s,2, o.w,o.s,2,0, + kva.w,kva.s,2,0,kvaln.data(), kvb.w,kvb.s,2,0, o.w,o.s,2,0, Lc,Rc,S,pos_base,0,eps,theta,ascale,got.data()); double ma=0,ym=0; for(size_t i=0;i1 (non- + // degenerate softmax), so the qabs (NOPE-side) dequant is exercised too, not just + // the always-live a_ctx (VH-side) dequant that a T=1 row alone would cover. + fail |= run_attn(1, 0, 128, "attn S=1 pos=0 kvb-fmt4-g128"); + fail |= run_attn(1, 37, 128, "attn S=1 pos=37 kvb-fmt4-g128"); + fail |= run_attn(4, 12, 128, "attn S=4 pos=12 kvb-fmt4-g128 (MTP)"); + fail |= run_attn(3, 0, 64, "attn S=3 pos=0 kvb-fmt4-g64"); printf("Metal fused attention tests (fmt=4 grouped q_a, proves bind_gemv gs plumbing):\n"); // pos_base=37 (not 0): at T=1 softmax is identically 1.0 regardless of q_a's output, // so an S=1 pos=0 case cannot catch ANY q_a defect (review round 1, auditor -- see diff --git a/c/tests/test_kvb_notice.c b/c/tests/test_kvb_notice.c new file mode 100644 index 000000000..f21e0d817 --- /dev/null +++ b/c/tests/test_kvb_notice.c @@ -0,0 +1,188 @@ +/* kv_b fmt-gate NOTICE test (see kvb_fmt_gate_notice in colibri.c, called once from + * model_init after all layer tensors are resolved). + * + * The trap: a v1-class mixed-precision container can mint kv_b_proj at a format outside + * the pair the fused Metal attention kernels actually serve. BOTH fused entries + * (attention_rows' fused decode and layer_forward_rows' fused layer decode) hard-gate on + * `l->kv_b.fmt==2||l->kv_b.fmt==4` (per-row or grouped int4 -- fmt=4 became a valid + * GPU-served kv_b configuration once #587's grouped-int4 kv_b support landed), so a kv_b + * outside {2,4} silently forces that layer's decode attention onto the CPU absorb path -- + * with no notice at all if this fires. This test proves the fix (a stderr notice, NOT a + * behavior change): it fires exactly once when it should, and stays silent for BOTH valid + * formats (fmt=2 and fmt=4) as well as for the MTP head (excluded by design: the notice + * only scans m->L[0..n_layers), never m->mtpL). It calls kvb_fmt_gate_notice() directly + * against a hand-built Model, so it needs no real snapshot on disk and no Metal backend + * link -- g_metal_enabled is declared unconditionally in colibri.c precisely so this + * stays portable. + * + * Portable stderr capture: freopen(stderr) onto a temp file, dup()/dup2() to save and + * restore the real fd 2 around it. Deliberately no child-process spawn/reap primitives + * anywhere in this file -- those are POSIX-only and the earlier version of this kind of + * test broke Windows CI. */ +#define main coli_glm_main_unused +#include "../colibri.c" +#undef main + +#include +#include +#include +#include + +static int fails = 0; +#define CHECK(c) do{ if(!(c)){ printf("FAIL %s:%d: %s\n", __FILE__, __LINE__, #c); fails++; } }while(0) + +/* Redirects stderr to `path` (truncated, read+write) and returns a dup of the original + * fd 2 so restore_stderr() can put it back. */ +static int redirect_stderr(const char *path){ + fflush(stderr); + int saved = dup(fileno(stderr)); + if(saved<0){ printf("FAIL: dup(stderr) failed\n"); fails++; } + if(!freopen(path, "w+", stderr)){ printf("FAIL: freopen(%s) failed\n", path); fails++; } + return saved; +} +/* Reads back everything written to stderr since redirect_stderr(), restores the real + * stderr fd, and NUL-terminates the captured text into buf. */ +static void restore_stderr(int saved, char *buf, size_t bufsz){ + fflush(stderr); + long n = ftell(stderr); + if(n<0) n=0; + rewind(stderr); + size_t want = (size_t)n < bufsz-1 ? (size_t)n : bufsz-1; + size_t got = fread(buf, 1, want, stderr); + buf[got] = 0; + fflush(stderr); + dup2(saved, fileno(stderr)); /* fd 2 back to the real stderr */ + close(saved); +} + +static Model make_model(int n_layers){ + Model m; memset(&m,0,sizeof m); + m.c.n_layers = n_layers; + m.L = calloc((size_t)n_layers, sizeof(Layer)); + return m; +} + +int main(void){ + char buf[4096]; + + /* (a) fmt=1 kv_b on 2/3 layers (the other is fmt=2) + Metal enabled -> notice fires + * exactly once, for a format OUTSIDE {2,4} -- fmt=1 (v1-class INT8 container) is + * chosen specifically because it's neither of the two now-valid GPU-served formats. */ + { + Model m = make_model(3); + m.L[0].kv_b.fmt = 1; /* v1-class INT8 container: outside {2,4} */ + m.L[1].kv_b.fmt = 2; + m.L[2].kv_b.fmt = 1; + g_metal_enabled = 1; + + int saved = redirect_stderr("tests/tmp_kvb_notice_a.stderr"); + kvb_fmt_gate_notice(&m); + restore_stderr(saved, buf, sizeof buf); + remove("tests/tmp_kvb_notice_a.stderr"); + + int occurrences = 0; const char *p = buf; + while((p = strstr(p, "[METAL] kv_b_proj"))){ occurrences++; p++; } + CHECK(occurrences == 1); /* exactly one notice, not one per layer */ + CHECK(strstr(buf, "fmt=1") != NULL); /* which fmt kv_b actually has */ + CHECK(strstr(buf, "2/3 layer") != NULL); /* count of affected layers */ + CHECK(strstr(buf, "fmt=2") != NULL); /* mentions the per-row option */ + CHECK(strstr(buf, "fmt=4") != NULL); /* mentions the grouped option -- fmt=4 is now valid */ + CHECK(strstr(buf, "CPU absorb path") != NULL); /* attention runs the CPU path */ + CHECK(strstr(buf, "--kvb-bits 4") != NULL); /* remedy hint */ + /* the pre-#587 message's PER-ROW-only warning and fmt=4-trap sentence must be + * gone -- fmt=4 is a valid GPU-served kv_b now, not a second way to miss the gate. */ + CHECK(strstr(buf, "PER-ROW") == NULL); + CHECK(strstr(buf, "keeping kv_b ungrouped") == NULL); + CHECK(strstr(buf, "also misses this fmt=2 gate") == NULL); + free(m.L); + } + + /* (b) fmt=2 kv_b on every layer -> no notice, even with Metal enabled */ + { + Model m = make_model(3); + for(int i=0;i<3;i++) m.L[i].kv_b.fmt = 2; + g_metal_enabled = 1; + + int saved = redirect_stderr("tests/tmp_kvb_notice_b.stderr"); + kvb_fmt_gate_notice(&m); + restore_stderr(saved, buf, sizeof buf); + remove("tests/tmp_kvb_notice_b.stderr"); + + CHECK(buf[0] == 0); + free(m.L); + } + + /* (b2) fmt=4 (grouped int4) kv_b on every layer -> no notice either: fmt=4 is a valid + * GPU-served kv_b configuration post-#587, not a trap (this is the semantic change + * under test -- the pre-#587 notice fired here). */ + { + Model m = make_model(3); + for(int i=0;i<3;i++) m.L[i].kv_b.fmt = 4; + g_metal_enabled = 1; + + int saved = redirect_stderr("tests/tmp_kvb_notice_b2.stderr"); + kvb_fmt_gate_notice(&m); + restore_stderr(saved, buf, sizeof buf); + remove("tests/tmp_kvb_notice_b2.stderr"); + + CHECK(buf[0] == 0); + free(m.L); + } + + /* (b3) mixed fmt=2/fmt=4 across layers -> still no notice (the OR gate, not just + * either format alone, is what the fused path actually requires per-layer). */ + { + Model m = make_model(4); + m.L[0].kv_b.fmt = 2; m.L[1].kv_b.fmt = 4; m.L[2].kv_b.fmt = 4; m.L[3].kv_b.fmt = 2; + g_metal_enabled = 1; + + int saved = redirect_stderr("tests/tmp_kvb_notice_b3.stderr"); + kvb_fmt_gate_notice(&m); + restore_stderr(saved, buf, sizeof buf); + remove("tests/tmp_kvb_notice_b3.stderr"); + + CHECK(buf[0] == 0); + free(m.L); + } + + /* (c) fmt=1 kv_b but Metal not enabled -> no notice (the fused paths this warns + * about are Metal-only, so with Metal off there is nothing to warn about) */ + { + Model m = make_model(2); + m.L[0].kv_b.fmt = 1; m.L[1].kv_b.fmt = 1; + g_metal_enabled = 0; + + int saved = redirect_stderr("tests/tmp_kvb_notice_c.stderr"); + kvb_fmt_gate_notice(&m); + restore_stderr(saved, buf, sizeof buf); + remove("tests/tmp_kvb_notice_c.stderr"); + + CHECK(buf[0] == 0); + free(m.L); + } + + /* (d) MTP-exclusion: the notice scans m->L[0..n_layers) only. A bad format planted + * on m->mtpL (the separate MTP-head Layer, never part of m->L) must NOT be counted or + * reported, even with every real layer at a valid format and Metal enabled -- the + * fused layer-decode gate itself already excludes the MTP head (lish_gate.fmt==shf`), but for fmt=4 + * (grouped int4) format equality alone isn't enough: two fmt=4 tensors can still disagree + * on GROUP SIZE, and moe_submit has no per-member gs -- every expert in the batch is + * dequantized with the SAME qgs, so a silently-included gs mismatch would corrupt every + * row belonging to the odd one out. mb_gs_compat() is the guard: it decides whether a + * candidate (fmt,gs) is safe to fold into a batch already established at (est_fmt,est_gs). + * + * This test exercises mb_gs_compat() directly rather than moe() end-to-end: MB_BUILD is a + * macro #define'd (and #undef'd) entirely inside moe()'s own function body, so it isn't + * reachable from outside moe() even via the "#include colibri.c" seam this file uses -- + * and moe() itself needs a fully-populated Model/Layer/ESlot/ecache/router-output + * apparatus no existing test constructs, which would make a true end-to-end harness + * heavy and fragile relative to what it proves. mb_gs_compat() is not a re-implementation + * of MB_BUILD's guard logic -- it IS that logic (MB_BUILD calls it directly, at both the + * routed-expert and the shared-expert call sites) -- so exercising it directly is a + * faithful, low-risk proxy for the integration behavior. See WORKER_REPORT_integ_postwave.md + * UNCERTAINTIES for the case against the heavier alternative. + * + * Also documents, by assertion (not just prose), the two gaps this guard deliberately + * does NOT cover: per-expert fmt heterogeneity (g vs g of a different expert), and + * per-tensor-role gs heterogeneity within one expert/shared-expert triple (g vs u vs d). + * Both are pre-existing, out of scope for F2 -- see the guard's own comment in colibri.c. */ +#define main coli_glm_main_unused +#include "../colibri.c" +#undef main + +#include + +int main(void){ + /* --- nothing established yet (est_fmt<0): always compatible, any (fmt,gs) --- */ + assert(mb_gs_compat(-1,0, 4,128)==1 && "first member of a subset is always accepted"); + assert(mb_gs_compat(-1,0, 2,0)==1 && "first member, non-grouped fmt"); + + /* --- established fmt=4: SAME gs is compatible (the common, correct case) --- */ + assert(mb_gs_compat(4,128, 4,128)==1 && "matching gs within a fmt=4 subset: ok"); + assert(mb_gs_compat(4,64, 4,64)==1 && "matching gs (different value): ok"); + + /* --- established fmt=4: DIFFERENT gs is the F2 trap -- must be rejected --- */ + assert(mb_gs_compat(4,128, 4,64)==0 && "gs mismatch within a fmt=4 subset: must reject"); + assert(mb_gs_compat(4,64, 4,128)==0 && "gs mismatch, reversed order: must reject"); + assert(mb_gs_compat(4,128, 4,127)==0 && "off-by-one gs mismatch: must reject (no slack)"); + + /* --- established fmt!=4 (no group size to disagree on): always compatible, + * REGARDLESS of what the candidate claims -- fmt heterogeneity itself is a + * separate, pre-existing gap this guard does not cover (documented, not fixed: + * MB_BUILD has never checked per-expert fmt agreement, fmt=4 or otherwise). --- */ + assert(mb_gs_compat(2,0, 4,999)==1 && "est fmt=2 (ungrouped): gs never checked"); + assert(mb_gs_compat(1,0, 4,999)==1 && "est fmt=1 (int8): gs never checked"); + + /* --- candidate fmt!=4 against an established fmt=4 batch: also out of scope + * (the candidate's OWN fmt mismatch is what would corrupt it, not gs -- and + * fmt agreement for the shared expert is still checked separately by MB_BUILD's + * TRY_SH condition before mb_gs_compat is ever consulted). --- */ + assert(mb_gs_compat(4,128, 2,0)==1 && "candidate fmt=2 against fmt=4 batch: gs moot"); + + printf("OK test_moe_gs_guard: mb_gs_compat (review F2) accepts matching gs, " + "rejects mismatched gs, documents its fmt-heterogeneity non-scope\n"); + return 0; +}