Unify continuous batching + heterogeneous runtime: decode batching, physical-core planning, disjoint VRAM/RAM placement, topp-policy warning (CPU-validated, CUDA on 6x5090) (#68)
* 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
This commit is contained in:
ZacharyZcR
+341
-6
@@ -1,9 +1,12 @@
|
||||
#include "backend_cuda.h"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <mma.h>
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <mutex>
|
||||
|
||||
struct ColiCudaTensor {
|
||||
void *weights;
|
||||
@@ -15,13 +18,27 @@ struct ColiCudaTensor {
|
||||
|
||||
typedef struct {
|
||||
int device;
|
||||
float *x, *y;
|
||||
size_t x_cap, y_cap;
|
||||
float *x, *y, *gate, *up;
|
||||
size_t x_cap, y_cap, gate_cap, up_cap;
|
||||
uint8_t *qx; float *qscale;
|
||||
size_t qx_cap, qscale_cap;
|
||||
float *host_x,*host_y; size_t host_x_cap,host_y_cap;
|
||||
float *aq,*al,*ar,*ac; size_t aq_cap,al_cap,ar_cap,ac_cap;
|
||||
cudaStream_t stream;
|
||||
void *group_desc; size_t group_desc_cap;
|
||||
size_t tensor_count, tensor_bytes;
|
||||
} DeviceContext;
|
||||
|
||||
typedef struct {
|
||||
const void *g,*u,*d; const float *gs,*us,*ds;
|
||||
int gf,uf,df,rows,offset;
|
||||
} GroupDesc;
|
||||
|
||||
static DeviceContext g_ctx[COLI_CUDA_MAX_DEVICES];
|
||||
static int g_nctx;
|
||||
static uint64_t g_group_calls,g_group_experts,g_group_rows;
|
||||
static double g_group_h2d_ms,g_group_kernel_ms,g_group_d2h_ms;
|
||||
static std::mutex g_group_stats_mu;
|
||||
|
||||
static int cuda_ok(cudaError_t err, const char *what) {
|
||||
if (err == cudaSuccess) return 1;
|
||||
@@ -38,7 +55,7 @@ static int select_ctx(DeviceContext *ctx) {
|
||||
return ctx && cuda_ok(cudaSetDevice(ctx->device), "select device");
|
||||
}
|
||||
|
||||
static size_t row_bytes(int fmt, int I) {
|
||||
__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;
|
||||
@@ -53,12 +70,16 @@ __device__ static float weight_at(const void *weights, int fmt, size_t row, int
|
||||
const uint8_t *q = base;
|
||||
if (fmt == 2) {
|
||||
uint8_t v = q[i >> 1];
|
||||
return static_cast<float>(((i & 1) ? (v >> 4) : (v & 15)) - 8);
|
||||
int n=(i&1)?(v>>4):(v&15); return static_cast<float>(n&8?n-16:n);
|
||||
}
|
||||
uint8_t v = q[i >> 2];
|
||||
return static_cast<float>(((v >> ((i & 3) * 2)) & 3) - 2);
|
||||
}
|
||||
|
||||
__global__ static void offset_to_signed_s4(uint8_t *q,size_t n){
|
||||
size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x;if(i<n)q[i]^=0x88;
|
||||
}
|
||||
|
||||
__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) {
|
||||
@@ -81,6 +102,137 @@ __global__ static void quant_matmul(float *y, const float *x, const void *weight
|
||||
y[(size_t)s * O + o] = partial[0] * (fmt ? scales[o] : 1.0f);
|
||||
}
|
||||
|
||||
__global__ static void silu_mul(float *gate, const float *up, size_t n) {
|
||||
size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < n) {
|
||||
float v = gate[i];
|
||||
gate[i] = (v / (1.0f + expf(-v))) * up[i];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void quantize_s4_rows(uint8_t *q,float *scale,const float *x,int S,int K){
|
||||
int s=blockIdx.x; if(s>=S)return; const float *xs=x+(size_t)s*K;
|
||||
float v=0; for(int i=threadIdx.x;i<K;i+=blockDim.x)v=fmaxf(v,fabsf(xs[i]));
|
||||
__shared__ float m[256]; m[threadIdx.x]=v; __syncthreads();
|
||||
for(int n=128;n;n>>=1){if(threadIdx.x<n)m[threadIdx.x]=fmaxf(m[threadIdx.x],m[threadIdx.x+n]);__syncthreads();}
|
||||
float sc=m[0]>0?m[0]/7.f:1.f; if(!threadIdx.x)scale[s]=sc;
|
||||
uint8_t *dst=q+(size_t)s*((K+1)/2);
|
||||
for(int b=threadIdx.x;b<(K+1)/2;b+=blockDim.x){
|
||||
int i=b*2,a=__float2int_rn(xs[i]/sc),c=i+1<K?__float2int_rn(xs[i+1]/sc):0;
|
||||
a=max(-8,min(7,a)); c=max(-8,min(7,c)); dst[b]=(uint8_t)((a&15)|((c&15)<<4));
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void grouped_s4_wmma(float *y,const uint8_t *x,const float *xscale,
|
||||
const GroupDesc *desc,int K,int O,int which){
|
||||
#if __CUDA_ARCH__ >= 750
|
||||
using namespace nvcuda;
|
||||
int warp=threadIdx.x/32,lane=threadIdx.x%32,tile=blockIdx.x*8+warp,c=blockIdx.y;
|
||||
if(tile*8>=O)return; GroupDesc d=desc[c];
|
||||
const void *w=which==0?d.g:(which==1?d.u:d.d);
|
||||
const float *ws=which==0?d.gs:(which==1?d.us:d.ds);
|
||||
int fmt=which==0?d.gf:(which==1?d.uf:d.df);
|
||||
if(fmt!=2)return;
|
||||
wmma::fragment<wmma::accumulator,8,8,32,int> acc; wmma::fill_fragment(acc,0);
|
||||
const uint8_t *a=x+(size_t)d.offset*((K+1)/2);
|
||||
const uint8_t *b=(const uint8_t*)w+(size_t)(tile*8)*((K+1)/2);
|
||||
for(int k=0;k<K;k+=32){
|
||||
wmma::fragment<wmma::matrix_a,8,8,32,wmma::experimental::precision::s4,wmma::row_major> af;
|
||||
wmma::fragment<wmma::matrix_b,8,8,32,wmma::experimental::precision::s4,wmma::col_major> bf;
|
||||
wmma::load_matrix_sync(af,a+k/2,K);
|
||||
wmma::load_matrix_sync(bf,b+k/2,K);
|
||||
wmma::mma_sync(acc,af,bf,acc);
|
||||
}
|
||||
__shared__ int out[8][64]; wmma::store_matrix_sync(out[warp],acc,8,wmma::mem_row_major);
|
||||
for(int i=lane;i<64;i+=32){int s=i/8,o=tile*8+i%8;
|
||||
if(s<d.rows&&o<O)y[(size_t)(d.offset+s)*O+o]=(float)out[warp][i]*xscale[d.offset+s]*ws[o];}
|
||||
#endif
|
||||
}
|
||||
|
||||
__global__ static void grouped_hidden(float *y,const float *x,const GroupDesc *desc,
|
||||
int I,int D,int which){
|
||||
int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z; GroupDesc d=desc[c];
|
||||
if(s>=d.rows) return;
|
||||
const void *w=which?d.u:d.g; const float *sc=which?d.us:d.gs; int fmt=which?d.uf:d.gf;
|
||||
size_t rb=row_bytes(fmt,D),row=(size_t)o*rb; const float *xs=x+(size_t)(d.offset+s)*D;
|
||||
float sum=0; for(int i=threadIdx.x;i<D;i+=blockDim.x) sum+=xs[i]*weight_at(w,fmt,row,i);
|
||||
__shared__ float p[256]; p[threadIdx.x]=sum; __syncthreads();
|
||||
for(int n=128;n;n>>=1){ if(threadIdx.x<n)p[threadIdx.x]+=p[threadIdx.x+n]; __syncthreads(); }
|
||||
if(!threadIdx.x) y[(size_t)(d.offset+s)*I+o]=p[0]*(fmt?sc[o]:1.f);
|
||||
}
|
||||
|
||||
__global__ static void grouped_down(float *y,const float *x,const GroupDesc *desc,int D,int I){
|
||||
int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z; GroupDesc d=desc[c];
|
||||
if(s>=d.rows) return;
|
||||
size_t rb=row_bytes(d.df,I),row=(size_t)o*rb; const float *xs=x+(size_t)(d.offset+s)*I;
|
||||
float sum=0; for(int i=threadIdx.x;i<I;i+=blockDim.x) sum+=xs[i]*weight_at(d.d,d.df,row,i);
|
||||
__shared__ float p[256]; p[threadIdx.x]=sum; __syncthreads();
|
||||
for(int n=128;n;n>>=1){ if(threadIdx.x<n)p[threadIdx.x]+=p[threadIdx.x+n]; __syncthreads(); }
|
||||
if(!threadIdx.x) y[(size_t)(d.offset+s)*D+o]=p[0]*(d.df?d.ds[o]:1.f);
|
||||
}
|
||||
|
||||
__device__ static void unpack_s4(uint8_t v,float *lo,float *hi){
|
||||
int a=v&15,b=v>>4; *lo=(float)(a&8?a-16:a); *hi=(float)(b&8?b-16:b);
|
||||
}
|
||||
|
||||
/* Exact low-row W4A32 path. It consumes each packed weight byte once instead
|
||||
* of routing both nibbles through weight_at(), preserving FP32 activations. */
|
||||
__global__ static void grouped_hidden_w4(float *y,const float *x,const GroupDesc *desc,
|
||||
int I,int D,int which){
|
||||
int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z;GroupDesc d=desc[c];if(s>=d.rows)return;
|
||||
const uint8_t *w=(const uint8_t*)(which?d.u:d.g);const float *sc=which?d.us:d.gs;
|
||||
const uint8_t *row=w+(size_t)o*((D+1)/2);const float *xs=x+(size_t)(d.offset+s)*D;
|
||||
float sum=0;for(int b=threadIdx.x;b<(D+1)/2;b+=blockDim.x){float a,z;unpack_s4(row[b],&a,&z);
|
||||
int i=b*2;sum+=xs[i]*a;if(i+1<D)sum+=xs[i+1]*z;}
|
||||
__shared__ float p[256];p[threadIdx.x]=sum;__syncthreads();
|
||||
for(int n=128;n;n>>=1){if(threadIdx.x<n)p[threadIdx.x]+=p[threadIdx.x+n];__syncthreads();}
|
||||
if(!threadIdx.x)y[(size_t)(d.offset+s)*I+o]=p[0]*sc[o];
|
||||
}
|
||||
|
||||
__global__ static void grouped_hidden_w4_dual(float *gate,float *up,const float *x,
|
||||
const GroupDesc *desc,int I,int D){
|
||||
int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z;GroupDesc d=desc[c];if(s>=d.rows)return;
|
||||
const uint8_t *gr=(const uint8_t*)d.g+(size_t)o*((D+1)/2);
|
||||
const uint8_t *ur=(const uint8_t*)d.u+(size_t)o*((D+1)/2);
|
||||
const float *xs=x+(size_t)(d.offset+s)*D;float ga=0,ua=0;
|
||||
for(int b=threadIdx.x;b<(D+1)/2;b+=blockDim.x){float g0,g1,u0,u1;unpack_s4(gr[b],&g0,&g1);unpack_s4(ur[b],&u0,&u1);
|
||||
int i=b*2;ga+=xs[i]*g0;ua+=xs[i]*u0;if(i+1<D){ga+=xs[i+1]*g1;ua+=xs[i+1]*u1;}}
|
||||
__shared__ float gp[256],upv[256];gp[threadIdx.x]=ga;upv[threadIdx.x]=ua;__syncthreads();
|
||||
for(int n=128;n;n>>=1){if(threadIdx.x<n){gp[threadIdx.x]+=gp[threadIdx.x+n];upv[threadIdx.x]+=upv[threadIdx.x+n];}__syncthreads();}
|
||||
if(!threadIdx.x){size_t z=(size_t)(d.offset+s)*I+o;gate[z]=gp[0]*d.gs[o];up[z]=upv[0]*d.us[o];}
|
||||
}
|
||||
|
||||
__global__ static void grouped_down_w4(float *y,const float *x,const GroupDesc *desc,int D,int I){
|
||||
int o=blockIdx.x,s=blockIdx.y,c=blockIdx.z;GroupDesc d=desc[c];if(s>=d.rows)return;
|
||||
const uint8_t *row=(const uint8_t*)d.d+(size_t)o*((I+1)/2);
|
||||
const float *xs=x+(size_t)(d.offset+s)*I;float sum=0;
|
||||
for(int b=threadIdx.x;b<(I+1)/2;b+=blockDim.x){float a,z;unpack_s4(row[b],&a,&z);
|
||||
int i=b*2;sum+=xs[i]*a;if(i+1<I)sum+=xs[i+1]*z;}
|
||||
__shared__ float p[256];p[threadIdx.x]=sum;__syncthreads();
|
||||
for(int n=128;n;n>>=1){if(threadIdx.x<n)p[threadIdx.x]+=p[threadIdx.x+n];__syncthreads();}
|
||||
if(!threadIdx.x)y[(size_t)(d.offset+s)*D+o]=p[0]*d.ds[o];
|
||||
}
|
||||
|
||||
__global__ static void attention_absorb_kernel(float *ctx,const float *q,const float *latent,
|
||||
const float *rope,const void *weights,const float *wscale,
|
||||
int fmt,int H,int Q,int R,int V,int K,int T,float scale){
|
||||
int h=blockIdx.x,tid=threadIdx.x,rbase=h*(Q+V);extern __shared__ float sm[];
|
||||
float *qa=sm,*cl=qa+K,*scores=cl+K;
|
||||
for(int k=tid;k<K;k+=blockDim.x){float a=0;for(int d=0;d<Q;d++)
|
||||
a+=q[(size_t)h*(Q+R)+d]*weight_at(weights,fmt,(size_t)(rbase+d)*row_bytes(fmt,K),k)*(fmt?wscale[rbase+d]:1.f);qa[k]=a;}
|
||||
__syncthreads();
|
||||
for(int t=tid;t<T;t+=blockDim.x){float a=0;const float *lt=latent+(size_t)t*K,*rt=rope+(size_t)t*R;
|
||||
for(int k=0;k<K;k++)a+=qa[k]*lt[k];for(int d=0;d<R;d++)a+=q[(size_t)h*(Q+R)+Q+d]*rt[d];scores[t]=a*scale;}
|
||||
__syncthreads();
|
||||
if(!tid){float mx=scores[0];for(int t=1;t<T;t++)mx=fmaxf(mx,scores[t]);float z=0;
|
||||
for(int t=0;t<T;t++){scores[t]=expf(scores[t]-mx);z+=scores[t];}for(int t=0;t<T;t++)scores[t]/=z;}
|
||||
__syncthreads();
|
||||
for(int k=tid;k<K;k+=blockDim.x){float a=0;for(int t=0;t<T;t++)a+=scores[t]*latent[(size_t)t*K+k];cl[k]=a;}
|
||||
__syncthreads();
|
||||
for(int v=tid;v<V;v+=blockDim.x){int row=rbase+Q+v;float a=0;size_t rb=row_bytes(fmt,K);
|
||||
for(int k=0;k<K;k++)a+=cl[k]*weight_at(weights,fmt,(size_t)row*rb,k);ctx[(size_t)h*V+v]=a*(fmt?wscale[row]:1.f);}
|
||||
}
|
||||
|
||||
static int reserve(float **ptr, size_t *cap, size_t bytes) {
|
||||
if (*cap >= bytes) return 1;
|
||||
if (*ptr) cudaFree(*ptr);
|
||||
@@ -91,6 +243,16 @@ static int reserve(float **ptr, size_t *cap, size_t bytes) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
static int reserve_bytes(void **ptr,size_t *cap,size_t bytes){
|
||||
if(*cap>=bytes) return 1; if(*ptr) cudaFree(*ptr); *ptr=nullptr; *cap=0;
|
||||
if(!cuda_ok(cudaMalloc(ptr,bytes),"descriptor allocation")) return 0; *cap=bytes; return 1;
|
||||
}
|
||||
|
||||
static int reserve_pinned(float **ptr,size_t *cap,size_t bytes){
|
||||
if(*cap>=bytes)return 1;if(*ptr)cudaFreeHost(*ptr);*ptr=nullptr;*cap=0;
|
||||
if(!cuda_ok(cudaMallocHost(ptr,bytes),"pinned staging allocation"))return 0;*cap=bytes;return 1;
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_init(const int *devices, int count) {
|
||||
int available = 0;
|
||||
if (!devices || count < 1 || count > COLI_CUDA_MAX_DEVICES) return 0;
|
||||
@@ -114,6 +276,9 @@ extern "C" int coli_cuda_init(const int *devices, int count) {
|
||||
if (!select_ctx(ctx)) { g_nctx = 0; return 0; }
|
||||
cudaDeviceProp prop{};
|
||||
if (!cuda_ok(cudaGetDeviceProperties(&prop, device), "device properties")) { g_nctx = 0; return 0; }
|
||||
if(!cuda_ok(cudaStreamCreateWithFlags(&ctx->stream,cudaStreamNonBlocking),"stream creation")){
|
||||
g_nctx=0;return 0;
|
||||
}
|
||||
g_nctx++;
|
||||
std::fprintf(stderr, "[CUDA] device %d: %s, %.1f GB VRAM, sm_%d%d\n",
|
||||
device, prop.name, prop.totalGlobalMem / 1e9, prop.major, prop.minor);
|
||||
@@ -127,8 +292,24 @@ extern "C" void coli_cuda_shutdown(void) {
|
||||
if (!select_ctx(ctx)) continue;
|
||||
if (ctx->x) cudaFree(ctx->x);
|
||||
if (ctx->y) cudaFree(ctx->y);
|
||||
ctx->x = ctx->y = nullptr;
|
||||
ctx->x_cap = ctx->y_cap = 0;
|
||||
if (ctx->gate) cudaFree(ctx->gate);
|
||||
if (ctx->up) cudaFree(ctx->up);
|
||||
if (ctx->qx) cudaFree(ctx->qx);
|
||||
if (ctx->qscale) cudaFree(ctx->qscale);
|
||||
if(ctx->aq)cudaFree(ctx->aq);if(ctx->al)cudaFree(ctx->al);if(ctx->ar)cudaFree(ctx->ar);if(ctx->ac)cudaFree(ctx->ac);
|
||||
if (ctx->host_x) cudaFreeHost(ctx->host_x);
|
||||
if (ctx->host_y) cudaFreeHost(ctx->host_y);
|
||||
if (ctx->stream) cudaStreamDestroy(ctx->stream);
|
||||
if (ctx->group_desc) cudaFree(ctx->group_desc);
|
||||
ctx->x = ctx->y = ctx->gate = ctx->up = nullptr;
|
||||
ctx->qx=nullptr; ctx->qscale=nullptr;
|
||||
ctx->aq=ctx->al=ctx->ar=ctx->ac=nullptr;
|
||||
ctx->host_x=ctx->host_y=nullptr;ctx->stream=nullptr;
|
||||
ctx->x_cap = ctx->y_cap = ctx->gate_cap = ctx->up_cap = 0;
|
||||
ctx->qx_cap=ctx->qscale_cap=0;
|
||||
ctx->aq_cap=ctx->al_cap=ctx->ar_cap=ctx->ac_cap=0;
|
||||
ctx->host_x_cap=ctx->host_y_cap=0;
|
||||
ctx->group_desc=nullptr; ctx->group_desc_cap=0;
|
||||
}
|
||||
g_nctx = 0;
|
||||
}
|
||||
@@ -155,6 +336,13 @@ extern "C" void coli_cuda_stats(int device, size_t *tensor_count, size_t *tensor
|
||||
if (tensor_bytes) *tensor_bytes = bytes;
|
||||
}
|
||||
|
||||
extern "C" void coli_cuda_group_stats(uint64_t *calls, uint64_t *experts, uint64_t *rows,
|
||||
double *h2d_ms, double *kernel_ms, double *d2h_ms) {
|
||||
if(calls) *calls=g_group_calls; if(experts) *experts=g_group_experts; if(rows) *rows=g_group_rows;
|
||||
if(h2d_ms) *h2d_ms=g_group_h2d_ms; if(kernel_ms) *kernel_ms=g_group_kernel_ms;
|
||||
if(d2h_ms) *d2h_ms=g_group_d2h_ms;
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_tensor_upload(ColiCudaTensor **tensor,
|
||||
const void *weights, const float *scales,
|
||||
int fmt, int I, int O, int device) {
|
||||
@@ -174,6 +362,8 @@ extern "C" int coli_cuda_tensor_upload(ColiCudaTensor **tensor,
|
||||
coli_cuda_tensor_free(t);
|
||||
return 0;
|
||||
}
|
||||
if(fmt==2){offset_to_signed_s4<<<(unsigned)((t->weight_bytes+255)/256),256>>>((uint8_t*)t->weights,t->weight_bytes);
|
||||
if(!cuda_ok(cudaGetLastError(),"int4 weight conversion")){coli_cuda_tensor_free(t);return 0;}}
|
||||
if (fmt) {
|
||||
if (!cuda_ok(cudaMalloc(&t->scales, (size_t)O * sizeof(float)), "scale allocation") ||
|
||||
!cuda_ok(cudaMemcpy(t->scales, scales, (size_t)O * sizeof(float), cudaMemcpyHostToDevice), "scale upload")) {
|
||||
@@ -207,6 +397,151 @@ extern "C" int coli_cuda_matmul(ColiCudaTensor **tensor,
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_expert_mlp(ColiCudaTensor *gate, ColiCudaTensor *up,
|
||||
ColiCudaTensor *down, float *y,
|
||||
const float *x, int S) {
|
||||
if (!gate || !up || !down || !x || !y || S < 1 ||
|
||||
gate->device != up->device || gate->device != down->device ||
|
||||
gate->I != up->I || gate->O != up->O ||
|
||||
down->I != gate->O || down->O != gate->I) return 0;
|
||||
DeviceContext *ctx = find_ctx(gate->device);
|
||||
if (!select_ctx(ctx)) return 0;
|
||||
int D = gate->I, I = gate->O;
|
||||
size_t xb=(size_t)S*D*sizeof(float), ib=(size_t)S*I*sizeof(float);
|
||||
size_t yb=(size_t)S*D*sizeof(float);
|
||||
if (!reserve(&ctx->x,&ctx->x_cap,xb) || !reserve(&ctx->y,&ctx->y_cap,yb) ||
|
||||
!reserve(&ctx->gate,&ctx->gate_cap,ib) || !reserve(&ctx->up,&ctx->up_cap,ib)) return 0;
|
||||
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<<<hidden_grid,256>>>(ctx->gate,ctx->x,gate->weights,gate->scales,
|
||||
gate->fmt,S,D,I,row_bytes(gate->fmt,D));
|
||||
quant_matmul<<<hidden_grid,256>>>(ctx->up,ctx->x,up->weights,up->scales,
|
||||
up->fmt,S,D,I,row_bytes(up->fmt,D));
|
||||
size_t n=(size_t)S*I;
|
||||
silu_mul<<<(unsigned)((n+255)/256),256>>>(ctx->gate,ctx->up,n);
|
||||
quant_matmul<<<output_grid,256>>>(ctx->y,ctx->gate,down->weights,down->scales,
|
||||
down->fmt,S,I,D,row_bytes(down->fmt,I));
|
||||
if (!cuda_ok(cudaGetLastError(),"expert MLP launch") ||
|
||||
!cuda_ok(cudaMemcpy(y,ctx->y,yb,cudaMemcpyDeviceToHost),"expert output download")) return 0;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
|
||||
ColiCudaTensor *const *ups,
|
||||
ColiCudaTensor *const *downs,
|
||||
const int *rows, int count,
|
||||
float *y, const float *x) {
|
||||
if (!gates || !ups || !downs || !rows || !x || !y || count < 1) return 0;
|
||||
ColiCudaTensor *first=gates[0];
|
||||
if (!first) return 0;
|
||||
int device=first->device,D=first->I,I=first->O,total=0,max_rows=0;
|
||||
GroupDesc host[64]; if(count>64) return 0;
|
||||
int all_s4=1;
|
||||
for(int c=0;c<count;c++){
|
||||
ColiCudaTensor *g=gates[c],*u=ups[c],*d=downs[c];
|
||||
if(!g||!u||!d||rows[c]<1||g->device!=device||u->device!=device||d->device!=device||
|
||||
g->I!=D||u->I!=D||g->O!=I||u->O!=I||d->I!=I||d->O!=D) return 0;
|
||||
host[c]={g->weights,u->weights,d->weights,g->scales,u->scales,d->scales,
|
||||
g->fmt,u->fmt,d->fmt,rows[c],total};
|
||||
all_s4&=g->fmt==2&&u->fmt==2&&d->fmt==2;
|
||||
total+=rows[c]; if(rows[c]>max_rows) max_rows=rows[c];
|
||||
}
|
||||
DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0;
|
||||
size_t xb=(size_t)total*D*sizeof(float), ib=(size_t)total*I*sizeof(float);
|
||||
if(!reserve(&ctx->x,&ctx->x_cap,xb)||!reserve(&ctx->y,&ctx->y_cap,xb)||
|
||||
!reserve(&ctx->gate,&ctx->gate_cap,ib)||!reserve(&ctx->up,&ctx->up_cap,ib)||
|
||||
!reserve_bytes(&ctx->group_desc,&ctx->group_desc_cap,(size_t)count*sizeof(GroupDesc))) return 0;
|
||||
int async=!getenv("COLI_CUDA_ASYNC")||atoi(getenv("COLI_CUDA_ASYNC"));
|
||||
if(async&&(!reserve_pinned(&ctx->host_x,&ctx->host_x_cap,xb)||
|
||||
!reserve_pinned(&ctx->host_y,&ctx->host_y_cap,xb)))return 0;
|
||||
cudaError_t copy_desc=async?cudaMemcpyAsync(ctx->group_desc,host,(size_t)count*sizeof(GroupDesc),
|
||||
cudaMemcpyHostToDevice,ctx->stream)
|
||||
:cudaMemcpy(ctx->group_desc,host,(size_t)count*sizeof(GroupDesc),cudaMemcpyHostToDevice);
|
||||
if(!cuda_ok(copy_desc,"expert group descriptors"))return 0;
|
||||
int profile=getenv("COLI_CUDA_PROFILE")&&atoi(getenv("COLI_CUDA_PROFILE"));
|
||||
cudaEvent_t ev[4]={};
|
||||
if(profile) for(int i=0;i<4;i++) if(!cuda_ok(cudaEventCreate(&ev[i]),"profile event")) profile=0;
|
||||
if(profile) cudaEventRecord(ev[0],ctx->stream);
|
||||
if(async)std::memcpy(ctx->host_x,x,xb);
|
||||
cudaError_t copy_x=async?cudaMemcpyAsync(ctx->x,ctx->host_x,xb,cudaMemcpyHostToDevice,ctx->stream)
|
||||
:cudaMemcpy(ctx->x,x,xb,cudaMemcpyHostToDevice);
|
||||
if(!cuda_ok(copy_x,"expert group input upload")) return 0;
|
||||
if(profile) cudaEventRecord(ev[1],ctx->stream);
|
||||
GroupDesc *dev=(GroupDesc*)ctx->group_desc;
|
||||
int tc=getenv("COLI_CUDA_TC_INT4")&&atoi(getenv("COLI_CUDA_TC_INT4"));
|
||||
tc=tc&&all_s4&&D%32==0&&I%32==0&&D%8==0&&I%8==0;
|
||||
int tc_min=getenv("COLI_CUDA_TC_MIN_ROWS")?atoi(getenv("COLI_CUDA_TC_MIN_ROWS")):8;
|
||||
for(int c=0;c<count&&tc;c++)tc=rows[c]>=tc_min;
|
||||
if(tc){
|
||||
size_t qb=(size_t)(total+7)*(size_t)(D>I?D:I)/2;
|
||||
if(!reserve_bytes((void**)&ctx->qx,&ctx->qx_cap,qb)||
|
||||
!reserve(&ctx->qscale,&ctx->qscale_cap,(size_t)(total+7)*sizeof(float)))return 0;
|
||||
cudaMemsetAsync(ctx->qx,0,qb,ctx->stream);
|
||||
quantize_s4_rows<<<total,256,0,ctx->stream>>>(ctx->qx,ctx->qscale,ctx->x,total,D);
|
||||
grouped_s4_wmma<<<dim3((unsigned)((I+63)/64),(unsigned)count),256,0,ctx->stream>>>(ctx->gate,ctx->qx,ctx->qscale,dev,D,I,0);
|
||||
grouped_s4_wmma<<<dim3((unsigned)((I+63)/64),(unsigned)count),256,0,ctx->stream>>>(ctx->up,ctx->qx,ctx->qscale,dev,D,I,1);
|
||||
silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I);
|
||||
quantize_s4_rows<<<total,256,0,ctx->stream>>>(ctx->qx,ctx->qscale,ctx->gate,total,I);
|
||||
grouped_s4_wmma<<<dim3((unsigned)((D+63)/64),(unsigned)count),256,0,ctx->stream>>>(ctx->y,ctx->qx,ctx->qscale,dev,I,D,2);
|
||||
}else if(all_s4&&(!getenv("COLI_CUDA_W4_PACKED")||atoi(getenv("COLI_CUDA_W4_PACKED")))){
|
||||
dim3 hg((unsigned)I,(unsigned)max_rows,(unsigned)count),og((unsigned)D,(unsigned)max_rows,(unsigned)count);
|
||||
int dual=!getenv("COLI_CUDA_DUAL_PROJ")||atoi(getenv("COLI_CUDA_DUAL_PROJ"));
|
||||
if(dual)grouped_hidden_w4_dual<<<hg,256,0,ctx->stream>>>(ctx->gate,ctx->up,ctx->x,dev,I,D);
|
||||
else{
|
||||
grouped_hidden_w4<<<hg,256,0,ctx->stream>>>(ctx->gate,ctx->x,dev,I,D,0);
|
||||
grouped_hidden_w4<<<hg,256,0,ctx->stream>>>(ctx->up,ctx->x,dev,I,D,1);
|
||||
}
|
||||
silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I);
|
||||
grouped_down_w4<<<og,256,0,ctx->stream>>>(ctx->y,ctx->gate,dev,D,I);
|
||||
}else{
|
||||
dim3 hg((unsigned)I,(unsigned)max_rows,(unsigned)count),og((unsigned)D,(unsigned)max_rows,(unsigned)count);
|
||||
grouped_hidden<<<hg,256,0,ctx->stream>>>(ctx->gate,ctx->x,dev,I,D,0);
|
||||
grouped_hidden<<<hg,256,0,ctx->stream>>>(ctx->up,ctx->x,dev,I,D,1);
|
||||
silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I);
|
||||
grouped_down<<<og,256,0,ctx->stream>>>(ctx->y,ctx->gate,dev,D,I);
|
||||
}
|
||||
if(profile) cudaEventRecord(ev[2],ctx->stream);
|
||||
if(!async&&!cuda_ok(cudaStreamSynchronize(ctx->stream),"expert group synchronize"))return 0;
|
||||
cudaError_t copy_y=async?cudaMemcpyAsync(ctx->host_y,ctx->y,xb,cudaMemcpyDeviceToHost,ctx->stream)
|
||||
:cudaMemcpy(y,ctx->y,xb,cudaMemcpyDeviceToHost);
|
||||
if(!cuda_ok(cudaGetLastError(),"expert group launch")||!cuda_ok(copy_y,"expert group output download"))return 0;
|
||||
if(async){if(!cuda_ok(cudaStreamSynchronize(ctx->stream),"expert group synchronize"))return 0;
|
||||
std::memcpy(y,ctx->host_y,xb);}
|
||||
if(profile){
|
||||
cudaEventRecord(ev[3],ctx->stream); cudaEventSynchronize(ev[3]); float a=0,b=0,c=0;
|
||||
cudaEventElapsedTime(&a,ev[0],ev[1]); cudaEventElapsedTime(&b,ev[1],ev[2]);
|
||||
cudaEventElapsedTime(&c,ev[2],ev[3]);
|
||||
{ std::lock_guard<std::mutex> lock(g_group_stats_mu);
|
||||
g_group_h2d_ms+=a; g_group_kernel_ms+=b; g_group_d2h_ms+=c; }
|
||||
for(int i=0;i<4;i++) cudaEventDestroy(ev[i]);
|
||||
}
|
||||
{ std::lock_guard<std::mutex> lock(g_group_stats_mu);
|
||||
g_group_calls++; g_group_experts+=(uint64_t)count; g_group_rows+=(uint64_t)total; }
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int coli_cuda_attention_absorb(ColiCudaTensor *w,float *ctx,const float *q,
|
||||
const float *latent,const float *rope,int H,int Q,
|
||||
int R,int V,int K,int T,float scale){
|
||||
if(!w||!ctx||!q||!latent||!rope||H<1||Q<1||R<1||V<1||K<1||K>512||T<1||T>4096||
|
||||
w->I!=K||w->O!=H*(Q+V))return 0;
|
||||
DeviceContext *dc=find_ctx(w->device);if(!select_ctx(dc))return 0;
|
||||
size_t qb=(size_t)H*(Q+R)*sizeof(float),lb=(size_t)T*K*sizeof(float);
|
||||
size_t rb=(size_t)T*R*sizeof(float),cb=(size_t)H*V*sizeof(float);
|
||||
if(!reserve(&dc->aq,&dc->aq_cap,qb)||!reserve(&dc->al,&dc->al_cap,lb)||
|
||||
!reserve(&dc->ar,&dc->ar_cap,rb)||!reserve(&dc->ac,&dc->ac_cap,cb))return 0;
|
||||
if(!cuda_ok(cudaMemcpyAsync(dc->aq,q,qb,cudaMemcpyHostToDevice,dc->stream),"attention q upload")||
|
||||
!cuda_ok(cudaMemcpyAsync(dc->al,latent,lb,cudaMemcpyHostToDevice,dc->stream),"attention latent upload")||
|
||||
!cuda_ok(cudaMemcpyAsync(dc->ar,rope,rb,cudaMemcpyHostToDevice,dc->stream),"attention rope upload"))return 0;
|
||||
size_t shared=(size_t)(2*K+T)*sizeof(float);
|
||||
attention_absorb_kernel<<<H,256,shared,dc->stream>>>(dc->ac,dc->aq,dc->al,dc->ar,w->weights,w->scales,
|
||||
w->fmt,H,Q,R,V,K,T,scale);
|
||||
if(!cuda_ok(cudaGetLastError(),"attention absorb launch")||
|
||||
!cuda_ok(cudaMemcpyAsync(ctx,dc->ac,cb,cudaMemcpyDeviceToHost,dc->stream),"attention context download")||
|
||||
!cuda_ok(cudaStreamSynchronize(dc->stream),"attention synchronize"))return 0;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) {
|
||||
if (!tensor) return;
|
||||
DeviceContext *ctx = find_ctx(tensor->device);
|
||||
|
||||
Reference in New Issue
Block a user