GPU resident pipeline: batch CUDA attention, head-sharded kv_b, prefill expert groups, W4A16 mixed dispatch (#111)
* Fuse CUDA expert MLP execution * Group CUDA expert transfers by device * Instrument grouped CUDA expert execution * Bound grouped CUDA decode scratch * Execute expert groups across GPUs in parallel * Release host backing for multi-GPU experts * Define quality-preserving memory policies * Overlap cold expert loading with resident compute * Adapt expert placement with session LFRU * Fuse q4 expert gate and up dispatch * Plan CPU work on physical cores * Batch grouped expert CUDA kernels * Separate VRAM and RAM expert placement * Add ragged multi-sequence decode forward * feat(runtime): add continuous decode scheduler * Route concurrent API requests through batch scheduler * Harden multiplex request lifecycle and framing * Cancel disconnected multiplex requests * Bind API port before starting the engine * fix automatic KV slot allocation * add native int4 Tensor Core grouped GEMM * add Tensor Core throughput benchmark * optimize packed int4 low-row kernels * add asynchronous CUDA staging streams * document validated six-GPU dense acceleration * tune six-GPU expert hot set * raise validated expert hot-set target * add CUDA MLA absorption core * fuse grouped expert gate and up projections * Warn for explicit lossy routing flags * Add full-resident expert placement mode * Adapt VRAM expert slots to live routes * Accelerate int4 matvec on AVX-512 * Reduce AVX-512 and RoPE decode overhead * Seed every GPU expert layer after prefill * Limit live GPU swaps during decode * CUDA batch MLA attention, kv_b head-sharding, fused o_proj, expert-group dispatch, W4A16 kernels Lab-qualified on the 6x RTX 5090 machine (914-token request benchmark): - batch MLA absorption kernel (COLI_CUDA_ATTN=1): whole-batch attention on device, 154.8s -> 102.4s - attention -> o_proj fusion on the layer device: -> 97.4s - kv_b head-sharding across cards (COLI_CUDA_ATTN_SHARD=1), no weight duplication: -> 94.05s - per-device expert-group dispatch with pinned-buffer async transfers, W4A16 tensor-core kernels for the shared expert, OMP hot-thread tuning Negative results (reverted, kept out): GPU-side weighted scatter-add (atomics + per-layer D2H lose 43.8%), shared-expert fused small-batch kernel (-38.8%), W4A4 grouped tensor cores (int4 activations corrupt output). Details in the lab research log. * GPU resident pipeline: device-resident prefill attention chain, GPU expert groups in prefill, batched router, W4A16 mixed dispatch COLI_CUDA_PIPE=1 keeps the prefill data plane on the layer home device; control flow (routing, cache/pin management) stays on CPU. Any CUDA failure falls back to the unchanged CPU path. - Device primitives + unit tests (tests/test_pipe_cuda.cu): rmsnorm (strided), interleaved RoPE, silu-mul, residual add, fixed-order row merge (no atomics), device-input GEMM, persistent per-device scratch. All verified against the engine's CPU math on SM120 (worst 1.2e-5). - attn_pipe_prefill: q_a -> norm -> q_b -> rope -> kv_a -> norm -> rope -> batch attention -> o_proj in one device chain (q_a/q_b/kv_a colocated with kv_b); only the final [S,D] and the new KV rows return to host. Attention 41.2s -> 30.8s on the 1571-token benchmark. - Prefill batch-union now uses the GPU expert groups (previously gated to S<=64, leaving all VRAM-resident experts idle during prefill - measured 21ms of GPU expert time in a 148s prefill). Expert phase 78.9s -> 69.0s. - Router computed as one batched matmul instead of S sequential rows (bit-identical math). - W4A16 tensor-core path for expert groups (COLI_CUDA_TC_W4A16=1) with row-count mixed dispatch: >=16 rows per expert use tensor cores, smaller batches keep the naive kernel (tensor cores measured negative below ~16 rows). Expert phase 69.0s -> 64.3s, decode unaffected. Net on the 1571-token prefill benchmark: 148.8s -> 114.3-126.8s (component timings stable across runs; wall drifts +-3-5s because .coli_usage placement learning shifts the expert tiers between runs). PROFILO now also prints the prefill-phase breakdown. * Skip OMP hot-thread tuning when CUDA is enabled The active-spin worker team measured 66.9s->20.9s on the CPU-only Zen5 build, but on the six-GPU full-residency workload the spinning workers contend with the CUDA dispatch threads: ~4x slower prefill with the process stuck near 1.8 cores. Gate the tuning on COLI_CUDA so each configuration keeps the behavior it was measured to prefer. * Inc.2a: sparse layers fully resident on the layer device, residual hops cards at layer boundaries COLI_CUDA_PIPE=2 keeps the residual stream on the layer home device for consecutive sparse layers (cudaMemcpyPeer at boundaries): in/post norms, attention chain, both residual adds and the shared-expert MLP run on device. Per layer only the post-norm activations (router + CPU-tier experts + group gather), the new KV rows and, on DSA indexer layers, the pre-attention norm leave the card. Per-layer transfers drop from ~130MB to ~70MB. A device-side snapshot at layer entry makes any mid-layer CUDA failure fall back to the unchanged CPU path idempotently. 1571-token prefill: 127.1s (PIPE=1 control) -> 117.6/118.9s, components attention 30.8->26.1, other 31.8->22.5-24.5; output verified coherent against the control. * Head-sharded attention inside the pipe: negative on PCIe star topology, gated opt-in Slicing q per card from the home device and collecting ctx back serializes ~95MB/layer through the home card's PCIe link: attention 26.1s -> 41.4/44.4s on the 1571-token benchmark (two repeats), wall 117.6 -> 135-138s. The standalone host-path sharding won because six cards uploaded from host RAM in parallel; a home-device star has no such parallelism without NVLink. Kept behind COLI_CUDA_PIPE_SHARD=1 for interconnects where peer bandwidth does not share one root port. * Inc.3: device-resident KV shadow for decode attention Decode re-uploaded the whole latent+rope window per layer per token (~300MB/token at 1571 context). Each layer now keeps a device shadow of the compressed KV on its kv_b card, bulk-synced when behind and appended incrementally; the host cache stays canonical. Invalidation on kv_bind (slot switch), kv_alloc (resize) and on any overwrite of mirrored rows, with the legacy full-upload path as fallback. Measured (COLI_CUDA_PIPE gate): short-context decode 5.48 -> 5.59/5.87 tok/s, 1571-context decode 4.14 -> 4.22 tok/s. Decode remains CPU-expert bound; the shadow removes the transfer tax, not the compute. * tools: unified user-experience benchmark (bench_ux.sh) Two fixed scenarios (short chat, long-document QA), TTFT + decode tok/s + first-line drift check, TEMP=0 DRAFT=0 enforced, medians over REPS runs. Encodes the measurement discipline from the lab record: same binary per comparison, judge medians because .coli_usage placement learning drifts wall times between runs. * tools: bench_ux.sh executable bit * gitignore compiled test binaries * tools: expert_atlas.py — measure per-expert topic affinity (#175) Diffs .coli_usage across 10 themed probe batches (code/math/chinese/ prose/science/law/poetry/structured/translation/casual, 3 prompts each) driven through a running API server — one engine load total. Every touched expert gets a topic-affinity vector, entropy, and a specialist/ generalist label; output experts.json feeds the Brain page hover. * serve: persist .coli_usage after every turn in mux mode, not only at exit run_serve_mux saved the learning cache once at shutdown; a crash lost the whole session's routing history, and live consumers of the file (expert_atlas.py diffs it between probe batches) saw a frozen snapshot. Now saved per turn like the interactive path (165KB write, negligible). * web: Brain hover shows measured expert atlas when published If /experts.json (from tools/expert_atlas.py, #175) is served next to the app, the tooltip upgrades from the depth heuristic to measured data: specialist/generalist label, entropy, and the top-3 topic affinities. Row index maps to real layer (row+3, last row = MTP 78). Falls back to the heuristic when no atlas is published. --------- Co-authored-by: JustVugg <JustVugg@users.noreply.github.com>
This commit is contained in:
ZacharyZcR
@@ -18,12 +18,14 @@ struct ColiCudaTensor {
|
||||
|
||||
typedef struct {
|
||||
int device;
|
||||
int compute_major,compute_minor;
|
||||
float *x, *y, *gate, *up;
|
||||
size_t x_cap, y_cap, gate_cap, up_cap;
|
||||
uint8_t *qx; float *qscale;
|
||||
size_t qx_cap, qscale_cap;
|
||||
float *host_x,*host_y; size_t host_x_cap,host_y_cap;
|
||||
float *aq,*al,*ar,*ac; size_t aq_cap,al_cap,ar_cap,ac_cap;
|
||||
float *pipe_buf[24]; size_t pipe_cap[24]; /* scratch persistenti del resident pipeline */
|
||||
cudaStream_t stream;
|
||||
void *group_desc; size_t group_desc_cap;
|
||||
size_t tensor_count, tensor_bytes;
|
||||
@@ -120,6 +122,68 @@ __global__ static void silu_mul(float *gate, const float *up, size_t n) {
|
||||
}
|
||||
}
|
||||
|
||||
/* Four warps share one A tile and compute 16x64 outputs. This matters for
|
||||
* prefill: the first prototype reloaded/converter A once per 16 output cols. */
|
||||
__global__ static void w4a16_matmul(float *y,const float *x,const uint8_t *w,
|
||||
const float *scale,int M,int K,int N){
|
||||
#if __CUDA_ARCH__ >= 700
|
||||
using namespace nvcuda;int warp=threadIdx.x>>5,lane=threadIdx.x&31;
|
||||
int m0=blockIdx.y*16,n0=blockIdx.x*64+warp*16;
|
||||
__shared__ __half ah[256],bh[4][256];
|
||||
wmma::fragment<wmma::accumulator,16,16,16,float> acc;wmma::fill_fragment(acc,0.f);
|
||||
size_t rb=(size_t)(K+1)/2;
|
||||
for(int k0=0;k0<K;k0+=16){
|
||||
for(int z=threadIdx.x;z<256;z+=blockDim.x){
|
||||
int m=z/16,k=z%16,gm=m0+m,gk=k0+k;
|
||||
ah[z]=(gm<M&&gk<K)?__float2half(x[(size_t)gm*K+gk]):__float2half(0.f);
|
||||
}
|
||||
for(int z=lane;z<256;z+=32){
|
||||
int n=z/16,gk=k0+(z%16),gn=n0+n;float v=0.f;
|
||||
if(gn<N&&gk<K){uint8_t q=w[(size_t)gn*rb+(gk>>1)];int a=(gk&1)?q>>4:q&15;
|
||||
v=(float)(a&8?a-16:a)*scale[gn];}
|
||||
bh[warp][z]=__float2half(v); /* [Ntile,Ktile] == B col-major */
|
||||
}
|
||||
__syncthreads();
|
||||
wmma::fragment<wmma::matrix_a,16,16,16,__half,wmma::row_major> af;
|
||||
wmma::fragment<wmma::matrix_b,16,16,16,__half,wmma::col_major> bf;
|
||||
wmma::load_matrix_sync(af,ah,16);wmma::load_matrix_sync(bf,bh[warp],16);
|
||||
wmma::mma_sync(acc,af,bf,acc);__syncthreads();
|
||||
}
|
||||
__shared__ float out[4][256];wmma::store_matrix_sync(out[warp],acc,16,wmma::mem_row_major);__syncwarp();
|
||||
for(int z=lane;z<256;z+=32){int m=z/16,n=z%16;
|
||||
if(m0+m<M&&n0+n<N)y[(size_t)(m0+m)*N+n0+n]=out[warp][z];}
|
||||
#endif
|
||||
}
|
||||
|
||||
/* Gate and up use the same input. Eight warps compute both 16x64 projections
|
||||
* while sharing the FP32->FP16 conversion of A. */
|
||||
__global__ static void w4a16_gate_up(float *gate,float *up,const float *x,
|
||||
const uint8_t *gw,const uint8_t *uw,const float *gs,const float *us,
|
||||
int M,int K,int N){
|
||||
#if __CUDA_ARCH__ >= 700
|
||||
using namespace nvcuda;int warp=threadIdx.x>>5,lane=threadIdx.x&31,which=warp&1,tile=warp>>1;
|
||||
int m0=blockIdx.y*16,n0=blockIdx.x*64+tile*16;const uint8_t *w=which?uw:gw;
|
||||
const float *scale=which?us:gs;float *y=which?up:gate;size_t rb=(size_t)(K+1)/2;
|
||||
__shared__ __half ah[256],bh[8][256];
|
||||
wmma::fragment<wmma::accumulator,16,16,16,float> acc;wmma::fill_fragment(acc,0.f);
|
||||
for(int k0=0;k0<K;k0+=16){
|
||||
for(int z=threadIdx.x;z<256;z+=blockDim.x){int m=z/16,k=z%16,gm=m0+m,gk=k0+k;
|
||||
ah[z]=(gm<M&&gk<K)?__float2half(x[(size_t)gm*K+gk]):__float2half(0.f);}
|
||||
for(int z=lane;z<256;z+=32){int n=z/16,gk=k0+(z%16),gn=n0+n;float v=0.f;
|
||||
if(gn<N&&gk<K){uint8_t q=w[(size_t)gn*rb+(gk>>1)];int a=(gk&1)?q>>4:q&15;
|
||||
v=(float)(a&8?a-16:a)*scale[gn];}bh[warp][z]=__float2half(v);}
|
||||
__syncthreads();
|
||||
wmma::fragment<wmma::matrix_a,16,16,16,__half,wmma::row_major> af;
|
||||
wmma::fragment<wmma::matrix_b,16,16,16,__half,wmma::col_major> bf;
|
||||
wmma::load_matrix_sync(af,ah,16);wmma::load_matrix_sync(bf,bh[warp],16);
|
||||
wmma::mma_sync(acc,af,bf,acc);__syncthreads();
|
||||
}
|
||||
__shared__ float out[8][256];wmma::store_matrix_sync(out[warp],acc,16,wmma::mem_row_major);__syncwarp();
|
||||
for(int z=lane;z<256;z+=32){int m=z/16,n=z%16;
|
||||
if(m0+m<M&&n0+n<N)y[(size_t)(m0+m)*N+n0+n]=out[warp][z];}
|
||||
#endif
|
||||
}
|
||||
|
||||
__global__ static void quantize_s4_rows(uint8_t *q,float *scale,const float *x,int S,int K){
|
||||
int s=blockIdx.x; if(s>=S)return; const float *xs=x+(size_t)s*K;
|
||||
float v=0; for(int i=threadIdx.x;i<K;i+=blockDim.x)v=fmaxf(v,fabsf(xs[i]));
|
||||
@@ -243,6 +307,37 @@ __global__ static void attention_absorb_kernel(float *ctx,const float *q,const f
|
||||
for(int k=0;k<K;k++)a+=cl[k]*weight_at(weights,fmt,(size_t)row*rb,k);ctx[(size_t)h*V+v]=a*(fmt?wscale[row]:1.f);}
|
||||
}
|
||||
|
||||
__global__ static void attention_absorb_batch_kernel(float *ctx,const float *q,
|
||||
const float *latent,const float *rope,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=T-S+s+1,rbase=h*(Q+V);
|
||||
if(s>=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;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=latent+(size_t)t*K;
|
||||
const float *rt=rope+(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]*latent[(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);
|
||||
@@ -286,6 +381,7 @@ extern "C" int coli_cuda_init(const int *devices, int count) {
|
||||
if (!select_ctx(ctx)) { g_nctx = 0; return 0; }
|
||||
cudaDeviceProp prop{};
|
||||
if (!cuda_ok(cudaGetDeviceProperties(&prop, device), "device properties")) { g_nctx = 0; return 0; }
|
||||
ctx->compute_major=prop.major;ctx->compute_minor=prop.minor;
|
||||
if(!cuda_ok(cudaStreamCreateWithFlags(&ctx->stream,cudaStreamNonBlocking),"stream creation")){
|
||||
g_nctx=0;return 0;
|
||||
}
|
||||
@@ -307,6 +403,7 @@ extern "C" void coli_cuda_shutdown(void) {
|
||||
if (ctx->qx) cudaFree(ctx->qx);
|
||||
if (ctx->qscale) cudaFree(ctx->qscale);
|
||||
if(ctx->aq)cudaFree(ctx->aq);if(ctx->al)cudaFree(ctx->al);if(ctx->ar)cudaFree(ctx->ar);if(ctx->ac)cudaFree(ctx->ac);
|
||||
for(int b=0;b<24;b++) if(ctx->pipe_buf[b]) cudaFree(ctx->pipe_buf[b]);
|
||||
if (ctx->host_x) cudaFreeHost(ctx->host_x);
|
||||
if (ctx->host_y) cudaFreeHost(ctx->host_y);
|
||||
if (ctx->stream) cudaStreamDestroy(ctx->stream);
|
||||
@@ -453,6 +550,34 @@ extern "C" int coli_cuda_expert_mlp(ColiCudaTensor *gate, ColiCudaTensor *up,
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_shared_mlp_w4a16(ColiCudaTensor *gate,ColiCudaTensor *up,
|
||||
ColiCudaTensor *down,float *y,const float *x,int S){
|
||||
if(!gate||!up||!down||!x||!y||S<1||gate->fmt!=2||up->fmt!=2||down->fmt!=2||
|
||||
gate->device!=up->device||gate->device!=down->device||gate->I!=up->I||
|
||||
gate->O!=up->O||down->I!=gate->O||down->O!=gate->I)return 0;
|
||||
DeviceContext *ctx=find_ctx(gate->device);if(!select_ctx(ctx)||ctx->compute_major<7)return 0;
|
||||
int D=gate->I,I=gate->O;size_t xb=(size_t)S*D*sizeof(float),ib=(size_t)S*I*sizeof(float);
|
||||
if(!reserve(&ctx->x,&ctx->x_cap,xb)||!reserve(&ctx->gate,&ctx->gate_cap,ib)||
|
||||
!reserve(&ctx->up,&ctx->up_cap,ib)||!reserve(&ctx->y,&ctx->y_cap,xb)||
|
||||
!reserve_pinned(&ctx->host_x,&ctx->host_x_cap,xb)||
|
||||
!reserve_pinned(&ctx->host_y,&ctx->host_y_cap,xb))return 0;
|
||||
std::memcpy(ctx->host_x,x,xb);
|
||||
if(!cuda_ok(cudaMemcpyAsync(ctx->x,ctx->host_x,xb,cudaMemcpyHostToDevice,ctx->stream),
|
||||
"shared w4a16 input upload"))return 0;
|
||||
dim3 hidden((unsigned)((I+63)/64),(unsigned)((S+15)/16));
|
||||
dim3 output((unsigned)((D+63)/64),(unsigned)((S+15)/16));
|
||||
w4a16_gate_up<<<hidden,256,0,ctx->stream>>>(ctx->gate,ctx->up,ctx->x,
|
||||
(const uint8_t*)gate->weights,(const uint8_t*)up->weights,gate->scales,up->scales,S,D,I);
|
||||
silu_mul<<<(unsigned)(((size_t)S*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)S*I);
|
||||
w4a16_matmul<<<output,128,0,ctx->stream>>>(ctx->y,ctx->gate,(const uint8_t*)down->weights,down->scales,S,I,D);
|
||||
if(!cuda_ok(cudaGetLastError(),"shared w4a16 launch")||
|
||||
!cuda_ok(cudaMemcpyAsync(ctx->host_y,ctx->y,xb,cudaMemcpyDeviceToHost,ctx->stream),
|
||||
"shared w4a16 output download")||
|
||||
!cuda_ok(cudaStreamSynchronize(ctx->stream),"shared w4a16 synchronize"))return 0;
|
||||
std::memcpy(y,ctx->host_y,xb);
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
|
||||
ColiCudaTensor *const *ups,
|
||||
ColiCudaTensor *const *downs,
|
||||
@@ -510,6 +635,38 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
|
||||
silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I);
|
||||
quantize_s4_rows<<<total,256,0,ctx->stream>>>(ctx->qx,ctx->qscale,ctx->gate,total,I);
|
||||
grouped_s4_wmma<<<dim3((unsigned)((D+63)/64),(unsigned)count),256,0,ctx->stream>>>(ctx->y,ctx->qx,ctx->qscale,dev,I,D,2);
|
||||
}else if(all_s4&&ctx->compute_major>=7&&getenv("COLI_CUDA_TC_W4A16")&&
|
||||
atoi(getenv("COLI_CUDA_TC_W4A16"))){
|
||||
/* W4A16 Tensor Core per gruppo: attivazioni fp16 per tile (lossless al
|
||||
* contrario del path W4A4), un lancio per expert dentro lo stream —
|
||||
* l'overhead di lancio e' trascurabile rispetto ai GEMM. */
|
||||
int tc16_min=getenv("COLI_CUDA_TC_W4A16_MIN")?atoi(getenv("COLI_CUDA_TC_W4A16_MIN")):16;
|
||||
int off16=0;
|
||||
for(int c=0;c<count;c++){
|
||||
int r=rows[c];
|
||||
float *g16=ctx->gate+(size_t)off16*I,*u16=ctx->up+(size_t)off16*I;
|
||||
float *x16=ctx->x+(size_t)off16*D,*y16=ctx->y+(size_t)off16*D;
|
||||
if(r>=tc16_min){
|
||||
dim3 hg16((unsigned)((I+63)/64),(unsigned)((r+15)/16));
|
||||
dim3 og16((unsigned)((D+63)/64),(unsigned)((r+15)/16));
|
||||
w4a16_gate_up<<<hg16,256,0,ctx->stream>>>(g16,u16,x16,
|
||||
(const uint8_t*)host[c].g,(const uint8_t*)host[c].u,host[c].gs,host[c].us,r,D,I);
|
||||
silu_mul<<<(unsigned)(((size_t)r*I+255)/256),256,0,ctx->stream>>>(g16,u16,(size_t)r*I);
|
||||
w4a16_matmul<<<og16,128,0,ctx->stream>>>(y16,g16,
|
||||
(const uint8_t*)host[c].d,host[c].ds,r,I,D);
|
||||
}else{
|
||||
/* piccoli batch: tile TC quasi vuoti + overhead di lancio — il
|
||||
* kernel naive per-elemento resta piu' veloce (misurato in decode) */
|
||||
quant_matmul<<<dim3((unsigned)I,(unsigned)r),256,0,ctx->stream>>>(g16,x16,
|
||||
host[c].g,host[c].gs,host[c].gf,r,D,I,row_bytes(host[c].gf,D));
|
||||
quant_matmul<<<dim3((unsigned)I,(unsigned)r),256,0,ctx->stream>>>(u16,x16,
|
||||
host[c].u,host[c].us,host[c].uf,r,D,I,row_bytes(host[c].uf,D));
|
||||
silu_mul<<<(unsigned)(((size_t)r*I+255)/256),256,0,ctx->stream>>>(g16,u16,(size_t)r*I);
|
||||
quant_matmul<<<dim3((unsigned)D,(unsigned)r),256,0,ctx->stream>>>(y16,g16,
|
||||
host[c].d,host[c].ds,host[c].df,r,I,D,row_bytes(host[c].df,I));
|
||||
}
|
||||
off16+=r;
|
||||
}
|
||||
}else if(all_s4&&(!getenv("COLI_CUDA_W4_PACKED")||atoi(getenv("COLI_CUDA_W4_PACKED")))){
|
||||
dim3 hg((unsigned)I,(unsigned)max_rows,(unsigned)count),og((unsigned)D,(unsigned)max_rows,(unsigned)count);
|
||||
int dual=!getenv("COLI_CUDA_DUAL_PROJ")||atoi(getenv("COLI_CUDA_DUAL_PROJ"));
|
||||
@@ -547,6 +704,7 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
|
||||
return 1;
|
||||
}
|
||||
|
||||
|
||||
extern "C" int coli_cuda_attention_absorb(ColiCudaTensor *w,float *ctx,const float *q,
|
||||
const float *latent,const float *rope,int H,int Q,
|
||||
int R,int V,int K,int T,float scale){
|
||||
@@ -569,6 +727,49 @@ extern "C" int coli_cuda_attention_absorb(ColiCudaTensor *w,float *ctx,const flo
|
||||
return 1;
|
||||
}
|
||||
|
||||
static int attention_absorb_batch_run(ColiCudaTensor *w,ColiCudaTensor *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 scale){
|
||||
if(!w||!out||!q||!latent||!rope||S<1||H<1||Q<1||R<1||V<1||K<1||K>512||
|
||||
T<S||T>8192||w->I!=K||w->O!=H*(Q+V))return 0;
|
||||
if(proj&&(proj->device!=w->device||proj->I!=H*V))return 0;
|
||||
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
|
||||
size_t qb=(size_t)S*H*(Q+R)*sizeof(float),lb=(size_t)T*K*sizeof(float);
|
||||
size_t rb=(size_t)T*R*sizeof(float),cb=(size_t)S*H*V*sizeof(float);
|
||||
if(!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))return 0;
|
||||
if(!cuda_ok(cudaMemcpyAsync(dc->aq,q,qb,cudaMemcpyHostToDevice,dc->stream),"attention batch q upload")||
|
||||
!cuda_ok(cudaMemcpyAsync(dc->al,latent,lb,cudaMemcpyHostToDevice,dc->stream),"attention batch latent upload")||
|
||||
!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<<<dim3(H,S),256,shared,dc->stream>>>(dc->ac,dc->aq,dc->al,
|
||||
dc->ar,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
|
||||
if(!cuda_ok(cudaGetLastError(),"attention batch launch"))return 0;
|
||||
const float *src=dc->ac;size_t ob=cb;
|
||||
if(proj){
|
||||
ob=(size_t)S*proj->O*sizeof(float);if(!reserve(&dc->y,&dc->y_cap,ob))return 0;
|
||||
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));
|
||||
if(!cuda_ok(cudaGetLastError(),"attention o_proj launch"))return 0;src=dc->y;
|
||||
}
|
||||
if(!cuda_ok(cudaMemcpyAsync(out,src,ob,cudaMemcpyDeviceToHost,dc->stream),
|
||||
proj?"attention projected output download":"attention batch context download")||
|
||||
!cuda_ok(cudaStreamSynchronize(dc->stream),"attention batch synchronize"))return 0;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_attention_absorb_batch(ColiCudaTensor *w,float *ctx,const float *q,
|
||||
const float *latent,const float *rope,int S,int H,int Q,int R,int V,int K,int T,
|
||||
float scale){
|
||||
return attention_absorb_batch_run(w,nullptr,ctx,q,latent,rope,S,H,Q,R,V,K,T,scale);
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_attention_project_batch(ColiCudaTensor *w,ColiCudaTensor *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 scale){
|
||||
return attention_absorb_batch_run(w,proj,out,q,latent,rope,S,H,Q,R,V,K,T,scale);
|
||||
}
|
||||
|
||||
extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) {
|
||||
if (!tensor) return;
|
||||
DeviceContext *ctx = find_ctx(tensor->device);
|
||||
@@ -590,3 +791,234 @@ extern "C" size_t coli_cuda_tensor_bytes(const ColiCudaTensor *tensor) {
|
||||
extern "C" int coli_cuda_tensor_device(const ColiCudaTensor *tensor) {
|
||||
return tensor ? tensor->device : -1;
|
||||
}
|
||||
|
||||
/* ==== resident-pipeline primitives (Inc.0, 2026-07-13) ====
|
||||
* Device-side building blocks so the residual stream can stay on the layer's
|
||||
* home device across a whole layer. Control flow stays on CPU; only the data
|
||||
* plane lives here. All entry points take DEVICE pointers (no transfers) —
|
||||
* the caller owns staging via the pipe buffer API below. */
|
||||
|
||||
__global__ static void pipe_rmsnorm_rows(float *y,const float *x,const float *w,
|
||||
int D,float eps,int xstride,int ystride){
|
||||
const float *xr=x+(size_t)blockIdx.x*xstride; float *yr=y+(size_t)blockIdx.x*ystride;
|
||||
__shared__ double sh[256];
|
||||
double a=0; for(int i=threadIdx.x;i<D;i+=blockDim.x){ double v=xr[i]; a+=v*v; }
|
||||
sh[threadIdx.x]=a; __syncthreads();
|
||||
for(int s=blockDim.x/2;s>0;s>>=1){ if(threadIdx.x<s) sh[threadIdx.x]+=sh[threadIdx.x+s]; __syncthreads(); }
|
||||
float r=rsqrtf((float)(sh[0]/D)+eps);
|
||||
for(int i=threadIdx.x;i<D;i+=blockDim.x) yr[i]=xr[i]*r*w[i];
|
||||
}
|
||||
|
||||
/* RoPE interleaved, identical math to glm.c rope_interleave. One block per row;
|
||||
* row layout: v + row*stride + offset holds R floats. pos index = row/heads
|
||||
* (heads=1 for k_rot rows, heads=H for [S,H,qh] query rows). */
|
||||
__global__ static void pipe_rope_rows(float *v,const int *pos,int pos_base,int stride,
|
||||
int offset,int R,int heads,float theta){
|
||||
float *p=v+(size_t)blockIdx.x*stride+offset;
|
||||
int half=R/2, ps=pos?pos[blockIdx.x/heads]:pos_base+(int)(blockIdx.x/heads);
|
||||
__shared__ float in[256];
|
||||
for(int j=threadIdx.x;j<R;j+=blockDim.x) in[j]=p[j];
|
||||
__syncthreads();
|
||||
for(int j=threadIdx.x;j<half;j+=blockDim.x){
|
||||
float inv=__powf(theta,-2.0f*j/R);
|
||||
float ang=ps*inv, cs=__cosf(ang), sn=__sinf(ang);
|
||||
float a=in[2*j], b=in[2*j+1];
|
||||
p[j]=a*cs-b*sn; p[half+j]=b*cs+a*sn;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void pipe_add_n(float *x,const float *t,size_t n){
|
||||
size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x;
|
||||
if(i<n) x[i]+=t[i];
|
||||
}
|
||||
|
||||
/* Fixed-order partial merge: block b adds partial row b into x row rows[b].
|
||||
* Target rows are unique by construction (CPU pre-sums per token), so no
|
||||
* atomics — the 9.20.7 lesson. */
|
||||
__global__ static void pipe_rows_add(float *x,const float *partial,const int *rows,
|
||||
int D){
|
||||
float *xr=x+(size_t)rows[blockIdx.x]*D;
|
||||
const float *pr=partial+(size_t)blockIdx.x*D;
|
||||
for(int i=threadIdx.x;i<D;i+=blockDim.x) xr[i]+=pr[i];
|
||||
}
|
||||
|
||||
/* scratch persistente per (device,slot): cresce e resta — niente cudaMalloc/Free
|
||||
* per layer (78 x ~10 alloc/richiesta erano puro churn). */
|
||||
extern "C" float *coli_cuda_pipe_scratch(int device,int slot,size_t bytes){
|
||||
DeviceContext *ctx=find_ctx(device);
|
||||
if(slot<0||slot>=24||!select_ctx(ctx)) return NULL;
|
||||
if(!reserve(&ctx->pipe_buf[slot],&ctx->pipe_cap[slot],bytes)) return NULL;
|
||||
return ctx->pipe_buf[slot];
|
||||
}
|
||||
extern "C" void *coli_cuda_pipe_alloc(int device,size_t bytes){
|
||||
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return NULL;
|
||||
void *p=NULL;
|
||||
if(!cuda_ok(cudaMalloc(&p,bytes),"pipe alloc")) return NULL;
|
||||
return p;
|
||||
}
|
||||
extern "C" void coli_cuda_pipe_free(int device,void *p){
|
||||
DeviceContext *ctx=find_ctx(device); if(!p||!select_ctx(ctx)) return;
|
||||
cudaFree(p);
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_upload(int device,void *dst,const void *src,size_t bytes){
|
||||
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
|
||||
return cuda_ok(cudaMemcpy(dst,src,bytes,cudaMemcpyHostToDevice),"pipe upload");
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_download(int device,const void *src,void *dst,size_t bytes){
|
||||
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
|
||||
return cuda_ok(cudaMemcpy(dst,src,bytes,cudaMemcpyDeviceToHost),"pipe download");
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_rmsnorm(int device,float *y_dev,const float *x_dev,
|
||||
const float *w_dev,int S,int D,float eps){
|
||||
DeviceContext *ctx=find_ctx(device);
|
||||
if(S<1||D<1||!select_ctx(ctx)) return 0;
|
||||
pipe_rmsnorm_rows<<<S,256>>>(y_dev,x_dev,w_dev,D,eps,D,D);
|
||||
return cuda_ok(cudaGetLastError(),"pipe rmsnorm");
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_rmsnorm_s(int device,float *y_dev,const float *x_dev,
|
||||
const float *w_dev,int S,int D,float eps,
|
||||
int xstride,int ystride){
|
||||
DeviceContext *ctx=find_ctx(device);
|
||||
if(S<1||D<1||xstride<D||ystride<D||!select_ctx(ctx)) return 0;
|
||||
pipe_rmsnorm_rows<<<S,256>>>(y_dev,x_dev,w_dev,D,eps,xstride,ystride);
|
||||
return cuda_ok(cudaGetLastError(),"pipe rmsnorm strided");
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_rope(int device,float *v_dev,const int *pos_dev,
|
||||
int rows,int stride,int offset,int R,int heads,
|
||||
float theta){
|
||||
DeviceContext *ctx=find_ctx(device);
|
||||
if(rows<1||R<2||R>256||heads<1||!select_ctx(ctx)) return 0;
|
||||
pipe_rope_rows<<<rows,128>>>(v_dev,pos_dev,0,stride,offset,R,heads,theta);
|
||||
return cuda_ok(cudaGetLastError(),"pipe rope");
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_rope_base(int device,float *v_dev,int pos_base,int rows,
|
||||
int stride,int offset,int R,int heads,float theta){
|
||||
DeviceContext *ctx=find_ctx(device);
|
||||
if(rows<1||R<2||R>256||heads<1||!select_ctx(ctx)) return 0;
|
||||
pipe_rope_rows<<<rows,128>>>(v_dev,NULL,pos_base,stride,offset,R,heads,theta);
|
||||
return cuda_ok(cudaGetLastError(),"pipe rope base");
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_copy2d(int device,float *dst,int dpitch,const float *src,
|
||||
int spitch,int width,int height){
|
||||
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
|
||||
return cuda_ok(cudaMemcpy2D(dst,(size_t)dpitch*4,src,(size_t)spitch*4,
|
||||
(size_t)width*4,height,cudaMemcpyDeviceToDevice),"pipe copy2d");
|
||||
}
|
||||
/* attention batch + fused o_proj with DEVICE-resident q/latent/rope: the whole
|
||||
* upstream projection chain stayed on this device, so nothing is uploaded here.
|
||||
* Only the final [S,O] projection is downloaded to host. */
|
||||
extern "C" int coli_cuda_attention_project_batch_dev(ColiCudaTensor *w,ColiCudaTensor *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(!w||!proj||!out||!q_dev||!latent_dev||!rope_dev||S<1||H<1||Q<1||R<1||V<1||
|
||||
K<1||K>512||T<S||T>8192||w->I!=K||w->O!=H*(Q+V)||
|
||||
proj->device!=w->device||proj->I!=H*V)return 0;
|
||||
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
|
||||
size_t cb=(size_t)S*H*V*sizeof(float);
|
||||
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<<<dim3(H,S),256,shared,dc->stream>>>(dc->ac,q_dev,latent_dev,
|
||||
rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
|
||||
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;
|
||||
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));
|
||||
if(!cuda_ok(cudaGetLastError(),"pipe o_proj launch"))return 0;
|
||||
if(!cuda_ok(cudaMemcpyAsync(out,dc->y,ob,cudaMemcpyDeviceToHost,dc->stream),"pipe attention download")||
|
||||
!cuda_ok(cudaStreamSynchronize(dc->stream),"pipe attention sync"))return 0;
|
||||
return 1;
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_silu_mul(int device,float *gate_dev,const float *up_dev,
|
||||
size_t n){
|
||||
DeviceContext *ctx=find_ctx(device); if(!n||!select_ctx(ctx)) return 0;
|
||||
silu_mul<<<(unsigned)((n+255)/256),256>>>(gate_dev,up_dev,n);
|
||||
return cuda_ok(cudaGetLastError(),"pipe silu mul");
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_add(int device,float *x_dev,const float *t_dev,size_t n){
|
||||
DeviceContext *ctx=find_ctx(device); if(!n||!select_ctx(ctx)) return 0;
|
||||
pipe_add_n<<<(unsigned)((n+255)/256),256>>>(x_dev,t_dev,n);
|
||||
return cuda_ok(cudaGetLastError(),"pipe add");
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_rows_add(int device,float *x_dev,const float *partial_dev,
|
||||
const int *rows_dev,int nrows,int D){
|
||||
DeviceContext *ctx=find_ctx(device); if(nrows<1||D<1||!select_ctx(ctx)) return 0;
|
||||
pipe_rows_add<<<nrows,256>>>(x_dev,partial_dev,rows_dev,D);
|
||||
return cuda_ok(cudaGetLastError(),"pipe rows add");
|
||||
}
|
||||
/* GEMM with device-resident activations: same quant_matmul kernel as
|
||||
* coli_cuda_matmul, zero host transfers. */
|
||||
extern "C" int coli_cuda_pipe_gemm(ColiCudaTensor *t,float *y_dev,const float *x_dev,
|
||||
int S){
|
||||
if(!t||S<1) return 0;
|
||||
DeviceContext *ctx=find_ctx(t->device); if(!select_ctx(ctx)) return 0;
|
||||
dim3 grid((unsigned)t->O,(unsigned)S);
|
||||
quant_matmul<<<grid,256>>>(y_dev,x_dev,t->weights,t->scales,t->fmt,S,t->I,t->O,
|
||||
row_bytes(t->fmt,t->I));
|
||||
return cuda_ok(cudaGetLastError(),"pipe gemm");
|
||||
}
|
||||
/* copia diretta scheda->scheda (P2P se disponibile, altrimenti staging driver) */
|
||||
extern "C" int coli_cuda_pipe_peer_copy(int dst_dev,float *dst,int src_dev,
|
||||
const float *src,size_t bytes){
|
||||
if(!dst||!src) return 0;
|
||||
if(dst_dev==src_dev){ DeviceContext *c=find_ctx(dst_dev); if(!select_ctx(c)) return 0;
|
||||
return cuda_ok(cudaMemcpy(dst,src,bytes,cudaMemcpyDeviceToDevice),"pipe intra copy"); }
|
||||
return cuda_ok(cudaMemcpyPeer(dst,dst_dev,src,src_dev,bytes),"pipe peer copy");
|
||||
}
|
||||
/* come attention_project_batch_dev ma l'uscita di o_proj RESTA sul device (out_dev). */
|
||||
extern "C" int coli_cuda_attention_project_batch_dev_out(ColiCudaTensor *w,ColiCudaTensor *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){
|
||||
if(!w||!proj||!out_dev||!q_dev||!latent_dev||!rope_dev||S<1||H<1||Q<1||R<1||V<1||
|
||||
K<1||K>512||T<S||T>8192||w->I!=K||w->O!=H*(Q+V)||
|
||||
proj->device!=w->device||proj->I!=H*V)return 0;
|
||||
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
|
||||
size_t cb=(size_t)S*H*V*sizeof(float);
|
||||
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<<<dim3(H,S),256,shared,dc->stream>>>(dc->ac,q_dev,latent_dev,
|
||||
rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
|
||||
if(!cuda_ok(cudaGetLastError(),"pipe attention launch (dev out)"))return 0;
|
||||
quant_matmul<<<dim3(proj->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));
|
||||
if(!cuda_ok(cudaGetLastError(),"pipe o_proj launch (dev out)"))return 0;
|
||||
return cuda_ok(cudaStreamSynchronize(dc->stream),"pipe attention sync (dev out)");
|
||||
}
|
||||
/* absorb batch con TUTTO su device (q/latent/rope gia' residenti sulla scheda
|
||||
* dello shard, ctx resta sul device): il cuore della attention head-shardata
|
||||
* dentro il pipeline. Nessun trasferimento host. */
|
||||
extern "C" int coli_cuda_attention_absorb_batch_dev(ColiCudaTensor *w,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){
|
||||
if(!w||!ctx_dev||!q_dev||!latent_dev||!rope_dev||S<1||H<1||Q<1||R<1||V<1||
|
||||
K<1||K>512||T<S||T>8192||w->I!=K||w->O!=H*(Q+V))return 0;
|
||||
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<<<dim3(H,S),256,shared,dc->stream>>>(ctx_dev,q_dev,latent_dev,
|
||||
rope_dev,w->weights,w->scales,w->fmt,S,H,Q,R,V,K,T,scale);
|
||||
if(!cuda_ok(cudaGetLastError(),"pipe shard attention launch"))return 0;
|
||||
return cuda_ok(cudaStreamSynchronize(dc->stream),"pipe shard attention sync");
|
||||
}
|
||||
/* absorb per il DECODE con KV gia' residente: carica solo q (poche KB),
|
||||
* latent/rope arrivano dall'ombra device. ctx torna a host (S piccolo). */
|
||||
extern "C" int coli_cuda_attention_absorb_kvdev(ColiCudaTensor *w,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){
|
||||
if(!w||!ctx||!q||!latent_dev||!rope_dev||H<1||Q<1||R<1||V<1||K<1||K>512||T<1||T>8192||
|
||||
w->I!=K||w->O!=H*(Q+V))return 0;
|
||||
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
|
||||
size_t qb=(size_t)H*(Q+R)*sizeof(float),cb=(size_t)H*V*sizeof(float);
|
||||
if(!reserve(&dc->aq,&dc->aq_cap,qb)||!reserve(&dc->ac,&dc->ac_cap,cb))return 0;
|
||||
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<<<dim3(H,1),256,shared,dc->stream>>>(dc->ac,dc->aq,latent_dev,
|
||||
rope_dev,w->weights,w->scales,w->fmt,1,H,Q,R,V,K,T,scale);
|
||||
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;
|
||||
return 1;
|
||||
}
|
||||
extern "C" int coli_cuda_pipe_sync(int device){
|
||||
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
|
||||
return cuda_ok(cudaDeviceSynchronize(),"pipe sync");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user