From d96e4254a4ab666513e435fccca9c63b7a71aa1c Mon Sep 17 00:00:00 2001 From: woolcoxm <13604288+woolcoxm@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:07:54 -0400 Subject: [PATCH] fix: apply per-group scales in CUDA attention kernels for fmt=4 (kv_b crash) PR #298 added fmt=4 (grouped int4, gs=64) support to quant_matmul but the two MLA attention absorb kernels kept the per-row scale semantic (wscale[row]) from fmt=2. The g64 model's kv_b is fmt=4 (ng=8 groups/row), so COLI_CUDA_ATTN=1 / COLI_CUDA_PIPE=2 routed it through attention_absorb* which indexed the O*ng scale array with a row index -> wrong stride -> GPU memory fault -> bugcheck 0x116 VIDEO_TDR_FAILURE -> reboot. Add absorb_scale() (mirrors quant_matmul's fmt==4 branch: wscale[row*ng+k/gs] for fmt=4, wscale[row] otherwise) and apply it inside the Q- and V-projection accumulation loops of both attention_absorb_kernel and attention_absorb_batch_kernel. Thread w->gs/w->ng through all six launch sites. No extern-C signature or header changes; the tensor already carries gs/ng from tensor_upload. For fmt!=4 ng==1 so k/gs==0 and the result is bit-identical to before. Validated: COLI_CUDA_ATTN=1 and COLI_CUDA_PIPE=2 (long-prompt prefill, the exact crash config) now run clean; base path unchanged. --- c/backend_cuda.cu | 40 +++++++++++++++++++++++++++------------- 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/c/backend_cuda.cu b/c/backend_cuda.cu index 30a1f4a..84c04de 100644 --- a/c/backend_cuda.cu +++ b/c/backend_cuda.cu @@ -105,6 +105,18 @@ __device__ static float weight_at(const void *weights, int fmt, size_t row, int return static_cast(((v >> ((i & 3) * 2)) & 3) - 2); } +/* Scale for output `row`, input element `k`. fmt=4 (grouped int4) stores ng + * scales per row at scales[row*ng + k/gs]; every other quantized format has + * one scale per row at scales[row]. Mirrors quant_matmul's fmt==4 branch so the + * attention absorb kernels apply per-group scales instead of the per-row + * (fmt=2) semantic that crashed #298's g64 kv_b. */ +__device__ static float absorb_scale(const float *wscale, int fmt, int gs, int ng, int row, int k) { + if (!fmt) return 1.f; + if (fmt != 4) return wscale[row]; + int g = k / gs; if (g >= ng) g = ng - 1; /* tail of the last (partial) group */ + return wscale[(size_t)row * ng + g]; +} + __global__ static void offset_to_signed_s4(uint8_t *q,size_t n){ size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x;if(i=S||nt<1)return; extern __shared__ float sm[];float *qa=sm,*cl=qa+K,*scores=cl+K,*red=scores+T; const float *qs=q+((size_t)s*H+h)*(Q+R); for(int k=tid;kar,rope,rb,cudaMemcpyHostToDevice,dc->stream),"attention rope upload"))return 0; size_t shared=(size_t)(2*K+T)*sizeof(float); attention_absorb_kernel<<stream>>>(dc->ac,dc->aq,dc->al,dc->ar,w->weights,w->scales, - w->fmt,H,Q,R,V,K,T,scale); + w->fmt,H,Q,R,V,K,T,scale,w->gs,w->ng); if(!cuda_ok(cudaGetLastError(),"attention absorb launch")|| !cuda_ok(cudaMemcpyAsync(ctx,dc->ac,cb,cudaMemcpyDeviceToHost,dc->stream),"attention context download")|| !cuda_ok(cudaStreamSynchronize(dc->stream),"attention synchronize"))return 0; @@ -1002,7 +1016,7 @@ static int attention_absorb_batch_run(ColiCudaTensor *w,ColiCudaTensor *proj,flo !cuda_ok(cudaMemcpyAsync(dc->ar,rope,rb,cudaMemcpyHostToDevice,dc->stream),"attention batch rope upload"))return 0; size_t shared=(size_t)(2*K+T+256)*sizeof(float); attention_absorb_batch_kernel<<stream>>>(dc->ac,dc->aq,dc->al, - dc->ar,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale); + dc->ar,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale,w->gs,w->ng); if(!cuda_ok(cudaGetLastError(),"attention batch launch"))return 0; const float *src=dc->ac;size_t ob=cb; if(proj){ @@ -1350,7 +1364,7 @@ extern "C" int coli_cuda_attention_project_batch_dev(ColiCudaTensor *w,ColiCudaT if(!reserve(&dc->ac,&dc->ac_cap,cb))return 0; size_t shared=(size_t)(2*K+T+256)*sizeof(float); attention_absorb_batch_kernel<<stream>>>(dc->ac,q_dev,latent_dev, - rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale); + rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale,w->gs,w->ng); if(!cuda_ok(cudaGetLastError(),"pipe attention launch"))return 0; size_t ob=(size_t)S*proj->O*sizeof(float); if(!reserve(&dc->y,&dc->y_cap,ob))return 0; @@ -1414,7 +1428,7 @@ extern "C" int coli_cuda_attention_project_batch_dev_out(ColiCudaTensor *w,ColiC if(!reserve(&dc->ac,&dc->ac_cap,cb))return 0; size_t shared=(size_t)(2*K+T+256)*sizeof(float); attention_absorb_batch_kernel<<stream>>>(dc->ac,q_dev,latent_dev, - rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale); + rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale,w->gs,w->ng); if(!cuda_ok(cudaGetLastError(),"pipe attention launch (dev out)"))return 0; quant_matmul<<O,S),256,0,dc->stream>>>(out_dev,dc->ac,proj->weights, proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I),proj->gs,proj->ng); @@ -1433,7 +1447,7 @@ extern "C" int coli_cuda_attention_absorb_batch_dev(ColiCudaTensor *w,float *ctx DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0; size_t shared=(size_t)(2*K+T+256)*sizeof(float); attention_absorb_batch_kernel<<stream>>>(ctx_dev,q_dev,latent_dev, - rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale); + rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale,w->gs,w->ng); if(!cuda_ok(cudaGetLastError(),"pipe shard attention launch"))return 0; return cuda_ok(cudaStreamSynchronize(dc->stream),"pipe shard attention sync"); } @@ -1451,7 +1465,7 @@ extern "C" int coli_cuda_attention_absorb_kvdev(ColiCudaTensor *w,float *ctx,con if(!cuda_ok(cudaMemcpyAsync(dc->aq,q,qb,cudaMemcpyHostToDevice,dc->stream),"kvdev q upload"))return 0; size_t shared=(size_t)(2*K+T+256)*sizeof(float); attention_absorb_batch_kernel<<stream>>>(dc->ac,dc->aq,latent_dev, - rope_dev,w->weights,w->scales,w->fmt,1,H,Q,R,V,K,T,scale); + rope_dev,w->weights,w->scales,w->fmt,1,H,Q,R,V,K,T,scale,w->gs,w->ng); if(!cuda_ok(cudaGetLastError(),"kvdev absorb launch")|| !cuda_ok(cudaMemcpyAsync(ctx,dc->ac,cb,cudaMemcpyDeviceToHost,dc->stream),"kvdev ctx download")|| !cuda_ok(cudaStreamSynchronize(dc->stream),"kvdev absorb sync"))return 0;