Skip to content

feat(qwen36): CUDA VRAM expert tier — heat-based placement across GPUs via the shared CUDA backend #713

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JustVugg merged 1 commit into from
Aug 18, 2026
Merged

feat(qwen36): CUDA VRAM expert tier — heat-based placement across GPUs via the shared CUDA backend #713

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 25 additions & 10 deletions c/Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -897,19 +897,34 @@ olmoe$(EXE): olmoe.c st.h json.h compat.h sample.h tok.h tok_unicode.h tok_unico
$(CC) $(NOCUDA_CFLAGS) olmoe.c -o olmoe$(EXE) $(NOCUDA_LDFLAGS)

# Qwen3.6-35B-A3B engine (hybrid Gated Attention + Gated DeltaNet + streaming
# MoE). CPU-only in this target; the optional CUDA expert tier is a separate
# follow-up (see docs/qwen36-phase01.md).
# NOCUDA_*: this file contains no CUDA, so building it with -DCOLI_CUDA and
# linking -lcudart only made `make qwen36 CUDA=1` depend on a toolkit it
# never calls. Same shape as olmoe. (The CUDA expert tier is #713, which
# adds its own sources and takes the CUDA flags back.)
qwen36$(EXE): qwen36.c st.h json.h compat.h
$(CC) $(NOCUDA_CFLAGS) qwen36.c -o qwen36$(EXE) $(NOCUDA_LDFLAGS)
# MoE). With CUDA=1 the optional VRAM expert tier (qwen36_tier.c) is compiled
# in and reuses the shared CUDA backend; without it the engine is CPU-only
# (qwen36_tier.h provides inline stubs).
#
# The flags follow the same switch. Without CUDA=1 this target keeps the
# NOCUDA_* flags it got when the engine had no CUDA at all: compiling
# -DCOLI_CUDA and linking -lcudart for a build that calls neither only made
# `make qwen36` depend on a toolkit it never touches. With CUDA=1 the tier is
# real CUDA and takes the normal flags back.
ifeq ($(CUDA),1)
QWEN36_TIER_SRC = qwen36_tier.c
QWEN36_CFLAGS = $(CFLAGS)
QWEN36_LDFLAGS = $(LDFLAGS)
else
QWEN36_TIER_SRC =
QWEN36_CFLAGS = $(NOCUDA_CFLAGS)
QWEN36_LDFLAGS = $(NOCUDA_LDFLAGS)
endif
qwen36$(EXE): qwen36.c qwen36_tier.h st.h json.h compat.h $(QWEN36_TIER_SRC) $(CUDA_OBJ)
$(CC) $(QWEN36_CFLAGS) qwen36.c $(QWEN36_TIER_SRC) $(CUDA_OBJ) -o qwen36$(EXE) $(QWEN36_LDFLAGS)

# Context-size gates: KV layout, growth across requests, and the attention
# capacity. Includes qwen36.c directly, so no model file is needed.
tests/test_qwen36_ctx$(EXE): tests/test_qwen36_ctx.c qwen36.c st.h json.h compat.h
$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
# Same tier sources as the engine: the test includes qwen36.c, so with CUDA=1
# it needs qwen36_tier.c and the backend object too (without CUDA, the header's
# inline stubs cover it and both are empty).
tests/test_qwen36_ctx$(EXE): tests/test_qwen36_ctx.c qwen36.c qwen36_tier.h st.h json.h compat.h $(QWEN36_TIER_SRC) $(CUDA_OBJ)
$(CC) $(CFLAGS) $< $(QWEN36_TIER_SRC) $(CUDA_OBJ) -o $@ $(LDFLAGS)

