Initial: Colibri upstream + HIP/ROCm port (Strix Halo / Radeon 8060S)

This commit is contained in:
2026-07-14 23:53:20 +02:00
parent 6d3ed7e62b
commit 53c37ec918
25 changed files with 567 additions and 63 deletions
Binary file not shown.
Binary file not shown.
Binary file not shown.
+565
View File
@@ -0,0 +1,565 @@
#include "backend_cuda.h"
#include <hip/hip_runtime.h>
// HIP: mma.h not available, TC_INT4 disabled via env
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <mutex>
struct ColiCudaTensor {
void *weights;
float *scales;
size_t weight_bytes;
int fmt, I, O, device;
int tracked;
};
typedef struct {
int device;
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;
hipStream_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(hipError_t err, const char *what) {
if (err == hipSuccess) return 1;
std::fprintf(stderr, "[CUDA] %s: %s\n", what, hipGetErrorString(err));
return 0;
}
static DeviceContext *find_ctx(int device) {
for (int i = 0; i < g_nctx; i++) if (g_ctx[i].device == device) return &g_ctx[i];
return nullptr;
}
static int select_ctx(DeviceContext *ctx) {
return ctx && cuda_ok(hipSetDevice(ctx->device), "select device");
}
__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 == 3) return (size_t)(I + 3) / 4;
return 0;
}
__device__ static float weight_at(const void *weights, int fmt, size_t row, int i) {
const uint8_t *base = static_cast<const uint8_t *>(weights) + row;
if (fmt == 0) return reinterpret_cast<const float *>(base)[i];
if (fmt == 1) return static_cast<float>(reinterpret_cast<const int8_t *>(base)[i]);
const uint8_t *q = base;
if (fmt == 2) {
uint8_t v = q[i >> 1];
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) {
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);
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (int n = blockDim.x >> 1; n; n >>= 1) {
if (threadIdx.x < n) partial[threadIdx.x] += partial[threadIdx.x + n];
__syncthreads();
}
if (!threadIdx.x)
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
// HIP: nvcuda removed
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) hipFree(*ptr);
*ptr = nullptr;
*cap = 0;
if (!cuda_ok(hipMalloc(ptr, bytes), "scratch allocation")) return 0;
*cap = bytes;
return 1;
}
static int reserve_bytes(void **ptr,size_t *cap,size_t bytes){
if(*cap>=bytes) return 1; if(*ptr) hipFree(*ptr); *ptr=nullptr; *cap=0;
if(!cuda_ok(hipMalloc(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)hipFreeHost(*ptr);*ptr=nullptr;*cap=0;
if(!cuda_ok(hipMallocHost((void**)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;
if (!cuda_ok(hipGetDeviceCount(&available), "device discovery")) return 0;
g_nctx = 0;
for (int i = 0; i < count; i++) {
int device = devices[i];
if (device < 0 || device >= available) {
std::fprintf(stderr, "[CUDA] invalid device %d (available: 0..%d)\n", device, available - 1);
g_nctx = 0;
return 0;
}
if (find_ctx(device)) {
std::fprintf(stderr, "[CUDA] duplicate device %d\n", device);
g_nctx = 0;
return 0;
}
DeviceContext *ctx = &g_ctx[g_nctx];
*ctx = {};
ctx->device = device;
if (!select_ctx(ctx)) { g_nctx = 0; return 0; }
hipDeviceProp_t prop{};
if (!cuda_ok(hipGetDeviceProperties(&prop, device), "device properties")) { g_nctx = 0; return 0; }
if(!cuda_ok(hipStreamCreateWithFlags(&ctx->stream,hipStreamNonBlocking),"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);
}
return 1;
}
extern "C" void coli_cuda_shutdown(void) {
for (int i = 0; i < g_nctx; i++) {
DeviceContext *ctx = &g_ctx[i];
if (!select_ctx(ctx)) continue;
if (ctx->x) hipFree(ctx->x);
if (ctx->y) hipFree(ctx->y);
if (ctx->gate) hipFree(ctx->gate);
if (ctx->up) hipFree(ctx->up);
if (ctx->qx) hipFree(ctx->qx);
if (ctx->qscale) hipFree(ctx->qscale);
if(ctx->aq)hipFree(ctx->aq);if(ctx->al)hipFree(ctx->al);if(ctx->ar)hipFree(ctx->ar);if(ctx->ac)hipFree(ctx->ac);
if (ctx->host_x) hipFreeHost(ctx->host_x);
if (ctx->host_y) hipFreeHost(ctx->host_y);
if (ctx->stream) hipStreamDestroy(ctx->stream);
if (ctx->group_desc) hipFree(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;
}
extern "C" int coli_cuda_device_count(void) { return g_nctx; }
extern "C" int coli_cuda_device_at(int index) {
return index >= 0 && index < g_nctx ? g_ctx[index].device : -1;
}
extern "C" int coli_cuda_mem_info(int device, size_t *free_bytes, size_t *total_bytes) {
DeviceContext *ctx = find_ctx(device);
if (!free_bytes || !total_bytes || !select_ctx(ctx)) return 0;
return cuda_ok(hipMemGetInfo(free_bytes, total_bytes), "memory info");
}
extern "C" void coli_cuda_stats(int device, size_t *tensor_count, size_t *tensor_bytes) {
size_t count = 0, bytes = 0;
for (int i = 0; i < g_nctx; i++) if (device < 0 || g_ctx[i].device == device) {
count += g_ctx[i].tensor_count;
bytes += g_ctx[i].tensor_bytes;
}
if (tensor_count) *tensor_count = count;
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) {
DeviceContext *ctx = find_ctx(device);
if (!tensor || !weights || I < 1 || O < 1 || !select_ctx(ctx)) return 0;
size_t rb = row_bytes(fmt, I);
if (!rb || (fmt && !scales)) return 0;
if (*tensor) {
ColiCudaTensor *t = *tensor;
return t->fmt == fmt && t->I == I && t->O == O && t->device == device;
}
ColiCudaTensor *t = static_cast<ColiCudaTensor *>(std::calloc(1, sizeof(*t)));
if (!t) return 0;
t->fmt = fmt; t->I = I; t->O = O; t->device = device; t->weight_bytes = rb * (size_t)O;
if (!cuda_ok(hipMalloc(&t->weights, t->weight_bytes), "tensor allocation") ||
!cuda_ok(hipMemcpy(t->weights, weights, t->weight_bytes, hipMemcpyHostToDevice), "tensor upload")) {
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(hipGetLastError(),"int4 weight conversion")){coli_cuda_tensor_free(t);return 0;}}
if (fmt) {
if (!cuda_ok(hipMalloc(&t->scales, (size_t)O * sizeof(float)), "scale allocation") ||
!cuda_ok(hipMemcpy(t->scales, scales, (size_t)O * sizeof(float), hipMemcpyHostToDevice), "scale upload")) {
coli_cuda_tensor_free(t);
return 0;
}
}
t->tracked = 1;
ctx->tensor_count++;
ctx->tensor_bytes += t->weight_bytes + (fmt ? (size_t)O * sizeof(float) : 0);
*tensor = t;
return 1;
}
extern "C" 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) {
if (S < 1 || !coli_cuda_tensor_upload(tensor, weights, scales, fmt, I, O, device)) return 0;
ColiCudaTensor *t = *tensor;
DeviceContext *ctx = find_ctx(t->device);
if (!select_ctx(ctx)) return 0;
size_t rb = row_bytes(fmt, I);
size_t xb = (size_t)S * I * sizeof(float), yb = (size_t)S * O * sizeof(float);
if (!reserve(&ctx->x, &ctx->x_cap, xb) || !reserve(&ctx->y, &ctx->y_cap, yb)) return 0;
if (!cuda_ok(hipMemcpy(ctx->x, x, xb, hipMemcpyHostToDevice), "input upload")) return 0;
dim3 grid((unsigned)O, (unsigned)S);
quant_matmul<<<grid, 256>>>(ctx->y, ctx->x, t->weights, t->scales, fmt, S, I, O, rb);
if (!cuda_ok(hipGetLastError(), "matmul launch") ||
!cuda_ok(hipMemcpy(y, ctx->y, yb, hipMemcpyDeviceToHost), "output download")) return 0;
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(hipMemcpy(ctx->x,x,xb,hipMemcpyHostToDevice),"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(hipGetLastError(),"expert MLP launch") ||
!cuda_ok(hipMemcpy(y,ctx->y,yb,hipMemcpyDeviceToHost),"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;
hipError_t copy_desc=async?hipMemcpyAsync(ctx->group_desc,host,(size_t)count*sizeof(GroupDesc),
hipMemcpyHostToDevice,ctx->stream)
:hipMemcpy(ctx->group_desc,host,(size_t)count*sizeof(GroupDesc),hipMemcpyHostToDevice);
if(!cuda_ok(copy_desc,"expert group descriptors"))return 0;
int profile=getenv("COLI_CUDA_PROFILE")&&atoi(getenv("COLI_CUDA_PROFILE"));
hipEvent_t ev[4]={};
if(profile) for(int i=0;i<4;i++) if(!cuda_ok(hipEventCreate(&ev[i]),"profile event")) profile=0;
if(profile) hipEventRecord(ev[0],ctx->stream);
if(async)std::memcpy(ctx->host_x,x,xb);
hipError_t copy_x=async?hipMemcpyAsync(ctx->x,ctx->host_x,xb,hipMemcpyHostToDevice,ctx->stream)
:hipMemcpy(ctx->x,x,xb,hipMemcpyHostToDevice);
if(!cuda_ok(copy_x,"expert group input upload")) return 0;
if(profile) hipEventRecord(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;
hipMemsetAsync(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) hipEventRecord(ev[2],ctx->stream);
if(!async&&!cuda_ok(hipStreamSynchronize(ctx->stream),"expert group synchronize"))return 0;
hipError_t copy_y=async?hipMemcpyAsync(ctx->host_y,ctx->y,xb,hipMemcpyDeviceToHost,ctx->stream)
:hipMemcpy(y,ctx->y,xb,hipMemcpyDeviceToHost);
if(!cuda_ok(hipGetLastError(),"expert group launch")||!cuda_ok(copy_y,"expert group output download"))return 0;
if(async){if(!cuda_ok(hipStreamSynchronize(ctx->stream),"expert group synchronize"))return 0;
std::memcpy(y,ctx->host_y,xb);}
if(profile){
hipEventRecord(ev[3],ctx->stream); hipEventSynchronize(ev[3]); float a=0,b=0,c=0;
hipEventElapsedTime(&a,ev[0],ev[1]); hipEventElapsedTime(&b,ev[1],ev[2]);
hipEventElapsedTime(&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++) hipEventDestroy(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(hipMemcpyAsync(dc->aq,q,qb,hipMemcpyHostToDevice,dc->stream),"attention q upload")||
!cuda_ok(hipMemcpyAsync(dc->al,latent,lb,hipMemcpyHostToDevice,dc->stream),"attention latent upload")||
!cuda_ok(hipMemcpyAsync(dc->ar,rope,rb,hipMemcpyHostToDevice,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(hipGetLastError(),"attention absorb launch")||
!cuda_ok(hipMemcpyAsync(ctx,dc->ac,cb,hipMemcpyDeviceToHost,dc->stream),"attention context download")||
!cuda_ok(hipStreamSynchronize(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);
if (ctx) select_ctx(ctx);
if (tensor->tracked && ctx) {
size_t bytes = tensor->weight_bytes + (tensor->fmt ? (size_t)tensor->O * sizeof(float) : 0);
if (ctx->tensor_count) ctx->tensor_count--;
if (ctx->tensor_bytes >= bytes) ctx->tensor_bytes -= bytes;
}
if (tensor->weights) hipFree(tensor->weights);
if (tensor->scales) hipFree(tensor->scales);
std::free(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;
}
extern "C" int coli_cuda_tensor_device(const ColiCudaTensor *tensor) {
return tensor ? tensor->device : -1;
}
BIN
View File
Binary file not shown.
+1 -1
View File
@@ -110,7 +110,7 @@ def cuda_binary():
if not os.path.exists(GLM) or sys.platform != "linux": return False
try:
linked=subprocess.run(["ldd",GLM],capture_output=True,text=True,timeout=3)
return any("libcudart" in line and "not found" not in line
return any("libamdhip64" in line and "not found" not in line
for line in linked.stdout.splitlines())
except (OSError,subprocess.SubprocessError): return False
Executable
BIN
View File
Binary file not shown.
BIN
View File
Binary file not shown.
Executable
BIN
View File
Binary file not shown.
Executable
BIN
View File
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
View File
Binary file not shown.
Binary file not shown.
BIN
View File
Binary file not shown.
BIN
View File
Binary file not shown.
BIN
View File
Binary file not shown.
BIN
View File
Binary file not shown.
Binary file not shown.