url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #365 from ZacharyZcR/feat/ragged-batch-attention-upstream

Browse Source 
cuda: batch ragged attention across independent streams
This commit is contained in:
Vincenzo Fornaro
2026-07-18 02:04:43 +02:00
committed by GitHub
parent b09273f73d 786f98d471
commit 605b0e28f9
5 changed files with 145 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/backend_cuda.cu
+71
View File
@@ -338,6 +338,40 @@ __global__ static void attention_absorb_batch_kernel(float *ctx,const float *q,
ctx[((size_t)s*H+h)*V+v]=a*(fmt?wscale[row]:1.f);}
}
/* Independent KV sequence per row. latent/rope are packed as [S,T,*], while
* lengths selects the valid prefix for each row. */
__global__ static void attention_absorb_ragged_kernel(float *ctx,const float *q,
const float *latent,const float *rope,const int *lengths,
const void *weights,const float *wscale,int fmt,int S,int H,int Q,int R,
int V,int K,int T,float scale){
int s=blockIdx.y,h=blockIdx.x,tid=threadIdx.x,nt=lengths[s],rbase=h*(Q+V);
if(s>=S||nt<1||nt>T)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);
const float *ls=latent+(size_t)s*T*K,*rs=rope+(size_t)s*T*R;
for(int k=tid;k<K;k+=blockDim.x){float a=0;for(int d=0;d<Q;d++)
a+=qs[d]*weight_at(weights,fmt,(size_t)(rbase+d)*row_bytes(fmt,K),k)*
(fmt?wscale[rbase+d]:1.f);qa[k]=a;}
__syncthreads();
for(int t=tid;t<nt;t+=blockDim.x){float a=0;const float *lt=ls+(size_t)t*K;
const float *rt=rs+(size_t)t*R;for(int k=0;k<K;k++)a+=qa[k]*lt[k];
for(int d=0;d<R;d++)a+=qs[Q+d]*rt[d];scores[t]=a*scale;}
__syncthreads();
float local=-3.402823466e+38F;for(int t=tid;t<nt;t+=blockDim.x)local=fmaxf(local,scores[t]);
red[tid]=local;__syncthreads();
for(int n=blockDim.x>>1;n;n>>=1){if(tid<n)red[tid]=fmaxf(red[tid],red[tid+n]);__syncthreads();}
float mx=red[0];local=0;for(int t=tid;t<nt;t+=blockDim.x){float e=expf(scores[t]-mx);scores[t]=e;local+=e;}
red[tid]=local;__syncthreads();
for(int n=blockDim.x>>1;n;n>>=1){if(tid<n)red[tid]+=red[tid+n];__syncthreads();}
float inv=1.f/red[0];for(int t=tid;t<nt;t+=blockDim.x)scores[t]*=inv;
__syncthreads();
for(int k=tid;k<K;k+=blockDim.x){float a=0;for(int t=0;t<nt;t++)a+=scores[t]*ls[(size_t)t*K+k];cl[k]=a;}
__syncthreads();
for(int v=tid;v<V;v+=blockDim.x){int row=rbase+Q+v;float a=0;size_t rb=row_bytes(fmt,K);
for(int k=0;k<K;k++)a+=cl[k]*weight_at(weights,fmt,(size_t)row*rb,k);
ctx[((size_t)s*H+h)*V+v]=a*(fmt?wscale[row]:1.f);}
}
static int reserve(float **ptr, size_t *cap, size_t bytes) {
if (*cap >= bytes) return 1;
if (*ptr) cudaFree(*ptr);
@@ -770,6 +804,43 @@ extern "C" int coli_cuda_attention_project_batch(ColiCudaTensor *w,ColiCudaTenso
return attention_absorb_batch_run(w,proj,out,q,latent,rope,S,H,Q,R,V,K,T,scale);
}
extern "C" int coli_cuda_attention_project_ragged(ColiCudaTensor *w,ColiCudaTensor *proj,
float *out,const float *q,const float *const *latent,const float *const *rope,
const int *lengths,int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!w||!proj||!out||!q||!latent||!rope||!lengths||S<1||S>512||T<1||T>512||
H<1||Q<1||R<1||V<1||K<1||K>512||w->I!=K||w->O!=H*(Q+V)||
proj->device!=w->device||proj->I!=H*V)return 0;
size_t ln=(size_t)S*T*K,rn=(size_t)S*T*R;
float *lh=(float*)std::calloc(ln,sizeof(float)),*rh=(float*)std::calloc(rn,sizeof(float));
if(!lh||!rh){std::free(lh);std::free(rh);return 0;}
for(int s=0;s<S;s++){
if(lengths[s]<1||lengths[s]>T){std::free(lh);std::free(rh);return 0;}
std::memcpy(lh+(size_t)s*T*K,latent[s],(size_t)lengths[s]*K*sizeof(float));
std::memcpy(rh+(size_t)s*T*R,rope[s],(size_t)lengths[s]*R*sizeof(float));
}
DeviceContext *dc=find_ctx(w->device);
if(!select_ctx(dc)){std::free(lh);std::free(rh);return 0;}
size_t qb=(size_t)S*H*(Q+R)*sizeof(float),lb=ln*sizeof(float),rb=rn*sizeof(float);
size_t cb=(size_t)S*H*V*sizeof(float),ob=(size_t)S*proj->O*sizeof(float);
int ok=reserve(&dc->aq,&dc->aq_cap,qb)&&reserve(&dc->al,&dc->al_cap,lb)&&
reserve(&dc->ar,&dc->ar_cap,rb)&&reserve(&dc->ac,&dc->ac_cap,cb)&&
reserve(&dc->y,&dc->y_cap,ob)&&
reserve_bytes(&dc->group_desc,&dc->group_desc_cap,(size_t)S*sizeof(int));
if(ok)ok=cuda_ok(cudaMemcpyAsync(dc->aq,q,qb,cudaMemcpyHostToDevice,dc->stream),"ragged q upload")&&
cuda_ok(cudaMemcpyAsync(dc->al,lh,lb,cudaMemcpyHostToDevice,dc->stream),"ragged latent upload")&&
cuda_ok(cudaMemcpyAsync(dc->ar,rh,rb,cudaMemcpyHostToDevice,dc->stream),"ragged rope upload")&&
cuda_ok(cudaMemcpyAsync(dc->group_desc,lengths,(size_t)S*sizeof(int),cudaMemcpyHostToDevice,dc->stream),"ragged lengths upload");
std::free(lh);std::free(rh);if(!ok)return 0;
size_t shared=(size_t)(2*K+T+256)*sizeof(float);
attention_absorb_ragged_kernel<<<dim3(H,S),256,shared,dc->stream>>>(dc->ac,dc->aq,dc->al,dc->ar,
(const int*)dc->group_desc,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
quant_matmul<<<dim3(proj->O,S),256,0,dc->stream>>>(dc->y,dc->ac,proj->weights,
proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I));
return cuda_ok(cudaGetLastError(),"ragged attention launch")&&
cuda_ok(cudaMemcpyAsync(out,dc->y,ob,cudaMemcpyDeviceToHost,dc->stream),"ragged output download")&&
cuda_ok(cudaStreamSynchronize(dc->stream),"ragged attention synchronize");
}
extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) {
if (!tensor) return;
DeviceContext *ctx = find_ctx(tensor->device);
c/backend_cuda.h
+5 -1
View File
@@ -14,6 +14,7 @@
#define COLI_CUDA_DLLEXPORT
#endif
#ifdef __cplusplus
extern "C" {
#endif
@@ -92,6 +93,10 @@ COLI_CUDA_DLLEXPORT int coli_cuda_attention_project_batch(ColiCudaTensor *kv_b,C
const float *rope,int S,int H,int Q,int R,
int V,int K,int T,float attention_scale);
COLI_CUDA_DLLEXPORT int coli_cuda_attention_project_ragged(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out,const float *q,const float *const *latent,const float *const *rope,
const int *lengths,int S,int H,int Q,int R,int V,int K,int max_t,float attention_scale);
COLI_CUDA_DLLEXPORT void coli_cuda_tensor_free(ColiCudaTensor *tensor);
COLI_CUDA_DLLEXPORT size_t coli_cuda_tensor_bytes(const ColiCudaTensor *tensor);
COLI_CUDA_DLLEXPORT int coli_cuda_tensor_device(const ColiCudaTensor *tensor);
@@ -143,4 +148,3 @@ COLI_CUDA_DLLEXPORT int coli_cuda_pipe_sync(int device);
#endif
#endif
c/backend_loader.c
+13
View File
@@ -61,6 +61,9 @@ typedef int (*fn_attention_absorb_batch)(ColiCudaTensor *kv_b,float *ctx,const f
typedef int (*fn_attention_absorb_batch_dev)(ColiCudaTensor *kv_b_shard,float *ctx_dev, const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale);
typedef int (*fn_attention_absorb_kvdev)(ColiCudaTensor *kv_b,float *ctx,const float *q, const float *latent_dev,const float *rope_dev,int H,int Q,int R,int V,int K,int T, float scale);
typedef int (*fn_attention_project_batch)(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out,const float *q,const float *latent, const float *rope,int S,int H,int Q,int R, int V,int K,int T,float attention_scale);
typedef int (*fn_attention_project_ragged)(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out,const float *q,const float *const *latent,const float *const *rope,
const int *lengths,int S,int H,int Q,int R,int V,int K,int max_t,float attention_scale);
typedef int (*fn_attention_project_batch_dev)(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out,const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale);
typedef int (*fn_attention_project_batch_dev_out)(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out_dev,const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale);
typedef int (*fn_pipe_add)(int device,float *x_dev,const float *t_dev,size_t n);
@@ -107,6 +110,7 @@ static struct {
fn_attention_absorb_batch_dev attention_absorb_batch_dev;
fn_attention_absorb_kvdev attention_absorb_kvdev;
fn_attention_project_batch attention_project_batch;
fn_attention_project_ragged attention_project_ragged;
fn_attention_project_batch_dev attention_project_batch_dev;
fn_attention_project_batch_dev_out attention_project_batch_dev_out;
fn_pipe_add pipe_add;
@@ -200,6 +204,7 @@ static int coli_cuda_load(void){
RESOLVE(attention_absorb_batch_dev, fn_attention_absorb_batch_dev)
RESOLVE(attention_absorb_kvdev, fn_attention_absorb_kvdev)
RESOLVE(attention_project_batch, fn_attention_project_batch)
RESOLVE(attention_project_ragged, fn_attention_project_ragged)
RESOLVE(attention_project_batch_dev, fn_attention_project_batch_dev)
RESOLVE(attention_project_batch_dev_out, fn_attention_project_batch_dev_out)
RESOLVE(pipe_add, fn_pipe_add)
@@ -342,6 +347,14 @@ int coli_cuda_attention_project_batch(ColiCudaTensor *kv_b,ColiCudaTensor *o_pro
return g_cuda.attention_project_batch(kv_b, o_proj, out, q, latent, rope, S, H, Q, R, V, K, T, attention_scale);
}
int coli_cuda_attention_project_ragged(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj,
float *out,const float *q,const float *const *latent,const float *const *rope,
const int *lengths,int S,int H,int Q,int R,int V,int K,int max_t,float attention_scale){
if(!coli_cuda_load()) return 0;
return g_cuda.attention_project_ragged(kv_b,o_proj,out,q,latent,rope,lengths,
S,H,Q,R,V,K,max_t,attention_scale);
}
int coli_cuda_attention_project_batch_dev(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, float *out,const float *q_dev,const float *latent_dev,const float *rope_dev, int S,int H,int Q,int R,int V,int K,int T,float scale){
if(!g_cuda.available){ return 0; }
return g_cuda.attention_project_batch_dev(kv_b, o_proj, out, q_dev, latent_dev, rope_dev, S, H, Q, R, V, K, T, scale);
c/glm.c
+21 -5
View File
@@ -1657,7 +1657,7 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits
c->index_topk, c->index_topk);
}
}
m->hlast=falloc(D); m->h_all=falloc((int64_t)64*D);
m->hlast=falloc(D); m->h_all=falloc((int64_t)512*D);
/* byte della parte DENSA residente (embed+lm_head+attn+mlp densa+shared+norme) */
int64_t rb=qt_bytes(&m->embed)+qt_bytes(&m->lm_head);
@@ -2713,7 +2713,23 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
float *sc_all = falloc((int64_t)omp_get_max_threads()*sc_cap);
int cuda_core=0,cuda_projected=0;
#ifdef COLI_CUDA
if(cuda_absorb&&l->n_kv_b_shard>1){
if(kvs&&g_cuda_enabled&&getenv("COLI_CUDA_ATTN")&&atoi(getenv("COLI_CUDA_ATTN"))&&
!dnsel&&l->kv_b.cuda_eligible&&l->o.cuda_eligible&&
qt_cuda_upload(&l->kv_b)&&qt_cuda_upload(&l->o)){
const float **rl=malloc((size_t)S*sizeof(*rl)),**rr=malloc((size_t)S*sizeof(*rr));
int *rn=malloc((size_t)S*sizeof(*rn)); int mt=0;
if(rl&&rr&&rn){
for(int s=0;s<S;s++){
int pos=positions[s],st0=kvs[s]->kv_start[layer]; rn[s]=pos+1-st0;
rl[s]=coli_kv_row(kvs[s]->Lc[layer],st0,kvl);
rr[s]=coli_kv_row(kvs[s]->Rc[layer],st0,c->qk_rope);
if(rn[s]>mt)mt=rn[s];
}
cuda_core=cuda_projected=coli_cuda_attention_project_ragged(l->kv_b.cuda,l->o.cuda,
out,Q,rl,rr,rn,S,H,c->qk_nope,c->qk_rope,vh,kvl,mt,c->attn_scale);
}
free(rl);free(rr);free(rn);
} else if(cuda_absorb&&l->n_kv_b_shard>1){
int n=l->n_kv_b_shard,st0=m->kv_start[layer],nt=pos_base+S-st0,ok=1;
float *qs=falloc((int64_t)S*H*qh),*cs=falloc((int64_t)S*H*vh);
for(int d=0;d<n;d++)for(int s=0;s<S;s++)memcpy(
@@ -4065,7 +4081,7 @@ static float *step_all(Model *m, const int *ids, int S, int pos_base){
float *x=falloc((int64_t)S*D);
for(int s=0;s<S;s++) embed_row(m, ids[s], x+(int64_t)s*D);
layers_forward(m,x,S,pos_base);
if(m->h_all) memcpy(m->h_all, x, (int64_t)S*D*sizeof(float)); /* hidden di TUTTE le pos (S<=64) */
if(m->h_all) memcpy(m->h_all, x, (int64_t)S*D*sizeof(float)); /* hidden di TUTTE le pos (S<=512) */
if(m->hlast) memcpy(m->hlast, x+(int64_t)(S-1)*D, D*sizeof(float));
float *lo=falloc((int64_t)S*c->vocab), *row=falloc(D);
for(int s=0;s<S;s++){ rmsnorm(row, x+(int64_t)s*D, m->final_norm, D, c->eps);
@@ -4078,8 +4094,8 @@ static float *step_all(Model *m, const int *ids, int S, int pos_base){
static float *step_decode_batch(Model *m, const DecodeRow *rows, int S){
Cfg *c=&m->c; int D=c->hidden;
/* Ragged KV currently uses MLA absorption; the stack kernel is sized to 512. */
if(!rows || S<1 || S>64 || c->kv_lora>512) return NULL;
KVState *kvs[64]; int positions[64];
if(!rows || S<1 || S>512 || c->kv_lora>512) return NULL;
KVState *kvs[512]; int positions[512];
float *x=falloc((int64_t)S*D);
for(int s=0;s<S;s++){
if(!rows[s].kv || !rows[s].kv->Lc || !rows[s].kv->Rc || !rows[s].kv->kv_start ||
c/tests/test_ragged_attention.cu
+35
View File
@@ -0,0 +1,35 @@
#include "../backend_cuda.h"
#include <cmath>
#include <cstdio>
#include <vector>
int main(){
int dev=0;if(!coli_cuda_init(&dev,1))return 77;
constexpr int S=3,H=2,Q=2,R=1,V=2,K=3,D=H*V,O=3,T=3;
std::vector<float> w(H*(Q+V)*K),p(O*D),q(S*H*(Q+R));
for(size_t i=0;i<w.size();i++)w[i]=((int)(i%11)-5)*.07f;
for(size_t i=0;i<p.size();i++)p[i]=((int)(i%7)-3)*.09f;
for(size_t i=0;i<q.size();i++)q[i]=((int)(i%13)-6)*.05f;
ColiCudaTensor *tw=nullptr,*tp=nullptr;
if(!coli_cuda_tensor_upload(&tw,w.data(),nullptr,0,K,H*(Q+V),dev)||
!coli_cuda_tensor_upload(&tp,p.data(),nullptr,0,D,O,dev))return 1;
int n[S]={1,2,3};std::vector<std::vector<float>> l(S),r(S);
const float *lp[S],*rp[S];
for(int s=0;s<S;s++){
l[s].resize(n[s]*K);r[s].resize(n[s]*R);
for(size_t i=0;i<l[s].size();i++)l[s][i]=((int)((i+s*3)%9)-4)*.08f;
for(size_t i=0;i<r[s].size();i++)r[s][i]=((int)((i+s)%5)-2)*.06f;
lp[s]=l[s].data();rp[s]=r[s].data();
}
float got[S*O],ref[S*O];
if(!coli_cuda_attention_project_ragged(tw,tp,got,q.data(),lp,rp,n,S,H,Q,R,V,K,T,.2f))return 2;
for(int s=0;s<S;s++)if(!coli_cuda_attention_project_batch(tw,tp,ref+s*O,
q.data()+s*H*(Q+R),lp[s],rp[s],1,H,Q,R,V,K,n[s],.2f))return 3;
double e=0,z=0;for(int i=0;i<S*O;i++){
double d=got[i]-ref[i];e+=d*d;z+=(double)ref[i]*ref[i];
}
double rms=std::sqrt(e/(z+1e-30));std::printf("ragged_relative_rms=%.9g\n",rms);
coli_cuda_tensor_free(tw);coli_cuda_tensor_free(tp);coli_cuda_shutdown();
return rms<1e-6?0:4;
}