Skip to content

metal: GPU attention for S>4 prefill (single command buffer) #763

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JustVugg merged 1 commit into from
Aug 2, 2026
+83 −22
Merged

metal: GPU attention for S>4 prefill (single command buffer) #763

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 67 additions & 20 deletions c/backend_metal.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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]; };
Expand All @@ -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;
}
Expand All @@ -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; }
Expand Down Expand Up @@ -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];
Expand Down
7 changes: 5 additions & 2 deletions c/colibri.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 &&
Expand Down Expand Up @@ -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.
Expand Down
11 changes: 11 additions & 0 deletions docs/metal.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading