diff --git a/c/backend_cuda.cu b/c/backend_cuda.cu index a0083f4..30a1f4a 100644 --- a/c/backend_cuda.cu +++ b/c/backend_cuda.cu @@ -21,6 +21,7 @@ struct ColiCudaTensor { size_t weight_bytes; int fmt, I, O, device; int gs; /* quant group size; 0 = per-row scales (#334) */ + int ng; /* number of scale groups per row = ceil(I/gs) for fmt=4 */ size_t scale_count; /* floats in `scales`: O per-row, O*ng grouped */ int tracked; RaggedKVEntry ragged[512]; @@ -85,7 +86,7 @@ static int select_ctx(DeviceContext *ctx) { __host__ __device__ static size_t row_bytes(int fmt, int I) { if (fmt == 0) return (size_t)I * sizeof(float); if (fmt == 1) return (size_t)I; - if (fmt == 2) return (size_t)(I + 1) / 2; + if (fmt == 2 || fmt == 4) return (size_t)(I + 1) / 2; /* fmt=4: same packed int4 */ if (fmt == 3) return (size_t)(I + 3) / 4; if (fmt == 4) return (size_t)(I + 1) / 2; /* grouped int4: nibbles like fmt 2 */ return 0; @@ -96,7 +97,7 @@ __device__ static float weight_at(const void *weights, int fmt, size_t row, int if (fmt == 0) return reinterpret_cast(base)[i]; if (fmt == 1) return static_cast(reinterpret_cast(base)[i]); const uint8_t *q = base; - if (fmt == 2) { + if (fmt == 2 || fmt == 4) { /* fmt=4: same nibble layout */ uint8_t v = q[i >> 1]; int n=(i&1)?(v>>4):(v&15); return static_cast(n&8?n-16:n); } @@ -110,14 +111,27 @@ __global__ static void offset_to_signed_s4(uint8_t *q,size_t n){ __global__ static void quant_matmul(float *y, const float *x, const void *weights, const float *scales, int fmt, int S, int I, int O, - size_t rb) { + size_t rb, int gs, int ng) { int o = blockIdx.x; int s = blockIdx.y; float sum = 0.0f; size_t row = (size_t)o * rb; const float *xs = x + (size_t)s * I; - for (int i = threadIdx.x; i < I; i += blockDim.x) - sum += xs[i] * weight_at(weights, fmt, row, i); + if (fmt == 4) { + /* Grouped int4: one f32 scale per gs elements along I (ng groups per row). + * Scale layout: scales[o*ng + g]. Each thread strides through I, applying + * the appropriate group scale as it crosses group boundaries. This matches + * the CPU matmul_i4_grouped accumulation exactly. */ + const float *scl = scales + (size_t)o * ng; + for (int i = threadIdx.x; i < I; i += blockDim.x) { + int g = i / gs; + if (g >= ng) g = ng - 1; /* tail elements in the last (partial) group */ + sum += xs[i] * weight_at(weights, fmt, row, i) * scl[g]; + } + } else { + for (int i = threadIdx.x; i < I; i += blockDim.x) + sum += xs[i] * weight_at(weights, fmt, row, i); + } __shared__ float partial[256]; partial[threadIdx.x] = sum; @@ -127,7 +141,7 @@ __global__ static void quant_matmul(float *y, const float *x, const void *weight __syncthreads(); } if (!threadIdx.x) - y[(size_t)s * O + o] = partial[0] * (fmt ? scales[o] : 1.0f); + y[(size_t)s * O + o] = (fmt && fmt != 4) ? partial[0] * scales[o] : partial[0]; } __global__ static void silu_mul(float *gate, const float *up, size_t n) { @@ -564,7 +578,7 @@ extern "C" int coli_cuda_tensor_upload(ColiCudaTensor **tensor, * on such a slot failed here — the GPU tier silently never computed * for host-released slab experts. */ ColiCudaTensor *t = *tensor; - return t->fmt == fmt && t->I == I && t->O == O && t->device == device; + return t->fmt == fmt && t->I == I && t->O == O && t->device == device && t->gs == gs; } DeviceContext *ctx = find_ctx(device); if (!weights || I < 1 || O < 1 || !select_ctx(ctx)) return 0; @@ -574,7 +588,8 @@ extern "C" int coli_cuda_tensor_upload(ColiCudaTensor **tensor, if (!t) return 0; t->fmt = fmt; t->I = I; t->O = O; t->device = device; t->weight_bytes = rb * (size_t)O; t->gs = (fmt==4 && g_upload_gs>0) ? g_upload_gs : 0; - t->scale_count = t->gs ? (size_t)O * (size_t)((I + t->gs - 1) / t->gs) : (size_t)O; + t->ng = t->gs ? (I + t->gs - 1) / t->gs : 1; + t->scale_count = t->gs ? (size_t)O * (size_t)t->ng : (size_t)O; if (!cuda_ok(cudaMalloc(&t->weights, t->weight_bytes), "tensor allocation") || !cuda_ok(cudaMemcpy(t->weights, weights, t->weight_bytes, cudaMemcpyHostToDevice), "tensor upload")) { coli_cuda_tensor_free(t); @@ -618,7 +633,9 @@ extern "C" int coli_cuda_tensor_update(ColiCudaTensor *tensor, (uint8_t*)tensor->weights,tensor->weight_bytes); if(!cuda_ok(cudaGetLastError(),"int4 weight refresh")) return 0; } + int ng = tensor->ng > 0 ? tensor->ng : 1; return !tensor->fmt || cuda_ok(cudaMemcpy(tensor->scales,scales, +<<<<<<< HEAD (tensor->scale_count?tensor->scale_count:(size_t)tensor->O)*sizeof(float), cudaMemcpyHostToDevice),"scale refresh"); } @@ -631,14 +648,22 @@ static long g_gpu_calls; static int fault_injected(void) { const char *fa = std::getenv("COLI_GPU_FAIL_AFTER"); return fa && g_gpu_calls++ >= std::atol(fa); +======= + (size_t)tensor->O*ng*sizeof(float),cudaMemcpyHostToDevice),"scale refresh"); +>>>>>>> 1a69113 (cuda+engine: full fmt=4 (grouped int4 gs=64) support + diagnostic harness) } extern "C" int coli_cuda_matmul(ColiCudaTensor **tensor, float *y, const float *x, const void *weights, const float *scales, +<<<<<<< HEAD int fmt, int S, int I, int O, int device) { if (fault_injected()) return 0; if (S < 1 || !coli_cuda_tensor_upload(tensor, weights, scales, fmt, I, O, device)) return 0; +======= + int fmt, int S, int I, int O, int device, int gs) { + if (S < 1 || !coli_cuda_tensor_upload(tensor, weights, scales, fmt, I, O, device, gs)) return 0; +>>>>>>> 1a69113 (cuda+engine: full fmt=4 (grouped int4 gs=64) support + diagnostic harness) ColiCudaTensor *t = *tensor; DeviceContext *ctx = find_ctx(t->device); if (!select_ctx(ctx)) return 0; @@ -647,7 +672,7 @@ extern "C" int coli_cuda_matmul(ColiCudaTensor **tensor, if (!reserve(&ctx->x, &ctx->x_cap, xb) || !reserve(&ctx->y, &ctx->y_cap, yb)) return 0; if (!cuda_ok(cudaMemcpy(ctx->x, x, xb, cudaMemcpyHostToDevice), "input upload")) return 0; dim3 grid((unsigned)O, (unsigned)S); - quant_matmul<<>>(ctx->y, ctx->x, t->weights, t->scales, fmt, S, I, O, rb); + quant_matmul<<>>(ctx->y, ctx->x, t->weights, t->scales, fmt, S, I, O, rb, t->gs, t->ng); if (!cuda_ok(cudaGetLastError(), "matmul launch") || !cuda_ok(cudaMemcpy(y, ctx->y, yb, cudaMemcpyDeviceToHost), "output download")) return 0; return 1; @@ -671,13 +696,13 @@ extern "C" int coli_cuda_expert_mlp(ColiCudaTensor *gate, ColiCudaTensor *up, if (!cuda_ok(cudaMemcpy(ctx->x,x,xb,cudaMemcpyHostToDevice),"expert input upload")) return 0; dim3 hidden_grid((unsigned)I,(unsigned)S), output_grid((unsigned)D,(unsigned)S); quant_matmul<<>>(ctx->gate,ctx->x,gate->weights,gate->scales, - gate->fmt,S,D,I,row_bytes(gate->fmt,D)); + gate->fmt,S,D,I,row_bytes(gate->fmt,D),gate->gs,gate->ng); quant_matmul<<>>(ctx->up,ctx->x,up->weights,up->scales, - up->fmt,S,D,I,row_bytes(up->fmt,D)); + up->fmt,S,D,I,row_bytes(up->fmt,D),up->gs,up->ng); size_t n=(size_t)S*I; silu_mul<<<(unsigned)((n+255)/256),256>>>(ctx->gate,ctx->up,n); quant_matmul<<>>(ctx->y,ctx->gate,down->weights,down->scales, - down->fmt,S,I,D,row_bytes(down->fmt,I)); + down->fmt,S,I,D,row_bytes(down->fmt,I),down->gs,down->ng); if (!cuda_ok(cudaGetLastError(),"expert MLP launch") || !cuda_ok(cudaMemcpy(y,ctx->y,yb,cudaMemcpyDeviceToHost),"expert output download")) return 0; return 1; @@ -732,9 +757,16 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, g->fmt,u->fmt,d->fmt,rows[c],total, g->gs,u->gs,d->gs}; all_s4&=g->fmt==2&&u->fmt==2&&d->fmt==2; +<<<<<<< HEAD all_q4&=(g->fmt==2||g->fmt==4)&&(u->fmt==2||u->fmt==4)&&(d->fmt==2||d->fmt==4)&& !(g->gs&1)&&!(u->gs&1)&&!(d->gs&1); /* even gs: a packed byte never straddles groups */ any_g4|=g->fmt==4||u->fmt==4||d->fmt==4; +======= + /* fmt=4 (grouped int4) experts have per-group scales that the grouped_hidden/ + * grouped_down kernels don't handle. Fall back to the per-expert path + * (coli_cuda_expert_mlp), which uses quant_matmul with correct gs/ng. */ + if(g->fmt==4||u->fmt==4||d->fmt==4) return 0; +>>>>>>> 1a69113 (cuda+engine: full fmt=4 (grouped int4 gs=64) support + diagnostic harness) total+=rows[c]; if(rows[c]>max_rows) max_rows=rows[c]; } DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0; @@ -805,12 +837,12 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, /* piccoli batch: tile TC quasi vuoti + overhead di lancio — il * kernel naive per-elemento resta piu' veloce (misurato in decode) */ quant_matmul<<stream>>>(g16,x16, - host[c].g,host[c].gs,host[c].gf,r,D,I,row_bytes(host[c].gf,D)); + host[c].g,host[c].gs,host[c].gf,r,D,I,row_bytes(host[c].gf,D),0,1); quant_matmul<<stream>>>(u16,x16, - host[c].u,host[c].us,host[c].uf,r,D,I,row_bytes(host[c].uf,D)); + host[c].u,host[c].us,host[c].uf,r,D,I,row_bytes(host[c].uf,D),0,1); silu_mul<<<(unsigned)(((size_t)r*I+255)/256),256,0,ctx->stream>>>(g16,u16,(size_t)r*I); quant_matmul<<stream>>>(y16,g16, - host[c].d,host[c].ds,host[c].df,r,I,D,row_bytes(host[c].df,I)); + host[c].d,host[c].ds,host[c].df,r,I,D,row_bytes(host[c].df,I),0,1); } off16+=r; } @@ -976,7 +1008,7 @@ static int attention_absorb_batch_run(ColiCudaTensor *w,ColiCudaTensor *proj,flo if(proj){ ob=(size_t)S*proj->O*sizeof(float);if(!reserve(&dc->y,&dc->y_cap,ob))return 0; 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)); + proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I),proj->gs,proj->ng); if(!cuda_ok(cudaGetLastError(),"attention o_proj launch"))return 0;src=dc->y; } if(!cuda_ok(cudaMemcpyAsync(out,src,ob,cudaMemcpyDeviceToHost,dc->stream), @@ -1087,7 +1119,8 @@ extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) { DeviceContext *ctx = find_ctx(tensor->device); if (ctx) select_ctx(ctx); if (tensor->tracked && ctx) { - size_t bytes = tensor->weight_bytes + (tensor->fmt ? (size_t)tensor->O * sizeof(float) : 0); + int ng = tensor->ng > 0 ? tensor->ng : 1; + size_t bytes = tensor->weight_bytes + (tensor->fmt ? (size_t)tensor->O * ng * sizeof(float) : 0); if (ctx->tensor_count) ctx->tensor_count--; if (ctx->tensor_bytes >= bytes) ctx->tensor_bytes -= bytes; } @@ -1101,7 +1134,9 @@ extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) { } extern "C" size_t coli_cuda_tensor_bytes(const ColiCudaTensor *tensor) { - return tensor ? tensor->weight_bytes + (tensor->fmt ? (size_t)tensor->O * sizeof(float) : 0) : 0; + if (!tensor) return 0; + int ng = tensor->ng > 0 ? tensor->ng : 1; + return tensor->weight_bytes + (tensor->fmt ? (size_t)tensor->O * ng * sizeof(float) : 0); } extern "C" int coli_cuda_tensor_device(const ColiCudaTensor *tensor) { @@ -1320,7 +1355,7 @@ extern "C" int coli_cuda_attention_project_batch_dev(ColiCudaTensor *w,ColiCudaT size_t ob=(size_t)S*proj->O*sizeof(float); if(!reserve(&dc->y,&dc->y_cap,ob))return 0; 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)); + proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I),proj->gs,proj->ng); 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; @@ -1355,7 +1390,7 @@ extern "C" int coli_cuda_pipe_gemm(ColiCudaTensor *t,float *y_dev,const float *x DeviceContext *ctx=find_ctx(t->device); if(!select_ctx(ctx)) return 0; dim3 grid((unsigned)t->O,(unsigned)S); quant_matmul<<>>(y_dev,x_dev,t->weights,t->scales,t->fmt,S,t->I,t->O, - row_bytes(t->fmt,t->I)); + row_bytes(t->fmt,t->I),t->gs,t->ng); return cuda_ok(cudaGetLastError(),"pipe gemm"); } /* copia diretta scheda->scheda (P2P se disponibile, altrimenti staging driver) */ @@ -1382,7 +1417,7 @@ extern "C" int coli_cuda_attention_project_batch_dev_out(ColiCudaTensor *w,ColiC 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<<O,S),256,0,dc->stream>>>(out_dev,dc->ac,proj->weights, - proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I)); + proj->scales,proj->fmt,S,proj->I,proj->O,row_bytes(proj->fmt,proj->I),proj->gs,proj->ng); if(!cuda_ok(cudaGetLastError(),"pipe o_proj launch (dev out)"))return 0; return cuda_ok(cudaStreamSynchronize(dc->stream),"pipe attention sync (dev out)"); } diff --git a/c/backend_cuda.h b/c/backend_cuda.h index dfe3809..3af966d 100644 --- a/c/backend_cuda.h +++ b/c/backend_cuda.h @@ -41,18 +41,19 @@ COLI_CUDA_DLLEXPORT int coli_cuda_tensor_upload_g(ColiCudaTensor **tensor, int fmt, int I, int O, int device, int gs); COLI_CUDA_DLLEXPORT int coli_cuda_tensor_upload(ColiCudaTensor **tensor, const void *weights, const float *scales, - int fmt, int I, int O, int device); + int fmt, int I, int O, int device, int gs); /* * y[S,O] = x[S,I] @ W[O,I]^T. - * fmt matches QT in glm.c: 0=f32, 1=int8, 2=int4, 3=int2. - * The first successful call uploads W and its row scales; later calls reuse it. + * fmt matches QT in glm.c: 0=f32, 1=int8, 2=int4, 3=int2, 4=grouped int4. + * gs is the group size for fmt=4 (0 for all other formats). + * The first successful call uploads W and its scales; later calls reuse it. * Returns 1 on success and 0 when CUDA is not initialized or the format is invalid. */ COLI_CUDA_DLLEXPORT int coli_cuda_matmul(ColiCudaTensor **tensor, float *y, const float *x, const void *weights, const float *scales, - int fmt, int S, int I, int O, int device); + int fmt, int S, int I, int O, int device, int gs); /* Fused expert pipeline: y = down(silu(gate(x)) * up(x)). All three tensors * must already be resident on one device. Activations cross PCIe once in diff --git a/c/backend_loader.c b/c/backend_loader.c index feb908e..bf0dd9a 100644 --- a/c/backend_loader.c +++ b/c/backend_loader.c @@ -54,7 +54,7 @@ typedef int (*fn_tensor_upload)(ColiCudaTensor **tensor, const void * typedef int (*fn_tensor_upload_g)(ColiCudaTensor **tensor, const void *weights, const float *scales, int fmt, int I, int O, int device, int gs); typedef int (*fn_matmul)(ColiCudaTensor **tensor, float *y, const float *x, const void *weights, const float *scales, - int fmt, int S, int I, int O, int device); + int fmt, int S, int I, int O, int device, int gs); typedef void (*fn_tensor_free)(ColiCudaTensor *tensor); typedef size_t (*fn_tensor_bytes)(const ColiCudaTensor *tensor); typedef int (*fn_tensor_device)(const ColiCudaTensor *tensor); @@ -325,9 +325,9 @@ int coli_cuda_attention_absorb(ColiCudaTensor *kv_b, float *ctx, const float *q, } int coli_cuda_tensor_upload(ColiCudaTensor **tensor, const void *weights, - const float *scales, int fmt, int I, int O, int device){ + const float *scales, int fmt, int I, int O, int device, int gs){ if(!g_cuda.available) return 0; - return g_cuda.tensor_upload(tensor, weights, scales, fmt, I, O, device); + return g_cuda.tensor_upload(tensor, weights, scales, fmt, I, O, device, gs); } int coli_cuda_tensor_upload_g(ColiCudaTensor **tensor, const void *weights, const float *scales, int fmt, int I, int O, int device, int gs){ @@ -337,9 +337,9 @@ int coli_cuda_tensor_upload_g(ColiCudaTensor **tensor, const void *weights, cons int coli_cuda_matmul(ColiCudaTensor **tensor, float *y, const float *x, const void *weights, const float *scales, - int fmt, int S, int I, int O, int device){ + int fmt, int S, int I, int O, int device, int gs){ if(!g_cuda.available) return 0; - return g_cuda.matmul(tensor, y, x, weights, scales, fmt, S, I, O, device); + return g_cuda.matmul(tensor, y, x, weights, scales, fmt, S, I, O, device, gs); } void coli_cuda_tensor_free(ColiCudaTensor *tensor){ diff --git a/c/colibri.c b/c/colibri.c index ef32b01..d8108be 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -448,6 +448,713 @@ static float *falloc(int64_t n){ float *p=malloc((size_t)n*sizeof(float)); if(!p){fprintf(stderr,"OOM\n");exit(1);} return p; } +/* y[S,O] = x[S,I] @ W^T, W[O,I] f32 */ +static void matmul(float *y, const float *x, const float *W, int S, int I, int O){ + #pragma omp parallel for schedule(static) + for (int o=0;o>1))); /* 8 byte=16 nibble */ + __m128i lo=_mm_and_si128(by,m4), hi=_mm_and_si128(_mm_srli_epi16(by,4),m4); + __m128i nib=_mm_unpacklo_epi8(lo,hi); /* nibble in ordine */ + __m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b8)); + __m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b8)); + acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0, acc); + acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1, acc); } + a=hsum256(acc); +#elif defined(__ARM_NEON) + const uint8x8_t m4=vdup_n_u8(0x0F); const int8x8_t b8=vdup_n_s8(8); + float32x4_t ac0=vdupq_n_f32(0), ac1=vdupq_n_f32(0); + for(;i+16<=I;i+=16){ uint8x8_t by=vld1_u8(w+(i>>1)); /* 8 byte=16 nibble */ + uint8x8x2_t z=vzip_u8(vand_u8(by,m4), vshr_n_u8(by,4)); /* nibble in ordine */ + int16x8_t w0=vmovl_s8(vsub_s8(vreinterpret_s8_u8(z.val[0]),b8)); + int16x8_t w1=vmovl_s8(vsub_s8(vreinterpret_s8_u8(z.val[1]),b8)); + ac0=vfmaq_f32(ac0, vld1q_f32(xs+i), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w0)))); + ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+4), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w0)))); + ac0=vfmaq_f32(ac0, vld1q_f32(xs+i+8), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w1)))); + ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+12), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w1)))); } + a=vaddvq_f32(vaddq_f32(ac0,ac1)); +#endif +#if defined(__AVX512F__) && defined(__AVX512BW__) + } +#endif + for(;i+1>1]; int lo=(int)(byte&0xF)-8, hi=(int)(byte>>4)-8; + a += xs[i]*(float)lo + xs[i+1]*(float)hi; } + if(i>1]; int lo=(int)(byte&0xF)-8; a += xs[i]*(float)lo; } + y[(int64_t)s*O+o]=a*sc; } } +} +/* y[S,O] = x[S,I] @ W^T with W int4 packed (2/byte) + per-GROUP scales (fmt=4). + * Same nibble math as matmul_i4, but the scale changes every `gs` elements along I. + * The accumulator resets at each group boundary: dot(x[grp], w[grp]) * scale[grp]. + * gs MUST be a multiple of 16 (the AVX2 vector width). */ +static void matmul_i4_grouped(float *y, const float *x, const uint8_t *q4, const float *scale, + int S, int I, int O, int gs){ + int rb=(I+1)/2; int ng=(I+gs-1)/gs; + #pragma omp parallel for schedule(static) + for(int o=0;oI) glen=I-base; + float sc=scl[g]; + int i=base; +#ifdef __AVX2__ + const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi32(8); + __m256 acc=_mm256_setzero_ps(); + for(; i+16<=base+glen; i+=16){ __m128i by=_mm_loadl_epi64((const __m128i*)(w+(i>>1))); + __m128i lo=_mm_and_si128(by,m4),hi=_mm_and_si128(_mm_srli_epi16(by,4),m4); + __m128i nib=_mm_unpacklo_epi8(lo,hi); + __m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b8)); + __m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b8)); + acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0, acc); + acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1, acc); } + a+=hsum256(acc)*sc; +#endif + /* scalar tail for the group remainder */ + for(; i>1]; + a+=(xs[i]*(float)((int)(byte&0xF)-8)+xs[i+1]*(float)((int)(byte>>4)-8))*sc; } + else { uint8_t byte=w[i>>1]; a+=xs[i]*(float)((int)(byte&0xF)-8)*sc; } + } + } + y[(int64_t)s*O+o]=a; + } + } +} +/* Fused gate+up for grouped int4 (fmt=4): computes both yg[S,O] and yu[S,O] from + * the same x[S,I], with per-group scales. One OpenMP dispatch covers both matrices, + * reading x once instead of twice — saves ~33% of expert-matmul time at decode. + * The per-group scale logic matches matmul_i4_grouped exactly. */ +static void matmul_i4_grouped_pair(float *yg, float *yu, const float *x, + const uint8_t *qg, const float *sg, + const uint8_t *qu, const float *su, + int S, int I, int O, int gs){ + int rb=(I+1)/2; int ng=(I+gs-1)/gs; + #pragma omp parallel for schedule(static) + for(int o=0;oI) glen=I-base; + float scg=sgl[g], scu=sul[g]; + int i=base; +#ifdef __AVX2__ + const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi32(8); + __m256 accg=_mm256_setzero_ps(), accu=_mm256_setzero_ps(); + for(; i+16<=base+glen; i+=16){ + __m128i byg=_mm_loadl_epi64((const __m128i*)(wg+(i>>1))); + __m128i log=_mm_and_si128(byg,m4),hig=_mm_and_si128(_mm_srli_epi16(byg,4),m4); + __m128i nibg=_mm_unpacklo_epi8(log,hig); + __m256 w0g=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nibg),b8)); + __m256 w1g=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nibg,8)),b8)); + accg=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0g, accg); + accg=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1g, accg); + __m128i byu=_mm_loadl_epi64((const __m128i*)(wu2+(i>>1))); + __m128i lou=_mm_and_si128(byu,m4),hiu=_mm_and_si128(_mm_srli_epi16(byu,4),m4); + __m128i nibu=_mm_unpacklo_epi8(lou,hiu); + __m256 w0u=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nibu),b8)); + __m256 w1u=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nibu,8)),b8)); + accu=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0u, accu); + accu=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1u, accu); + } + ag+=hsum256(accg)*scg; au+=hsum256(accu)*scu; +#endif + for(; i+1>1], bu=wu2[i>>1]; + ag+=(xs[i]*(float)((int)(bg&0xF)-8)+xs[i+1]*(float)((int)(bg>>4)-8))*scg; + au+=(xs[i]*(float)((int)(bu&0xF)-8)+xs[i+1]*(float)((int)(bu>>4)-8))*scu; + } + if(i>1], bu=wu2[i>>1]; + ag+=xs[i]*(float)((int)(bg&0xF)-8)*scg; + au+=xs[i]*(float)((int)(bu&0xF)-8)*scu; + } + } + yg[(int64_t)s*O+o]=ag; yu[(int64_t)s*O+o]=au; + } + } +} +/* Decode hot path for gate+up: same exact q4 dot products as matmul_i4, but one + * OpenMP dispatch covers both matrices. KTransformers uses persistent pools; + * this keeps colibri dependency-free while removing one team launch/expert. */ +static void matmul_i4_pair(float *yg, float *yu, const float *x, + const uint8_t *qg, const float *sg, + const uint8_t *qu, const float *su, int I, int O){ + int rb=(I+1)/2; + #pragma omp parallel for schedule(static) + for(int z=0;z<2*O;z++){ + int o=z>1))); + __m128i lo=_mm_and_si128(by,m4),hi=_mm_and_si128(_mm_srli_epi16(by,4),m4); + __m128i nib=_mm_unpacklo_epi8(lo,hi); + __m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b8)); + __m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b8)); + acc=_mm256_fmadd_ps(_mm256_loadu_ps(x+i),w0,acc); + acc=_mm256_fmadd_ps(_mm256_loadu_ps(x+i+8),w1,acc); } + a=hsum256(acc); +#elif defined(__ARM_NEON) + const uint8x8_t m4=vdup_n_u8(0x0F); const int8x8_t b8=vdup_n_s8(8); + float32x4_t ac0=vdupq_n_f32(0),ac1=vdupq_n_f32(0); + for(;i+16<=I;i+=16){ uint8x8_t by=vld1_u8(w+(i>>1)); + uint8x8x2_t n=vzip_u8(vand_u8(by,m4),vshr_n_u8(by,4)); + int16x8_t w0=vmovl_s8(vsub_s8(vreinterpret_s8_u8(n.val[0]),b8)); + int16x8_t w1=vmovl_s8(vsub_s8(vreinterpret_s8_u8(n.val[1]),b8)); + ac0=vfmaq_f32(ac0,vld1q_f32(x+i),vcvtq_f32_s32(vmovl_s16(vget_low_s16(w0)))); + ac1=vfmaq_f32(ac1,vld1q_f32(x+i+4),vcvtq_f32_s32(vmovl_s16(vget_high_s16(w0)))); + ac0=vfmaq_f32(ac0,vld1q_f32(x+i+8),vcvtq_f32_s32(vmovl_s16(vget_low_s16(w1)))); + ac1=vfmaq_f32(ac1,vld1q_f32(x+i+12),vcvtq_f32_s32(vmovl_s16(vget_high_s16(w1)))); } + a=vaddvq_f32(vaddq_f32(ac0,ac1)); +#endif +#if defined(__AVX512F__) && defined(__AVX512BW__) + } +#endif + for(;i+1>1]; a+=x[i]*(float)((b&15)-8)+x[i+1]*(float)((b>>4)-8); } + if(i>1]&15)-8); + (z=2) non calcolano la STESSA funzione. Tre interruttori dipendono da S: il gate + * int4-IDOT (S>=g_i4s — asimmetrico proprio dove g_i4s>1), la fusione gate+up solo-S==1, + * e la soglia righe del GEMM Metal. Con SPEC_PIN=1 (default) ogni forward emesso mentre + * i draft del modello sono attivi resta sulla famiglia di kernel di S=1: draft e verifica + * coincidono per costruzione. Prefill e decode non speculativo sono intoccati. + * EN: MTP acceptance collapses when the draft (S=1) and verify (S>=2) forwards do not + * compute the SAME function. Three switches are S-dependent: the int4 IDOT gate + * (S>=g_i4s — asymmetric exactly on ISAs where g_i4s>1), the S==1-only gate+up fusion, + * and the Metal GEMM row threshold. SPEC_PIN=1 (default) pins every forward issued + * while model drafts are live to the platform's S=1 kernel family, so draft and verify + * agree by construction; prefill and non-speculative decode are untouched. + * SPEC_PIN=0 restores the S-dependent gates (A/B). */ +static int g_spec_pin=1; +static int g_spec_live=0; /* set by spec_decode while drafts are live */ +static inline int spec_pinned(void){ return g_spec_pin && g_spec_live; } +static void expert_gate_up(float *g,float *u,const float *x,QT *wg,QT *wu,int S){ + if(!g_no_fused_pair&&!spec_pinned()&&S==1&&wg->fmt==2&&wu->fmt==2&&wg->I==wu->I&&wg->O==wu->O) + matmul_i4_pair(g,u,x,wg->q4,wg->s,wu->q4,wu->s,wg->I,wg->O); + else if(!g_no_fused_pair&&S==1&&wg->fmt==4&&wu->fmt==4&&wg->I==wu->I&&wg->O==wu->O&&wg->gs==wu->gs) + matmul_i4_grouped_pair(g,u,x,wg->q4,wg->s,wu->q4,wu->s,S,wg->I,wg->O,wg->gs); + else { matmul_qt(g,x,wg,S); matmul_qt(u,x,wu,S); } +} +/* y[S,O] = x[S,I] @ W^T con W int2 impacchettato (4 valori/byte) + scala[O]. nibble 2-bit -> [-2,1]. */ +static void matmul_i2(float *y, const float *x, const uint8_t *q2, const float *scale, int S, int I, int O){ + int rb=(I+3)/4; + #pragma omp parallel for schedule(static) + for (int o=0;o>2))); /* 4 byte=16 valori */ + __m128i p0=_mm_and_si128(by,m2), p1=_mm_and_si128(_mm_srli_epi16(by,2),m2); + __m128i p2=_mm_and_si128(_mm_srli_epi16(by,4),m2), p3=_mm_and_si128(_mm_srli_epi16(by,6),m2); + __m128i lo=_mm_unpacklo_epi8(p0,p1), hi=_mm_unpacklo_epi8(p2,p3); + __m128i nib=_mm_unpacklo_epi16(lo,hi); /* 16 valori in ordine */ + __m256 w0=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(nib),b2)); + __m256 w1=_mm256_cvtepi32_ps(_mm256_sub_epi32(_mm256_cvtepu8_epi32(_mm_srli_si128(nib,8)),b2)); + acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i), w0, acc); + acc=_mm256_fmadd_ps(_mm256_loadu_ps(xs+i+8), w1, acc); } + a=hsum256(acc); +#elif defined(__ARM_NEON) + const uint8x8_t m2v=vdup_n_u8(3); const int8x8_t b2v=vdup_n_s8(2); + float32x4_t ac0=vdupq_n_f32(0), ac1=vdupq_n_f32(0); + for(;i+16<=I;i+=16){ uint32_t wd; memcpy(&wd, w+(i>>2), 4); /* 4 byte=16 valori */ + uint8x8_t by=vreinterpret_u8_u32(vdup_n_u32(wd)); + uint8x8x2_t z01=vzip_u8(vand_u8(by,m2v), vand_u8(vshr_n_u8(by,2),m2v)); + uint8x8x2_t z23=vzip_u8(vand_u8(vshr_n_u8(by,4),m2v), vshr_n_u8(by,6)); + uint16x4x2_t zz=vzip_u16(vreinterpret_u16_u8(z01.val[0]), vreinterpret_u16_u8(z23.val[0])); + int16x8_t w0=vmovl_s8(vsub_s8(vreinterpret_s8_u16(zz.val[0]),b2v)); /* 16 valori in ordine */ + int16x8_t w1=vmovl_s8(vsub_s8(vreinterpret_s8_u16(zz.val[1]),b2v)); + ac0=vfmaq_f32(ac0, vld1q_f32(xs+i), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w0)))); + ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+4), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w0)))); + ac0=vfmaq_f32(ac0, vld1q_f32(xs+i+8), vcvtq_f32_s32(vmovl_s16(vget_low_s16(w1)))); + ac1=vfmaq_f32(ac1, vld1q_f32(xs+i+12), vcvtq_f32_s32(vmovl_s16(vget_high_s16(w1)))); } + a=vaddvq_f32(vaddq_f32(ac0,ac1)); +#endif + for(;i>2]; int sh=(i&3)*2; a += xs[i]*(float)((int)((byte>>sh)&3)-2); } + y[(int64_t)s*O+o]=a*sc; } } +} +/* ---- KERNEL INTERI (IDOT): attivazioni quantizzate a int8 per riga (absmax/127, + * stile Q8_0), prodotto scalare INTERO via maddubs/madd AVX2 — niente conversione + * f32 dei pesi nel ciclo caldo. ~2-3x sui matmul quantizzati; errore aggiunto ~0.3% + * RMS per matmul (attivazione int8), IDOT=0 torna al percorso f32 esatto. */ +#if defined(__AVX512VNNI__) && defined(__AVX512BW__) +#define IDOT_KERNEL "avx512-vnni" +#elif defined(__AVXVNNI__) && defined(__AVX2__) +#define IDOT_KERNEL "avx-vnni" +#elif defined(__AVX2__) +#define IDOT_KERNEL "avx2" +#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8) +#define IDOT_KERNEL "neon-i8mm" +#elif defined(__ARM_NEON) +#define IDOT_KERNEL "neon" +#elif defined(__VSX__) +#define IDOT_KERNEL "vsx" +#else +#define IDOT_KERNEL "scalar" +#endif +static int g_idot=1; +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD) +static int g_i4s=1; /* SDOT presente: int4 IDOT conviene anche a S=1 (decode). Misurato + * su Apple M-series: +14%%, expert-matmul -16%%. EN: with SDOT, int4 + * IDOT pays even at S=1 (decode); measured on Apple M-series. */ +#elif defined(__VSX__) +static int g_i4s=1; /* POWER8 vec_msum: qui il fallback f32 e' SCALARE, quindi l'IDOT + * int4 conviene anche a S=1. Misurato su POWER8 S824 (vedi PR). + * EN: on VSX the f32 fallback is plain scalar C, so int4 IDOT + * pays even at S=1. Measured on a POWER8 S824 (see PR). */ +#else +static int g_i4s=2; /* senza SDOT / altrove: soglia originale (misura AVX2 dell'autore). + * EN: without SDOT / elsewhere: original threshold (author's AVX2). */ +#endif +static inline float qrow_i8(const float *x, int8_t *q, int I){ + float amax=0; for(int i=0;iamax)amax=a; } + float s=amax/127.f; if(s<1e-12f) s=1e-12f; float inv=1.f/s; + for(int i=0;i s32 directly, 64 bytes/iter, no 16-bit intermediate. + * AVX-512 has no vpsignb: |w| via abs, sign folded into x with a mask-negate + * (w==0 -> product 0 either way). |x|<=127 (qrow_i8), |w|<=128 as u8: each + * s32 lane adds <= 4*128*127, safe up to I=16384 like the AVX2 bound. */ + __m512i acc=_mm512_setzero_si512(); + for(;i+64<=I;i+=64){ + __m512i wv=_mm512_loadu_si512((const void*)(w+i)); + __m512i xv=_mm512_loadu_si512((const void*)(x+i)); + __mmask64 neg=_mm512_movepi8_mask(wv); + __m512i xs=_mm512_mask_sub_epi8(xv,neg,_mm512_setzero_si512(),xv); + acc=_mm512_dpbusd_epi32(acc,_mm512_abs_epi8(wv),xs); + } + sum=_mm512_reduce_add_epi32(acc); +#elif defined(__AVXVNNI__) && defined(__AVX2__) + /* AVX-VNNI 128-bit: vpdpbusd u8*s8 -> s32, 16 byte/iter. Stesso trucco del + * segno della variante 512-bit: |w| via abs, segno piegato in x con maschera + * (w==0 -> product 0). __AVX2__ serve per _mm_sign_epi8 / abs. */ + __m128i acc=_mm_setzero_si128(); + for(;i+16<=I;i+=16){ + __m128i wv=_mm_loadu_si128((const __m128i*)(w+i)); + __m128i xv=_mm_loadu_si128((const __m128i*)(x+i)); + __m128i xs=_mm_sign_epi8(xv,wv); /* x * sign(w); _mm_sign zona __AVX2__ */ + acc=_mm_dpbusd_epi32(acc,_mm_abs_epi8(wv),xs); + } + sum=hsum128_i32(acc); +#elif defined(__AVX2__) + __m256i acc=_mm256_setzero_si256(); const __m256i ones=_mm256_set1_epi16(1); + for(;i+32<=I;i+=32){ + __m256i wv=_mm256_loadu_si256((const __m256i*)(w+i)); + __m256i xv=_mm256_loadu_si256((const __m256i*)(x+i)); + __m256i p=_mm256_maddubs_epi16(_mm256_sign_epi8(wv,wv),_mm256_sign_epi8(xv,wv)); + acc=_mm256_add_epi32(acc,_mm256_madd_epi16(p,ones)); + } + sum=hsum256_i32(acc); +#elif defined(__ARM_NEON) + /* ARM: SDOT nativo se disponibile (Apple Silicon: sempre); altrimenti vmull/vpadal. + * Stesso bound anti-overflow del trucco AVX2: coppie <= 128*127*2 = 32512 < 32767. */ +#if defined(__ARM_FEATURE_DOTPROD) + /* 4 accumulatori indipendenti: SDOT ha latenza ~3-4 cicli, con un solo acc la + * catena seriale strozza il core a ~26 GB/s di pesi; con 4 lane indipendenti il + * dot diventa memory-bound (misurato su M4: 26 -> 63 GB/s per core, 2.4x). */ + int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0); + for(;i+64<=I;i+=64){ + a0=vdotq_s32(a0,vld1q_s8(w+i), vld1q_s8(x+i)); + a1=vdotq_s32(a1,vld1q_s8(w+i+16),vld1q_s8(x+i+16)); + a2=vdotq_s32(a2,vld1q_s8(w+i+32),vld1q_s8(x+i+32)); + a3=vdotq_s32(a3,vld1q_s8(w+i+48),vld1q_s8(x+i+48)); + } + int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3)); + for(;i+16<=I;i+=16) acc=vdotq_s32(acc,vld1q_s8(w+i),vld1q_s8(x+i)); + sum=vaddvq_s32(acc); +#else + int32x4_t acc=vdupq_n_s32(0); + for(;i+16<=I;i+=16){ + int8x16_t wv=vld1q_s8(w+i), xv=vld1q_s8(x+i); + int16x8_t p=vmull_s8(vget_low_s8(wv),vget_low_s8(xv)); + p=vmlal_s8(p,vget_high_s8(wv),vget_high_s8(xv)); + acc=vpadalq_s16(acc,p); + } + sum=vaddvq_s32(acc); +#endif +#elif defined(__VSX__) + /* POWER8: vec_msum (s8 x u8 -> s32) somma i prodotti byte DIRETTAMENTE in lane + * s32, 16 byte/iter: il bound anti-saturazione a 16 bit di maddubs qui non serve. + * Stesso trucco del segno (|w| u8 per x*sign(w) s8), ma |w| via select+sub MODULO + * e non vec_abs: -128 deve diventare 128 u8, non saturare a 127. + * EN: vec_msum accumulates byte products straight into s32 lanes; |w| is built + * with a modulo subtract select instead of vec_abs so w=-128 wraps to 128 (u8) + * rather than saturating to 127. |x|<=127 from qrow_i8, so x negation is safe. */ + __vector signed int acc=vec_splats(0); + const __vector signed char vz=vec_splats((signed char)0); + for(;i+16<=I;i+=16){ + __vector signed char wv=vec_xl(0,(const signed char*)(w+i)); + __vector signed char xv=vec_xl(0,(const signed char*)(x+i)); + __vector __bool char neg=vec_cmplt(wv,vz); + __vector signed char xs=vec_sel(xv,vec_sub(vz,xv),neg); + __vector unsigned char wa=(__vector unsigned char)vec_sel(wv,vec_sub(vz,wv),neg); + acc=vec_msum(xs,wa,acc); + } + sum=vec_extract(acc,0)+vec_extract(acc,1)+vec_extract(acc,2)+vec_extract(acc,3); +#endif + for(;i int8 [-8,7] al volo, poi stesso trucco */ +static inline int32_t dot_i4i8(const uint8_t *w4, const int8_t *x, int I){ + int32_t sum=0; int i=0; +#if defined(__AVX512VNNI__) && defined(__AVX512BW__) + /* 32 bytes = 64 nibbles -> int8 in [-8,7], one vpdpbusd per 64 values. + * 256-bit unpack leaves values in per-128-lane order [0-15][32-47]/[16-31][48-63]; + * dot pairing is order-invariant, so permute x's 128-bit blocks to match + * instead of re-ordering w (one vpermq per iter, off the critical unpack path). */ + const __m256i m4v=_mm256_set1_epi8(0x0F); + const __m512i b8v=_mm512_set1_epi8(8); + const __m512i xidx=_mm512_setr_epi64(0,1,4,5,2,3,6,7); + __m512i acc=_mm512_setzero_si512(); + for(;i+64<=I;i+=64){ + __m256i by=_mm256_loadu_si256((const __m256i*)(w4+(i>>1))); + __m256i lo=_mm256_and_si256(by,m4v), hi=_mm256_and_si256(_mm256_srli_epi16(by,4),m4v); + __m256i z0=_mm256_unpacklo_epi8(lo,hi), z1=_mm256_unpackhi_epi8(lo,hi); + __m512i wv=_mm512_sub_epi8(_mm512_inserti64x4(_mm512_castsi256_si512(z0),z1,1),b8v); + __m512i xv=_mm512_permutexvar_epi64(xidx,_mm512_loadu_si512((const void*)(x+i))); + __mmask64 neg=_mm512_movepi8_mask(wv); + __m512i xs=_mm512_mask_sub_epi8(xv,neg,_mm512_setzero_si512(),xv); + acc=_mm512_dpbusd_epi32(acc,_mm512_abs_epi8(wv),xs); + } + sum=_mm512_reduce_add_epi32(acc); +#elif defined(__AVXVNNI__) && defined(__AVX2__) + /* AVX-VNNI 128-bit, int4: 16 byte = 32 nibble -> int8 [-8,7] in due half + * (n0/n1), ciascuno alimentato a un vpdpbusd da 16 byte. Stesso unpack + * 128-bit del ramo AVX2 sotto; 32 elementi/iter come li. */ + const __m128i m4=_mm_set1_epi8(0x0F); const __m128i b8=_mm_set1_epi8(8); + __m128i acc=_mm_setzero_si128(); + for(;i+32<=I;i+=32){ + __m128i by=_mm_loadu_si128((const __m128i*)(w4+(i>>1))); /* 16 byte = 32 nibble */ + __m128i lo=_mm_and_si128(by,m4), hi=_mm_and_si128(_mm_srli_epi16(by,4),m4); + __m128i n0=_mm_unpacklo_epi8(lo,hi), n1=_mm_unpackhi_epi8(lo,hi); /* nibble in ordine */ + __m128i w0=_mm_sub_epi8(n0,b8), w1=_mm_sub_epi8(n1,b8); + __m128i x0=_mm_loadu_si128((const __m128i*)(x+i)); + __m128i x1=_mm_loadu_si128((const __m128i*)(x+i+16)); + acc=_mm_dpbusd_epi32(acc,_mm_abs_epi8(w0),_mm_sign_epi8(x0,w0)); + acc=_mm_dpbusd_epi32(acc,_mm_abs_epi8(w1),_mm_sign_epi8(x1,w1)); + } + sum=hsum128_i32(acc); +#elif defined(__AVX2__) + const __m128i m4=_mm_set1_epi8(0x0F); const __m256i b8=_mm256_set1_epi8(8); + const __m256i ones=_mm256_set1_epi16(1); + __m256i acc=_mm256_setzero_si256(); + for(;i+32<=I;i+=32){ + __m128i by=_mm_loadu_si128((const __m128i*)(w4+(i>>1))); /* 16 byte = 32 nibble */ + __m128i lo=_mm_and_si128(by,m4), hi=_mm_and_si128(_mm_srli_epi16(by,4),m4); + __m128i n0=_mm_unpacklo_epi8(lo,hi), n1=_mm_unpackhi_epi8(lo,hi); /* in ordine */ + __m256i wv=_mm256_sub_epi8(_mm256_set_m128i(n1,n0),b8); + __m256i xv=_mm256_loadu_si256((const __m256i*)(x+i)); + __m256i p=_mm256_maddubs_epi16(_mm256_sign_epi8(wv,wv),_mm256_sign_epi8(xv,wv)); + acc=_mm256_add_epi32(acc,_mm256_madd_epi16(p,ones)); + } + sum=hsum256_i32(acc); +#elif defined(__ARM_NEON) + const uint8x16_t m4q=vdupq_n_u8(0x0F); const int8x16_t b8q=vdupq_n_s8(8); +#if defined(__ARM_FEATURE_DOTPROD) + /* 4 accumulatori indipendenti (vedi dot_i8i8): spezza la catena seriale su acc. + * Misurato su M4: 12.4 -> 29.9 GB/s di pesi per core (2.4x). */ + int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0); + for(;i+64<=I;i+=64){ + uint8x16_t byA=vld1q_u8(w4+(i>>1)), byB=vld1q_u8(w4+(i>>1)+16); + uint8x16x2_t zA=vzipq_u8(vandq_u8(byA,m4q), vshrq_n_u8(byA,4)); /* nibble in ordine */ + uint8x16x2_t zB=vzipq_u8(vandq_u8(byB,m4q), vshrq_n_u8(byB,4)); + a0=vdotq_s32(a0,vsubq_s8(vreinterpretq_s8_u8(zA.val[0]),b8q),vld1q_s8(x+i)); + a1=vdotq_s32(a1,vsubq_s8(vreinterpretq_s8_u8(zA.val[1]),b8q),vld1q_s8(x+i+16)); + a2=vdotq_s32(a2,vsubq_s8(vreinterpretq_s8_u8(zB.val[0]),b8q),vld1q_s8(x+i+32)); + a3=vdotq_s32(a3,vsubq_s8(vreinterpretq_s8_u8(zB.val[1]),b8q),vld1q_s8(x+i+48)); + } + int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3)); + for(;i+32<=I;i+=32){ + uint8x16_t by=vld1q_u8(w4+(i>>1)); /* 16 byte = 32 nibble */ + uint8x16x2_t z=vzipq_u8(vandq_u8(by,m4q), vshrq_n_u8(by,4)); /* nibble in ordine */ + acc=vdotq_s32(acc,vsubq_s8(vreinterpretq_s8_u8(z.val[0]),b8q),vld1q_s8(x+i)); + acc=vdotq_s32(acc,vsubq_s8(vreinterpretq_s8_u8(z.val[1]),b8q),vld1q_s8(x+i+16)); + } + sum=vaddvq_s32(acc); +#else + int32x4_t acc=vdupq_n_s32(0); + for(;i+32<=I;i+=32){ + uint8x16_t by=vld1q_u8(w4+(i>>1)); /* 16 byte = 32 nibble */ + uint8x16x2_t z=vzipq_u8(vandq_u8(by,m4q), vshrq_n_u8(by,4)); /* nibble in ordine */ + int8x16_t w0=vsubq_s8(vreinterpretq_s8_u8(z.val[0]),b8q); + int8x16_t w1=vsubq_s8(vreinterpretq_s8_u8(z.val[1]),b8q); + int8x16_t x0=vld1q_s8(x+i), x1=vld1q_s8(x+i+16); + int16x8_t p=vmull_s8(vget_low_s8(w0),vget_low_s8(x0)); /* |w|<=8: nessun overflow */ + p=vmlal_s8(p,vget_high_s8(w0),vget_high_s8(x0)); + acc=vpadalq_s16(acc,p); + p=vmull_s8(vget_low_s8(w1),vget_low_s8(x1)); + p=vmlal_s8(p,vget_high_s8(w1),vget_high_s8(x1)); + acc=vpadalq_s16(acc,p); + } + sum=vaddvq_s32(acc); +#endif +#elif defined(__VSX__) + /* 16 byte = 32 nibble. vec_mergeh/vec_mergel su ppc64le (GCC) interallacciano come + * unpacklo/unpackhi x86 (verificato empiricamente su POWER8): i nibble escono in + * ordine di memoria. |w|<=8 dopo il -8, quindi stesso trucco del segno di dot_i8i8. + * EN: vec_mergeh/l on ppc64le interleave like x86 unpacklo/hi (verified on POWER8), + * so nibbles come out in memory order; then the same sign trick as dot_i8i8. */ + const __vector unsigned char m4v=vec_splats((unsigned char)0x0F); + const __vector unsigned char sh4=vec_splats((unsigned char)4); + const __vector signed char b8v=vec_splats((signed char)8); + const __vector signed char vz=vec_splats((signed char)0); + __vector signed int acc=vec_splats(0); + for(;i+32<=I;i+=32){ + __vector unsigned char by=vec_xl(0,w4+(i>>1)); /* 16 byte = 32 nibble */ + __vector unsigned char lo=vec_and(by,m4v), hi=vec_sr(by,sh4); + __vector signed char w0=vec_sub((__vector signed char)vec_mergeh(lo,hi),b8v); + __vector signed char w1=vec_sub((__vector signed char)vec_mergel(lo,hi),b8v); + __vector signed char x0=vec_xl(0,(const signed char*)(x+i)); + __vector signed char x1=vec_xl(0,(const signed char*)(x+i+16)); + __vector __bool char n0=vec_cmplt(w0,vz), n1=vec_cmplt(w1,vz); + acc=vec_msum(vec_sel(x0,vec_sub(vz,x0),n0), + (__vector unsigned char)vec_sel(w0,vec_sub(vz,w0),n0),acc); + acc=vec_msum(vec_sel(x1,vec_sub(vz,x1),n1), + (__vector unsigned char)vec_sel(w1,vec_sub(vz,w1),n1),acc); + } + sum=vec_extract(acc,0)+vec_extract(acc,1)+vec_extract(acc,2)+vec_extract(acc,3); +#endif + for(;i+1>1]; sum+=((int)(b&0xF)-8)*x[i]+((int)(b>>4)-8)*x[i+1]; } + if(i>1]; sum+=((int)(b&0xF)-8)*x[i]; } + return sum; +} +#if defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8) +/* SMMLA (i8mm): vmmlaq_s32 vede ogni int8x16_t come matrice 2x8 row-major (byte 0-7 = + * riga 0, byte 8-15 = riga 1) e accumula C += A*B^T nel 2x2 int32: lane0=a0.b0, + * lane1=a0.b1, lane2=a1.b0, lane3=a1.b1. vcombine di due mezze-righe costruisce la + * matrice: A = due righe di peso (o,o+1), B = due righe di attivazione (s,s+1), quindi + * meta' traffico pesi e doppio lavoro per istruzione a S>=2. EN: vmmlaq_s32 treats each + * int8x16_t as a 2x8 row-major matrix and does C += A*B^T on a 2x2 int32 tile; vcombine + * of vget_low/high halves builds the 2-row register from two weight/activation rows. */ +static inline int32x4_t mm_tile16(int32x4_t acc, int8x16_t wo, int8x16_t wo1, + int8x16_t xs, int8x16_t xs1){ + acc=vmmlaq_s32(acc, vcombine_s8(vget_low_s8(wo), vget_low_s8(wo1)), + vcombine_s8(vget_low_s8(xs), vget_low_s8(xs1))); + return vmmlaq_s32(acc, vcombine_s8(vget_high_s8(wo), vget_high_s8(wo1)), + vcombine_s8(vget_high_s8(xs), vget_high_s8(xs1))); +} +static void matmul_q_idot_mm(float *y, const int8_t *xq, const float *sx, const int8_t *q, + const float *scale, int S, int I, int O){ + #pragma omp parallel for schedule(static) + for(int o=0;o<(O&~1);o+=2){ + const int8_t *wo=q+(int64_t)o*I, *wo1=q+(int64_t)(o+1)*I; + float sc0=scale[o], sc1=scale[o+1]; + for(int s=0;s<(S&~1);s+=2){ + const int8_t *xs=xq+(int64_t)s*I, *xs1=xq+(int64_t)(s+1)*I; + /* 4 accumulatori indipendenti: una sola catena vmmla e' latency-bound. + * EN: 4 independent accumulators; a single vmmla chain is latency-bound. */ + int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0); int i=0; + for(;i+64<=I;i+=64){ + a0=mm_tile16(a0,vld1q_s8(wo+i), vld1q_s8(wo1+i), vld1q_s8(xs+i), vld1q_s8(xs1+i)); + a1=mm_tile16(a1,vld1q_s8(wo+i+16),vld1q_s8(wo1+i+16),vld1q_s8(xs+i+16),vld1q_s8(xs1+i+16)); + a2=mm_tile16(a2,vld1q_s8(wo+i+32),vld1q_s8(wo1+i+32),vld1q_s8(xs+i+32),vld1q_s8(xs1+i+32)); + a3=mm_tile16(a3,vld1q_s8(wo+i+48),vld1q_s8(wo1+i+48),vld1q_s8(xs+i+48),vld1q_s8(xs1+i+48)); + } + for(;i+16<=I;i+=16) + a0=mm_tile16(a0,vld1q_s8(wo+i),vld1q_s8(wo1+i),vld1q_s8(xs+i),vld1q_s8(xs1+i)); + int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3)); + int32_t d00=vgetq_lane_s32(acc,0), d01=vgetq_lane_s32(acc,1); + int32_t d10=vgetq_lane_s32(acc,2), d11=vgetq_lane_s32(acc,3); + for(;i>1)), byo1=vld1q_u8(wo1+(i>>1)); + uint8x16_t cyo=vld1q_u8(wo+(i>>1)+16), cyo1=vld1q_u8(wo1+(i>>1)+16); + uint8x16x2_t zo =vzipq_u8(vandq_u8(byo, m4q), vshrq_n_u8(byo, 4)); + uint8x16x2_t zo1=vzipq_u8(vandq_u8(byo1,m4q), vshrq_n_u8(byo1,4)); + uint8x16x2_t ko =vzipq_u8(vandq_u8(cyo, m4q), vshrq_n_u8(cyo, 4)); + uint8x16x2_t ko1=vzipq_u8(vandq_u8(cyo1,m4q), vshrq_n_u8(cyo1,4)); + a0=mm_tile16(a0, vsubq_s8(vreinterpretq_s8_u8(zo.val[0]),b8q), + vsubq_s8(vreinterpretq_s8_u8(zo1.val[0]),b8q), + vld1q_s8(xs+i), vld1q_s8(xs1+i)); + a1=mm_tile16(a1, vsubq_s8(vreinterpretq_s8_u8(zo.val[1]),b8q), + vsubq_s8(vreinterpretq_s8_u8(zo1.val[1]),b8q), + vld1q_s8(xs+i+16), vld1q_s8(xs1+i+16)); + a2=mm_tile16(a2, vsubq_s8(vreinterpretq_s8_u8(ko.val[0]),b8q), + vsubq_s8(vreinterpretq_s8_u8(ko1.val[0]),b8q), + vld1q_s8(xs+i+32), vld1q_s8(xs1+i+32)); + a3=mm_tile16(a3, vsubq_s8(vreinterpretq_s8_u8(ko.val[1]),b8q), + vsubq_s8(vreinterpretq_s8_u8(ko1.val[1]),b8q), + vld1q_s8(xs+i+48), vld1q_s8(xs1+i+48)); + } + for(;i+32<=I;i+=32){ + uint8x16_t byo=vld1q_u8(wo+(i>>1)), byo1=vld1q_u8(wo1+(i>>1)); + uint8x16x2_t zo =vzipq_u8(vandq_u8(byo, m4q), vshrq_n_u8(byo, 4)); + uint8x16x2_t zo1=vzipq_u8(vandq_u8(byo1,m4q), vshrq_n_u8(byo1,4)); + a0=mm_tile16(a0, vsubq_s8(vreinterpretq_s8_u8(zo.val[0]),b8q), + vsubq_s8(vreinterpretq_s8_u8(zo1.val[0]),b8q), + vld1q_s8(xs+i), vld1q_s8(xs1+i)); + a1=mm_tile16(a1, vsubq_s8(vreinterpretq_s8_u8(zo.val[1]),b8q), + vsubq_s8(vreinterpretq_s8_u8(zo1.val[1]),b8q), + vld1q_s8(xs+i+16), vld1q_s8(xs1+i+16)); + } + int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3)); + int32_t d00=vgetq_lane_s32(acc,0), d01=vgetq_lane_s32(acc,1); + int32_t d10=vgetq_lane_s32(acc,2), d11=vgetq_lane_s32(acc,3); + for(;i+1>1], bo1=wo1[i>>1]; + int a0=(int)(bo&0xF)-8, a1=(int)(bo>>4)-8, b0=(int)(bo1&0xF)-8, b1=(int)(bo1>>4)-8; + int u0=xs[i],u1=xs[i+1],v0=xs1[i],v1=xs1[i+1]; + d00+=a0*u0+a1*u1; d01+=a0*v0+a1*v1; d10+=b0*u0+b1*u1; d11+=b0*v0+b1*v1; } + if(i>1], bo1=wo1[i>>1]; + int a0=(int)(bo&0xF)-8, b0=(int)(bo1&0xF)-8; + d00+=a0*xs[i]; d01+=a0*xs1[i]; d10+=b0*xs[i]; d11+=b0*xs1[i]; } + y[(int64_t)s*O+o] =(float)d00*sc0*sx[s]; + y[(int64_t)s*O+(o+1)] =(float)d10*sc1*sx[s]; + y[(int64_t)(s+1)*O+o] =(float)d01*sc0*sx[s+1]; + y[(int64_t)(s+1)*O+(o+1)]=(float)d11*sc1*sx[s+1]; + } + if(S&1){ int s=S-1; const int8_t *xs=xq+(int64_t)s*I; + y[(int64_t)s*O+o] =(float)dot_i4i8(wo, xs,I)*sc0*sx[s]; + y[(int64_t)s*O+(o+1)]=(float)dot_i4i8(wo1,xs,I)*sc1*sx[s]; } + } + if(O&1){ int o=O-1; const uint8_t *w=q4+(int64_t)o*rb; float sc=scale[o]; + #pragma omp parallel for schedule(static) + for(int s=0;s=2){ matmul_q_idot_mm(y,xq,sx,q,scale,S,I,O); return; } +#endif + #pragma omp parallel for schedule(static) + for(int o=0;o=2){ matmul_i4_idot_mm(y,xq,sx,q4,scale,S,I,O); return; } +#endif + #pragma omp parallel for schedule(static) + for(int o=0;og_qscratch.xq_cap){ + int8_t *p=realloc(g_qscratch.xq,xn); + if(!p){ fprintf(stderr,"OOM quant scratch\n"); exit(1); } + g_qscratch.xq=p; g_qscratch.xq_cap=xn; + } + if(sn>g_qscratch.sx_cap){ + float *p=realloc(g_qscratch.sx,sn*sizeof(float)); + if(!p){ fprintf(stderr,"OOM quant scales\n"); exit(1); } + g_qscratch.sx=p; g_qscratch.sx_cap=sn; + } + *xq=g_qscratch.xq; *sx=g_qscratch.sx; +} + +/* allow_idot=0: forza il kernel int4/int8 ESATTO (attivazioni f32). Serve alle proiezioni di + * attenzione: sono sensibili alla quantizzazione int8 delle attivazioni dell'IDOT. Misurato su + * GLM-5.2 int4, 1023 token, log-lik -5040.33 (esatto) -> -5160.47 (IDOT) = +0.117 nat/token, + * ~+12% perplexity. Gli altri matmul del prefill (o_proj, kv_b, expert) tengono l'IDOT. + * EN: allow_idot=0 forces the EXACT int4/int8 kernel (f32 activations). The attention + * projections need it: IDOT's int8 activation quantization costs +0.117 nats/token there + * (~+12% perplexity), measured. Every other prefill matmul keeps IDOT as before. */ +static void matmul_qt_ex(float *y, const float *x, QT *w, int S, int allow_idot); +static void matmul_qt(float *y, const float *x, QT *w, int S){ matmul_qt_ex(y,x,w,S,1); } static void matmul_qt_ex(float *y, const float *x, QT *w, int S, int allow_idot){ #ifdef COLI_METAL if(g_metal_enabled && S>=g_metal_gemm_min && !spec_pinned() && (w->fmt==1||w->fmt==2) && !omp_in_parallel()){ @@ -459,7 +1166,7 @@ static void matmul_qt_ex(float *y, const float *x, QT *w, int S, int allow_idot) if(g_cuda_enabled && w->cuda_eligible && !w->cuda_failed && w->fmt!=5 && !omp_in_parallel()){ const void *weights = w->fmt==0 ? (const void*)w->qf : w->fmt==1 ? (const void*)w->q8 : (const void*)w->q4; - if(coli_cuda_matmul(&w->cuda,y,x,weights,w->s,w->fmt,S,w->I,w->O,w->cuda_device)) return; + if(coli_cuda_matmul(&w->cuda,y,x,weights,w->s,w->fmt,S,w->I,w->O,w->cuda_device,w->gs)) return; w->cuda_failed=1; fprintf(stderr,"[CUDA] tensor [%d,%d] on device %d disabled after an error; falling back to CPU\n", w->O,w->I,w->cuda_device); @@ -934,13 +1641,14 @@ static void qt_cuda_colocate(QT *dst,const QT *src){ } 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); + int rb=l->kv_b.fmt==1?l->kv_b.I: + (l->kv_b.fmt==2||l->kv_b.fmt==4)?(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;dkv_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; + 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],l->kv_b.gs))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;ikv_b.cuda_device)old=i; @@ -1122,6 +1830,14 @@ static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits static void embed_row(Model *m, int tok, float *x){ int D=m->c.hidden; QT *e=&m->embed; if(e->fmt==0){ memcpy(x, e->qf+(int64_t)tok*D, D*sizeof(float)); return; } + if(e->fmt==4){ /* grouped int4: per-group scale (embed/lm_head at io_bits, usually fmt 0/1) */ + const uint8_t *q=e->q4+(int64_t)tok*((D+1)/2); int gs=e->gs,ng=(D+gs-1)/gs; + const float *scl=e->s+(int64_t)tok*ng; + for(int g=0;g*gsD)glen=D-base; float s=scl[g]; + for(int i=base;i+1>1]; x[i]=(float)((int)(byte&0xF)-8)*s; + x[i+1]=(float)((int)(byte>>4)-8)*s; } + if(glen&1){ uint8_t byte=q[(base+glen-1)>>1]; x[base+glen-1]=(float)((int)(byte&0xF)-8)*s; } } + return; } if(e->fmt==1){ const int8_t *q=e->q8+(int64_t)tok*D; float s=e->s[tok]; for(int i=0;ifmt==2){ const uint8_t *q=e->q4+(int64_t)tok*((D+1)/2); float s=e->s[tok]; /* int4 */ @@ -1885,6 +2601,13 @@ static void qt_matvec_rows(const QT *t, int r0, int n, const float *x, float *y) int I=t->I; for(int j=0;jfmt==0){ const float *w=t->qf+(int64_t)row*I; for(int i=0;ifmt==4){ /* grouped int4: per-group scale */ + const uint8_t *w=t->q4+(int64_t)row*((I+1)/2); int gs=t->gs,ng=(I+gs-1)/gs; + const float *scl=t->s+(int64_t)row*ng; + for(int g=0;g*gsI)glen=I-base; float sc=scl[g]; float acc=0; + for(int i=base;i+1>1]; acc+=((int)(b&0xF)-8)*x[i]+((int)(b>>4)-8)*x[i+1]; } + if(glen&1){ uint8_t b=w[(base+glen-1)>>1]; acc+=((int)(b&0xF)-8)*x[base+glen-1]; } + a+=acc*sc; } } else if(t->fmt==1){ const int8_t *w=t->q8+(int64_t)row*I; float s=t->s[row]; float acc=0; for(int i=0;ifmt==2){ const uint8_t *w=t->q4+(int64_t)row*((I+1)/2); float s=t->s[row]; float acc=0; @@ -4479,9 +5202,21 @@ static void run_text(Model *m, const char *snap, const char *prompt, int ngen){ grammar_setup(&T); /* metodo F: GRAMMAR=file.gbnf (#48) */ if(g_temp<0) g_temp=0.7f; /* auto: 0.7, NON l'1.0 ufficiale — la coda della * distribuzione int4 e' rumore di quantizzazione */ - int cap=(int)strlen(prompt)+16; int *pids=malloc(cap*sizeof(int)); + int cap=(int)strlen(prompt)+16; int *pids=malloc((cap+2)*sizeof(int)); int np=tok_encode(&T,prompt,(int)strlen(prompt),pids,cap); if(np<1){ fprintf(stderr,"prompt is empty after tokenization\n"); return; } + /* GLM prefix (#108): every GLM training sequence starts [gMASK]; without it the + * model is out-of-distribution and generates garbage. Serve/SCORE modes prepend it; + * run_text must too. CHAT_TEMPLATE=0 disables (raw prompt, for debugging/non-GLM). + * A prompt that already starts with the prefix passes through intact. */ + int templ=getenv("CHAT_TEMPLATE")?atoi(getenv("CHAT_TEMPLATE")):1; + if(templ){ + int gmask=tok_id_of(&T,"[gMASK]"), sop=tok_id_of(&T,""); + if(gmask>=0 && sop>=0 && (np<2 || pids[0]!=gmask || pids[1]!=sop)){ + memmove(pids+2,pids,np*sizeof(int)); + pids[0]=gmask; pids[1]=sop; np+=2; + } + } printf("prompt: %d tokens | generating up to %d (EOS stop=%d) | n-gram draft=%d\n", np, ngen, eos, g_draft); fputs(prompt,stdout); fflush(stdout); kv_alloc(m, np+ngen+g_draft+2); diff --git a/c/tests/bench_tensor_core.cu b/c/tests/bench_tensor_core.cu index c8c4362..3666eb4 100644 --- a/c/tests/bench_tensor_core.cu +++ b/c/tests/bench_tensor_core.cu @@ -30,9 +30,9 @@ int main(){ for(int i=0;imx?score[t]:mx; for(int t=0;t<3;t++){score[t]=std::exp(score[t]-mx);z+=score[t];}for(int t=0;t<3;t++)score[t]/=z; @@ -163,9 +163,9 @@ int main(int argc, char **argv) { for(int i=0;i<32;i++)ws4[i]=0.01f+(i%5)*0.002f; for(int i=0;i<64;i++)gx4[i]=std::sin((float)(i+1)*0.17f)*2.f; ColiCudaTensor *g4=nullptr,*u4=nullptr,*d4=nullptr; - if(!coli_cuda_tensor_upload(&g4,w4,ws4,2,32,32,d0)|| - !coli_cuda_tensor_upload(&u4,w4,ws4,2,32,32,d0)|| - !coli_cuda_tensor_upload(&d4,w4,ws4,2,32,32,d0))return 1; + if(!coli_cuda_tensor_upload(&g4,w4,ws4,2,32,32,d0,0)|| + !coli_cuda_tensor_upload(&u4,w4,ws4,2,32,32,d0,0)|| + !coli_cuda_tensor_upload(&d4,w4,ws4,2,32,32,d0,0))return 1; ColiCudaTensor *gg4[2]={g4,g4},*ug4[2]={u4,u4},*dg4[2]={d4,d4}; if(!coli_cuda_expert_group(gg4,ug4,dg4,group_rows,2,scalar4,gx4))return 1; setenv("COLI_CUDA_TC_INT4","1",1); diff --git a/c/tests/test_pipe_cuda.cu b/c/tests/test_pipe_cuda.cu index f20a9e5..fc66dbd 100644 --- a/c/tests/test_pipe_cuda.cu +++ b/c/tests/test_pipe_cuda.cu @@ -119,7 +119,7 @@ int main(void){ for(int i=0;i/). Each phase writes a raw log +(_.log) and all metrics are collected into report.json + report.md. + +Design notes: + - stdout and stderr are captured separately via subprocess.PIPE. The engine streams + generated text to stdout (interleaved with prompt + PROFILE stats), but TOKENS=1 dumps + clean token-id lists to stderr — that is the primary text-capture path. + - Every regex is anchored to the exact printf format strings in glm.c (verified against + profile_print line 3853, run_text line 3948, the banner line 5299, etc.). + - Subprocess calls have a hard timeout (default 600s) and are killed cleanly on expiry. + - A single phase can be run standalone; results accumulate in the output dir. +""" +import os, sys, re, json, time, argparse, subprocess, signal, traceback +from datetime import datetime +from pathlib import Path + +# --------------------------------------------------------------------------- +# PROMPT SUITE — curated across categories. "expect" is a case-insensitive +# substring checked against the generated text (None = coherence-only check). +# --------------------------------------------------------------------------- +PROMPTS = [ + {"id":"fact_capital", "cat":"factual", "prompt":"The capital of France is", + "expect":"Paris", "note":"basic world knowledge"}, + {"id":"fact_boiling", "cat":"factual", "prompt":"What is the boiling point of water in Celsius?", + "expect":"100", "note":"basic science"}, + {"id":"fact_planet", "cat":"factual", "prompt":"What planet is closest to the Sun?", + "expect":"Mercury","note":"astronomy fact"}, + {"id":"math_mult", "cat":"math", "prompt":"What is 15 times 12?", + "expect":"180", "note":"2-digit multiplication"}, + {"id":"math_add", "cat":"math", "prompt":"What is 847 plus 153?", + "expect":"1000", "note":"3-digit addition"}, + {"id":"reason_train", "cat":"reasoning", "prompt":"If a train travels 60 mph for 2.5 hours, how far does it go?", + "expect":"150", "note":"rate-time-distance"}, + {"id":"code_factorial","cat":"code", "prompt":"Write a Python function that computes the factorial of a number.", + "expect":"def", "note":"code generation"}, + {"id":"explain_nn", "cat":"explanation","prompt":"Explain what a neural network is in one sentence.", + "expect":None, "note":"coherence check"}, + {"id":"creative_story","cat":"creative", "prompt":"Write a one-sentence story about a lighthouse.", + "expect":None, "note":"creative coherence"}, + {"id":"edge_hello", "cat":"edge", "prompt":"Hello, how are you today?", + "expect":None, "note":"conversational opener"}, + {"id":"edge_the", "cat":"edge", "prompt":"The weather today is", + "expect":None, "note":"simple continuation"}, + {"id":"edge_repeat", "cat":"edge", "prompt":"The quick brown fox jumps over the lazy dog. The quick brown fox", + "expect":None, "note":"repetition-bait"}, +] + +# --------------------------------------------------------------------------- +# METRIC EXTRACTION — each parser is (name, compiled_regex, group_index, cast). +# Regexes match the EXACT printf format strings in glm.c. +# --------------------------------------------------------------------------- +def _f1(g): return float(g) +def _i1(g): return int(g) + +# Build parsers as (regex, lambda(match)->value) for clarity +RX = { + # stdout: run_text summary line (line 3948) + "decode_toks": (re.compile(r"decode (\d+) tokens in"), lambda m: int(m.group(1))), + "decode_secs": (re.compile(r"decode \d+ tokens in ([\d.]+)s"), lambda m: float(m.group(1))), + "decode_tps": (re.compile(r"decode \d+ tokens in [\d.]+s \(([\d.]+) tok/s\)"), lambda m: float(m.group(1))), + "prefill_toks": (re.compile(r"prefill (\d+) tokens in"), lambda m: int(m.group(1))), + "prefill_secs": (re.compile(r"prefill \d+ tokens in ([\d.]+)s"), lambda m: float(m.group(1))), + "hit_rate": (re.compile(r"expert hit rate ([\d.]+)%"), lambda m: float(m.group(1))), + "rss_gb": (re.compile(r"RSS ([\d.]+) GB"), lambda m: float(m.group(1))), + "experts_per_tok":(re.compile(r"experts loaded/token: ([\d.]+)"), lambda m: float(m.group(1))), + "mtp_accept": (re.compile(r"MTP acceptance ([\d.]+)%"), lambda m: float(m.group(1))), + "mtp_acc_cnt": (re.compile(r"MTP acceptance \d+% \((\d+)/(\d+)\)"), lambda m: (int(m.group(1)), int(m.group(2)))), + "spec_tok_per_fw":(re.compile(r"speculation: ([\d.]+) tokens/forward"),lambda m: float(m.group(1))), + # stdout: PROFILE lines (profile_print, line 3853-3864) + "prof_expert_disk": (re.compile(r"expert-disk ([\d.]+)s service"), lambda m: float(m.group(1))), + "prof_expert_wait": (re.compile(r"service / ([\d.]+)s wait"), lambda m: float(m.group(1))), + "prof_expert_mm": (re.compile(r"expert-matmul ([\d.]+)s"), lambda m: float(m.group(1))), + "prof_attention": (re.compile(r"\| attention ([\d.]+)s"), lambda m: float(m.group(1))), + "prof_kvb": (re.compile(r"including kvb ([\d.]+)s"), lambda m: float(m.group(1))), + "prof_lm_head": (re.compile(r"lm_head ([\d.]+)s"), lambda m: float(m.group(1))), + "prof_other": (re.compile(r"\| other ([\d.]+)s"), lambda m: float(m.group(1))), + # stdout: banner (line 5299) + "load_secs": (re.compile(r"loaded in ([\d.]+)s"), lambda m: float(m.group(1))), + "resident_mb": (re.compile(r"resident dense: ([\d.]+) MB"), lambda m: float(m.group(1))), + "mtp_status": (re.compile(r"MTP (ACTIVE|absent)"), lambda m: m.group(1)), + "n_layers": (re.compile(r"layers=(\d+) experts=(\d+)"), lambda m: int(m.group(1))), + "n_experts": (re.compile(r"layers=(\d+) experts=(\d+)"), lambda m: int(m.group(2))), + # stdout: banner idot kernel + "idot_kernel": (re.compile(r"idot: (\S+) =="), lambda m: m.group(1)), + "cache_cap": (re.compile(r"cache=(\d+) experts/layer"), lambda m: int(m.group(1))), + # stderr: RAM_GB (line 5086-5112) + "ram_budget": (re.compile(r"\[RAM_GB=([\d.]+)"), lambda m: float(m.group(1))), + "cap_lowered": (re.compile(r"cap lowered (\d+)->(\d+)"), lambda m: (int(m.group(1)), int(m.group(2)))), + "cap_raised": (re.compile(r"cap raised (\d+)->(\d+)"), lambda m: (int(m.group(1)), int(m.group(2)))), + "cap_ok": (re.compile(r"cap=(\d+) ok"), lambda m: int(m.group(1))), + # stderr: CUDA (backend_cuda.cu:389) + "cuda_device": (re.compile(r"\[CUDA\] device (\d+): (.*?), ([\d.]+) GB VRAM, sm_(\d)(\d)"), + lambda m: {"id":int(m.group(1)),"name":m.group(2).strip(),"vram_gb":float(m.group(3)), + "sm":f"{m.group(4)}.{m.group(5)}"}), + "cuda_mode": (re.compile(r"\[CUDA\] mode: (.+)"), lambda m: m.group(1)), + "cuda_tier": (re.compile(r"CUDA expert tier: (\d+) resident experts \(([\d.]+) GB\)"), + lambda m: {"resident":int(m.group(1)),"vram_gb":float(m.group(2))}), + # stderr: TOKENS dump (line 4010-4012) + "tokens_dump": (re.compile(r"^\[TOKENS\] (\d+) generated:(.*)$", re.MULTILINE), + lambda m: [int(x) for x in m.group(2).split()]), + # stderr: per-16-token progress (emit_stream line 3742) + "progress_tps": (re.compile(r"t=(\d+)\s+RSS ([\d.]+) GB\s+hit ([\d.]+)%\s+([\d.]+) tok/s\s+([\d.]+) tok/fw"), + lambda m: {"tok":int(m.group(1)),"rss":float(m.group(2)),"hit":float(m.group(3)), + "tps":float(m.group(4)),"tpf":float(m.group(5))}), + # stderr: DSA, USAGE, KV startup lines + "usage_loaded": (re.compile(r"\[USAGE\].*?(\d+) selections"), lambda m: int(m.group(1))), + "kv_slots": (re.compile(r"\[KV\].*?(\d+) context slots"), lambda m: int(m.group(1))), +} + +def extract_metrics(stdout: str, stderr: str) -> dict: + """Parse all metrics from the engine's stdout+stderr output.""" + text_out = stdout or "" + text_err = stderr or "" + metrics = {} + # For most parsers we search BOTH streams (engine is inconsistent about which + # channel a given line lands on). Some are stream-specific (noted below). + combined = text_out + "\n" + text_err + for name, (rx, fn) in RX.items(): + m = rx.search(combined) + if m: + try: metrics[name] = fn(m) + except (ValueError, IndexError): pass + # TOKENS dump is stderr-only and may appear once; grab it explicitly + m = RX["tokens_dump"][0].search(text_err) + if m: + try: metrics["tokens_dump"] = RX["tokens_dump"][1](m) + except: pass + # Multiple CUDA devices — collect all + cuda_devs = [] + for m in re.finditer(r"\[CUDA\] device (\d+): (.*?), ([\d.]+) GB VRAM, sm_(\d)(\d)", text_err): + cuda_devs.append({"id":int(m.group(1)),"name":m.group(2).strip(), + "vram_gb":float(m.group(3)),"sm":f"{m.group(4)}.{m.group(5)}"}) + if cuda_devs: metrics["cuda_devices"] = cuda_devs + # Multiple progress checkpoints — collect the full curve + progress = [] + for m in RX["progress_tps"][0].finditer(text_err): + try: + progress.append(RX["progress_tps"][1](m)) + except: pass + if progress: metrics["progress_curve"] = progress + # Multiple PROFILE lines — there are two (prefill + decode). Keep both. + prof_lines = [(i, line) for i, line in enumerate(text_out.splitlines()) if line.startswith("PROFILE:")] + for idx, (line_no, line) in enumerate(prof_lines): + label = "prefill_profile" if idx == 0 else "decode_profile" + p = {} + for pname, (rx, fn) in RX.items(): + if not pname.startswith("prof_"): continue + mm = rx.search(line) + if mm: + try: p[pname] = fn(mm) + except: pass + if p: metrics[label] = p + return metrics + + +# --------------------------------------------------------------------------- +# ENGINE RUNNER — subprocess wrapper with timeout, signal handling, logging. +# --------------------------------------------------------------------------- +class EngineRunner: + def __init__(self, glm_path, snap, out_dir, default_env=None, timeout=600): + self.glm = str(glm_path) + self.snap = str(snap) + self.out_dir = Path(out_dir) + self.out_dir.mkdir(parents=True, exist_ok=True) + self.default_env = default_env or {} + self.timeout = timeout + self.default_cap = 75 # production default (matches bench_full.sh) + + def run_prompt(self, prompt, ngen=64, env_extra=None, log_name=None, cap=None): + """Run the engine in PROMPT mode and return (stdout, stderr, returncode, elapsed).""" + env = dict(os.environ, SNAP=self.snap, PROMPT=prompt, NGEN=str(ngen)) + env.update(self.default_env) + if env_extra: env.update(env_extra) + env["TOKENS"] = "1" # always capture token ids for reliable text extraction + cmd = [self.glm, str(cap if cap is not None else self.default_cap)] + return self._exec(cmd, env, log_name) + + def run_score(self, score_file, cap=None, env_extra=None, log_name=None): + """Run the engine in SCORE (log-likelihood) mode.""" + env = dict(os.environ, SNAP=self.snap, SCORE=score_file) + env.update(self.default_env) + if env_extra: env.update(env_extra) + cmd = [self.glm, str(cap if cap is not None else self.default_cap)] + return self._exec(cmd, env, log_name) + + def _exec(self, cmd, env, log_name): + t0 = time.time() + log_path = self.out_dir / (log_name or f"run_{int(t0)}.log") + try: + proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + text=True, encoding="utf-8", errors="replace") + try: + stdout, stderr = proc.communicate(timeout=self.timeout) + rc = proc.returncode + except subprocess.TimeoutExpired: + proc.kill() + stdout, stderr = proc.communicate() + rc = -1 + stderr = (stderr or "") + f"\n[TIMEOUT after {self.timeout}s]\n" + except Exception as e: + stdout, stderr, rc = "", f"[EXCEPTION] {e}\n{traceback.format_exc()}", -2 + elapsed = time.time() - t0 + # Write raw log (both streams, clearly delimited) + with open(log_path, "w", encoding="utf-8") as f: + f.write(f"=== CMD: {' '.join(cmd)}\n=== ELAPSED: {elapsed:.1f}s\n=== RC: {rc}\n\n") + f.write("--- STDOUT ---\n"); f.write(stdout or ""); f.write("\n") + f.write("--- STDERR ---\n"); f.write(stderr or ""); f.write("\n") + return stdout or "", stderr or "", rc, elapsed + + +# --------------------------------------------------------------------------- +# TEXT RECOVERY — decode the [TOKENS] ids back to text using the tokenizer. +# Falls back to stdout extraction if tokenizer/tokens unavailable. +# --------------------------------------------------------------------------- +def recover_text(stdout, stderr, prompt, tokenizer=None): + """Extract generated text. Primary: TOKENS dump decoded. Fallback: stdout parse.""" + # Method 1: decode TOKENS ids via tokenizer + if tokenizer: + m = re.search(r"^\[TOKENS\] \d+ generated:(.*)$", stderr, re.MULTILINE) + if m: + ids = [int(x) for x in m.group(1).split()] + if ids: + try: + return tokenizer.decode(ids), "tokens" + except Exception: pass + # Method 2: stdout — text is between the prompt string and "PROFILO" or "\n---" + text = stdout or "" + # Find the prompt in stdout, take everything after it up to PROFILO + idx = text.find(prompt) + if idx >= 0: + after = text[idx + len(prompt):] + cut = after.find("PROFILO") + if cut >= 0: + return after[:cut].strip(), "stdout" + cut = after.find("\n---") + if cut >= 0: + return after[:cut].strip(), "stdout" + return "", "none" + + +# --------------------------------------------------------------------------- +# CORRUPTION / COHERENCE CHECKS +# --------------------------------------------------------------------------- +def check_repetition(token_ids): + """Detect degenerate repetition: same token 3+ consecutive times.""" + if not token_ids or len(token_ids) < 6: + return False, 0 + max_run = 1; cur = 1 + for i in range(1, len(token_ids)): + if token_ids[i] == token_ids[i-1]: cur += 1; max_run = max(max_run, cur) + else: cur = 1 + return max_run >= 3, max_run + +def check_expected(text, expect): + """Check if expected substring appears case-insensitively in the first 200 chars.""" + if not expect or not text: return None + return expect.lower() in text[:200].lower() + + +# --------------------------------------------------------------------------- +# PHASES +# --------------------------------------------------------------------------- +class DiagnosticHarness: + def __init__(self, args): + self.args = args + self.snap = args.snap + self.glm = args.glm + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + self.out_dir = Path(args.out or f"./diag_results/{ts}") + self.out_dir.mkdir(parents=True, exist_ok=True) + # Base env for all runs + base_env = {"TEMP": "0", "PIPE": "1", "PIPE_WORKERS": "8", "DIRECT": "1"} + if args.ram: base_env["RAM_GB"] = str(args.ram) + if args.cuda: + base_env.update({ + "COLI_CUDA": "1", + "COLI_GPU": str(args.gpu) if args.gpu is not None else "0", + "CUDA_DENSE": "1", + "COLI_CUDA_ATTN": "1", + "COLI_CUDA_PIPE": "2", + "COLI_CUDA_PIPE_S_MIN": "1", + }) + self.runner = EngineRunner(self.glm, self.snap, self.out_dir, base_env, timeout=args.timeout) + self.runner.default_cap = args.cap + # Load tokenizer for text recovery + self.tokenizer = None + tok_path = os.path.join(self.snap, "tokenizer.json") + if os.path.exists(tok_path): + try: + from tokenizers import Tokenizer + self.tokenizer = Tokenizer.from_file(tok_path) + except Exception as e: + print(f"[warn] could not load tokenizer: {e}", file=sys.stderr) + self.results = {"meta": {"snap": self.snap, "glm": self.glm, + "timestamp": ts, "args": vars(args)}, + "phases": {}} + + def phase_system(self): + """Phase 0: minimal run to capture startup telemetry.""" + print("\n" + "="*60 + "\nPHASE 0: SYSTEM PROBE\n" + "="*60) + stdout, stderr, rc, elapsed = self.runner.run_prompt( + "Hello", ngen=1, log_name="system_probe.log") + metrics = extract_metrics(stdout, stderr) + result = { + "rc": rc, "elapsed": elapsed, + "load_secs": metrics.get("load_secs"), + "resident_mb": metrics.get("resident_mb"), + "n_layers": metrics.get("n_layers"), + "n_experts": metrics.get("n_experts"), + "mtp_status": metrics.get("mtp_status"), + "idot_kernel": metrics.get("idot_kernel"), + "cache_cap_requested": metrics.get("cache_cap"), + "ram_budget_gb": metrics.get("ram_budget"), + "cap_lowered": metrics.get("cap_lowered"), + "cap_raised": metrics.get("cap_raised"), + "cap_final": metrics.get("cap_ok"), + "cuda_devices": metrics.get("cuda_devices", []), + "cuda_mode": metrics.get("cuda_mode"), + } + # Print summary + print(f" load time: {result['load_secs']:.2f}s" if result['load_secs'] else " load time: FAILED") + print(f" resident: {result['resident_mb']:.0f} MB" if result['resident_mb'] else "") + print(f" layers/exp: {result['n_layers']}/{result['n_experts']}" if result['n_layers'] else "") + print(f" MTP: {result['mtp_status']}") + print(f" idot kernel: {result['idot_kernel']}") + print(f" RAM budget: {result['ram_budget_gb']} GB" if result['ram_budget_gb'] else " RAM budget: auto") + if result['cap_lowered']: + print(f" cache cap: {result['cap_lowered'][0]} -> {result['cap_lowered'][1]} (RAM-lowered)") + elif result['cap_final']: + print(f" cache cap: {result['cap_final']} (ok)") + else: + print(f" cache cap: {result['cache_cap_requested']}") + for d in result['cuda_devices']: + print(f" GPU {d['id']}: {d['name']}, {d['vram_gb']:.1f} GB, sm_{d['sm']}") + if rc != 0 and rc != -1: + print(f" WARNING: engine returned rc={rc}") + self.results["phases"]["system"] = result + return result + + def phase_smoke(self): + """Phase 1: correctness smoke test across curated prompts.""" + print("\n" + "="*60 + "\nPHASE 1: CORRECTNESS SMOKE TEST\n" + "="*60) + ngen = self.args.ngen + prompt_results = [] + pass_count = 0; total = len(PROMPTS) + for p in PROMPTS: + stdout, stderr, rc, elapsed = self.runner.run_prompt( + p["prompt"], ngen=ngen, log_name=f"smoke_{p['id']}.log") + metrics = extract_metrics(stdout, stderr) + token_ids = metrics.get("tokens_dump", []) + text, method = recover_text(stdout, stderr, p["prompt"], self.tokenizer) + is_rep, max_run = check_repetition(token_ids) + expect_ok = check_expected(text, p.get("expect")) + # Verdict: pass if text is non-empty, no severe repetition, and (if factual) expected found + has_text = bool(text.strip()) + ok = has_text and not is_rep + if p.get("expect") and expect_ok is False: ok = False + if ok: pass_count += 1 + # Diagnose failure mode for reporting + if not has_text: + fail_reason = "no tokens generated (immediate EOS)" + elif is_rep: + fail_reason = f"repetition loop (max run={max_run})" + elif p.get("expect") and expect_ok is False: + fail_reason = f"expected '{p['expect']}' not found" + else: + fail_reason = "" + prompt_results.append({ + "id": p["id"], "cat": p["cat"], "prompt": p["prompt"], + "expect": p.get("expect"), "generated": text[:300], + "expect_match": expect_ok, "repetition": is_rep, "max_run": max_run, + "toks_generated": len(token_ids), "decode_tps": metrics.get("decode_tps"), + "hit_rate": metrics.get("hit_rate"), "rc": rc, "elapsed": elapsed, + "text_method": method, "pass": ok, "fail_reason": fail_reason, + }) + status = "PASS" if ok else "FAIL" + extra = f" expect={'Y' if expect_ok else ('N' if expect_ok is False else '-')}" if p.get("expect") else "" + tps = f" {metrics.get('decode_tps',0):.2f}t/s" if metrics.get("decode_tps") else "" + print(f" [{status}] {p['id']:<22} ({p['cat']:<11}) rep={'Y' if is_rep else 'N'}{extra}{tps}") + if text: + preview = text.replace("\n", " ")[:80] + print(f" -> {preview}") + elif rc != 0: + print(f" -> [engine rc={rc}]") + else: + print(f" -> [{fail_reason}]") + result = {"prompts": prompt_results, "pass": pass_count, "total": total, + "pass_rate": 100.0 * pass_count / total if total else 0} + print(f"\n SMOKE SUMMARY: {pass_count}/{total} passed ({result['pass_rate']:.0f}%)") + self.results["phases"]["smoke"] = result + return result + + def phase_diagnostic(self): + """Phase 2: single deep-instrumented run — full PROFILE + routing + MTP.""" + print("\n" + "="*60 + "\nPHASE 2: FULL SYSTEM DIAGNOSTIC\n" + "="*60) + diag_env = { + "LOOKA": "1", "DISK_SPLIT": "1", "ROUTE_AGREE": "1", + } + if self.args.cuda: + diag_env["COLI_CUDA_PROFILE"] = "1" + prompt = ("Write a short paragraph explaining how photosynthesis works. " + "Include the roles of sunlight, water, and carbon dioxide.") + stdout, stderr, rc, elapsed = self.runner.run_prompt( + prompt, ngen=self.args.ngen, env_extra=diag_env, log_name="diagnostic.log") + metrics = extract_metrics(stdout, stderr) + text, _ = recover_text(stdout, stderr, prompt, self.tokenizer) + result = { + "rc": rc, "elapsed": elapsed, + "generated_text": text[:500], + "prefill_profile": metrics.get("prefill_profile", {}), + "decode_profile": metrics.get("decode_profile", {}), + "decode_tps": metrics.get("decode_tps"), + "prefill_secs": metrics.get("prefill_secs"), + "hit_rate": metrics.get("hit_rate"), + "rss_gb": metrics.get("rss_gb"), + "experts_per_tok": metrics.get("experts_per_tok"), + "mtp_accept": metrics.get("mtp_accept"), + "mtp_counts": metrics.get("mtp_acc_cnt"), + "spec_tok_per_fw": metrics.get("spec_tok_per_fw"), + "cuda_tier": metrics.get("cuda_tier"), + "progress_curve": metrics.get("progress_curve", []), + } + # Print the PROFILE breakdown + dp = result["decode_profile"] + print(f"\n DECODE TIMING BREAKDOWN (per-bucket, seconds):") + print(f" expert-disk: {dp.get('prof_expert_disk', '?'):>8}") + print(f" expert-matmul: {dp.get('prof_expert_mm', '?'):>8}") + print(f" attention: {dp.get('prof_attention', '?'):>8}") + print(f" lm_head: {dp.get('prof_lm_head', '?'):>8}") + print(f" other: {dp.get('prof_other', '?'):>8}") + print(f"\n PERFORMANCE:") + print(f" prefill: {result['prefill_secs']:.2f}s" if result['prefill_secs'] else " prefill: ?") + print(f" decode: {result['decode_tps']:.3f} tok/s" if result['decode_tps'] else " decode: ?") + print(f" hit rate: {result['hit_rate']:.1f}%" if result.get('hit_rate') is not None else " hit rate: ?") + print(f" RSS: {result['rss_gb']:.1f} GB" if result.get('rss_gb') else " RSS: ?") + print(f" exp/tok: {result['experts_per_tok']:.1f}" if result.get('experts_per_tok') else " exp/tok: ?") + print(f" MTP: {result['mtp_accept']:.0f}% accept" if result.get('mtp_accept') is not None else " MTP: ?") + if result.get('cuda_tier'): + ct = result['cuda_tier'] + print(f" CUDA tier: {ct['resident']} experts, {ct['vram_gb']:.1f} GB VRAM") + # Show generated text preview + if text: + print(f"\n GENERATED TEXT (first 200 chars):") + print(f" {text[:200].replace(chr(10), ' ')}") + else: + print(f"\n GENERATED TEXT: [none recovered]") + self.results["phases"]["diagnostic"] = result + return result + + def phase_quality(self): + """Phase 3: benchmark accuracy via eval_glm.py SCORE mode.""" + print("\n" + "="*60 + "\nPHASE 3: QUALITY BENCHMARKS\n" + "="*60) + eval_script = os.path.join(os.path.dirname(__file__), "eval_glm.py") + bench_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "bench") + if not os.path.exists(eval_script): + print(f" [SKIP] eval_glm.py not found at {eval_script}") + self.results["phases"]["quality"] = {"error": "eval_glm.py not found"} + return None + tasks = ["hellaswag", "arc_challenge", "mmlu"] + missing = [t for t in tasks if not os.path.exists(os.path.join(bench_dir, f"{t}.jsonl"))] + if missing: + print(f" [SKIP] benchmark data missing: {missing}") + print(f" run: python tools/fetch_benchmarks.py --out {bench_dir}") + self.results["phases"]["quality"] = {"error": f"missing benchmark data: {missing}"} + return None + limit = self.args.quality_limit + py = sys.executable + cmd = [py, eval_script, "--snap", self.snap, "--glm", self.glm, + "--data", bench_dir, "--tasks", ",".join(tasks), "--limit", str(limit)] + env = dict(os.environ) + if self.args.ram: env["RAM_GB"] = str(self.args.ram) + print(f" Running eval_glm.py (tasks={tasks}, n={limit})...") + print(f" This takes ~{limit*3*4/0.05:.0f}s at 0.05 tok/s (worst case)...") + t0 = time.time() + log_path = self.out_dir / "quality_eval.log" + try: + with open(log_path, "w", encoding="utf-8") as logf: + proc = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=self.args.timeout*3, + encoding="utf-8", errors="replace") + logf.write(proc.stdout); logf.write("\n--- STDERR ---\n"); logf.write(proc.stderr) + elapsed = time.time() - t0 + # Parse the acc/acc_norm table from eval_glm.py output + scores = {} + for line in proc.stdout.splitlines(): + # lines like: "hellaswag 40 45.0% 50.0%" + m = re.match(r"(\w+)\s+(\d+)\s+([\d.]+)%\s+([\d.]+)%", line.strip()) + if m: + scores[m.group(1)] = {"n": int(m.group(2)), "acc": float(m.group(3)), + "acc_norm": float(m.group(4))} + mean_m = re.search(r"MEAN acc_norm:\s*([\d.]+)%", proc.stdout) + result = {"scores": scores, "mean_acc_norm": float(mean_m.group(1)) if mean_m else None, + "elapsed": elapsed, "rc": proc.returncode} + print(f" Completed in {elapsed:.0f}s\n") + print(f" {'task':<18} {'n':>4} {'acc':>7} {'acc_norm':>9}") + for t, s in scores.items(): + print(f" {t:<18} {s['n']:>4} {s['acc']:>6.1f}% {s['acc_norm']:>8.1f}%") + if result["mean_acc_norm"] is not None: + print(f"\n MEAN acc_norm: {result['mean_acc_norm']:.1f}%") + except subprocess.TimeoutExpired: + result = {"error": f"eval timed out after {self.args.timeout*3}s"} + print(f" [TIMEOUT] eval_glm.py exceeded {self.args.timeout*3}s") + except Exception as e: + result = {"error": str(e)} + print(f" [ERROR] {e}") + self.results["phases"]["quality"] = result + return result + + def phase_throughput(self): + """Phase 4: tok/s comparison — MTP on vs off.""" + print("\n" + "="*60 + "\nPHASE 4: THROUGHPUT BENCHMARK\n" + "="*60) + prompt = "Summarize the plot of Romeo and Juliet in three sentences." + ngen = self.args.ngen + results = {} + # Run 1: with MTP (default) + print(f" [1/2] MTP ON (draft=3)...") + stdout, stderr, rc, t = self.runner.run_prompt( + prompt, ngen=ngen, log_name="throughput_mtp_on.log") + m_on = extract_metrics(stdout, stderr) + results["mtp_on"] = {"tps": m_on.get("decode_tps"), "hit": m_on.get("hit_rate"), + "mtp_accept": m_on.get("mtp_accept"), "elapsed": t} + print(f" {m_on.get('decode_tps',0):.3f} tok/s | hit {m_on.get('hit_rate',0):.1f}% | " + f"MTP {m_on.get('mtp_accept',0):.0f}%") + # Run 2: MTP off + print(f" [2/2] MTP OFF (MTP=0)...") + stdout, stderr, rc, t = self.runner.run_prompt( + prompt, ngen=ngen, env_extra={"MTP": "0"}, log_name="throughput_mtp_off.log") + m_off = extract_metrics(stdout, stderr) + results["mtp_off"] = {"tps": m_off.get("decode_tps"), "hit": m_off.get("hit_rate"), + "elapsed": t} + print(f" {m_off.get('decode_tps',0):.3f} tok/s | hit {m_off.get('hit_rate',0):.1f}%") + # Compute speedup + if results["mtp_on"]["tps"] and results["mtp_off"]["tps"] and results["mtp_off"]["tps"] > 0: + sp = results["mtp_on"]["tps"] / results["mtp_off"]["tps"] + results["mtp_speedup"] = sp + print(f"\n MTP speedup: {sp:.2f}x ({'MTP helps' if sp > 1.05 else 'MTP hurts' if sp < 0.95 else 'no effect'})") + else: + print(f"\n MTP speedup: (insufficient data)") + self.results["phases"]["throughput"] = results + return results + + def write_report(self): + """Write report.json and report.md.""" + ts = self.results["meta"]["timestamp"] + json_path = self.out_dir / "report.json" + md_path = self.out_dir / "report.md" + # JSON + with open(json_path, "w", encoding="utf-8") as f: + json.dump(self.results, f, indent=2, default=str) + # Markdown + lines = [ + f"# Diagnostic Report — {ts}", + f"", + f"**Model:** `{self.snap}`", + f"**Engine:** `{self.glm}`", + f"**Date:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", + f"", + ] + phases = self.results["phases"] + # System + if "system" in phases: + s = phases["system"] + lines += ["## Phase 0: System Probe", ""] + lines += [f"- Load time: **{s.get('load_secs','?')}s**"] + lines += [f"- Layers/experts: {s.get('n_layers','?')}/{s.get('n_experts','?')}"] + lines += [f"- MTP: {s.get('mtp_status','?')}"] + lines += [f"- idot kernel: `{s.get('idot_kernel','?')}`"] + lines += [f"- RAM budget: {s.get('ram_budget_gb','auto')} GB"] + if s.get("cap_lowered"): + lines += [f"- Cache cap: {s['cap_lowered'][0]}→{s['cap_lowered'][1]} (RAM-lowered)"] + elif s.get("cap_final"): + lines += [f"- Cache cap: {s['cap_final']}"] + for d in s.get("cuda_devices", []): + lines += [f"- GPU {d['id']}: {d['name']}, {d['vram_gb']:.1f} GB, sm_{d['sm']}"] + lines.append("") + # Smoke + if "smoke" in phases: + sm = phases["smoke"] + lines += [f"## Phase 1: Correctness Smoke", "", + f"**{sm['pass']}/{sm['total']} prompts passed ({sm['pass_rate']:.0f}%)**", "", + "| ID | Category | Pass | Expect | Repetition | tok/s | Generated (first 80 chars) |", + "|---|---|---|---|---|---|---|"] + for p in sm["prompts"]: + gen = p["generated"][:80].replace("|", "\\|").replace("\n", " ") if p["generated"] else "" + exp = "Y" if p.get("expect_match") else ("N" if p.get("expect_match") is False else "-") + tps = f"{p.get('decode_tps',0):.2f}" if p.get("decode_tps") else "-" + lines.append(f"| {p['id']} | {p['cat']} | {'✅' if p['pass'] else '❌'} | {exp} | " + f"{'⚠️' if p['repetition'] else '-'} (run={p.get('max_run',0)}) | {tps} | {gen} |") + lines.append("") + # Diagnostic + if "diagnostic" in phases: + d = phases["diagnostic"] + lines += ["## Phase 2: Full System Diagnostic", ""] + dp = d.get("decode_profile", {}) + lines += ["### Decode Timing Breakdown", "", + "| Bucket | Seconds |", "|---|---|"] + for k, label in [("prof_expert_disk","expert-disk"),("prof_expert_mm","expert-matmul"), + ("prof_attention","attention"),("prof_lm_head","lm_head"), + ("prof_other","other")]: + v = dp.get(k, "?") + lines.append(f"| {label} | {v} |") + lines += [f"", f"### Performance", f"- Decode: **{d.get('decode_tps','?')} tok/s**", + f"- Prefill: {d.get('prefill_secs','?')}s", + f"- Expert hit rate: {d.get('hit_rate','?')}%", + f"- RSS: {d.get('rss_gb','?')} GB", + f"- Experts/token: {d.get('experts_per_tok','?')}", + f"- MTP acceptance: {d.get('mtp_accept','?')}%", ""] + if d.get("generated_text"): + lines += [f"### Generated Text", f"```\n{d['generated_text'][:300]}\n```", ""] + # Quality + if "quality" in phases: + q = phases["quality"] + lines += ["## Phase 3: Quality Benchmarks", ""] + if "error" in q: + lines += [f"⚠️ {q['error']}", ""] + else: + lines += [f"**Mean acc_norm: {q.get('mean_acc_norm','?')}%**", "", + "| Task | n | acc | acc_norm |", "|---|---|---|---|"] + for t, s in q.get("scores", {}).items(): + lines.append(f"| {t} | {s['n']} | {s['acc']:.1f}% | {s['acc_norm']:.1f}% |") + lines.append("") + # Throughput + if "throughput" in phases: + th = phases["throughput"] + lines += ["## Phase 4: Throughput", "", + "| Mode | tok/s | hit% | MTP accept |", "|---|---|---|---|"] + on = th.get("mtp_on", {}); off = th.get("mtp_off", {}) + lines.append(f"| MTP ON | {on.get('tps','?')} | {on.get('hit','?')}% | {on.get('mtp_accept','?')}% |") + lines.append(f"| MTP OFF | {off.get('tps','?')} | {off.get('hit','?')}% | — |") + if th.get("mtp_speedup"): + lines.append(f"\n**MTP speedup: {th['mtp_speedup']:.2f}x**") + lines.append("") + with open(md_path, "w", encoding="utf-8") as f: + f.write("\n".join(lines)) + print(f"\n{'='*60}") + print(f"Report written:") + print(f" JSON: {json_path}") + print(f" MD: {md_path}") + print(f" Logs: {self.out_dir}/*.log") + print(f"{'='*60}") + + def run(self): + phase = self.args.phase + if phase == "all": + self.phase_system() + self.phase_smoke() + self.phase_diagnostic() + self.phase_quality() + self.phase_throughput() + elif phase == "system": self.phase_system() + elif phase == "smoke": self.phase_smoke() + elif phase == "diagnostic": self.phase_diagnostic() + elif phase == "quality": self.phase_quality() + elif phase == "throughput": self.phase_throughput() + else: + print(f"Unknown phase: {phase}", file=sys.stderr); sys.exit(1) + self.write_report() + + +def main(): + ap = argparse.ArgumentParser(description="Comprehensive model diagnostic harness for colibri GLM-5.2") + ap.add_argument("--snap", required=True, help="model snapshot directory") + ap.add_argument("--glm", default=None, help="engine binary path (default: ./glm.exe or ./glm)") + ap.add_argument("--phase", default="all", + choices=["all","system","smoke","diagnostic","quality","throughput"], + help="which test phase to run") + ap.add_argument("--out", default=None, help="output directory (default: ./diag_results/)") + ap.add_argument("--ngen", type=int, default=64, help="generation length for smoke/throughput (default 64)") + ap.add_argument("--quality-limit", type=int, default=40, help="questions per benchmark task (default 40)") + ap.add_argument("--ram", type=float, default=0, help="RAM_GB override (0=auto)") + ap.add_argument("--cuda", action="store_true", help="enable COLI_CUDA GPU tier") + ap.add_argument("--gpu", type=int, default=None, help="GPU device ordinal (with --cuda)") + ap.add_argument("--cap", type=int, default=75, help="experts-per-layer cache cap (default 75)") + ap.add_argument("--timeout", type=int, default=600, help="per-run timeout in seconds (default 600)") + a = ap.parse_args() + # Auto-detect engine binary + if a.glm is None: + for cand in ["./glm.exe", "./glm", "../glm.exe", os.path.join(os.path.dirname(__file__), "..", "glm.exe")]: + if os.path.exists(cand): a.glm = os.path.abspath(cand); break + if a.glm is None: + print("ERROR: could not find glm/glm.exe. Specify with --glm", file=sys.stderr); sys.exit(1) + if not os.path.isdir(a.snap): + print(f"ERROR: snapshot dir does not exist: {a.snap}", file=sys.stderr); sys.exit(1) + harness = DiagnosticHarness(a) + harness.run() + +if __name__ == "__main__": + main()