From 1223f9ca2e12a5bfeaccf1f1eca90462fbd0ad97 Mon Sep 17 00:00:00 2001 From: ZacharyZcR Date: Sat, 18 Jul 2026 01:54:31 +0800 Subject: [PATCH 1/2] cuda: batch ragged attention across independent streams --- c/backend_cuda.cu | 71 ++++++++++++++++++++++++++++++++ c/backend_cuda.h | 6 ++- c/glm.c | 26 +++++++++--- c/tests/test_ragged_attention.cu | 35 ++++++++++++++++ 4 files changed, 132 insertions(+), 6 deletions(-) create mode 100644 c/tests/test_ragged_attention.cu diff --git a/c/backend_cuda.cu b/c/backend_cuda.cu index 9ce142f..188330d 100644 --- a/c/backend_cuda.cu +++ b/c/backend_cuda.cu @@ -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>1;n;n>>=1){if(tid>1;n;n>>=1){if(tid= 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;sT){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<<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<<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); diff --git a/c/backend_cuda.h b/c/backend_cuda.h index acbe4bc..edddbfe 100644 --- a/c/backend_cuda.h +++ b/c/backend_cuda.h @@ -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 - diff --git a/c/glm.c b/c/glm.c index 003a42d..298065b 100644 --- a/c/glm.c +++ b/c/glm.c @@ -1644,7 +1644,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); @@ -2640,7 +2640,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;skv_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;dh_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;sfinal_norm, D, c->eps); @@ -3988,8 +4004,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;sLc || !rows[s].kv->Rc || !rows[s].kv->kv_start || diff --git a/c/tests/test_ragged_attention.cu b/c/tests/test_ragged_attention.cu new file mode 100644 index 0000000..c00a8eb --- /dev/null +++ b/c/tests/test_ragged_attention.cu @@ -0,0 +1,35 @@ +#include "../backend_cuda.h" + +#include +#include +#include + +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 w(H*(Q+V)*K),p(O*D),q(S*H*(Q+R)); + for(size_t i=0;i> l(S),r(S); + const float *lp[S],*rp[S]; + for(int s=0;s Date: Sat, 18 Jul 2026 03:50:48 +0800 Subject: [PATCH 2/2] cuda: load ragged attention entry point on Windows --- c/backend_loader.c | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/c/backend_loader.c b/c/backend_loader.c index b743d1b..bbcb8ca 100644 --- a/c/backend_loader.c +++ b/c/backend_loader.c @@ -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);