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

GPU resident pipeline: batch CUDA attention, head-sharded kv_b, prefill expert groups, W4A16 mixed dispatch (#111)

Browse Source 
* 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
2026-07-15 00:18:05 +08:00
committed by GitHub
parent 2ead86a27f
commit ec89136029
13 changed files with 1267 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/glm.c
+411 -29
View File
@@ -113,6 +113,11 @@ typedef struct {
float *in_ln, *post_ln;
/* MLA (densa, quantizzata) */
QT q_a, q_b, kv_a, kv_b, o; float *q_a_ln, *kv_a_ln;
#ifdef COLI_CUDA
ColiCudaTensor *kv_b_shard[COLI_CUDA_MAX_DEVICES];
int shard_h0[COLI_CUDA_MAX_DEVICES],shard_hn[COLI_CUDA_MAX_DEVICES],n_kv_b_shard;
int shared_w4a16_failed;
#endif
int sparse;
/* dense mlp (sparse==0) */
QT gate_proj, up_proj, down_proj;
@@ -152,6 +157,7 @@ typedef struct {
int *kv_start; /* prima pos valida nella KV del layer (MTP: parziale) */
KVState *kv;
ESlot **ecache; int *ecn; int ecap; /* LRU expert per-layer */
float **kv_dev_L, **kv_dev_R; int *kv_dev_valid; /* ombra KV su device (decode) */
ESlot ws[64]; /* working set del layer corrente (load paralleli) */
ESlot **pin; int *npin; /* HOT-STORE: expert pinnati in RAM (mai evicted) */
uint32_t **eusage; /* contatori persistenti (per STATS/PIN) */
@@ -1065,6 +1071,33 @@ static float *ld(Model *m, const char *name){ /* tensore 1D f32 residente (nor
float *p=(float*)qalloc((size_t)n*sizeof(float)); /* registrato per la GPU sotto METAL */
st_read_f32(&m->S,name,p,0); return p;
}
#ifdef COLI_CUDA
static void qt_cuda_colocate(QT *dst,const QT *src){
if(!g_cuda_enabled||!g_cuda_dense||!dst->cuda_eligible||!src->cuda_eligible||
dst->cuda_device==src->cuda_device)return;
int old=-1,now=-1;for(int i=0;i<g_cuda_ndev;i++){
if(g_cuda_devices[i]==dst->cuda_device)old=i;if(g_cuda_devices[i]==src->cuda_device)now=i;
}
if(old>=0)g_cuda_dense_projected[old]-=qt_bytes(dst);
if(now>=0)g_cuda_dense_projected[now]+=qt_bytes(dst);
dst->cuda_device=src->cuda_device;
}
static void layer_cuda_shard_kvb(Layer *l,int H,int Q,int V){
if(!g_cuda_enabled||!g_cuda_dense||g_cuda_ndev<2||l->kv_b.fmt==0)return;
int rb=l->kv_b.fmt==1?l->kv_b.I:(l->kv_b.fmt==2?(l->kv_b.I+1)/2:(l->kv_b.I+3)/4);
const uint8_t *weights=l->kv_b.fmt==1?(const uint8_t*)l->kv_b.q8:l->kv_b.q4;
for(int d=0,h0=0;d<g_cuda_ndev;d++){
int hn=H/g_cuda_ndev+(d<H%g_cuda_ndev),rows=hn*(Q+V);
const void *part=weights+(int64_t)h0*(Q+V)*rb;
const float *scale=l->kv_b.s+(int64_t)h0*(Q+V);
if(!coli_cuda_tensor_upload(&l->kv_b_shard[d],part,scale,l->kv_b.fmt,l->kv_b.I,rows,g_cuda_devices[d]))return;
l->shard_h0[d]=h0;l->shard_hn[d]=hn;l->n_kv_b_shard++;h0+=hn;
}
int old=-1;for(int i=0;i<g_cuda_ndev;i++)if(g_cuda_devices[i]==l->kv_b.cuda_device)old=i;
if(old>=0)g_cuda_dense_projected[old]-=qt_bytes(&l->kv_b);
l->kv_b.cuda_eligible=0;
}
#endif
static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits){
memset(m,0,sizeof(*m)); m->ebits=ebits; m->dbits=dbits;
@@ -1079,6 +1112,8 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits
m->L=calloc(c->n_layers,sizeof(Layer));
int NR=c->n_layers+1; /* +1: riga del layer MTP */
m->ecap=cap; m->ecache=calloc(NR,sizeof(ESlot*)); m->ecn=calloc(NR,sizeof(int));
m->kv_dev_L=calloc(NR,sizeof(float*)); m->kv_dev_R=calloc(NR,sizeof(float*));
m->kv_dev_valid=calloc(NR,sizeof(int));
m->eroute=calloc(NR,sizeof(int*)); m->enr=calloc(NR,sizeof(int));
m->pin=calloc(NR,sizeof(ESlot*)); m->npin=calloc(NR,sizeof(int));
m->eusage=calloc(NR,sizeof(uint32_t*)); m->eheat=calloc(NR,sizeof(uint32_t*));
@@ -1097,6 +1132,14 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits
l->kv_a_ln= ld(m,P("self_attn.kv_a_layernorm.weight"));
l->kv_b = qt_load(m,P("self_attn.kv_b_proj.weight"), H*(c->qk_nope+c->v_head), c->kv_lora, dbits);
l->o = qt_load(m,P("self_attn.o_proj.weight"), D, H*c->v_head, dbits);
#ifdef COLI_CUDA
qt_cuda_colocate(&l->o,&l->kv_b);
qt_cuda_colocate(&l->q_a,&l->kv_b); /* PIPE: intera catena attention sulla */
qt_cuda_colocate(&l->q_b,&l->kv_b); /* stessa scheda / whole attention chain */
qt_cuda_colocate(&l->kv_a,&l->kv_b); /* on the layer home device */
if(getenv("COLI_CUDA_ATTN_SHARD")&&atoi(getenv("COLI_CUDA_ATTN_SHARD")))
layer_cuda_shard_kvb(l,H,c->qk_nope,c->v_head);
#endif
l->sparse = (i >= c->first_dense);
if(!l->sparse){
l->gate_proj = qt_load(m,P("mlp.gate_proj.weight"), c->dense_inter, D, dbits);
@@ -1109,6 +1152,11 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits
l->sh_gate = qt_load(m,P("mlp.shared_experts.gate_proj.weight"), sI, D, dbits);
l->sh_up = qt_load(m,P("mlp.shared_experts.up_proj.weight"), sI, D, dbits);
l->sh_down = qt_load(m,P("mlp.shared_experts.down_proj.weight"), D, sI, dbits);
#ifdef COLI_CUDA
qt_cuda_colocate(&l->sh_gate,&l->kv_b); /* PIPE2: shared chain on the layer home device */
qt_cuda_colocate(&l->sh_up,&l->sh_gate);
qt_cuda_colocate(&l->sh_down,&l->sh_gate);
#endif
m->ecache[i]=calloc(cap,sizeof(ESlot));
m->eroute[i]=calloc(c->topk,sizeof(int)); /* metodo C: ultimo routing del layer */
m->eusage[i]=calloc(c->n_experts,sizeof(uint32_t));
@@ -1573,7 +1621,10 @@ static void qt_matvec_rows(const QT *t, int r0, int n, const float *x, float *y)
y[j]=(float)a;
}
}
static int g_absorb=-1; /* ABSORB: -1 auto (decode S<=4), 0 mai, 1 sempre (test) */
static int g_absorb=-1;
#ifdef COLI_CUDA
static int g_cuda_pipe=0; /* COLI_CUDA_PIPE=1: prefill attention chain resident on the layer home device */
#endif /* ABSORB: -1 auto (decode S<=4), 0 mai, 1 sempre (test) */
static int g_dsa_force=0; /* DSA_FORCE=1: selezione sempre attiva (test: top-min(k,T)=denso) */
static int cmp_fdesc(const void *a,const void *b){
float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
@@ -1581,6 +1632,141 @@ static int cmp_fdesc(const void *a,const void *b){
/* attenzione MLA con KV-cache compressa, su token nuovi x[S,hidden], pos_base = pos del primo */
/* kvs/pos describe a ragged decode batch: each row may belong to a different
* sequence. NULL keeps the original contiguous, currently-bound KV path. */
#ifdef COLI_CUDA
/* Ombra KV su device per il DECODE: righe [0,upto) valide sulla scheda di kv_b.
* L'host resta canonico; l'ombra si riallinea in blocco quando resta indietro e
* viene invalidata da kv_bind / dalla riscrittura di righe gia' specchiate. */
static int kv_dev_sync(Model *m, Layer *l, int layer, int upto){
Cfg *c=&m->c; int kvl=c->kv_lora, R=c->qk_rope, dev=l->kv_b.cuda_device;
if(upto>m->max_t) return 0;
if(!m->kv_dev_L[layer]){
m->kv_dev_L[layer]=(float*)coli_cuda_pipe_alloc(dev,(size_t)m->max_t*kvl*4);
m->kv_dev_R[layer]=(float*)coli_cuda_pipe_alloc(dev,(size_t)m->max_t*R*4);
m->kv_dev_valid[layer]=0;
if(!m->kv_dev_L[layer]||!m->kv_dev_R[layer]) return 0;
}
int v=m->kv_dev_valid[layer];
if(v<upto){
if(!coli_cuda_pipe_upload(dev,m->kv_dev_L[layer]+(size_t)v*kvl,
coli_kv_row(m->kv->Lc[layer],v,kvl),(size_t)(upto-v)*kvl*4)||
!coli_cuda_pipe_upload(dev,m->kv_dev_R[layer]+(size_t)v*R,
coli_kv_row(m->kv->Rc[layer],v,R),(size_t)(upto-v)*R*4)) return 0;
m->kv_dev_valid[layer]=upto;
}
return 1;
}
/* Inc.1a — catena attention residente sul device del layer / attention chain
* resident on the layer home device. Proiezioni q/kv, norme, RoPE, batch
* attention e o_proj girano sulla scheda di kv_b; scaricano solo out [S,D],
* i nuovi record KV [S,kvl+R] e nulla altro. Ritorna 0 su qualsiasi errore:
* il chiamante riesegue il percorso CPU (idempotente). */
static int attn_pipe_prefill(Model *m, Layer *l, int layer, const float *x, int x_is_dev,
int S, int pos_base, float *out, float *out_dev){
Cfg *c=&m->c; int H=c->n_heads, D=c->hidden, qh=c->qk_head;
int kvl=c->kv_lora, R=c->qk_rope, ql=c->q_lora;
int dev=l->kv_b.cuda_device;
if(l->q_a.cuda_device!=dev||l->q_b.cuda_device!=dev||
l->kv_a.cuda_device!=dev||l->o.cuda_device!=dev) return 0;
int st0=m->kv_start[layer], T=pos_base+S-st0, old=pos_base-st0;
if(T<S||T>8192) return 0;
double t0=now_s();
size_t xb=(size_t)S*D*4, qrb=(size_t)S*ql*4, qb=(size_t)S*H*qh*4;
size_t cb=(size_t)S*(kvl+R)*4, lb=(size_t)T*kvl*4, rb=(size_t)T*R*4;
float *chost=NULL; int ok=0;
/* scratch persistenti (slot fissi per device): zero churn di cudaMalloc */
float *xd =x_is_dev?(float*)x:coli_cuda_pipe_scratch(dev,0,xb);
float *qrd=coli_cuda_pipe_scratch(dev,1,qrb);
float *qd =coli_cuda_pipe_scratch(dev,2,qb), *cd =coli_cuda_pipe_scratch(dev,3,cb);
float *ld_=coli_cuda_pipe_scratch(dev,4,lb), *rd =coli_cuda_pipe_scratch(dev,5,rb);
float *w1 =coli_cuda_pipe_scratch(dev,6,(size_t)ql*4);
float *w2 =coli_cuda_pipe_scratch(dev,7,(size_t)kvl*4);
chost=(float*)malloc(cb);
if(!xd||!qrd||!qd||!cd||!ld_||!rd||!w1||!w2||!chost) goto done;
if((!x_is_dev&&!coli_cuda_pipe_upload(dev,xd,x,xb))||
!coli_cuda_pipe_upload(dev,w1,l->q_a_ln,(size_t)ql*4)||
!coli_cuda_pipe_upload(dev,w2,l->kv_a_ln,(size_t)kvl*4)) goto done;
/* proiezioni + norme + rope, tutto sul device */
if(!coli_cuda_pipe_gemm(l->q_a.cuda,qrd,xd,S)) goto done;
if(!coli_cuda_pipe_rmsnorm(dev,qrd,qrd,w1,S,ql,c->eps)) goto done;
if(!coli_cuda_pipe_gemm(l->q_b.cuda,qd,qrd,S)) goto done;
if(!coli_cuda_pipe_rope_base(dev,qd,pos_base,S*H,qh,c->qk_nope,R,H,c->theta)) goto done;
if(!coli_cuda_pipe_gemm(l->kv_a.cuda,cd,xd,S)) goto done;
if(!coli_cuda_pipe_rmsnorm_s(dev,cd,cd,w2,S,kvl,c->eps,kvl+R,kvl+R)) goto done;
if(!coli_cuda_pipe_rope_base(dev,cd,pos_base,S,kvl+R,kvl,R,1,c->theta)) goto done;
/* cache latente [T,kvl] + rot [T,R] contigue: righe vecchie da host, nuove da cd */
if(old>0){
if(!coli_cuda_pipe_upload(dev,ld_,coli_kv_row(m->Lc[layer],st0,kvl),(size_t)old*kvl*4)||
!coli_cuda_pipe_upload(dev,rd,coli_kv_row(m->Rc[layer],st0,R),(size_t)old*R*4)) goto done;
}
if(!coli_cuda_pipe_copy2d(dev,ld_+(size_t)old*kvl,kvl,cd,kvl+R,kvl,S)) goto done;
if(!coli_cuda_pipe_copy2d(dev,rd+(size_t)old*R,R,cd+kvl,kvl+R,R,S)) goto done;
/* KV host resta canonica: scarica i record nuovi (gia' normati+ropati) */
if(!coli_cuda_pipe_download(dev,cd,chost,cb)) goto done;
for(int s=0;s<S;s++){
memcpy(coli_kv_row(m->Lc[layer],pos_base+s,kvl),chost+(size_t)s*(kvl+R),kvl*4);
memcpy(coli_kv_row(m->Rc[layer],pos_base+s,R),chost+(size_t)s*(kvl+R)+kvl,R*4);
}
if(m->kv_dev_valid[layer]>pos_base) m->kv_dev_valid[layer]=pos_base;
m->t_aproj+=now_s()-t0; t0=now_s();
#ifdef COLI_CUDA
/* Negativo (2026-07-13): P2P a stella dal device di casa serializza ~95MB/layer
* sul suo link PCIe — attention 26->41-44s. Resta opt-in per topologie NVLink. */
if(out_dev && l->n_kv_b_shard>1 &&
getenv("COLI_CUDA_PIPE_SHARD") && atoi(getenv("COLI_CUDA_PIPE_SHARD"))){
/* head-shard nel pipeline: q gia' sul device di casa. Per ogni scheda:
* slice di q (repack strided->contiguo), broadcast latent+rope via P2P,
* score parallelo sui rispettivi head, ctx slice riportata a casa e
* ricomposta, poi o_proj residente. */
int n=l->n_kv_b_shard, vh=c->v_head;
size_t ctxb=(size_t)S*H*vh*4;
size_t stage_one=(size_t)S*H*(size_t)(c->qk_head>vh?c->qk_head:vh)*4;
float *ctx_full=coli_cuda_pipe_scratch(dev,16,ctxb);
float *stage=coli_cuda_pipe_scratch(dev,17,stage_one*n);
int ok_sh=(ctx_full&&stage)?1:0;
if(ok_sh){
#pragma omp parallel for schedule(static) reduction(&:ok_sh)
for(int d2=0;d2<n;d2++){
int hn=l->shard_hn[d2], h0=l->shard_h0[d2];
int sdev=coli_cuda_tensor_device(l->kv_b_shard[d2]);
float *st=stage+(size_t)d2*(stage_one/4);
size_t qsb=(size_t)S*hn*qh*4, csb=(size_t)S*hn*vh*4;
float *qs_r=coli_cuda_pipe_scratch(sdev,18,qsb);
float *ld_r=coli_cuda_pipe_scratch(sdev,19,(size_t)T*kvl*4);
float *rr_r=coli_cuda_pipe_scratch(sdev,20,(size_t)T*R*4);
float *cx_r=coli_cuda_pipe_scratch(sdev,21,csb);
int okd=qs_r&&ld_r&&rr_r&&cx_r;
/* slice di q: [S,H,qh] -> [S,hn,qh] contigua sul device di casa */
okd=okd&&coli_cuda_pipe_copy2d(dev,st,hn*qh,qd+(size_t)h0*qh,H*qh,hn*qh,S);
okd=okd&&coli_cuda_pipe_peer_copy(sdev,qs_r,dev,st,qsb);
okd=okd&&coli_cuda_pipe_peer_copy(sdev,ld_r,dev,ld_,(size_t)T*kvl*4);
okd=okd&&coli_cuda_pipe_peer_copy(sdev,rr_r,dev,rd,(size_t)T*R*4);
okd=okd&&coli_cuda_attention_absorb_batch_dev(l->kv_b_shard[d2],cx_r,qs_r,ld_r,rr_r,
S,hn,c->qk_nope,R,vh,kvl,T,c->attn_scale);
okd=okd&&coli_cuda_pipe_peer_copy(dev,st,sdev,cx_r,csb);
okd=okd&&coli_cuda_pipe_copy2d(dev,ctx_full+(size_t)h0*vh,H*vh,st,hn*vh,hn*vh,S);
ok_sh&=okd;
}
}
if(ok_sh){
ok=coli_cuda_pipe_gemm(l->o.cuda,out_dev,ctx_full,S)&&coli_cuda_pipe_sync(dev);
} else ok=0;
if(!ok)
ok=coli_cuda_attention_project_batch_dev_out(l->kv_b.cuda,l->o.cuda,out_dev,qd,ld_,rd,
S,H,c->qk_nope,R,c->v_head,kvl,T,c->attn_scale);
} else
#endif
ok=out_dev?coli_cuda_attention_project_batch_dev_out(l->kv_b.cuda,l->o.cuda,out_dev,qd,ld_,rd,
S,H,c->qk_nope,R,c->v_head,kvl,T,c->attn_scale)
:coli_cuda_attention_project_batch_dev(l->kv_b.cuda,l->o.cuda,out,qd,ld_,rd,
S,H,c->qk_nope,R,c->v_head,kvl,T,c->attn_scale);
m->t_acore+=now_s()-t0;
done:
free(chost); /* gli scratch device restano al contesto */
return ok;
}
#endif
static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int pos_base,
KVState *const *kvs, const int *positions, float *out){
Cfg *c=&m->c; int H=c->n_heads, D=c->hidden, qh=c->qk_head, vh=c->v_head;
@@ -1631,12 +1817,24 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
* batch-union. matmul_qt_ex(...,0) keeps them on the EXACT int4 kernel: letting S>1 pull
* them into IDOT is much faster but costs ~12% perplexity (measured). Batching alone is
* bit-identical to upstream; the kernel switch is not. */
matmul_qt_ex(QR, x, &l->q_a, S, 0);
for(int s=0;s<S;s++){ float *qr=QR+(int64_t)s*c->q_lora;
rmsnorm(qr, qr, l->q_a_ln, c->q_lora, c->eps); } /* q_b legge il residuo NORMATO */
matmul_qt_ex(Q, QR, &l->q_b, S, 0);
matmul_qt_ex(comp, x, &l->kv_a, S, 0);
for(int s=0;s<S;s++){
int pipe_done=0;
#ifdef COLI_CUDA
if(g_cuda_pipe&&!kvs&&S>=8&&layer<c->n_layers&&g_cuda_enabled&&c->kv_lora<=512&&
!(m->has_dsa&&pos_base+S>c->index_topk)&&
l->q_a.cuda_eligible&&l->q_b.cuda_eligible&&l->kv_a.cuda_eligible&&
l->kv_b.cuda_eligible&&l->o.cuda_eligible&&
qt_cuda_upload(&l->q_a)&&qt_cuda_upload(&l->q_b)&&qt_cuda_upload(&l->kv_a)&&
qt_cuda_upload(&l->kv_b)&&qt_cuda_upload(&l->o))
pipe_done=attn_pipe_prefill(m,l,layer,x,0,S,pos_base,out,NULL);
#endif
if(!pipe_done){
matmul_qt_ex(QR, x, &l->q_a, S, 0);
for(int s=0;s<S;s++){ float *qr=QR+(int64_t)s*c->q_lora;
rmsnorm(qr, qr, l->q_a_ln, c->q_lora, c->eps); } /* q_b legge il residuo NORMATO */
matmul_qt_ex(Q, QR, &l->q_b, S, 0);
matmul_qt_ex(comp, x, &l->kv_a, S, 0);
}
if(!pipe_done) for(int s=0;s<S;s++){
KVState *ks=kvs?kvs[s]:m->kv;
int pos=positions?positions[s]:pos_base+s;
float *qfull=Q+(int64_t)s*H*qh;
@@ -1644,6 +1842,10 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
const float *cs=comp+(int64_t)s*cw;
float *Ldst=coli_kv_row(ks->Lc[layer],pos,c->kv_lora);
float *Rdst=coli_kv_row(ks->Rc[layer],pos,c->qk_rope);
#ifdef COLI_CUDA
if(ks==m->kv&&m->kv_dev_valid&&layer<=c->n_layers&&m->kv_dev_valid[layer]>pos)
m->kv_dev_valid[layer]=pos; /* riga riscritta: l'ombra si accorcia */
#endif
memcpy(Ldst, cs, c->kv_lora*sizeof(float));
rmsnorm(Ldst, Ldst, l->kv_a_ln, c->kv_lora, c->eps); /* latente normato */
memcpy(Rdst, cs+c->kv_lora, c->qk_rope*sizeof(float));
@@ -1720,7 +1922,17 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
* k/v per ogni token del contesto. Per linearita':
* q·k_nope_t = (W_K^hT q_nope)·L_t ctx^h = W_V^h (Σ_t a_t L_t)
* costo per step ~O(T·kv_lora) invece di O(T·H·(nope+vh)) del matmul kvb_all. */
int absorb = kvs || g_absorb==1 || (g_absorb<0 && S<=4);
if(pipe_done){
free(ctx); free(Q); free(QR); free(comp);
m->t_attn += now_s()-ta0;
return;
}
int cuda_absorb=0;
#ifdef COLI_CUDA
cuda_absorb=layer<c->n_layers&&!kvs&&g_cuda_enabled&&getenv("COLI_CUDA_ATTN")&&
atoi(getenv("COLI_CUDA_ATTN"))&&c->kv_lora<=512;
#endif
int absorb = kvs || g_absorb==1 || (g_absorb<0 && S<=4) || cuda_absorb;
if(absorb && c->kv_lora<=512){
m->t_aproj+=now_s()-ta0; double tac=now_s();
int kvl=c->kv_lora, r0v=c->qk_nope; /* offset righe V dentro il blocco di testa */
@@ -1739,19 +1951,48 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
if(nt>sc_cap) sc_cap=nt;
}
float *sc_all = falloc((int64_t)omp_get_max_threads()*sc_cap);
int cuda_core=0;
int cuda_core=0,cuda_projected=0;
#ifdef COLI_CUDA
if(S<=4&&g_cuda_enabled&&getenv("COLI_CUDA_ATTN")&&atoi(getenv("COLI_CUDA_ATTN"))&&
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(
qs+(int64_t)l->shard_h0[d]*S*qh+(int64_t)s*l->shard_hn[d]*qh,
Q+((int64_t)s*H+l->shard_h0[d])*qh,(size_t)l->shard_hn[d]*qh*sizeof(float));
#pragma omp parallel for schedule(static) reduction(&:ok)
for(int d=0;d<n;d++)ok&=coli_cuda_attention_absorb_batch(l->kv_b_shard[d],
cs+(int64_t)l->shard_h0[d]*S*vh,qs+(int64_t)l->shard_h0[d]*S*qh,
coli_kv_row(m->Lc[layer],st0,kvl),coli_kv_row(m->Rc[layer],st0,c->qk_rope),
S,l->shard_hn[d],c->qk_nope,c->qk_rope,vh,kvl,nt,c->attn_scale);
if(ok)for(int d=0;d<n;d++)for(int s=0;s<S;s++)memcpy(
ctx+((int64_t)s*H+l->shard_h0[d])*vh,
cs+(int64_t)l->shard_h0[d]*S*vh+(int64_t)s*l->shard_hn[d]*vh,
(size_t)l->shard_hn[d]*vh*sizeof(float));
free(qs);free(cs);cuda_core=ok;
} else if(cuda_absorb&&l->kv_b.cuda_eligible&&l->o.cuda_eligible&&
qt_cuda_upload(&l->kv_b)&&qt_cuda_upload(&l->o)){
int st0=m->kv_start[layer],nt=pos_base+S-st0;
cuda_core=cuda_projected=coli_cuda_attention_project_batch(l->kv_b.cuda,l->o.cuda,out,Q,
coli_kv_row(m->Lc[layer],st0,kvl),coli_kv_row(m->Rc[layer],st0,c->qk_rope),
S,H,c->qk_nope,c->qk_rope,vh,kvl,nt,c->attn_scale);
} else if(S<=4&&g_cuda_enabled&&getenv("COLI_CUDA_ATTN")&&atoi(getenv("COLI_CUDA_ATTN"))&&
l->kv_b.cuda_eligible&&qt_cuda_upload(&l->kv_b)){
cuda_core=1;
for(int s=0;s<S&&cuda_core;s++){
KVState *ks=kvs?kvs[s]:m->kv;int pos=positions?positions[s]:pos_base+s;
int st0=ks->kv_start[layer],nt=pos+1-st0;
if(dnsel&&dnsel[s]>0){cuda_core=0;break;}
cuda_core=coli_cuda_attention_absorb(l->kv_b.cuda,ctx+(int64_t)s*H*vh,
Q+(int64_t)s*H*qh,coli_kv_row(ks->Lc[layer],st0,kvl),
coli_kv_row(ks->Rc[layer],st0,c->qk_rope),H,c->qk_nope,c->qk_rope,
vh,kvl,nt,c->attn_scale);
cuda_core=0;
if(g_cuda_pipe&&ks==m->kv&&layer<c->n_layers&&kv_dev_sync(m,l,layer,pos+1))
cuda_core=coli_cuda_attention_absorb_kvdev(l->kv_b.cuda,ctx+(int64_t)s*H*vh,
Q+(int64_t)s*H*qh,m->kv_dev_L[layer]+(size_t)st0*kvl,
m->kv_dev_R[layer]+(size_t)st0*c->qk_rope,H,c->qk_nope,c->qk_rope,
vh,kvl,nt,c->attn_scale);
if(!cuda_core)
cuda_core=coli_cuda_attention_absorb(l->kv_b.cuda,ctx+(int64_t)s*H*vh,
Q+(int64_t)s*H*qh,coli_kv_row(ks->Lc[layer],st0,kvl),
coli_kv_row(ks->Rc[layer],st0,c->qk_rope),H,c->qk_nope,c->qk_rope,
vh,kvl,nt,c->attn_scale);
}
}
#endif
@@ -1786,7 +2027,7 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
}
}
m->t_acore+=now_s()-tac; double tao=now_s();
matmul_qt(out, ctx, &l->o, S); m->t_aout+=now_s()-tao;
if(!cuda_projected){matmul_qt(out, ctx, &l->o, S);} m->t_aout+=now_s()-tao;
free(ctx); free(Q); free(QR); free(comp); free(sc_all);
m->t_attn += now_s()-ta0;
return;
@@ -1840,7 +2081,7 @@ static void attention(Model *m, Layer *l, int layer, float *x, int S, int pos_ba
* una volta sola e moltiplicato per tutte le posizioni che lo usano (pesi letti 1 volta);
* lo shared expert e' un unico matmul a S righe. Per posizione l'accumulo resta
* nell'ordine (routed nel loro ordine di union, poi shared). */
static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int with_shared){
if(g_pilot_real){ /* barriera cross-layer: prendi possesso di QUESTO layer e aspetta
* l'eventuale load-pilota in volo sullo stesso layer (dopodiche' il
* worker droppa ogni nuovo load <= layer -> ecache[layer] e' stabile
@@ -1855,11 +2096,14 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
}
}
Cfg *c=&m->c; int D=c->hidden, E=c->n_experts, K=c->topk, I=c->moe_inter;
float *logit=falloc(E), *choice=falloc(E);
float *choice=falloc(E);
int sI=c->moe_inter*c->n_shared;
/* ---- FASE A: routing di tutte le S posizioni ---- */
int *idxs=malloc((size_t)S*K*sizeof(int)); float *ws=malloc((size_t)S*K*sizeof(float));
int *keff=malloc(S*sizeof(int));
/* router in UN matmul batch: stessa matematica, via le S chiamate S=1 */
float *logits_all=falloc((int64_t)S*E);
int pre_routed=0; (void)pre_routed;
#ifdef COLI_METAL
if(g_pre_idx){ /* routing gia' calcolata dal layer CB (GPU) */
memcpy(idxs,g_pre_idx,(size_t)S*K*sizeof(int));
@@ -1874,11 +2118,13 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
}
for(int d=0;d<D;d++) out[(int64_t)s*D+d]=0;
}
} else
pre_routed=1;
}
#endif
if(!pre_routed) matmul(logits_all, x, l->router, S, D, E);
if(!pre_routed)
for(int s=0;s<S;s++){
const float *xs=x+(int64_t)s*D;
matmul(logit, xs, l->router, 1, D, E);
float *logit=logits_all+(int64_t)s*E;
for(int e=0;e<E;e++){ logit[e]=sigmoidf(logit[e]); choice[e]=logit[e]+l->router_bias[e]; }
int *idx=idxs+(int64_t)s*K; float *w=ws+(int64_t)s*K;
int Ksel = g_topk>0 ? (g_topk<K?g_topk:K) : K;
@@ -1938,7 +2184,10 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
float *xg=falloc((int64_t)S*D), *gg=falloc((int64_t)S*I), *uu=falloc((int64_t)S*I), *hh=falloc((int64_t)S*D);
int *rows=malloc(S*sizeof(int)); float *rw=malloc(S*sizeof(float));
#ifdef COLI_CUDA
int group_enabled=S<=64;
/* PIPE Inc.1b: il batch-union del prefill passa dai gruppi GPU — prima di
* questo, 9343 expert in VRAM restavano INUTILIZZATI durante il prefill
* (misurato: 81s di expert-matmul tutto su CPU, GPU groups 21ms totali). */
int group_enabled = S<=64 || (g_cuda_pipe && S<=4096);
float *group_x=group_enabled?falloc((int64_t)S*K*D):NULL;
float *group_y=group_enabled?falloc((int64_t)S*K*D):NULL;
int *group_row=group_enabled?malloc((size_t)64*S*sizeof(int)):NULL;
@@ -2174,22 +2423,43 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){
ESlot tmp=*dst; *dst=m->ws[q]; m->ws[q]=tmp; dst->used=(uint64_t)__atomic_add_fetch(&m->eclock,1,__ATOMIC_RELAXED); }
}
}
/* ---- FASE E: shared expert, un matmul a S righe (skipped se fuso nel blocco GPU) ---- */
float *sg=falloc((int64_t)S*sI), *su=falloc((int64_t)S*sI);
/* ---- FASE E: shared expert (PIPE2: gia' sul device; Metal CB: gia' sommata) ---- */
if(!with_shared) goto shared_done;
{
float *sg=NULL,*su=NULL;int shared_cuda=0;
#ifdef COLI_METAL
if(g_pre_sh){ for(int64_t z=0;z<(int64_t)S*D;z++) out[z]+=g_pre_sh[z]; shared_on_gpu=1; }
if(shared_on_gpu) shared_cuda=2; /* gia' sommato in out: salta calcolo e add */
#endif
if(!shared_on_gpu){
#ifdef COLI_CUDA
int shared_min=getenv("COLI_CUDA_SHARED_W4A16_MIN_ROWS")?
atoi(getenv("COLI_CUDA_SHARED_W4A16_MIN_ROWS")):32;
if(shared_min<16)shared_min=16;
if(shared_cuda==0&&S>=shared_min&&!l->shared_w4a16_failed&&!omp_in_parallel()&&g_cuda_enabled&&
l->sh_gate.fmt==2&&l->sh_up.fmt==2&&l->sh_down.fmt==2&&
getenv("COLI_CUDA_SHARED_W4A16")&&atoi(getenv("COLI_CUDA_SHARED_W4A16"))&&
qt_cuda_upload(&l->sh_gate)&&qt_cuda_upload(&l->sh_up)&&qt_cuda_upload(&l->sh_down)){
shared_cuda=coli_cuda_shared_mlp_w4a16(l->sh_gate.cuda,l->sh_up.cuda,
l->sh_down.cuda,hh,x,S);
if(!shared_cuda)l->shared_w4a16_failed=1;
}
#endif
if(!shared_cuda){
sg=falloc((int64_t)S*sI);su=falloc((int64_t)S*sI);
matmul_qt(sg, x, &l->sh_gate, S);
matmul_qt(su, x, &l->sh_up, S);
for(int64_t z=0;z<(int64_t)S*sI;z++) sg[z]=siluf(sg[z])*su[z];
matmul_qt(hh, sg, &l->sh_down, S);
for(int64_t z=0;z<(int64_t)S*D;z++) out[z]+=hh[z];
}
free(logit); free(choice); free(idxs); free(ws); free(keff); free(uniq);
free(xg); free(gg); free(uu); free(hh); free(rows); free(rw); free(sg); free(su);
if(shared_cuda!=2) for(int64_t z=0;z<(int64_t)S*D;z++) out[z]+=hh[z];
free(sg); free(su);
}
shared_done:
free(logits_all); free(choice); free(idxs); free(ws); free(keff); free(uniq);
free(xg); free(gg); free(uu); free(hh); free(rows); free(rw);
#ifdef COLI_CUDA
free(group_x); free(group_y); free(group_row); free(group_weight);
free(group_x);free(group_y);
free(group_row); free(group_weight);
#endif
}
@@ -2393,6 +2663,65 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S){
}
/* forward di UN layer (usato dai 78 principali e dal layer MTP) */
#ifdef COLI_CUDA
/* Inc.2a — intero layer SPARSO residente sul device del layer. x_dev entra e resta;
* lasciano il device solo: nrm post-attention (router + expert CPU + gather dei
* gruppi), i nuovi record KV, e la nrm pre-attention sui layer con indexer DSA.
* Ritorna 0 su errore: il chiamante ripristina lo snapshot e rifa' il layer su CPU. */
static int pipe_layer_sparse(Model *m, Layer *l, int li, float *x_dev, int S, int pos_base,
float *nrm_host, float *out_host){
Cfg *c=&m->c; int D=c->hidden, dev=l->kv_b.cuda_device;
int sI=c->moe_inter*c->n_shared;
size_t xb=(size_t)S*D*4;
if(!l->sh_gate.cuda_eligible||!l->sh_up.cuda_eligible||!l->sh_down.cuda_eligible||
!qt_cuda_upload(&l->sh_gate)||!qt_cuda_upload(&l->sh_up)||!qt_cuda_upload(&l->sh_down)||
l->sh_gate.cuda_device!=dev||l->sh_up.cuda_device!=dev||l->sh_down.cuda_device!=dev) return 0;
float *w_in =coli_cuda_pipe_scratch(dev,8,(size_t)D*4);
float *w_post=coli_cuda_pipe_scratch(dev,9,(size_t)D*4);
float *nrm_d=coli_cuda_pipe_scratch(dev,10,xb);
float *y_d =coli_cuda_pipe_scratch(dev,11,xb);
float *sg_d =coli_cuda_pipe_scratch(dev,12,(size_t)S*sI*4);
float *su_d =coli_cuda_pipe_scratch(dev,13,(size_t)S*sI*4);
float *snap =coli_cuda_pipe_scratch(dev,14,xb);
if(!w_in||!w_post||!nrm_d||!y_d||!sg_d||!su_d||!snap) return 0;
if(!coli_cuda_pipe_peer_copy(dev,snap,dev,x_dev,xb)) return 0; /* snapshot per il fallback */
if(!coli_cuda_pipe_upload(dev,w_in,l->in_ln,(size_t)D*4)||
!coli_cuda_pipe_upload(dev,w_post,l->post_ln,(size_t)D*4)) return 0;
double ta=now_s();
if(!coli_cuda_pipe_rmsnorm(dev,nrm_d,x_dev,w_in,S,D,c->eps)) return 0;
/* DSA: i layer con indexer FULL cachano k_idx dalla nrm pre-attention (CPU, piccolo) */
if(m->has_dsa && li<c->n_layers && m->kv_start[li]==0 && c->idx_type[li]){
if(!coli_cuda_pipe_download(dev,nrm_d,nrm_host,xb)) return 0;
int nh=c->index_nh, hd=c->index_hd; (void)nh;
for(int s=0;s<S;s++){
int pos=pos_base+s;
float *kd=coli_kv_row(m->kv->Ic[li],pos,hd);
matmul_qt(kd, nrm_host+(int64_t)s*D, &m->ix_wk[li], 1);
layernorm(kd, m->ix_knw[li], m->ix_knb[li], hd, 1e-6f);
rope_interleave(kd, pos, c);
}
}
if(!attn_pipe_prefill(m,l,li,nrm_d,1,S,pos_base,NULL,y_d)) return 0;
if(!coli_cuda_pipe_add(dev,x_dev,y_d,(size_t)S*D)) return 0; /* prima mutazione */
if(!coli_cuda_pipe_rmsnorm(dev,nrm_d,x_dev,w_post,S,D,c->eps)) return 0;
if(!coli_cuda_pipe_download(dev,nrm_d,nrm_host,xb)) return 0;
m->t_attn+=now_s()-ta;
/* expert routed su CPU/gruppi GPU come oggi (shared saltata: la fa il device) */
moe(m,l,li,nrm_host,S,out_host,0);
double te=now_s();
if(!coli_cuda_pipe_upload(dev,y_d,out_host,xb)) return 0;
if(!coli_cuda_pipe_add(dev,x_dev,y_d,(size_t)S*D)) return 0;
if(!coli_cuda_pipe_gemm(l->sh_gate.cuda,sg_d,nrm_d,S)) return 0;
if(!coli_cuda_pipe_gemm(l->sh_up.cuda,su_d,nrm_d,S)) return 0;
if(!coli_cuda_pipe_silu_mul(dev,sg_d,su_d,(size_t)S*sI)) return 0;
if(!coli_cuda_pipe_gemm(l->sh_down.cuda,y_d,sg_d,S)) return 0;
if(!coli_cuda_pipe_add(dev,x_dev,y_d,(size_t)S*D)) return 0;
if(!coli_cuda_pipe_sync(dev)) return 0;
m->t_emm+=now_s()-te;
return 1;
}
#endif
static void layer_forward_rows(Model *m, Layer *l, int li, float *x, int S, int pos_base,
KVState *const *kvs, const int *positions, float *nrm, float *tmp){
Cfg *c=&m->c; int D=c->hidden;
@@ -2459,7 +2788,7 @@ static void layer_forward_rows(Model *m, Layer *l, int li, float *x, int S, int
if(g_pilot && S<=8 && li+1<c->n_layers && m->L[li+1].sparse) pilot_prefetch(m,li+1,x,S);
if(g_looka && S==1 && li+1<c->n_layers && m->L[li+1].sparse) la_predict(m,li+1,x,1);
for(int s=0;s<S;s++) rmsnorm(nrm+(int64_t)s*D, x+(int64_t)s*D, l->post_ln, D, c->eps);
if(l->sparse) moe(m,l,li,nrm,S,tmp); else dense_mlp(l,nrm,S,D,c->dense_inter,tmp);
if(l->sparse) moe(m,l,li,nrm,S,tmp,1); else dense_mlp(l,nrm,S,D,c->dense_inter,tmp);
for(int64_t j=0;j<(int64_t)S*D;j++) x[j]+=tmp[j];
}
static void layer_forward(Model *m, Layer *l, int li, float *x, int S, int pos_base, float *nrm, float *tmp){
@@ -2474,13 +2803,53 @@ static void layers_forward_rows(Model *m, float *x, int S, int pos_base,
pthread_mutex_unlock(&g_pilot_mx);
}
float *nrm=falloc((int64_t)S*D), *tmp=falloc((int64_t)S*D);
#ifdef COLI_CUDA
/* PIPE2 (Inc.2a): il residuo resta sul device del layer, saltando tra le schede
* ai confini di layer. x host diventa STALE finche' la residenza e' attiva. */
float *x_dev=NULL; int x_dev_on=-1;
size_t xb=(size_t)S*(size_t)D*4;
int pipe2 = g_cuda_pipe>=2 && !kvs && S>=8 && g_cuda_enabled && c->kv_lora<=512 &&
!(m->has_dsa && pos_base+S>c->index_topk);
#endif
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);
#ifdef COLI_CUDA
Layer *l=&m->L[i];
if(pipe2 && l->sparse && i<c->n_layers &&
l->q_a.cuda_eligible&&l->q_b.cuda_eligible&&l->kv_a.cuda_eligible&&
l->kv_b.cuda_eligible&&l->o.cuda_eligible&&
qt_cuda_upload(&l->q_a)&&qt_cuda_upload(&l->q_b)&&qt_cuda_upload(&l->kv_a)&&
qt_cuda_upload(&l->kv_b)&&qt_cuda_upload(&l->o)&&
l->q_a.cuda_device==l->kv_b.cuda_device&&l->q_b.cuda_device==l->kv_b.cuda_device&&
l->kv_a.cuda_device==l->kv_b.cuda_device&&l->o.cuda_device==l->kv_b.cuda_device){
int dev=l->kv_b.cuda_device, ok=1;
float *dst=coli_cuda_pipe_scratch(dev,15,xb);
if(dst){
if(x_dev_on<0) ok=coli_cuda_pipe_upload(dev,dst,x,xb);
else if(x_dev_on!=dev){ ok=coli_cuda_pipe_peer_copy(dev,dst,x_dev_on,x_dev,xb); }
else dst=x_dev;
if(ok){
x_dev=dst; x_dev_on=dev;
if(pipe_layer_sparse(m,l,i,x_dev,S,pos_base,nrm,tmp)) continue;
/* fallback: snapshot -> host, layer rifatto sul percorso CPU */
coli_cuda_pipe_peer_copy(dev,x_dev,dev,coli_cuda_pipe_scratch(dev,14,xb),xb);
coli_cuda_pipe_download(dev,x_dev,x,xb);
x_dev_on=-1;
}else x_dev_on=-1;
}
} else if(x_dev_on>=0){ /* layer fuori pipe: il residuo torna a casa */
coli_cuda_pipe_download(x_dev_on,x_dev,x,xb);
x_dev_on=-1;
}
#endif
layer_forward_rows(m,&m->L[i],i,x,S,pos_base,kvs,positions,nrm,tmp);
}
#ifdef COLI_CUDA
if(x_dev_on>=0) coli_cuda_pipe_download(x_dev_on,x_dev,x,xb);
#endif
free(nrm); free(tmp);
}
static void layers_forward(Model *m, float *x, int S, int pos_base){
@@ -2490,6 +2859,14 @@ static void layers_forward(Model *m, float *x, int S, int pos_base){
static void kv_alloc(Model *m, int max_t){
Cfg *c=&m->c;
KVState *k=m->kv;
#ifdef COLI_CUDA
if(m->kv_dev_L) for(int i=0;i<c->n_layers+1;i++){ /* dimensioni cambiate: ombra da rifare */
if(m->kv_dev_L[i]){ coli_cuda_pipe_free(m->L[i<c->n_layers?i:0].kv_b.cuda_device,m->kv_dev_L[i]); m->kv_dev_L[i]=NULL; }
if(m->kv_dev_R[i]){ coli_cuda_pipe_free(m->L[i<c->n_layers?i:0].kv_b.cuda_device,m->kv_dev_R[i]); m->kv_dev_R[i]=NULL; }
m->kv_dev_valid[i]=0;
}
#endif
if(k->Lc){ for(int i=0;i<c->n_layers+1;i++){ free(k->Lc[i]); free(k->Rc[i]); } free(k->Lc); free(k->Rc); }
if(k->Lc){ for(int i=0;i<c->n_layers+1;i++){
#ifdef COLI_METAL
if(g_metal_enabled){ coli_metal_unregister(k->Lc[i]); coli_metal_unregister(k->Rc[i]); }
@@ -2522,6 +2899,8 @@ static void kv_alloc(Model *m, int max_t){
}
static void kv_bind(Model *m, KVState *k){
if(m->kv!=k && m->kv_dev_valid) /* ombra legata al KVState corrente */
for(int i=0;i<m->c.n_layers+1;i++) m->kv_dev_valid[i]=0;
m->kv=k; m->Lc=k->Lc; m->Rc=k->Rc; m->Ic=k->Ic;
m->max_t=k->max_t; m->kv_start=k->kv_start;
}
@@ -3034,6 +3413,7 @@ static void run_text(Model *m, const char *snap, const char *prompt, int ngen){
repin_pass_limit(m,limit); /* prompt routing seeds every GPU layer */
}
prefill_t=now_s()-prefill_t;
printf("PROFILO PREFILL (%.2fs):\n",prefill_t); profile_print(m,prefill_t);
m->hits=m->miss=m->ereq=m->gpu_expert_calls=0;
m->n_emit=m->n_fw=0;
g_last_repin=0;
@@ -3329,6 +3709,7 @@ static void mux_done(Model *m, ServeCtx *sc, ServeReq *r){
double dt=now_s()-r->started; if(dt<1e-6) dt=1e-6;
double dh=(double)(m->hits-r->hits0), dm=(double)(m->miss-r->miss0);
hwinfo_emit(m);
usage_save(m); /* la cache che impara non deve aspettare l'uscita */
tiers_emit(m);
emap_emit(m);
hits_emit(m);
@@ -4204,6 +4585,7 @@ int main(int argc, char **argv){
if(!g_cuda_enabled){ fprintf(stderr,"[CUDA] requested backend is unavailable\n"); return 2; }
}
g_cuda_dense=getenv("CUDA_DENSE")?atoi(getenv("CUDA_DENSE")):0;
g_cuda_pipe=getenv("COLI_CUDA_PIPE")?atoi(getenv("COLI_CUDA_PIPE")):0;
const char *cuda_expert=getenv("CUDA_EXPERT_GB");
g_cuda_expert_auto=cuda_expert&&!strcmp(cuda_expert,"auto");
g_cuda_expert_gb=cuda_expert&&!g_cuda_expert_auto?atof(cuda_expert):0;