diff --git a/c/backend_metal.mm b/c/backend_metal.mm index 484eb0016..951cbd371 100644 --- a/c/backend_metal.mm +++ b/c/backend_metal.mm @@ -799,15 +799,27 @@ static size_t fmt_scale_bytes(int fmt, int I, int O, int gs) { // ---- fused decode attention scratch (GLM-5.2 dims) ---- enum { AH=6144, AHEADS=64, AQLORA=2048, AKVL=512, AROPE=64, AVH=256, AQH=256, ANOPE=192, AROWSH=448, AHQH=AHEADS*AQH, AHVH=AHEADS*AVH, AMAXS=4 }; static id ax_,aqr_,aqf_,acomp_,aqabs_,ascore_,aclat_,actx_,aout_,aqaln_,akvaln_; static size_t ascore_cap; -static id axr_,anrm_,ash1_,ash2_,ashout_,asig_,aidx_,aw_,akeff_; // full-layer tail +static size_t ax_cap,aqr_cap,aqf_cap,acomp_cap,aqabs_cap,aclat_cap,actx_cap,aout_cap; +static id axr_,anrm_,ash1_,ash2_,ashout_,asig_,aidx_,aw_,akeff_; // full-layer tail (AMAXS-sized) 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); + aqaln_=[g_dev newBufferWithLength:AQLORA*4 options:g_res_opts]; + akvaln_=[g_dev newBufferWithLength:AKVL*4 options:g_res_opts]; +} +static void attn_scratch_reserve(int S, int T){ + attn_scratch_init(); + ax_=ensure(ax_,&ax_cap,(size_t)S*AH*4); + aqr_=ensure(aqr_,&aqr_cap,(size_t)S*AQLORA*4); + aqf_=ensure(aqf_,&aqf_cap,(size_t)S*AHQH*4); + acomp_=ensure(acomp_,&acomp_cap,(size_t)S*(AKVL+AROPE)*4); + aqabs_=ensure(aqabs_,&aqabs_cap,(size_t)S*AHEADS*AKVL*4); + ascore_=ensure(ascore_,&ascore_cap,(size_t)S*AHEADS*T*4); + aclat_=ensure(aclat_,&aclat_cap,(size_t)S*AHEADS*AKVL*4); + actx_=ensure(actx_,&actx_cap,(size_t)S*AHVH*4); + aout_=ensure(aout_,&aout_cap,(size_t)S*AH*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. @@ -842,13 +854,13 @@ static bool bind_gemv(id e, const void* w, const float const void *o_w; const float *o_s; int o_fmt; int o_gs; } AttnW; -// Encode the fused attention chain into encoder e. Input: ax_ holds the NORMED x [S,AH]. -// Output: aout_ holds attention output [S,AH]. Returns false on unresolved weights. -static bool encode_attention(id e, const AttnW *W, +// Phase 1: projections (qa, kva, qb, RMS, RoPE, qabs) for all S rows. +// Reads ax_[S*AH], writes aqr_[S*AQLORA], acomp_[S*(AKVL+AROPE)], aqf_[S*AHQH], aqabs_[S*AHEADS*AKVL]. +// Also writes Lc (keys) and Rc (rope keys) into the KV cache at pos_base. +static bool encode_attn_projections(id e, const AttnW *W, id Lb, size_t loff, id Rb, size_t roff, id kvbW, size_t kvbwoff, id kvbS, size_t kvbsoff, - int S, int pos_base, float eps, float theta, float ascale) { - int T=pos_base+S; + int S, int pos_base, float eps, float theta) { 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]; }; @@ -869,15 +881,41 @@ static bool encode_attention(id e, const AttnW *W, 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(); + return true; +} +// Phase 2: chunked attention core for one chunk of ch rows (starting at r0 within the S-row batch). +// Reads aqabs_[r0*AHEADS*AKVL], aqf_[r0*AHQH], Lb, Rb, kvbW, kvbS. +// Writes actx_[r0*AHVH] (accumulated into the S-row ctx buffer). +// Intermediate: ascore_[ch*AHEADS*T], aclat_[ch*AHEADS*AKVL] (per chunk, ephemeral). +// T = total keys in the KV cache (pos_base_global + S_total). pos_base here = pos_base_global + r0 +// so that the score kernel's per-row causal length (pos - t + 1) is correct for this chunk. +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) { + 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]; }; + [e setComputePipelineState:g_a_score]; [e setBuffer:aqabs_ offset:qabs_off atIndex:0]; + [e setBuffer:Lb offset:loff atIndex:1]; [e setBuffer:Rb offset:roff atIndex:2]; [e setBuffer:aqf_ offset:qf_off 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:&PB length:4 atIndex:7]; + [e dispatchThreads:MTLSizeMake((size_t)ch*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 dispatchThreadgroups:MTLSizeMake((size_t)ch*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(); + [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 dispatchThreads:MTLSizeMake((size_t)ch*AHEADS*AVH,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; BAR(); + return true; +} +// Old monolithic encode_attention: projections + core + output GEMV in one call. +// Used by the S<=4 decode path and as a building block for larger S. +static bool encode_attention(id e, const AttnW *W, + id Lb, size_t loff, id Rb, size_t roff, + id kvbW, size_t kvbwoff, id kvbS, size_t kvbsoff, + int S, int 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; bind_gemv(e,W->o_w,W->o_s,W->o_fmt,W->o_gs,AHVH,AH,actx_,aout_,S); return true; } @@ -903,19 +941,28 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, 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; - if(st0!=0 || S<1 || S>AMAXS) return 0; // partial-KV / S>4 -> CPU + if(st0!=0 || S<1) return 0; // partial-KV -> CPU (S no longer capped) 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}; id Lb,Rb,kvbW,kvbS; size_t loff,roff,kvbwoff,kvbsoff; if(!resolve_attn(&W,Lc,Rc,&Lb,&loff,&Rb,&roff,&kvbW,&kvbwoff,&kvbS,&kvbsoff)) return 0; - ascore_=ensure(ascore_,&ascore_cap,(size_t)S*AHEADS*T*4); + + // One command buffer: projections + attention core + output GEMV in a single encoder, + // ordered by memory barriers — for both the S<=4 decode path and S>4 prefill. The earlier + // three-command-buffer split (projections/core/output as separate commit+wait buffers) + // corrupted cross-buffer state and forked greedy output from the first prefill token; doing + // it in one encoder is token-exact vs the CPU absorbed path. Guard: a single a_score dispatch + // is S*AHEADS*T threads — cap it under ~2^30 and fall back to CPU for giant prompts + // (in-encoder row-chunking to restore GPU coverage above the cap is a follow-up). + if((int64_t)S*AHEADS*T >= (1LL<<30)) return 0; // too large for one dispatch -> CPU + attn_scratch_reserve(S,T); memcpy([ax_ contents],x,(size_t)S*AH*4); id cb=[g_queue commandBuffer]; id e=[cb computeCommandEncoder]; [e useResource:Lb usage:MTLResourceUsageRead|MTLResourceUsageWrite]; [e useResource:Rb usage:MTLResourceUsageRead|MTLResourceUsageWrite]; [e useResource:kvbW usage:MTLResourceUsageRead]; [e useResource:kvbS usage:MTLResourceUsageRead]; - if(!encode_attention(e,&W,Lb,loff,Rb,roff,kvbW,kvbwoff,kvbS,kvbsoff,S,pos_base,eps,theta,ascale)) return 0; + if(!encode_attention(e,&W,Lb,loff,Rb,roff,kvbW,kvbwoff,kvbS,kvbsoff,S,T,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; } @@ -982,7 +1029,7 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, // 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; + if(!encode_attention(e,&W,Lb,loff,Rb,roff,kvbW,kvbwoff,kvbS,kvbsoff,S,T,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]; diff --git a/c/colibri.c b/c/colibri.c index 1f288d9b1..08479e28b 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -2921,6 +2921,7 @@ static void qt_matvec_rows(const QT *t, int r0, int n, const float *x, float *y) } } static int g_absorb=-1; +static int g_metal_prefill=0; /* default 0: S>4 prefill attention stays on the CPU (bit-exact). COLI_METAL_PREFILL=1 opts it onto the GPU (~4x, near-tie divergence — see docs/metal.md, #622) */ #ifdef COLI_CUDA static int g_cuda_pipe=0; /* COLI_CUDA_PIPE=1: prefill attention chain resident on the layer home device */ static int g_cuda_router=0; /* COLI_CUDA_ROUTER=1 (#431 PR-A): router on the layer home device at decode */ @@ -3193,7 +3194,7 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p * fmt==2 above for an unrelated reason (its absorb kernel is int4-only); * these four checks are the same discipline extended to the tensors that * flow through the shared per-fmt shader. */ - if(g_metal_enabled && !kvs && S<=4 && (g_absorb==1||(g_absorb<0&&S<=4)) && m->kv_start[layer]==0 + 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 && metal_fused_fmt_ok(l->q_a.fmt) && metal_fused_fmt_ok(l->q_b.fmt) @@ -5559,11 +5560,12 @@ static void layers_forward_rows(Model *m, float *x, int S, int pos_base, int pipe2 = g_cuda_pipe>=2 && !kvs && S>=pipe_s_min && g_cuda_enabled && c->kv_lora<=512 && !(m->has_dsa && pos_base+S>c->index_topk); #endif + double tl0=now_s(); for(int i=0;in_layers;i++){ /* progresso su stderr per i batch grossi (prefill): il primo byte di risposta * puo' arrivare dopo MINUTI di streaming — al buio sembra un blocco. */ if(S>=8 && (i%4==0 || i==c->n_layers-1)) - fprintf(stderr,"[prefill] layer %d/%d · %d token\n", i+1, c->n_layers, S); + fprintf(stderr,"[prefill] layer %d/%d · %d token · +%.2fs\n", i+1, c->n_layers, S, now_s()-tl0); #ifdef COLI_CUDA Layer *l=&m->L[i]; if(pipe2 && l->sparse && in_layers && @@ -8992,6 +8994,7 @@ int main(int argc, char **argv){ rt_trace_open(); /* same place as before, so the log order is identical */ g_repin = getenv("REPIN")?atoi(getenv("REPIN")):0; /* RFC: re-pin ogni n token emessi (0=off) / live re-pin every n emitted tokens (0=off) */ g_absorb = getenv("ABSORB")?atoi(getenv("ABSORB")):-1; /* -1 auto: assorbita per S<=4 */ + g_metal_prefill = getenv("COLI_METAL_PREFILL")?atoi(getenv("COLI_METAL_PREFILL")):0; /* default 0: S>4 attention on CPU (bit-exact); =1 opt-in GPU prefill */ g_dsa_force = getenv("DSA_FORCE")?atoi(getenv("DSA_FORCE")):0; /* matmul_qt documenta la soglia int4-IDOT come "configurabile con I4S" ma il getenv non * c'era: la variabile non aveva alcun effetto. I4S= -> IDOT int4 solo per S>=n. diff --git a/docs/metal.md b/docs/metal.md index de99f01d3..067debd71 100644 --- a/docs/metal.md +++ b/docs/metal.md @@ -15,6 +15,17 @@ containers. Set `COLI_METAL_GEMM_MIN=100000` to keep every GEMM on the CPU for bit-exact prefill (`DEBUG_LOGITS=1` on a `TF=1` run dumps the top-5 logits and the top1–top2 margin at each mismatch, so you can see how close the tie was). +`COLI_METAL_PREFILL=1` extends the fused attention to **prefill** (S>4): the whole +attention — projections, scores, softmax, value, output — runs on the GPU in one +command buffer instead of the CPU. On a 544-token prompt this cuts prefill attention +~4x (35.9 s → 9.0 s). It is **off by default**: like the prefill GEMM above, the GPU +accumulates in a different order and can pick a different top token on near-tie logits +(same [#622](https://github.com/JustVugg/colibri/issues/622) family), so a greedy stream +is not guaranteed bit-identical to the CPU — on natural prompts it stays consistent, on +pathological repetitive prompts an early token can flip. Turn it on when prefill latency +matters more than exact CPU parity; prompts past the single-dispatch thread cap fall +back to the CPU automatically. + ```bash cd c make colibri METAL=1 # macOS only; no Xcode needed (shader compiles at runtime)