mirror of
https://github.com/JustVugg/colibri.git
synced 2026-10-02 02:54:37 +08:00
metal: GPU attention for S>4 prefill (single command buffer, opt-in)
Lift the S<=4 cap on coli_metal_attn_decode so large-S prefill attention can run on the GPU. Projections + attention core + output GEMV run in ONE command buffer (ordered by memory barriers), the same structure the S<=4 decode path already used. On a 544-token prompt this cuts prefill attention ~4x (35.9s -> 9.0s). Off by default (COLI_METAL_PREFILL=0): like the prefill GEMM, the GPU accumulates in a different order and can pick a different top token on near-tie logits (#622 family), so a greedy stream is not guaranteed bit-identical to the CPU. COLI_METAL_PREFILL=1 opts in. Documented in docs/metal.md. A single a_score dispatch is S*AHEADS*T threads; guarded under ~2^30 and falls back to CPU for giant prompts (in-encoder row-chunking is a follow-up). CPU and CUDA paths unchanged.
This commit is contained in:
Doug Sharp
+67
-20
@@ -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<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 size_t ax_cap,aqr_cap,aqf_cap,acomp_cap,aqabs_cap,aclat_cap,actx_cap,aout_cap;
|
||||
static id<MTLBuffer> 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<MTLComputeCommandEncoder> 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<MTLComputeCommandEncoder> 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<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;
|
||||
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<MTLComputeCommandEncoder> 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<MTLComputeCommandEncoder> e,
|
||||
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 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<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 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<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);
|
||||
|
||||
// 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<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;
|
||||
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];
|
||||
|
||||
+5
-2
@@ -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;i<c->n_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 && i<c->n_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=<n> -> IDOT int4 solo per S>=n.
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user