diff --git a/GPU_BACKENDS.md b/GPU_BACKENDS.md index be421bfb..7f877811 100644 --- a/GPU_BACKENDS.md +++ b/GPU_BACKENDS.md @@ -41,6 +41,10 @@ distributing or on machines with an unsupported iGPU visible to the runtime header) because gfx GPUs report `compute_major >= 7` and a runtime check alone would select empty kernel bodies under HIP. On AMD, all compute uses the portable kernels; rocWMMA matrix-core support is a possible follow-up. +- `COLI_CUDA_WARP_DECODE=1` — opt-in low-row INT4 decode kernel on NVIDIA. + Maps one output neuron to one 32-lane warp and eight adjacent outputs to + each block, while preserving the original 256-thread accumulation order. + This removes block-wide barriers without changing results bit-for-bit. ## Validation diff --git a/c/Makefile b/c/Makefile index f29e498c..08ae1282 100644 --- a/c/Makefile +++ b/c/Makefile @@ -130,6 +130,8 @@ INSTALL ?= install # clean. See backend_loader.c and README "cuda-dll" below. CUDA ?= 0 CUDA_DLL ?= 0 +COLI_ANS ?= 0 +DIETGPU_ROOT ?= ifneq ($(IS_WIN),) # the CUDA installer sets CUDA_PATH system-wide (e.g. C:\Program Files\NVIDIA # GPU Computing Toolkit\CUDA\v13.2); fall back to it before the POSIX default. @@ -228,6 +230,16 @@ CFLAGS += -DCOLI_CUDA LDFLAGS += -L$(CUDA_HOME)/lib64 -Wl,-rpath,$(CUDA_HOME)/lib64 -lcudart -lstdc++ CUDA_OBJ = backend_cuda.o INK_CUDA_OBJ = backend_cuda_ink.o +ifeq ($(COLI_ANS),1) +ifeq ($(strip $(DIETGPU_ROOT)),) +$(error DIETGPU_ROOT is required when COLI_ANS=1) +endif +CFLAGS += -DCOLI_ANS +GPUFLAGS += -DCOLI_ANS -I$(DIETGPU_ROOT) -I$(DIETGPU_ROOT)/build/third_party/glog -I$(DIETGPU_ROOT)/third_party/glog/src +ANS_LIBS = -L$(DIETGPU_ROOT)/build/lib -Wl,-rpath,$(DIETGPU_ROOT)/build/lib -lgpu_ans -ldietgpu_utils -lglog +ANS_NVCC_LIBS = -L$(DIETGPU_ROOT)/build/lib -Xlinker -rpath -Xlinker $(DIETGPU_ROOT)/build/lib -lgpu_ans -ldietgpu_utils -lglog +LDFLAGS += $(ANS_LIBS) +endif endif ifeq ($(HIP),1) @@ -324,7 +336,7 @@ gemm-test: tests/test_gemm_largebatch.mm backend_metal.mm backend_metal.h cuda-test: backend_cuda.cu backend_cuda.h backend_gpu_compat.h tests/test_backend_cuda.cu @command -v "$(GPUCC)" >/dev/null 2>&1 || { echo "$(GPUCC_NAME) not found" >&2; exit 1; } - "$(GPUCC)" $(GPUFLAGS) backend_cuda.cu tests/test_backend_cuda.cu -o backend_cuda_test$(EXE) + "$(GPUCC)" $(GPUFLAGS) backend_cuda.cu tests/test_backend_cuda.cu -o backend_cuda_test$(EXE) $(ANS_NVCC_LIBS) ./backend_cuda_test$(EXE) # convenience alias: kernel correctness on AMD (same test, hipcc toolchain) @@ -334,11 +346,11 @@ hip-test: # CI: compile the backend and its test binary WITHOUT executing them (for # runners with the toolchain but no GPU). Pass CUDA_ARCH/HIP=1+HIP_ARCH. gpu-compile: backend_cuda.o - "$(GPUCC)" $(GPUFLAGS) backend_cuda.cu tests/test_backend_cuda.cu -o backend_cuda_test$(EXE) + "$(GPUCC)" $(GPUFLAGS) backend_cuda.cu tests/test_backend_cuda.cu -o backend_cuda_test$(EXE) $(ANS_NVCC_LIBS) cuda-bench: backend_cuda.cu backend_cuda.h backend_gpu_compat.h tests/bench_tensor_core.cu @command -v "$(GPUCC)" >/dev/null 2>&1 || { echo "$(GPUCC_NAME) not found" >&2; exit 1; } - "$(GPUCC)" $(GPUFLAGS) backend_cuda.cu tests/bench_tensor_core.cu -o backend_cuda_bench$(EXE) + "$(GPUCC)" $(GPUFLAGS) backend_cuda.cu tests/bench_tensor_core.cu -o backend_cuda_bench$(EXE) $(ANS_NVCC_LIBS) ./backend_cuda_bench$(EXE) olmoe$(EXE): olmoe.c st.h json.h compat.h diff --git a/c/backend_cuda.cu b/c/backend_cuda.cu index a85c7adb..a59361d7 100644 --- a/c/backend_cuda.cu +++ b/c/backend_cuda.cu @@ -5,7 +5,20 @@ #include #include #include +#include +#include #include +#include + +#if defined(__linux__) +#include +#include +#endif + +#ifdef COLI_ANS +#include +#include +#endif #if defined(__HIP_PLATFORM_AMD__) || defined(__HIP__) #include @@ -27,10 +40,18 @@ struct ColiCudaTensor { 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; + int weights_owned; +#ifdef COLI_ANS + size_t archive_bytes; + int compressed; +#endif RaggedKVEntry ragged[512]; int ragged_count; }; +#ifdef COLI_ANS +struct AnsArenaChunk { uint8_t *p; size_t used,cap; }; +#endif typedef struct { int device; int compute_major,compute_minor; @@ -46,6 +67,13 @@ typedef struct { void *group_desc; size_t group_desc_cap; size_t tensor_count, tensor_bytes; int group_pending; size_t group_pending_bytes; /* async expert-group in flight (Inc.4) */ +#ifdef COLI_ANS + void *ans_raw; size_t ans_raw_cap; + void *ans_host; size_t ans_host_cap; + int ans_copy_pending; + dietgpu::StackDeviceMemory *ans_scratch; + std::vector *ans_chunks; +#endif } DeviceContext; typedef struct { @@ -59,6 +87,21 @@ 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; +#ifdef COLI_ANS +static FILE *g_ans_sidecar; +static int g_ans_sidecar_pack; +#if defined(__linux__) +static int g_ans_direct_fd=-1; +static off_t g_ans_direct_off; +#endif +static uint64_t g_ans_load_records; +static double g_ans_header_s,g_ans_read_s,g_ans_stage_s,g_ans_enqueue_s; +static int g_ans_profile_printed; +static double ans_now_s(){ + using clock=std::chrono::steady_clock; + return std::chrono::duration(clock::now().time_since_epoch()).count(); +} +#endif static int cuda_ok(cudaError_t err, const char *what) { if (err == cudaSuccess) return 1; @@ -407,6 +450,59 @@ __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); } +__device__ static float warp_sum(float v){ + for(int n=16;n;n>>=1)v+=__shfl_down_sync(0xffffffffu,v,n); + return v; +} + +/* Decode-oriented W4 path following WarpDecode's token-centric mapping: + * one warp owns one output neuron, so an 8-warp block produces 8 adjacent + * outputs without block-wide barriers. The resident row-major W4 layout is + * already contiguous for this access pattern; an offline permutation is not + * required for this first A/B. */ +__global__ static void grouped_hidden_w4_warp(float *gate,float *up,const float *x, + const GroupDesc *desc,int I,int D){ + int lane=threadIdx.x&31,warp=threadIdx.x>>5; + int o=(int)blockIdx.x*8+warp,s=blockIdx.y,c=blockIdx.z; + GroupDesc d=desc[c];if(o>=I||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[8]={},ua[8]={}; + /* These are the eight logical lanes lane+32*k of the old 256-thread + * block. Keeping their stride and reduction order makes the warp mapping + * bit-identical to grouped_hidden_w4_dual. */ + for(int k=0;k<8;k++)for(int b=lane+32*k;b<(D+1)/2;b+=256){float g0,g1,u0,u1; + unpack_s4(gr[b],&g0,&g1);unpack_s4(ur[b],&u0,&u1); + int i=b*2;ga[k]+=xs[i]*g0;ua[k]+=xs[i]*u0; + if(i+1>5; + int o=(int)blockIdx.x*8+warp,s=blockIdx.y,c=blockIdx.z; + GroupDesc d=desc[c];if(o>=D||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[8]={}; + for(int k=0;k<8;k++)for(int b=lane+32*k;b<(I+1)/2;b+=256){float a,z; + unpack_s4(row[b],&a,&z);int i=b*2; + sum[k]+=xs[i]*a;if(i+1ans_chunks)ctx->ans_chunks=new std::vector; + if(ctx->ans_chunks->empty()||ctx->ans_chunks->back().cap-ctx->ans_chunks->back().usedans_chunks->push_back({p,0,cap}); + } + AnsArenaChunk &c=ctx->ans_chunks->back();void *p=c.p+c.used;c.used+=bytes;return p; +} +static int ans_host_reserve(DeviceContext *ctx,size_t bytes){ + if(ctx->ans_copy_pending){ + if(!cuda_ok(cudaStreamSynchronize(ctx->stream),"ANS sidecar upload synchronize"))return 0; + ctx->ans_copy_pending=0; + } + if(ctx->ans_host_cap>=bytes)return 1; + if(ctx->ans_host)cudaFreeHost(ctx->ans_host); + ctx->ans_host=nullptr;ctx->ans_host_cap=0; + if(!cuda_ok(cudaMallocHost(&ctx->ans_host,bytes),"ANS pinned staging allocation"))return 0; + ctx->ans_host_cap=bytes;return 1; +} +static int prepare_group_weights(DeviceContext *ctx, + ColiCudaTensor *const *gates,ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs,int count,GroupDesc *host){ + int n=0; size_t total=0; + for(int c=0;ccompressed){n++;total+=q[k]->weight_bytes;} + } + if(!n) return 1; + if(!g_ans_profile_printed&&std::getenv("COLI_ANS_PROFILE")){ + g_ans_profile_printed=1; + std::fprintf(stderr, + "[ANS] load profile: %llu records | header %.2fs | read %.2fs | " + "staging/alloc %.2fs | enqueue %.2fs\n", + (unsigned long long)g_ans_load_records,g_ans_header_s,g_ans_read_s, + g_ans_stage_s,g_ans_enqueue_s); + } + if(!ctx->ans_scratch||!reserve_bytes(&ctx->ans_raw,&ctx->ans_raw_cap,total)) return 0; + std::vector in; in.reserve(n); + std::vector out; out.reserve(n); + std::vector cap; cap.reserve(n); + size_t off=0; + for(int c=0;ccompressed){ + void *raw=(uint8_t*)ctx->ans_raw+off; + in.push_back(q[k]->weights);out.push_back(raw);cap.push_back((uint32_t)q[k]->weight_bytes); + *dst[k]=raw;off+=q[k]->weight_bytes; + } + } + dietgpu::ANSCodecConfig config(11,false); + dietgpu::ansDecodeBatchPointer(*ctx->ans_scratch,config,(uint32_t)n,in.data(),out.data(), + cap.data(),nullptr,nullptr,ctx->stream); + return cuda_ok(cudaGetLastError(),"ANS expert decode launch"); +} +#else +static int prepare_group_weights(DeviceContext *,ColiCudaTensor *const *, + ColiCudaTensor *const *,ColiCudaTensor *const *,int,GroupDesc *){return 1;} +#endif + /* Publish quant.h's E8 codebook to every configured device. __constant__ memory * is per-device, so this walks the contexts; the engine calls it once after init * rather than the backend carrying a second copy of the table that could drift @@ -661,6 +820,12 @@ extern "C" int coli_cuda_init(const int *devices, int count) { if(!cuda_ok(cudaStreamCreateWithFlags(&ctx->stream,cudaStreamNonBlocking),"stream creation")){ g_nctx=0;return 0; } +#ifdef COLI_ANS + if(std::getenv("CUDA_RAW_EXPERTS")){ + ctx->ans_scratch = new dietgpu::StackDeviceMemory(device, 1ull << 30); + ctx->ans_chunks = new std::vector; + } +#endif 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); @@ -685,6 +850,15 @@ extern "C" void coli_cuda_shutdown(void) { if (ctx->host_kv) cudaFreeHost(ctx->host_kv); if (ctx->stream) cudaStreamDestroy(ctx->stream); if (ctx->group_desc) cudaFree(ctx->group_desc); +#ifdef COLI_ANS + if(ctx->ans_copy_pending)cudaStreamSynchronize(ctx->stream); + if(ctx->ans_host)cudaFreeHost(ctx->ans_host); + if (ctx->ans_raw) cudaFree(ctx->ans_raw); + if(ctx->ans_chunks){for(auto &c:*ctx->ans_chunks)cudaFree(c.p);delete ctx->ans_chunks;} + delete ctx->ans_scratch; + ctx->ans_scratch=nullptr;ctx->ans_chunks=nullptr;ctx->ans_raw=nullptr;ctx->ans_raw_cap=0; + ctx->ans_host=nullptr;ctx->ans_host_cap=0;ctx->ans_copy_pending=0; +#endif ctx->x = ctx->y = ctx->gate = ctx->up = nullptr; ctx->qx=nullptr; ctx->qscale=nullptr; ctx->aq=ctx->al=ctx->ar=ctx->ac=nullptr; @@ -696,6 +870,12 @@ extern "C" void coli_cuda_shutdown(void) { ctx->group_desc=nullptr; ctx->group_desc_cap=0; } g_nctx = 0; +#ifdef COLI_ANS + if(g_ans_sidecar){std::fclose(g_ans_sidecar);g_ans_sidecar=nullptr;} +#if defined(__linux__) + if(g_ans_direct_fd>=0){close(g_ans_direct_fd);g_ans_direct_fd=-1;g_ans_direct_off=0;} +#endif +#endif } extern "C" int coli_cuda_device_count(void) { return g_nctx; } @@ -765,6 +945,7 @@ extern "C" int coli_cuda_tensor_upload(ColiCudaTensor **tensor, coli_cuda_tensor_free(t); return 0; } + t->weights_owned=1; if(fmt==2||fmt==4){ /* same nibble layout: offset-binary -> signed in place */ 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;}} @@ -791,10 +972,170 @@ extern "C" int coli_cuda_tensor_upload_g(ColiCudaTensor **tensor, return r; } +#ifdef COLI_ANS +struct AnsSidecarHeader { + uint32_t magic,raw_bytes,archive_bytes,fmt,I,O; +}; +static FILE *ans_sidecar(void){ + if(g_ans_sidecar) return g_ans_sidecar; + const char *path=std::getenv("COLI_ANS_SIDECAR"); + if(!path||!*path) return nullptr; + g_ans_sidecar_pack=std::getenv("COLI_ANS_PACK")&&std::atoi(std::getenv("COLI_ANS_PACK")); + g_ans_sidecar=std::fopen(path,g_ans_sidecar_pack?"wb":"rb"); +#if defined(__linux__) + if(g_ans_sidecar&&!g_ans_sidecar_pack&&std::getenv("COLI_ANS_DIRECT")&& + std::atoi(std::getenv("COLI_ANS_DIRECT"))){ + g_ans_direct_fd=open(path,O_RDONLY|O_DIRECT); + if(g_ans_direct_fd<0)std::fprintf(stderr,"[ANS] O_DIRECT unavailable; using buffered sidecar\n"); + } + if(g_ans_sidecar&&!g_ans_sidecar_pack&&g_ans_direct_fd<0) + posix_fadvise(fileno(g_ans_sidecar),0,0,POSIX_FADV_SEQUENTIAL); +#endif + return g_ans_sidecar; +} +extern "C" int coli_cuda_tensor_upload_compressed(ColiCudaTensor **tensor, + const void *weights,const float *scales,int fmt,int I,int O,int device){ + if(fmt!=2 || !tensor || *tensor) return 0; /* prototype: per-row int4 experts only */ + FILE *sidecar=ans_sidecar(); + if(sidecar&&!g_ans_sidecar_pack){ + AnsSidecarHeader h{}; + size_t expected=((size_t)I+1)/2*(size_t)O; + double t0=ans_now_s(); +#if defined(__linux__) + size_t direct_delta=0,direct_bytes=0,direct_data=0; + if(g_ans_direct_fd>=0){ + alignas(4096) uint8_t first[8192]; + off_t start=g_ans_direct_off&~off_t(4095); + direct_delta=(size_t)(g_ans_direct_off-start); + ssize_t got=pread(g_ans_direct_fd,first,sizeof(first),start); + if(got<(ssize_t)(direct_delta+sizeof(h))){ + std::fprintf(stderr,"[ANS] direct header read failed at %lld: got %lld errno %d\n", + (long long)g_ans_direct_off,(long long)got,errno); + return 0; + } + std::memcpy(&h,first+direct_delta,sizeof(h)); + }else +#endif + if(std::fread(&h,sizeof(h),1,sidecar)!=1)return 0; + if(h.magic!=0x31534e41u|| + h.fmt!=(uint32_t)fmt||h.I!=(uint32_t)I||h.O!=(uint32_t)O|| + expected>UINT32_MAX||h.raw_bytes!=(uint32_t)expected|| + !h.archive_bytes||h.archive_bytes>dietgpu::getMaxCompressedSize(h.raw_bytes)){ + std::fprintf(stderr,"[ANS] invalid or mismatched sidecar record\n"); + return 0; + } + g_ans_header_s+=ans_now_s()-t0; + DeviceContext *ctx=find_ctx(device); if(!ctx||!select_ctx(ctx)) return 0; + ColiCudaTensor *t=(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=h.raw_bytes; + t->scale_count=(size_t)O;t->archive_bytes=h.archive_bytes;t->compressed=1; + t0=ans_now_s(); + t->weights=ans_arena_alloc(ctx,h.archive_bytes); + size_t scale_bytes=(size_t)O*sizeof(float),scale_off; +#if defined(__linux__) + if(g_ans_direct_fd>=0){ + size_t record_bytes=sizeof(h)+(size_t)h.archive_bytes; + direct_data=direct_delta+sizeof(h); + direct_bytes=(direct_delta+record_bytes+4095)&~size_t(4095); + scale_off=(direct_bytes+255)&~size_t(255); + }else +#endif + scale_off=(h.archive_bytes+255)&~size_t(255); + if(!t->weights||!ans_host_reserve(ctx,scale_off+scale_bytes)|| + !cuda_ok(cudaMalloc(&t->scales,scale_bytes),"ANS sidecar scales")){ + coli_cuda_tensor_free(t);return 0; + } + g_ans_stage_s+=ans_now_s()-t0; + t0=ans_now_s(); +#if defined(__linux__) + if(g_ans_direct_fd>=0){ + off_t start=g_ans_direct_off&~off_t(4095); + ssize_t got=pread(g_ans_direct_fd,ctx->ans_host,direct_bytes,start); + if(got<(ssize_t)(direct_data+h.archive_bytes)){ + std::fprintf(stderr, + "[ANS] direct record read failed at %lld: need %zu got %lld errno %d\n", + (long long)g_ans_direct_off,direct_data+h.archive_bytes, + (long long)got,errno); + coli_cuda_tensor_free(t);return 0; + } + g_ans_direct_off+=(off_t)sizeof(h)+(off_t)h.archive_bytes; + }else +#endif + if(std::fread(ctx->ans_host,h.archive_bytes,1,sidecar)!=1){ + std::fprintf(stderr,"[ANS] truncated sidecar record\n"); + coli_cuda_tensor_free(t);return 0; + } + g_ans_read_s+=ans_now_s()-t0; + std::memcpy((uint8_t*)ctx->ans_host+scale_off,scales,scale_bytes); + t0=ans_now_s(); + void *archive_src= +#if defined(__linux__) + g_ans_direct_fd>=0?(uint8_t*)ctx->ans_host+direct_data: +#endif + ctx->ans_host; + if(!cuda_ok(cudaMemcpyAsync(t->weights,archive_src,h.archive_bytes, + cudaMemcpyHostToDevice,ctx->stream),"ANS sidecar upload")|| + !cuda_ok(cudaMemcpyAsync(t->scales,(uint8_t*)ctx->ans_host+scale_off,scale_bytes, + cudaMemcpyHostToDevice,ctx->stream),"ANS sidecar scale upload")){ + coli_cuda_tensor_free(t);return 0; + } + g_ans_enqueue_s+=ans_now_s()-t0;g_ans_load_records++; + ctx->ans_copy_pending=1; + t->tracked=1;ctx->tensor_count++;ctx->tensor_bytes+=h.archive_bytes+(size_t)O*sizeof(float); + *tensor=t;return 1; + } + if(!sidecar||!g_ans_sidecar_pack) return 0; + if(!coli_cuda_tensor_upload(tensor,weights,scales,fmt,I,O,device)) return 0; + ColiCudaTensor *t=*tensor; + DeviceContext *ctx=find_ctx(device); + if(!ctx||!ctx->ans_scratch||!select_ctx(ctx)){ coli_cuda_tensor_free(t);*tensor=nullptr;return 0; } + uint32_t raw=(uint32_t)t->weight_bytes; + uint32_t bound=dietgpu::getMaxCompressedSize(raw), *dsize=nullptr; + void *tmp=nullptr; + if(!cuda_ok(cudaMalloc(&tmp,bound),"ANS archive allocation")|| + !cuda_ok(cudaMalloc(&dsize,sizeof(*dsize)),"ANS size allocation")){ + if(tmp)cudaFree(tmp);if(dsize)cudaFree(dsize);coli_cuda_tensor_free(t);*tensor=nullptr;return 0; + } + dietgpu::ANSCodecConfig config(11,false); + /* tensor_upload converted offset-binary nibbles on the legacy stream. + * ctx->stream is explicitly non-blocking, so it does not inherit the + * legacy-stream dependency. Finish that one-time conversion before the + * encoder reads the bytes. */ + if(!cuda_ok(cudaStreamSynchronize(0),"ANS source conversion synchronize")){ + cudaFree(tmp);cudaFree(dsize);coli_cuda_tensor_free(t);*tensor=nullptr;return 0; + } + dietgpu::ansEncodeBatchStride(*ctx->ans_scratch,config,1,t->weights,raw,raw,nullptr, + tmp,bound,dsize,ctx->stream); + uint32_t used=0; + int ok=cuda_ok(cudaMemcpyAsync(&used,dsize,sizeof(used),cudaMemcpyDeviceToHost,ctx->stream), + "ANS size download")&& + cuda_ok(cudaStreamSynchronize(ctx->stream),"ANS encode synchronize")&&used>0&&used archive(used); + AnsSidecarHeader h{0x31534e41u,raw,used,(uint32_t)fmt,(uint32_t)I,(uint32_t)O}; + ok=cuda_ok(cudaMemcpy(archive.data(),tmp,used,cudaMemcpyDeviceToHost),"ANS sidecar download")&& + std::fwrite(&h,sizeof(h),1,sidecar)==1&& + std::fwrite(archive.data(),archive.size(),1,sidecar)==1; + cudaFree(tmp);coli_cuda_tensor_free(t);*tensor=nullptr; + if(!ok)return 0; + t=(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=raw; + t->archive_bytes=used;t->compressed=1;*tensor=t; + return 1; + } + return 0; +} +#endif + extern "C" int coli_cuda_tensor_update(ColiCudaTensor *tensor, const void *weights, const float *scales) { if (!tensor || !weights || (tensor->fmt && tensor->fmt != 6 && !scales)) return 0; +#ifdef COLI_ANS + if(tensor->compressed) return 0; +#endif DeviceContext *ctx=find_ctx(tensor->device); if (!select_ctx(ctx)) return 0; if (!cuda_ok(cudaMemcpy(tensor->weights,weights,tensor->weight_bytes, @@ -804,7 +1145,6 @@ 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; /* fmt=6 has no scale buffer at all (scales live in-block, scale_count 0), and * the fallback below would otherwise copy O floats out of a NULL host pointer. */ return !tensor->fmt || tensor->fmt==6 || cuda_ok(cudaMemcpy(tensor->scales,scales, @@ -943,6 +1283,7 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, total+=rows[c]; if(rows[c]>max_rows) max_rows=rows[c]; } DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0; + if(!prepare_group_weights(ctx,gates,ups,downs,count,host)) 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)|| @@ -1021,15 +1362,20 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, off16+=r; } }else if(all_s4&&(!getenv("COLI_CUDA_W4_PACKED")||atoi(getenv("COLI_CUDA_W4_PACKED")))){ - dim3 hg((unsigned)I,(unsigned)max_rows,(unsigned)count),og((unsigned)D,(unsigned)max_rows,(unsigned)count); + int warp_decode=COLI_GPU_HAS_WARP32&&getenv("COLI_CUDA_WARP_DECODE")&& + atoi(getenv("COLI_CUDA_WARP_DECODE")); + dim3 hg((unsigned)(warp_decode?(I+7)/8:I),(unsigned)max_rows,(unsigned)count); + dim3 og((unsigned)(warp_decode?(D+7)/8: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<<stream>>>(ctx->gate,ctx->up,ctx->x,dev,I,D); + if(warp_decode)grouped_hidden_w4_warp<<stream>>>(ctx->gate,ctx->up,ctx->x,dev,I,D); + else if(dual)grouped_hidden_w4_dual<<stream>>>(ctx->gate,ctx->up,ctx->x,dev,I,D); else{ /* non-dual path has no fused epilogue: silu stays a kernel here */ grouped_hidden_w4<<stream>>>(ctx->gate,ctx->x,dev,I,D,0); grouped_hidden_w4<<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<<stream>>>(ctx->y,ctx->gate,dev,D,I); + if(warp_decode)grouped_down_w4_warp<<stream>>>(ctx->y,ctx->gate,dev,D,I); + else grouped_down_w4<<stream>>>(ctx->y,ctx->gate,dev,D,I); }else if(all_q4&&any_g4){ /* grouped-int4 (fmt=4) present: per-group scales (#334). fmt=2 members * ride along as the ng=1 special case. silu fused in the dual epilogue. */ @@ -1099,6 +1445,7 @@ extern "C" int coli_cuda_expert_group_issue(ColiCudaTensor *const *gates, } if(total>8) return 0; /* decode-scale only */ DeviceContext *ctx=find_ctx(device); if(!ctx||ctx->group_pending||!select_ctx(ctx)) return 0; + if(!prepare_group_weights(ctx,gates,ups,downs,count,host)) 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)|| @@ -1315,12 +1662,17 @@ extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) { /* Must mirror the upload's accounting exactly: fmt=6 never charged for a * scale buffer, and over-subtracting here trips the >= guard below, which * silently leaves the tensor's bytes on the device counter forever. */ - size_t bytes = tensor->weight_bytes + + size_t storage_bytes = +#ifdef COLI_ANS + tensor->compressed ? tensor->archive_bytes : +#endif + tensor->weight_bytes; + size_t bytes = storage_bytes + ((tensor->fmt && tensor->fmt != 6) ? (size_t)tensor->O * ng * sizeof(float) : 0); if (ctx->tensor_count) ctx->tensor_count--; if (ctx->tensor_bytes >= bytes) ctx->tensor_bytes -= bytes; } - if (tensor->weights) cudaFree(tensor->weights); + if (tensor->weights&&tensor->weights_owned) cudaFree(tensor->weights); if (tensor->scales) cudaFree(tensor->scales); for(int i=0;iragged_count;i++){ if(tensor->ragged[i].latent)cudaFree(tensor->ragged[i].latent); @@ -1332,7 +1684,12 @@ extern "C" void coli_cuda_tensor_free(ColiCudaTensor *tensor) { extern "C" size_t coli_cuda_tensor_bytes(const ColiCudaTensor *tensor) { 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); + size_t storage_bytes = +#ifdef COLI_ANS + tensor->compressed ? tensor->archive_bytes : +#endif + tensor->weight_bytes; + return storage_bytes + (tensor->fmt ? (size_t)tensor->O * ng * sizeof(float) : 0); } extern "C" int coli_cuda_tensor_device(const ColiCudaTensor *tensor) { @@ -1576,6 +1933,7 @@ extern "C" int coli_cuda_expert_group_resident_issue(ColiCudaTensor *const *gate } if(!all_s4) return 0; /* resident path: per-row int4 only */ DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0; + if(!prepare_group_weights(ctx,gates,ups,downs,count,host)) return 0; if(!ctx->ev_done_ok){ if(!cuda_ok(cudaEventCreateWithFlags(&ctx->ev_done,cudaEventDisableTiming), "resident group event")) return 0; diff --git a/c/backend_cuda.h b/c/backend_cuda.h index e3198368..c98eb6ad 100644 --- a/c/backend_cuda.h +++ b/c/backend_cuda.h @@ -47,6 +47,14 @@ COLI_CUDA_DLLEXPORT int coli_cuda_tensor_upload_g(ColiCudaTensor **tensor, COLI_CUDA_DLLEXPORT int coli_cuda_tensor_upload(ColiCudaTensor **tensor, const void *weights, const float *scales, int fmt, int I, int O, int device); +#ifdef COLI_ANS +/* Experimental Linux-only GPU-resident entropy tier. The archive remains in + * VRAM and is decoded into per-device scratch immediately before a grouped + * expert launch. */ +COLI_CUDA_DLLEXPORT int coli_cuda_tensor_upload_compressed(ColiCudaTensor **tensor, + const void *weights, const float *scales, + int fmt, int I, int O, int device); +#endif /* * y[S,O] = x[S,I] @ W[O,I]^T. diff --git a/c/backend_gpu_compat.h b/c/backend_gpu_compat.h index 74bbd578..c9ea4b03 100644 --- a/c/backend_gpu_compat.h +++ b/c/backend_gpu_compat.h @@ -27,10 +27,12 @@ #if __has_include() #include #define COLI_GPU_HAS_WMMA 1 +#define COLI_GPU_HAS_WARP32 0 #define __CUDA_ARCH__ 700 #define __half rocwmma::float16_t namespace nvcuda { namespace wmma = ::rocwmma; } #define __syncwarp() __syncthreads() +#define __shfl_down_sync(mask, value, delta) __shfl_down(value, delta) #else #error "rocWMMA headers not found. Install rocwmma-dev (or rocm-hip-runtime-dev) to build with HIP." #endif @@ -77,6 +79,7 @@ namespace nvcuda { namespace wmma = ::rocwmma; } #include #include #define COLI_GPU_HAS_WMMA 1 +#define COLI_GPU_HAS_WARP32 1 #endif #endif /* COLIBRI_BACKEND_GPU_COMPAT_H */ diff --git a/c/colibri.c b/c/colibri.c index 8a978767..bd755c63 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -290,6 +290,7 @@ static int g_cuda_expert_auto; static int g_cuda_dense; static int g_cuda_release_host; static double g_cuda_reserve_gb; /* CUDA_RESERVE_GB: VRAM headroom kept free of expert tier (default 2 GB) */ +static int g_cuda_raw_experts=-1; /* experimental ANS tier: keep this global hot prefix raw */ static int g_cuda_devices[COLI_CUDA_MAX_DEVICES], g_cuda_ndev, g_cuda_rr; static int64_t g_cuda_dense_projected[COLI_CUDA_MAX_DEVICES]; static void qt_cuda_reset(QT *t){ @@ -309,6 +310,12 @@ static int qt_cuda_upload(QT *t){ return coli_cuda_tensor_upload_g(&t->cuda,weights,t->s,t->fmt,t->I,t->O,t->cuda_device,t->gs); return coli_cuda_tensor_upload(&t->cuda,weights,t->s,t->fmt,t->I,t->O,t->cuda_device); } +#ifdef COLI_ANS +static int qt_cuda_upload_compressed(QT *t){ + if(t->fmt!=2) return 0; + return coli_cuda_tensor_upload_compressed(&t->cuda,t->q4,t->s,t->fmt,t->I,t->O,t->cuda_device); +} +#endif static int qt_cuda_update(QT *t){ const void *weights=t->fmt==0?(const void*)t->qf: t->fmt==1?(const void*)t->q8:(const void*)t->q4; @@ -2070,7 +2077,19 @@ static void expert_host_release(Model *m, ESlot *s){ * fixed at the original expert_load site. fslab is plain malloc/falloc * on the CPU path, so its free() stays plain (Metal path frees it before * re-alloc and never reaches here with an aligned fslab on _WIN32). */ - if(s->aslab){ s->slab=NULL; s->fslab=NULL; } /* arena slice (#419): detach, keep caps, never free */ + if(s->aslab){ +#ifdef __linux__ + /* The arena owns the virtual address, not the resident pages. A GPU + * slot no longer needs its host copy after the synchronous H2D upload. + * Keeping those anonymous pages resident made every live REPIN pass + * grow RSS by one swapped batch (~300 MB at 16 experts). DONTNEED + * preserves the reusable slice while returning its pages to Linux; + * expert_host_ensure reloads them before the slot can run on CPU. */ + madvise(s->slab,(size_t)s->slab_cap,MADV_DONTNEED); + madvise(s->fslab,(size_t)s->fslab_cap*sizeof(float),MADV_DONTNEED); +#endif + s->slab=NULL; s->fslab=NULL; /* detach, keep caps/arena ownership */ + } else { compat_aligned_free(s->slab); free(s->fslab); s->slab=NULL; s->fslab=NULL; s->slab_cap=s->fslab_cap=0; } QT *q[3]={&s->g,&s->u,&s->d}; for(int k=0;k<3;k++){ q[k]->qf=NULL; q[k]->q8=NULL; q[k]->q4=NULL; q[k]->s=NULL; } @@ -2371,6 +2390,32 @@ static int attn_pipe_prefill(Model *m, Layer *l, int layer, const float *x, int free(chost); /* gli scratch device restano al contesto */ return ok; } + +/* Decode-scale prefix only: upload x once and reuse it for q_a and kv_a. Keep + * the stock CPU RMSNorm between q_a and q_b so greedy output stays exact, then + * return Q/comp to the stock attention path. */ +static int attn_pipe_prefix(Model *m,Layer *l,const float *x,int S,float *QR,float *Q,float *comp){ + Cfg *c=&m->c; int D=c->hidden,ql=c->q_lora,kvl=c->kv_lora,R=c->qk_rope; + int dev=l->q_a.cuda_device; + if(S<1||S>4||dev<0||l->q_b.cuda_device!=dev||l->kv_a.cuda_device!=dev) return 0; + size_t xb=(size_t)S*D*4,qrb=(size_t)S*ql*4; + size_t qb=(size_t)S*c->n_heads*c->qk_head*4,cb=(size_t)S*(kvl+R)*4; + float *xd=coli_cuda_pipe_scratch(dev,0,xb); + float *qrd=coli_cuda_pipe_scratch(dev,1,qrb); + float *qd=coli_cuda_pipe_scratch(dev,2,qb); + float *cd=coli_cuda_pipe_scratch(dev,3,cb); + if(!xd||!qrd||!qd||!cd) return 0; + if(!coli_cuda_pipe_upload(dev,xd,x,xb)|| + !coli_cuda_pipe_gemm(l->q_a.cuda,qrd,xd,S)|| + !coli_cuda_pipe_download(dev,qrd,QR,qrb)) return 0; + for(int s=0;sq_a_ln,ql,c->eps); + if(!coli_cuda_pipe_upload(dev,qrd,QR,qrb)|| + !coli_cuda_pipe_gemm(l->q_b.cuda,qd,qrd,S)|| + !coli_cuda_pipe_gemm(l->kv_a.cuda,cd,xd,S)|| + !coli_cuda_pipe_download(dev,qd,Q,qb)|| + !coli_cuda_pipe_download(dev,cd,comp,cb)) return 0; + return 1; +} #endif static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int pos_base, @@ -2429,7 +2474,7 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p * batch-union. matmul_qt_ex(...,0) keeps them on the EXACT int4 kernel: letting S>1 pull * them into IDOT is much faster but costs ~12% perplexity (measured). Batching alone is * bit-identical to upstream; the kernel switch is not. */ - int pipe_done=0; + int pipe_done=0,prefix_done=0; #ifdef COLI_CUDA if(g_cuda_pipe&&!kvs&&S>=8&&layern_layers&&g_cuda_enabled&&c->kv_lora<=512&& !(m->has_dsa&&pos_base+S>c->index_topk)&& @@ -2438,8 +2483,14 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p qt_cuda_upload(&l->q_a)&&qt_cuda_upload(&l->q_b)&&qt_cuda_upload(&l->kv_a)&& qt_cuda_upload(&l->kv_b)&&qt_cuda_upload(&l->o)) pipe_done=attn_pipe_prefill(m,l,layer,x,0,S,pos_base,out,NULL); + if(!pipe_done&&getenv("COLI_CUDA_ATTN_PREFIX")&&atoi(getenv("COLI_CUDA_ATTN_PREFIX"))&& + !kvs&&S<=4&&g_cuda_enabled&& + (!m->has_dsa||pos_base+S<=c->index_topk)&& + l->q_a.cuda_eligible&&l->q_b.cuda_eligible&&l->kv_a.cuda_eligible&& + qt_cuda_upload(&l->q_a)&&qt_cuda_upload(&l->q_b)&&qt_cuda_upload(&l->kv_a)) + prefix_done=attn_pipe_prefix(m,l,x,S,QR,Q,comp); #endif - if(!pipe_done){ + if(!pipe_done&&!prefix_done){ matmul_qt_ex(QR, x, &l->q_a, S, 0); for(int s=0;sq_lora; rmsnorm(qr, qr, l->q_a_ln, c->q_lora, c->eps); } /* q_b legge il residuo NORMATO */ @@ -6321,6 +6372,13 @@ static void pin_load(Model *m, const char *statspath, double gb){ if(g_cuda_expert_auto) budget=safe_total; if(g_cuda_enabled&&g_cuda_release_host&&budget>0){ prefix_est=(int)(budget/eb)+g_cuda_ndev; +#ifdef COLI_ANS + if(g_cuda_raw_experts>=0){ + int raw=g_cuda_raw_experts; + if((double)raw*eb>budget) raw=(int)(budget/eb); + prefix_est=raw+(int)((budget-(double)raw*eb)/(0.80*eb))+g_cuda_ndev; + } +#endif npin+=prefix_est; /* additive: prefix RAM is returned after upload */ } #endif @@ -6362,17 +6420,34 @@ static void pin_load(Model *m, const char *statspath, double gb){ int li=r[a].l; { ESlot *s=&m->pin[li][slot_of[a]]; int64_t need=qt_bytes(&s->g)+qt_bytes(&s->u)+qt_bytes(&s->d); - if(m->gpu_expert_bytes+need>budget) break; + int compress=0; +#ifdef COLI_ANS + compress=g_cuda_raw_experts>=0&&a>=g_cuda_raw_experts; +#endif + int64_t projected=compress?need*80/100:need; + if(m->gpu_expert_bytes+projected>budget) break; int tried[COLI_CUDA_MAX_DEVICES]={0}, placed=0; for(int attempt=0;attempt=need && + for(int i=0;i=projected && (best<0||placed_b[i]g.cuda_device=s->u.cuda_device=s->d.cuda_device=g_cuda_devices[best]; s->g.cuda_eligible=s->u.cuda_eligible=s->d.cuda_eligible=1; - if(qt_cuda_upload(&s->g) && qt_cuda_upload(&s->u) && qt_cuda_upload(&s->d)){ + int uploaded= +#ifdef COLI_ANS + compress ? (qt_cuda_upload_compressed(&s->g)&&qt_cuda_upload_compressed(&s->u)&& + qt_cuda_upload_compressed(&s->d)) : +#endif + (qt_cuda_upload(&s->g) && qt_cuda_upload(&s->u) && qt_cuda_upload(&s->d)); +#ifdef COLI_ANS + if(compress&&!uploaded){ + fprintf(stderr,"[ANS] sidecar load failed; refusing partial VRAM placement\n"); + exit(2); + } +#endif + if(uploaded){ int64_t actual=(int64_t)coli_cuda_tensor_bytes(s->g.cuda) +(int64_t)coli_cuda_tensor_bytes(s->u.cuda) +(int64_t)coli_cuda_tensor_bytes(s->d.cuda); @@ -6839,8 +6914,22 @@ int main(int argc, char **argv){ g_cuda_expert_auto=cuda_expert&&!strcmp(cuda_expert,"auto"); g_cuda_expert_gb=cuda_expert&&!g_cuda_expert_auto?atof(cuda_expert):0; g_cuda_reserve_gb=getenv("CUDA_RESERVE_GB")?atof(getenv("CUDA_RESERVE_GB")):2.0; +#ifdef COLI_ANS + g_cuda_raw_experts=getenv("CUDA_RAW_EXPERTS")?atoi(getenv("CUDA_RAW_EXPERTS")):-1; + if(g_cuda_raw_experts>=0&&(!getenv("COLI_ANS_SIDECAR")||!*getenv("COLI_ANS_SIDECAR"))){ + fprintf(stderr,"COLI_ANS_SIDECAR is required when CUDA_RAW_EXPERTS is set\n"); + return 1; + } + if(g_cuda_raw_experts>=0&&g_repin){ + fprintf(stderr,"REPIN is incompatible with the fixed-order ANS sidecar\n"); + return 1; + } +#endif if(!getenv("REPIN")&&g_cuda_expert_auto&&getenv("PIN_GB")&& !strcmp(getenv("PIN_GB"),"all")) g_repin=16; +#ifdef COLI_ANS + if(g_cuda_raw_experts>=0) g_repin=0; +#endif g_cuda_release_host=getenv("CUDA_RELEASE_HOST")?atoi(getenv("CUDA_RELEASE_HOST")):(g_cuda_ndev>1); if((getenv("COLI_GPU")||getenv("COLI_GPUS"))&&!g_cuda_enabled){ fprintf(stderr,"COLI_GPU(S) requires COLI_CUDA=1\n"); return 2; } if(g_cuda_dense&&!g_cuda_enabled){ fprintf(stderr,"CUDA_DENSE requires COLI_CUDA=1\n"); return 2; } @@ -6969,6 +7058,13 @@ int main(int argc, char **argv){ pin_load(&m,pin,pin_gb&&!strcmp(pin_gb,"all")?-1.0:pin_gb?atof(pin_gb):10.0); /* PIN_GB=all (#80) */ } } +#if defined(COLI_CUDA) && defined(COLI_ANS) + if(getenv("COLI_ANS_PACK")&&atoi(getenv("COLI_ANS_PACK"))){ + fprintf(stderr,"[ANS] sidecar packing complete; exiting before inference\n"); + coli_cuda_shutdown(); + return 0; + } +#endif if(getenv("COUPLE")&&*getenv("COUPLE")){ /* coupling-scored cross-layer prefetch (#176) */ g_couple_k=getenv("COUPLE_K")?atoi(getenv("COUPLE_K")):8; if(g_couple_k<1)g_couple_k=1; if(g_couple_k>32)g_couple_k=32; diff --git a/c/tests/test_backend_cuda.cu b/c/tests/test_backend_cuda.cu index f1537750..188a3447 100644 --- a/c/tests/test_backend_cuda.cu +++ b/c/tests/test_backend_cuda.cu @@ -341,7 +341,7 @@ int main(int argc, char **argv) { /* Native s4 WMMA path: compare the quantized-activation result against the existing FP32-activation/s4-weight grouped implementation. */ - uint8_t w4[32*32/2]; float ws4[32], gx4[64], scalar4[64], async4[64], tensor4[64]; + uint8_t w4[32*32/2]; float ws4[32], gx4[64], scalar4[64], warp4[64], async4[64], tensor4[64]; for(int i=0;i<(int)sizeof(w4);i++){ int lo=((i%15)-7)&15,hi=(((i*3)%15)-7)&15; w4[i]=(uint8_t)(lo|(hi<<4)); @@ -354,6 +354,13 @@ int main(int argc, char **argv) { !coli_cuda_tensor_upload_g(&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_WARP_DECODE","1",1); + if(!coli_cuda_expert_group(gg4,ug4,dg4,group_rows,2,warp4,gx4))return 1; + unsetenv("COLI_CUDA_WARP_DECODE"); + if(std::memcmp(warp4,scalar4,sizeof(warp4))){ + std::fprintf(stderr,"warp-decode packed-s4 group differs from block path\n"); + return 1; + } if(!coli_cuda_expert_group_issue(gg4,ug4,dg4,group_rows,2,gx4))return 1; const float *async_result=coli_cuda_expert_group_take(d0); if(!async_result)return 1; @@ -371,7 +378,7 @@ int main(int argc, char **argv) { coli_cuda_tensor_free(g4);coli_cuda_tensor_free(u4);coli_cuda_tensor_free(d4); uint64_t group_calls=0,group_experts=0,group_total_rows=0; coli_cuda_group_stats(&group_calls,&group_experts,&group_total_rows,nullptr,nullptr,nullptr); - if(group_calls!=4||group_experts!=8||group_total_rows!=8) return 1; + if(group_calls!=5||group_experts!=10||group_total_rows!=10) return 1; coli_cuda_stats(-1, &count, &bytes); if (count != 7 || bytes != 166) { diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index 7854e195..3d0e4772 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -98,6 +98,7 @@ Per-drive byte counts are reported in a `MIRROR:` stats line. Combine with `DIRE | `CUDA_EXPERT_GB` | `0` | VRAM budget (GB) for caching experts on the GPU. | | `CUDA_RELEASE_HOST` | auto (`1` if >1 device) | Release host-side copies after upload. | | `COLI_CUDA_ATTN` | off | Run S≤4 attention on the GPU. | +| `COLI_CUDA_ATTN_PREFIX` | off | Reuse one uploaded decode activation across `q_a` and `kv_a` while preserving the stock CPU RMSNorm path. | | `COLI_CUDA_ATTN_SHARD` | off | `=1` splits KV-b heads across devices during attention load (multi-GPU). | | `COLI_CUDA_PROFILE` | off | Emit CUDA timing. | | `COLI_CUDA_PIPE` | `0` (off) | `1` engages the multi-step attention pipeline; `2` enables the pipe2 path. | @@ -113,6 +114,11 @@ Per-drive byte counts are reported in a `MIRROR:` stats line. Combine with `DIRE | `COLI_CUDA_TC_W4A16_MIN` | `16` | Per-expert row threshold above which W4A16 TC tiles dispatch (smaller batches fall back to the naive kernel). | | `COLI_CUDA_SHARED_W4A16` | off | `=1` uploads shared-expert weights and runs the shared-MLP W4A16 Tensor Core kernel. | | `COLI_CUDA_SHARED_W4A16_MIN_ROWS` | `32` | Min row count to engage the shared-MLP W4A16 kernel. | +| `CUDA_RAW_EXPERTS` | unset | Experimental ANS build only: keep this many hottest experts raw, then store subsequent VRAM experts losslessly compressed. Requires `COLI_ANS_SIDECAR`. | +| `COLI_ANS_SIDECAR` | unset | Experimental ANS build only: path to the sequential compressed-expert sidecar. | +| `COLI_ANS_PACK` | `0` | Experimental ANS build only: `=1` creates `COLI_ANS_SIDECAR` during pinning and exits before inference. | +| `COLI_ANS_DIRECT` | `0` | Experimental ANS build on Linux: `=1` reads the sidecar with aligned `O_DIRECT`, bypassing page-cache overhead. Falls back to buffered I/O if unavailable. | +| `COLI_ANS_PROFILE` | `0` | Experimental ANS build: print sidecar header, read, staging/allocation, and H2D enqueue timings on first use. | | `COLI_METAL_UNTRACKED` | off (Metal only) | `=1` sets `MTLResourceHazardTrackingModeUntracked` on Metal buffers (reduces hazard-tracking overhead). | > **Windows note.** On Windows, a bare `coli chat` / `coli run` / `coli serve` diff --git a/docs/cuda.md b/docs/cuda.md index 835ba971..4490a082 100644 --- a/docs/cuda.md +++ b/docs/cuda.md @@ -66,6 +66,49 @@ dispatch (`COLI_CUDA_TC_W4A16=1`), 96-token greedy decode measured Full experiment log: [experiments/glm52-6x5090-2026-07-12.md](experiments/glm52-6x5090-2026-07-12.md). These are host-specific capacity results, not portable defaults. +## Experimental lossless compressed expert tier + +An optional [DietGPU](https://github.com/facebookresearch/dietgpu) ANS tier keeps +warm int4 experts compressed in VRAM and decodes the selected tensors directly +to reusable device scratch before grouped expert matmul. It does not quantize, +prune, substitute, or move expert weights over PCIe during decode. + +This path is disabled in normal builds. DietGPU must first be built as shared +libraries, then supplied explicitly: + +```bash +make CUDA=1 COLI_ANS=1 DIETGPU_ROOT=/opt/dietgpu +``` + +Create the sidecar once using the same model, hot-expert profile, device set, +and placement budget intended for inference: + +```bash +COLI_CUDA=1 COLI_GPUS=0,1,2,3 CUDA_EXPERT_GB=auto \ +PIN=stats.txt PIN_GB=all CUDA_RAW_EXPERTS=2500 \ +COLI_ANS_SIDECAR=/models/glm52-experts.ans COLI_ANS_PACK=1 \ +SNAP=/models/glm52-int4 ./colibri 1 4 4 +``` + +Remove `COLI_ANS_PACK` to load the sidecar. Records are checked against the +expected format, dimensions, raw size, codec bound, and file length. A sidecar +is tied to its exact expert placement order; changing the model, profile, +budgets, raw-prefix size, or device set requires repacking it. +On Linux, `COLI_ANS_DIRECT=1` uses aligned direct reads while retaining the same +sidecar format. `COLI_ANS_PROFILE=1` prints a one-line load-time breakdown. + +On a 6× RTX 5090 host with GLM-5.2 int4, a 2,500-expert raw tier plus 8,128 +compressed experts increased VRAM capacity from 9,335 to 10,628 experts +(+13.9%). Greedy decode was byte-identical to the raw-tier baseline and measured +6.19→7.12 tok/s over 32 tokens (+15.0%), and 6.75→7.31 tok/s over 128 tokens +(+8.3%). Pinned asynchronous staging reduced placement from 197 to 157 seconds; +aligned direct reads reduced it further to 80–103 seconds across two runs, with +the 110 GB archive read accounting for 18.94–19.15 seconds. These results are a +general compression guarantee: incompressible records are rejected, and the +archived DietGPU dependency makes this an experimental build-time option. +Live `REPIN` is disabled because it would invalidate the sidecar's fixed expert +order. + ## The GPU-resident pipeline (`COLI_CUDA_PIPE`) `COLI_CUDA_PIPE=2` keeps the residual stream on-device across layers: rmsnorms,