inkling$(EXE): inkling.c st.h json.h compat.h $(INK_CUDA_OBJ) $(METAL_OBJ)
$(CC) $(CFLAGS) inkling.c $(INK_CUDA_OBJ) $(METAL_OBJ) -o inkling$(EXE) $(LDFLAGS)
Expand Down
171 changes: 161 additions & 10 deletions c/qwen36.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ static int qwen36_max_ctx(void) {
#endif
#include "st.h"
#include "json.h" /* tokenizer.json parsing (reuse minimal parser) */
#include "qwen36_tier.h" /* optional transparent Vulkan compute backend for MoE experts */

#ifdef _WIN32
#include <windows.h>
Expand Down Expand Up @@ -583,7 +584,7 @@ typedef struct {
} Layer;

/* ---------- LRU expert cache (int8 weights + per-row float scales) ---------- */
typedef struct { int eid; int pinned; int is_int4; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot;
typedef struct { int eid; int pinned; int is_int4; int8_t *g, *u, *d; uint8_t *g4, *u4, *d4; float *gs, *us, *ds; uint64_t used; } Slot;
typedef struct { Slot *slots; int n, cap; } LCache;

typedef struct {
Expand Down Expand Up @@ -651,6 +652,7 @@ static double g_tm_dec[6], g_tm_pre[6]; /* 0=deltanet 1=attention 2=moe_total
static long g_tm_dec_tokens = 0, g_tm_pre_tokens = 0;
static double tm_now(void){ struct timespec ts; clock_gettime(CLOCK_MONOTONIC,&ts); return ts.tv_sec*1e3 + ts.tv_nsec/1e6; }
static int tm_on(void){ if(g_timers<0){ const char *e=getenv("COLI_TIMERS"); g_timers = (e && *e=='1'); } return g_timers; }
double g_qt_iss=0, g_qt_cpu=0, g_qt_tak=0; /* QTIER-Phasen (Decode) */
double g_dn_sub[4]; /* DN: proj, conv+split, l2n+rec, norm+out */
double g_tm_step=0; /* step() total (decode) */
static double g_tm_win_moe=0; static int g_tm_win_n=0;
Expand Down Expand Up @@ -682,6 +684,9 @@ static void tm_report(void){
if(g_dn_sub[0]+g_dn_sub[1]+g_dn_sub[2]+g_dn_sub[3]>0)
fprintf(stderr,"[timers] dn-sub: proj %.1f | conv %.1f | l2n+rec %.1f | norm+out %.1f ms/token\n",
g_dn_sub[0]/g_tm_dec_tokens,g_dn_sub[1]/g_tm_dec_tokens,g_dn_sub[2]/g_tm_dec_tokens,g_dn_sub[3]/g_tm_dec_tokens);
if(g_qt_iss+g_qt_cpu+g_qt_tak>0)
fprintf(stderr,"[timers] qtier: issue %.2f | cpu-miss %.2f | take %.2f ms/token\n",
g_qt_iss/g_tm_dec_tokens, g_qt_cpu/g_tm_dec_tokens, g_qt_tak/g_tm_dec_tokens);
fprintf(stderr,"[timers] prefill: %ld tokens dn=%.0f attn=%.0f moe=%.0f(sh=%.0f rt=%.0f) head=%.0f ms\n",
g_tm_pre_tokens,g_tm_pre[0],g_tm_pre[1],g_tm_pre[2],g_tm_pre[3],g_tm_pre[4],g_tm_pre[5]);
}
Expand Down Expand Up @@ -1197,6 +1202,7 @@ static void slot_ensure_allocated(Model *m, Slot *s) {
s->ds = s_block + 2*scale_count_gu(c);
s->pinned = 0;
s->is_int4 = 0;
s->g4 = s->u4 = s->d4 = NULL; /* packed int4 (allocated on int4 load if GPU int4 active) */
}

static void load_expert_merged(Model *m, int layer, int eid, Slot *s) {
Expand Down Expand Up @@ -1234,14 +1240,29 @@ static void load_expert_merged(Model *m, int layer, int eid, Slot *s) {
s->g[i] = v;
}
s->is_int4 = 1;
/* The packed int4 bytes are NOT kept here. This engine only ever reads
* the unpacked int8 copy above, so retaining them doubled expert-cache
* RSS on the recommended gs64 container for nothing. They are the upload
* source for the CUDA expert tier and come back with it (#713), which is
* where they are actually read. */
/* Free any previous occupant first (LRU slot reuse). */
free(s->g4); free(s->u4); free(s->d4); s->g4 = s->u4 = s->d4 = NULL;
/* Keep the packed int4 bytes alongside the unpacked int8 copy only when
* the CUDA tier is actually running: they are its upload source, and
* they let slot_ensure_int8() rematerialize an evicted expert whose
* int8 copy the warmstart freed. Without the tier nothing ever reads
* them, and keeping them would add ~50% to expert-cache RSS on the
* recommended gs64 container -- so qt_ready() gates the allocation.
* Under CUDA=0 that is an inline `return 0` and this costs nothing. */
if (qt_ready()) {
int64_t gp = ng / 2, up = ng / 2, dp = nd / 2; /* gate/up/down packed sizes */
s->g4 = (uint8_t *)malloc((size_t)gp);
s->u4 = (uint8_t *)malloc((size_t)up);
s->d4 = (uint8_t *)malloc((size_t)dp);
if (!s->g4 || !s->u4 || !s->d4) { fprintf(stderr, "OOM int4-packed %s\n", nm); exit(1); }
memcpy(s->g4, raw, (size_t)gp);
memcpy(s->u4, raw + gp, (size_t)up);
memcpy(s->d4, raw + gp + up, (size_t)dp);
}
free(raw);
} else {
s->is_int4 = 0;
free(s->g4); free(s->u4); free(s->d4); s->g4 = s->u4 = s->d4 = NULL;
st_read_raw(&m->S, nm, s->g, 1);
}
st_read_f32(&m->S, qsnm, s->gs, 0);
Expand All @@ -1262,6 +1283,31 @@ static int container_is_int4(Model *m) {
return (tw->nbytes == want_w / 2) ? 1 : 0;
}

/* Rematerialize a slot's int8 block from its packed int4 copy on demand
* (~0.5 ms, no container access). Needed after the warmstart freed the int8
* copies of VRAM-resident experts and one of them got LFRU-evicted. */
static void slot_ensure_int8(Model *m, Slot *s) {
if (s->g || !s->g4) return;
Cfg *c = &m->c;
int64_t ng = (int64_t)c->inter * c->hidden, nd = (int64_t)c->hidden * c->inter;
int8_t *w = malloc((size_t)(ng + ng + nd));
if (!w) { fprintf(stderr, "OOM slot_ensure_int8\n"); exit(1); }
const uint8_t *src4[3] = { s->g4, s->u4, s->d4 };
int64_t lens[3] = { ng, ng, nd };
int8_t *dst = w;
for (int t = 0; t < 3; t++) {
const uint8_t *p = src4[t];
for (int64_t i = 0; i < lens[t]; i += 2) {
uint8_t b = p[i >> 1];
int8_t lo = (int8_t)(b & 0xF); if (lo & 8) lo -= 16;
int8_t hi = (int8_t)((b >> 4) & 0xF); if (hi & 8) hi -= 16;
dst[i] = lo; dst[i + 1] = hi;
}
dst += lens[t];
}
s->g = w; s->u = w + ng; s->d = w + ng + ng;
}

static void expert_get(Model *m, int layer, int eid, Slot **out) {
LCache *lc = &m->cache[layer];
pthread_mutex_lock(&g_pilot_mx);
Expand Down Expand Up @@ -1572,9 +1618,59 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
for (int kk = 0; kk < K; kk++) if (idx[kk] >= 0) freq_l[idx[kk]]++;
}
const float *xs = x + (int64_t)s*D;
{
int shared_done = 0;
if (qt_ready()) {
/* CUDA expert tier: run the resident experts as async groups on
* all devices, compute the misses on the CPU (overlapped), then
* collect the GPU results. */
for (int kk = 0; kk < K; kk++) {
Slot *e; expert_get(m, layer, idx[kk], &e);
if (e->g4) qt_note(layer, idx[kk], e->g4, e->u4, e->d4, e->gs, e->us, e->ds);
}
double _q0 = tm_on()? tm_now():0;
uint32_t qmask = qt_issue(layer, idx, K, xs);
double _q1 = tm_on()? tm_now():0;
for (int kk = 0; kk < K; kk++) {
if (qmask & (1u<<kk)) continue;
Slot *e; expert_get(m, layer, idx[kk], &e);
slot_ensure_int8(m, e);
matmul_qe(g, xs, e->g, e->gs, D, I);
matmul_qe(u, xs, e->u, e->us, D, I);
for (int i = 0; i < I; i++) { float gv = g[i]; g[i] = (gv / (1.f + expf(-gv))) * u[i]; }
matmul_qe(hh, g, e->d, e->ds, I, D);
float w = val[kk]; float *os = out + (int64_t)s*D;
for (int d = 0; d < D; d++) os[d] += w * hh[d];
}
/* Compute the shared expert NOW so it overlaps with the GPU
* groups; the common block below is skipped. */
{
double _ts2 = tm_on() ? tm_now() : 0.0;
int Ish = c->shared_inter;
matmul_d(sh, xs, l->sh_g, 1, D, Ish);
matmul_d(shu, xs, l->sh_u, 1, D, Ish);
for (int i = 0; i < Ish; i++) { float sv = sh[i]; sh[i] = (sv / (1.f + expf(-sv))) * shu[i]; }
matmul_d(shd, sh, l->sh_d, 1, Ish, D);
float sgate = 1.f;
if (l->sh_gate) {
float sg = 0.f; const float *wg = l->sh_gate;
for (int i = 0; i < D; i++) sg += xs[i] * wg[i];
sgate = 1.f / (1.f + expf(-sg));
}
float *os = out + (int64_t)s*D;
for (int d = 0; d < D; d++) os[d] += sgate * shd[d];
if (tm_on()) tm_add(S, 3, tm_now()-_ts2);
}
shared_done = 1;
double _q2 = tm_on()? tm_now():0;
qt_take(qmask, val, K, out + (int64_t)s*D);
if (tm_on() && S==1) {
extern double g_qt_iss, g_qt_cpu, g_qt_tak;
g_qt_iss += _q1-_q0; g_qt_cpu += _q2-_q1; g_qt_tak += tm_now()-_q2;
}
} else {
for (int kk = 0; kk < K; kk++) {
Slot *e; expert_get(m, layer, idx[kk], &e);
slot_ensure_int8(m, e);
matmul_qe(g, xs, e->g, e->gs, D, I);
matmul_qe(u, xs, e->u, e->us, D, I);
for (int i = 0; i < I; i++) { float gv = g[i]; g[i] = (gv / (1.f + expf(-gv))) * u[i]; }
Expand All @@ -1585,6 +1681,7 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
}
}
/* shared expert (SwiGLU), sigmoid-gated by shared_expert_gate */
if (shared_done) { if (0) goto shared_skip_dummy; shared_skip_dummy: continue; }
double _ts = tm_on() ? tm_now() : 0.0;
int Ish = c->shared_inter;
matmul_d(sh, xs, l->sh_g, 1, D, Ish);
Expand Down Expand Up @@ -1919,9 +2016,14 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) {
if (idx[b] >= 0 && (idx[a] < 0 || idx[a] > idx[b])) { int t = idx[a]; idx[a] = idx[b]; idx[b] = t; }
for (int kk = 0; kk < cand; kk++) {
int eid = idx[kk]; if (eid < 0) continue;
int found = 0; pthread_mutex_lock(&g_pilot_mx); LCache *lc = &m->cache[lnext];
for (int z = 0; z < lc->n; z++) if (lc->slots[z].eid == eid) { found = 1; break; }
int found = 0, fz = -1; pthread_mutex_lock(&g_pilot_mx); LCache *lc = &m->cache[lnext];
for (int z = 0; z < lc->n; z++) if (lc->slots[z].eid == eid) { found = 1; fz = z; break; }
pthread_mutex_unlock(&g_pilot_mx);
/* Lookahead: RAM-resident layer-L+1 candidates go to VRAM asynchronously */
if (found && fz >= 0 && qt_ready()) {
Slot *ps = &lc->slots[fz];
if (ps->g4) qt_note(lnext, eid, ps->g4, ps->u4, ps->d4, ps->gs, ps->us, ps->ds);
}
if (!found) {
int gidx = lnext*E + eid;
pthread_mutex_lock(&g_pilot_mx); int already_queued = m->is_queued[gidx];
Expand Down Expand Up @@ -2307,7 +2409,55 @@ int main(int argc, char **argv) {
g_qdw_n, now_s()-tq, freed/1073741824.0);
}

/* coli serve mode: speak the gateway wire protocol instead of argv generation */
/* Optional CUDA VRAM expert tier (COLI_CUDA=1): hot experts live in
* DEVICE_LOCAL memory across the configured GPUs, misses fall back to the
* CPU int8 path. See qwen36_tier.h. */
if (qt_init(m.c.n_layers, m.c.n_experts, m.c.hidden, m.c.inter, cap, m.c.topk, m.c.expert_gs)) {
fprintf(stderr, "[gpu] MoE experts -> CUDA VRAM tier\n");
atexit(qt_shutdown);
/* Warmstart: fill the VRAM budget BEFORE the first token (heat order
* when HEAT_FILE exists, natural order otherwise), loading all RAM
* slots along the way. */
const char *nws = getenv("QT_NO_WARMSTART");
if (!(nws && *nws=='1')) {
/* Plan the set (heat order), then load+stage IN PARALLEL. The
* load path is thread-safe: expert_get locks the layer cache
* (g_pilot_mx), st_read_raw uses pread; entries are unique. */
double t0 = now_s();
int cap_total = m.c.n_layers * m.c.n_experts;
int *wpl = malloc((size_t)cap_total*sizeof(int));
int *wpe = malloc((size_t)cap_total*sizeof(int));
int wn = qt_plan_fill(wpl, wpe, cap_total);
/* Load ALL experts into RAM, not just the planned (VRAM) set:
* otherwise the first touch of a CPU-fallback expert triggers a
* ~12 ms container read in the middle of decode (measured: 139
* ms/token on a single-GPU run). Planned ones also go to VRAM. */
uint8_t *planned = calloc((size_t)cap_total, 1);
for (int i = 0; i < wn; i++) planned[wpl[i]*m.c.n_experts + wpe[i]] = 1;
int keep8 = getenv("COLI_KEEP_INT8") != NULL;
#pragma omp parallel for schedule(dynamic, 16)
for (int gi = 0; gi < cap_total; gi++) {
int l = gi / m.c.n_experts, eidw = gi % m.c.n_experts;
Slot *e; expert_get(&m, l, eidw, &e);
if (planned[gi] && e->g4) {
qt_note_planned(l, eidw, e->g4, e->u4, e->d4, e->gs, e->us, e->ds);
/* The staging copy is done; free the int8 copy RIGHT AWAY
* so it never shows up in peak RSS. On LFRU eviction
* slot_ensure_int8() rematerializes from g4 (no container
* access). */
if (!keep8 && e->g) { free(e->g); e->g = e->u = e->d = NULL; }
}
}
qt_fill_wait();
free(wpl); free(wpe); free(planned);
fprintf(stderr, "[qtier] warmstart (parallel): all %d experts in RAM (int8 only for non-residents), %d in VRAM -- %.1f s\n",
cap_total, wn, now_s()-t0);
}
}

/* coli serve mode: speak the gateway wire protocol instead of argv
* generation. AFTER the tier init: serve sessions ride the VRAM experts
* exactly like argv runs, and serve_loop never returns. */
if (getenv("SERVE") && getenv("SERVE")[0] == '1') {
if (!g_tok) { fprintf(stderr, "[serve] tokenizer.json required (put in SNAP or set TOK)\n"); return 1; }
serve_loop(&m);
Expand Down Expand Up @@ -2383,6 +2533,7 @@ int main(int argc, char **argv) {
double tot = m.hits + m.miss;
if (g_ttft >= 0) fprintf(stderr, "TTFT: %.2f s (time to first token)\n", g_ttft);
tm_report();
qt_stats();
fprintf(stderr, "\nPEAK RSS: %.2f GB\n", rss_gb());
fprintf(stderr, "Expert cache hit rate: %.1f%% (hit=%llu miss=%llu)\n", tot?100.0*m.hits/tot:0.0,
(unsigned long long)m.hits, (unsigned long long)m.miss);
Expand Down
Loading
Loading