Something went wrong. Try again.
A single-model (Bonsai-27B) inference engine for Intel Arc GPUs
Something went wrong. Try again.
C++
at extra-debug
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477#include "qxmx_ref.h"
#include <algorithm>#include <cmath>#include <cstdio>#include <cstring>
namespace qx {
/* ---------- f16 ---------- */static float f16_to_f32(uint16_t bits) { uint32_t sign = (uint32_t)(bits & 0x8000) << 16; uint32_t exp = (bits >> 10) & 0x1f, mant = bits & 0x3ff; union { uint32_t u; float f; } cv; if (exp == 0) { cv.f = (float)mant * 0.00006103515625f; cv.u |= sign; return cv.f; } if (exp == 31) { cv.u = sign | 0x7f800000u | (mant ? 0x00400000u : 0); return cv.f; } cv.u = sign | ((exp + 112) << 23) | (mant << 13); return cv.f;}
/* ---------- Q2 dequant ---------- *//* GGUF Q2 2D tensor dims[0]=K (in, fast), dims[1]=M (out). * Supports both g64 (18 bytes/block) and g128 (34 bytes/block). * q in {0,1,2} -> (q-1) in {-1,0,+1}. */void ref_dequant_row(const qx_weight* w, int m_idx, float* out) { int K = (int)w->dims[0]; int qk2 = (w->nbytes * 64 / w->n_elem == 18) ? 64 : 128; int ngrp = K / qk2; if (qk2 == 64) { const qx_block_q2* row = (const qx_block_q2*)w->data + (size_t)m_idx * ngrp; for (int g = 0; g < ngrp; g++) { float d = f16_to_f32(row[g].d); for (int k = 0; k < 64; k++) { int q = (row[g].qs[k / 4] >> ((k % 4) * 2)) & 3; out[g * 64 + k] = d * (float)(q - 1); } } } else { const qx_block_q2_g128* row = (const qx_block_q2_g128*)w->data + (size_t)m_idx * ngrp; for (int g = 0; g < ngrp; g++) { float d = f16_to_f32(row[g].d); for (int k = 0; k < 128; k++) { int q = (row[g].qs[k / 4] >> ((k % 4) * 2)) & 3; out[g * 128 + k] = d * (float)(q - 1); } } }}
/* W(M,N) @ x(N) -> y(M), W stored as GGUF Q2 [N][M] (dims[0]=N, dims[1]=M). * Dequant+dot per group, no intermediate. Handles both g64 and g128. */static void q2_matvec(const qx_weight* w, const float* x, float* y) { int N = (int)w->dims[0], M = (int)w->dims[1]; int qk2 = (w->nbytes * 64 / w->n_elem == 18) ? 64 : 128; int ngrp = N / qk2; if (qk2 == 64) { const qx_block_q2* W = (const qx_block_q2*)w->data; for (int m = 0; m < M; m++) { float acc = 0.0f; const qx_block_q2* row = W + (size_t)m * ngrp; for (int g = 0; g < ngrp; g++) { float d = f16_to_f32(row[g].d); float p = 0.0f; for (int k = 0; k < 64; k++) { int q = (row[g].qs[k / 4] >> ((k % 4) * 2)) & 3; p += (float)(q - 1) * x[g * 64 + k]; } acc += d * p; } y[m] = acc; } } else { const qx_block_q2_g128* W = (const qx_block_q2_g128*)w->data; for (int m = 0; m < M; m++) { float acc = 0.0f; const qx_block_q2_g128* row = W + (size_t)m * ngrp; for (int g = 0; g < ngrp; g++) { float d = f16_to_f32(row[g].d); float p = 0.0f; for (int k = 0; k < 128; k++) { int q = (row[g].qs[k / 4] >> ((k % 4) * 2)) & 3; p += (float)(q - 1) * x[g * 128 + k]; } acc += d * p; } y[m] = acc; } }}
/* F32 matvec: y(M) = W(M,N) @ x(N), W row-major [M][N]. *//* F32 weight stored as GGUF tensor (1D or 2D). */
/* ---------- RMSNorm ---------- * llama.cpp uses a single plain RMSNorm everywhere: out = w * norm(x). The * GGUF converter bakes the +1 Gemma offset into the stored weight for every * norm EXCEPT the SSM post-delta norm (which is plain in HF too), so at runtime * all norms are identical. Verified on the real GGUF: input/post/final/q/k * norms read ~1.0, ssm_norm ~0.87 (no offset baked). */static void rmsnorm(const float* x, const float* w, float* out, int dim, float eps) { double ss = 0.0; for (int i = 0; i < dim; i++) ss += (double)x[i] * x[i]; float inv = 1.0f / sqrtf((float)(ss / dim) + eps); for (int i = 0; i < dim; i++) out[i] = x[i] * inv * w[i];}
/* ---------- activations ---------- */static inline float silu(float x) { return x / (1.0f + expf(-x)); }static inline float sigmoid(float x) { return 1.0f / (1.0f + expf(-x)); }
/* ---------- causal depthwise conv1d (Qwen3-Next / qwen35 convention) ---------- * Per Aila + Candle bug #2: PyTorch Conv1d left-pads. For kernel k=4, output at * step t = sum_{j=0..k-1} w[j] * x_hist[mod], where x_hist = [x[t-3], x[t-2], * x[t-1], x[t]]. GGUF conv1d.weight is stored channel-major [channels][k] * (dims[0]=k is the fast axis): w_conv[c*k+j], w[0]*oldest, w[k-1]*newest. * Conv history is [k-1][channels], most-recent-last. */static void conv1d_step(const float* w_conv /*[channels][k], channel-major*/, const float* x_hist /*[k-1][channels]*/, float* x_hist_new /*[k-1][channels]*/, const float* x_new /*[channels]*/, float* out /*[channels]*/, int k, int channels) { for (int c = 0; c < channels; c++) { float acc = 0.0f; /* history is x[t-(k-1)] .. x[t-1]; w[0] is oldest (channel-major weight) */ for (int j = 0; j < k - 1; j++) acc += w_conv[c * k + j] * x_hist[j * channels + c]; acc += w_conv[c * k + (k - 1)] * x_new[c]; out[c] = acc; } /* shift history: drop oldest, append x_new */ if (k > 1) { memmove(x_hist_new, x_hist + channels, sizeof(float) * channels * (k - 2)); memcpy(x_hist_new + channels * (k - 2), x_new, sizeof(float) * channels); }}
/* ---------- mRoPE with partial rotation (Candle bug #4) ---------- * Apply RoPE only to the first rot_dim=64 of each head_dim=256 head, per the * mRoPE section widths [t,h,w,extra]=[11,11,10,0]. For text (1D positions) * the three section position-ids are all = pos. Frequencies: inv_freq[i] = * 1 / (freq_base ^ (2i/rot_dim)), i in [0, rot_dim/2). */static void apply_rope(float* q /*[n_qk_heads][head_dim]*/, float* k /*[n_kv_heads][head_dim]*/, int n_qk, int n_kv, int head_dim, int rot_dim, int64_t pos, double freq_base) { int half = rot_dim / 2; /* precompute (cos, sin) for the rot_dim/2 freqs at this position */ std::vector<float> ctab(half), stab(half); for (int i = 0; i < half; i++) { double freq = 1.0 / pow(freq_base, (double)(2 * i) / rot_dim); double angle = pos * freq; ctab[i] = (float)cos(angle); stab[i] = (float)sin(angle); } auto rope_one = [&](float* v) { /* rotate pairs (v[i], v[i+half]) for i in [0,half) */ for (int i = 0; i < half; i++) { float a = v[i], b = v[i + half]; v[i] = a * ctab[i] - b * stab[i]; v[i+half] = a * stab[i] + b * ctab[i]; } }; for (int h = 0; h < n_qk; h++) rope_one(q + h * head_dim); for (int h = 0; h < n_kv; h++) rope_one(k + h * head_dim);}
/* ---------- Gated DeltaNet recurrence (decode, single token) ---------- * Inputs: q,k [NKH][HD=128], v [NVH][HD=128], alpha/beta [NVH], already conv1d'd * and (for q,k) about to be L2-normalized inside. State: [NVH][HD][HD]. * * Math (Aila fused kernel, matches CKE REF): * scale_qk = 1/sqrt(HD) * q <- L2(q), k <- L2(k) (per V-head's key group) * decay = exp( ssm_a * softplus(alpha + dt_bias) ) // ssm_a = -exp(A_log), baked * beta_s = sigmoid(beta) * S <- decay * S (per V-head) * kv_mem = S^T k * delta = (v - kv_mem) * beta_s * S <- S + k outer delta * out = scale_qk * (S^T q) * The dt_bias here is ssm_dt.bias; alpha,beta are per-V-head projections of x. * Note: in this model ssm_alpha/ssm_beta are PROJECTIONS (out=48=NVH), not the * raw gate scalars -- they are applied per V-head after the in-proj split. * Exported (non-static) for the M1.1 chunked-scan equivalence harness. */void deltanet_step(const float* q_in /*[NKH*HD]*/, const float* k_in, const float* v_in /*[NVH*HD]*/, const float* alpha /*[NVH]*/, const float* beta /*[NVH]*/, const float* a_log /*[NVH] = ssm_a*/, const float* dt_bias /*[NVH]*/, float* state_io /*[NVH*HD*HD]*/, float* out /*[NVH*HD]*/) { constexpr int HD = QX_SSM_STATE; /* 128 */ constexpr int NKH = QX_SSM_GROUPS; /* 16 */ constexpr int NVH = QX_SSM_DT_RANK; /* 48 */ const float scale_qk = 1.0f / sqrtf((float)HD);
/* working q,k copied per V-head's key group (3 V-heads share a key head) */ float q[NKH * HD], k[NKH * HD]; memcpy(q, q_in, sizeof(q)); memcpy(k, k_in, sizeof(k));
/* L2 normalize q,k per KEY head (shared across the 3 V-heads in its group) */ for (int hk = 0; hk < NKH; hk++) { float* qp = q + hk * HD; float* kp = k + hk * HD; float qs = 0, ks = 0; for (int i = 0; i < HD; i++) { qs += qp[i]*qp[i]; ks += kp[i]*kp[i]; } float qinv = 1.0f / sqrtf(qs + QX_RMS_EPS); float kinv = 1.0f / sqrtf(ks + QX_RMS_EPS); for (int i = 0; i < HD; i++) { qp[i] *= qinv; kp[i] *= kinv; } }
for (int hv = 0; hv < NVH; hv++) { int hk = hv % NKH; /* tiled V-head order: hv -> key head hv%NKH */ const float* qp = q + hk * HD; const float* kp = k + hk * HD; const float* vp = v_in + hv * HD; float* S = state_io + (size_t)hv * HD * HD; float* outp = out + hv * HD;
/* decay: HF g = -exp(A_log) * softplus(a + dt_bias). The GGUF converter * bakes -exp(A_log) into ssm_a (conversion/qwen.py), so g = ssm_a * sp; * ssm_a is already negative. decay = exp(g) in (0,1). */ float apb = alpha[hv] + dt_bias[hv]; float sp = apb > 20.0f ? apb : (apb < -20.0f ? expf(apb) : log1pf(expf(apb))); float g = a_log[hv] * sp; /* ssm_a = -exp(A_log), baked in */ float decay = expf(g); float beta_s = sigmoid(beta[hv]);
/* S <- decay*S; kv_mem = S^T k; delta = (v-kv_mem)*beta_s */ float kv_mem[HD], delta[HD]; for (int col = 0; col < HD; col++) { float acc = 0.0f; for (int row = 0; row < HD; row++) { float s = S[row * HD + col] * decay; S[row * HD + col] = s; acc += s * kp[row]; } kv_mem[col] = acc; delta[col] = (vp[col] - acc) * beta_s; } /* S <- S + k outer delta; out = scale_qk * (S^T q) */ for (int col = 0; col < HD; col++) { float acc = 0.0f; for (int row = 0; row < HD; row++) { float s = S[row * HD + col] + kp[row] * delta[col]; S[row * HD + col] = s; acc += s * qp[row]; } outp[col] = scale_qk * acc; } }}
/* ---------- session ---------- */void ref_session::init(const qx_model* m, int max_ctx) { int n_ssm = 0; for (int L = 0; L < QX_N_LAYER; L++) if (!m->blocks[L].is_attention) n_ssm++; int n_attn = QX_N_LAYER - n_ssm; ssm.assign(n_ssm, {}); for (auto& s : ssm) s.resize(QX_SSM_DT_RANK, QX_SSM_STATE, QX_SSM_CONVD, QX_SSM_CONV_K); /* KV cache: [n_attn][max_ctx][KV_DIM], one slab per attention layer. */ kv_k.assign((size_t)n_attn * max_ctx * QX_KV_DIM, 0.0f); kv_v.assign((size_t)n_attn * max_ctx * QX_KV_DIM, 0.0f); n_tokens = 0; this->max_ctx = max_ctx;}
/* ---------- block forwards ---------- */static void ffn_forward(const qx_ffn_block* f, const float* x, float* out) { std::vector<float> h(QX_D_MODEL), g(QX_FF_INTER), u(QX_FF_INTER), t(QX_FF_INTER); rmsnorm(x, (const float*)f->post_norm.data, h.data(), QX_D_MODEL, QX_RMS_EPS); q2_matvec(&f->gate, h.data(), g.data()); q2_matvec(&f->up, h.data(), u.data()); for (int i = 0; i < QX_FF_INTER; i++) t[i] = silu(g[i]) * u[i]; q2_matvec(&f->down, t.data(), out);}
/* Attention block forward (decode, batch=1). Updates KV cache for this layer * and this position. kv_k/kv_v are [max_ctx][KV_DIM]; head_dim=256, GQA 24Q/4KV. */static void attn_block_forward(const qx_attn_block* a, const float* x, float* out, float* kv_k_layer, float* kv_v_layer, int64_t pos) { std::vector<float> h(QX_D_MODEL); rmsnorm(x, (const float*)a->attn_norm.data, h.data(), QX_D_MODEL, QX_RMS_EPS);
/* GGUF attn_q [5120x12288] emits q AND the attention gate, laid out PER-HEAD * interleaved [h0_q(256)|h0_gate(256)|h1_q(256)|h1_gate(256)|...] (llama.cpp * ggml_view_3d with stride head_dim*2). De-interleave into q[6144]/gate[6144]. */ std::vector<float> qall(QX_Q_DIM * 2), q(QX_Q_DIM), gate(QX_Q_DIM); q2_matvec(&a->q, h.data(), qall.data()); for (int hh = 0; hh < QX_N_HEAD; hh++) { memcpy(q.data() + hh * QX_HEAD_DIM, qall.data() + hh * QX_HEAD_DIM * 2, sizeof(float) * QX_HEAD_DIM); memcpy(gate.data() + hh * QX_HEAD_DIM, qall.data() + hh * QX_HEAD_DIM * 2 + QX_HEAD_DIM, sizeof(float) * QX_HEAD_DIM); } std::vector<float> k(QX_KV_DIM), v(QX_KV_DIM); q2_matvec(&a->k, h.data(), k.data()); q2_matvec(&a->v, h.data(), v.data());
/* per-head q/k plain RMSNorm over head_dim=256 (q_norm/k_norm [256], +1 baked). */ const float* qn = (const float*)a->q_norm.data; const float* kn = (const float*)a->k_norm.data; for (int hh = 0; hh < QX_N_HEAD; hh++) rmsnorm(q.data() + hh * QX_HEAD_DIM, qn, q.data() + hh * QX_HEAD_DIM, QX_HEAD_DIM, QX_RMS_EPS); for (int hh = 0; hh < QX_N_HEAD_KV; hh++) rmsnorm(k.data() + hh * QX_HEAD_DIM, kn, k.data() + hh * QX_HEAD_DIM, QX_HEAD_DIM, QX_RMS_EPS);
/* Partial mRoPE (IMROPE, half-split pairs): rot_dim=64 only; text 1D pos. */ apply_rope(q.data(), k.data(), QX_N_HEAD, QX_N_HEAD_KV, QX_HEAD_DIM, QX_ROT_DIM, pos, QX_ROPE_FREQ_BASE);
/* append K,V to cache at this position */ memcpy(kv_k_layer + pos * QX_KV_DIM, k.data(), sizeof(float) * QX_KV_DIM); memcpy(kv_v_layer + pos * QX_KV_DIM, v.data(), sizeof(float) * QX_KV_DIM);
/* GQA attention: 24 Q heads, 4 KV heads (ratio 6). head_dim 256. */ std::vector<float> attn_out(QX_Q_DIM, 0.0f); float scale = 1.0f / sqrtf((float)QX_HEAD_DIM); for (int hh = 0; hh < QX_N_HEAD; hh++) { int hk = hh / (QX_N_HEAD / QX_N_HEAD_KV); const float* qph = q.data() + hh * QX_HEAD_DIM; std::vector<float> scores(pos + 1); float maxs = -1e30f; for (int64_t t = 0; t <= pos; t++) { const float* kt = kv_k_layer + t * QX_KV_DIM + hk * QX_HEAD_DIM; float s = 0.0f; for (int d = 0; d < QX_HEAD_DIM; d++) s += qph[d] * kt[d]; s *= scale; scores[t] = s; if (s > maxs) maxs = s; } float sum = 0.0f; for (int64_t t = 0; t <= pos; t++) { scores[t] = expf(scores[t] - maxs); sum += scores[t]; } float inv = 1.0f / sum; std::vector<float> ao(QX_HEAD_DIM, 0.0f); for (int64_t t = 0; t <= pos; t++) { float w = scores[t] * inv; const float* vt = kv_v_layer + t * QX_KV_DIM + hk * QX_HEAD_DIM; for (int d = 0; d < QX_HEAD_DIM; d++) ao[d] += w * vt[d]; } memcpy(attn_out.data() + hh * QX_HEAD_DIM, ao.data(), sizeof(float) * QX_HEAD_DIM); }
/* output gate: sigmoid(gate) * attn_out, then o_proj */ for (int i = 0; i < QX_Q_DIM; i++) attn_out[i] *= sigmoid(gate[i]); q2_matvec(&a->o, attn_out.data(), out);}
/* SSM block forward (decode, batch=1). */static void ssm_block_forward(const qx_ssm_block* s, ref_ssm_state* st, const float* x, float* out) { std::vector<float> h(QX_D_MODEL), qkv(QX_SSM_CONVD), conv_out(QX_SSM_CONVD); rmsnorm(x, (const float*)s->attn_norm.data, h.data(), QX_D_MODEL, QX_RMS_EPS);
/* fused qkv -> [q(2048) | k(2048) | v(6144)] = convd 10240 */ q2_matvec(&s->qkv, h.data(), qkv.data()); if (getenv("QXMX_DUMP_LAYERS")) { static bool ref_dbg = false; if (!ref_dbg) { ref_dbg = true; fprintf(stderr, " REF qkv[0..2]: %.4f %.4f %.4f qkv[2048..2050]: %.4f %.4f %.4f qkv[4096..4098]: %.4f %.4f %.4f\n", qkv[0],qkv[1],qkv[2], qkv[QX_SSM_KEYD],qkv[QX_SSM_KEYD+1],qkv[QX_SSM_KEYD+2], qkv[2*QX_SSM_KEYD],qkv[2*QX_SSM_KEYD+1],qkv[2*QX_SSM_KEYD+2]); } }
/* causal conv1d on the full 10240-wide qkv (depthwise, k=4) */ const float* wc = (const float*)s->conv1d.data; /* [4][10240] */ float hist_new[QX_SSM_CONV_K - 1][QX_SSM_CONVD]; conv1d_step(wc, st->conv.data(), (float*)hist_new, qkv.data(), conv_out.data(), QX_SSM_CONV_K, QX_SSM_CONVD); memcpy(st->conv.data(), hist_new, sizeof(hist_new));
/* SiLU on the conv output before splitting (Aila: silu_dev(v) on each convolved channel) */ for (int i = 0; i < QX_SSM_CONVD; i++) conv_out[i] = silu(conv_out[i]); if (getenv("QXMX_DUMP_LAYERS")) { static bool conv_dbg = false; if (!conv_dbg) { conv_dbg = true; fprintf(stderr, " REF post_conv[0..2]: %.4f %.4f %.4f [2048..2050]: %.4f %.4f %.4f [4096..4098]: %.4f %.4f %.4f\n", conv_out[0],conv_out[1],conv_out[2], conv_out[QX_SSM_KEYD],conv_out[QX_SSM_KEYD+1],conv_out[QX_SSM_KEYD+2], conv_out[2*QX_SSM_KEYD],conv_out[2*QX_SSM_KEYD+1],conv_out[2*QX_SSM_KEYD+2]); } }
const float* q = conv_out.data(); const float* k = conv_out.data() + QX_SSM_KEYD; const float* v = conv_out.data() + QX_SSM_KEYD * 2;
/* alpha, beta projections -> per-V-head scalars (NVH=48) */ std::vector<float> alpha(QX_SSM_DT_RANK), beta(QX_SSM_DT_RANK); q2_matvec(&s->alpha, h.data(), alpha.data()); /* uses pre-conv h (Aila) */ q2_matvec(&s->beta, h.data(), beta.data()); if (getenv("QXMX_DUMP_LAYERS")) { static bool ab_dbg = false; if (!ab_dbg) { ab_dbg = true; fprintf(stderr, " REF alpha[0..3]: %.4f %.4f %.4f %.4f beta[0..3]: %.4f %.4f %.4f %.4f\n", alpha[0],alpha[1],alpha[2],alpha[3], beta[0],beta[1],beta[2],beta[3]); } }
/* gated delta-rule recurrence; out is [NVH*HD] = [6144] */ std::vector<float> ssm_out(QX_SSM_VALD); deltanet_step(q, k, v, alpha.data(), beta.data(), (const float*)s->a.data, (const float*)s->dt_bias.data, st->state.data(), ssm_out.data()); if (getenv("QXMX_DUMP_LAYERS")) { static bool dn_dbg = false; if (!dn_dbg) { dn_dbg = true; fprintf(stderr, " REF ssm_out[0..2]: %.4f %.4f %.4f [128..130]: %.4f %.4f %.4f [6141..6143]: %.4f %.4f %.4f\n", ssm_out[0],ssm_out[1],ssm_out[2], ssm_out[128],ssm_out[129],ssm_out[130], ssm_out[6141],ssm_out[6142],ssm_out[6143]); } }
/* group RMSNorm over the state (ssm_norm is [128], applied per V-head's 128) * then output gate (sigmoid(attn_gate) * silu?), then ssm_out proj. Per Aila: * the gate is silu(z) where z comes from the attn_gate projection of h. */ const float* nw = (const float*)s->ssm_norm.data; /* [128] */ for (int hv = 0; hv < QX_SSM_DT_RANK; hv++) { rmsnorm(ssm_out.data() + hv * QX_SSM_STATE, nw, ssm_out.data() + hv * QX_SSM_STATE, QX_SSM_STATE, QX_RMS_EPS); }
/* output gate z = attn_gate(h) -> [6144], gate = silu(z); out *= gate */ std::vector<float> zg(QX_SSM_VALD); q2_matvec(&s->gate, h.data(), zg.data()); if (getenv("QXMX_DUMP_LAYERS")) { static bool g_dbg = false; if (!g_dbg) { g_dbg = true; fprintf(stderr, " REF zg[0..2]: %.4f %.4f %.4f ssm_out_gated[0..2]: %.4f %.4f %.4f\n", zg[0],zg[1],zg[2], ssm_out[0]*silu(zg[0]), ssm_out[1]*silu(zg[1]), ssm_out[2]*silu(zg[2])); } } for (int i = 0; i < QX_SSM_VALD; i++) ssm_out[i] *= silu(zg[i]);
q2_matvec(&s->out, ssm_out.data(), out);}
/* ---------- top-level decode ---------- */void ref_decode(const qx_model* m, ref_session* s, const float* x_in, float* logits_out) { std::vector<float> x(x_in, x_in + QX_D_MODEL), y(QX_D_MODEL); if (getenv("QXMX_DUMP_LAYERS")) { double norm = 0; for (int i = 0; i < QX_D_MODEL; i++) norm += (double)x[i]*x[i]; norm = sqrt(norm / QX_D_MODEL); fprintf(stderr, " EMB: x_rms=%.4f [%.4f %.4f %.4f]\n", norm, x[0], x[1], x[2]); } int64_t pos = s->n_tokens; int ssm_idx = 0, attn_idx = 0; size_t kv_stride = (size_t)s->max_ctx * QX_KV_DIM; std::vector<float> h1(QX_D_MODEL), f(QX_D_MODEL); for (int L = 0; L < QX_N_LAYER; L++) { const qx_block* b = &m->blocks[L]; if (b->is_attention) { attn_block_forward(&b->attn, x.data(), y.data(), s->kv_k.data() + (size_t)attn_idx * kv_stride, s->kv_v.data() + (size_t)attn_idx * kv_stride, pos); attn_idx++; } else { ssm_block_forward(&b->ssm, &s->ssm[ssm_idx], x.data(), y.data()); ssm_idx++; } /* residual after attention/SSM, then dense SwiGLU FFN with its own * post_norm + residual: out = (x + sublayer) + ffn(post_norm(x+sublayer)). */ for (int i = 0; i < QX_D_MODEL; i++) h1[i] = x[i] + y[i]; ffn_forward(&b->ffn, h1.data(), f.data()); for (int i = 0; i < QX_D_MODEL; i++) x[i] = h1[i] + f[i]; if (getenv("QXMX_DUMP_LAYERS")) { double norm = 0; for (int i = 0; i < QX_D_MODEL; i++) norm += (double)x[i]*x[i]; norm = sqrt(norm / QX_D_MODEL); fprintf(stderr, " L%02d: x_rms=%.4f [%.4f %.4f %.4f]\n", L, norm, x[0], x[1], x[2]); } } /* output norm + lm_head */ std::vector<float> h(QX_D_MODEL); rmsnorm(x.data(), (const float*)m->output_norm.data, h.data(), QX_D_MODEL, QX_RMS_EPS); q2_matvec(&m->output, h.data(), logits_out); s->n_tokens++;}
} // namespace qx