Merge pull request #272 from hexxyan/feat/cpu-q4k-routed-experts

CPU support for Q4_K routed experts (fixes #171)

Makefile

@@ -197,9 +197,13 @@ else
 	$(NVCC) $(NVCCFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(CUDA_LDLIBS)
 endif
 
-test: ds4_test ds4-eval
+test: ds4_test ds4-eval q4k-dot-test
 	./ds4-eval --self-test-extractors
 	./ds4_test
 
+q4k-dot-test: tests/test_q4k_dot.c
+	$(CC) -O2 -Wall -Wextra -std=c99 -o tests/test_q4k_dot tests/test_q4k_dot.c -lm -pthread
+	./tests/test_q4k_dot
+
 clean:
-	rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o
+	rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test tests/test_q4k_dot *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o

ds4.c

@@ -2244,6 +2244,108 @@ static void ds4_vec_dot_q2_K_q8_K(int n, float *s, const block_q2_K *x, const bl
 #endif
 }
 
+static inline void q4_k_get_scale_min(int j, const uint8_t *q, uint8_t *sc, uint8_t *m) {
+    if (j < 4) {
+        *sc = q[j] & 63;
+        *m  = q[j + 4] & 63;
+    } else {
+        *sc = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
+        *m  = (q[j + 4] >> 4)  | ((q[j - 0] >> 6) << 4);
+    }
+}
+
+static void ds4_vec_dot_q4_K_q8_K(int n, float *s, const block_q4_K *x, const block_q8_K *y) {
+    const int nb = n / QK_K;
+
+#if defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD)
+    const int32x4_t zero = vdupq_n_s32(0);
+    float sumf = 0.0f;
+
+    for (int i = 0; i < nb; i++) {
+        const float d  = y[i].d * f16_to_f32(x[i].d);
+        const float dm = -y[i].d * f16_to_f32(x[i].dmin);
+
+        const uint8_t *qs = x[i].qs;
+        const uint8_t *sc = x[i].scales;
+        const int8_t  *q8 = y[i].qs;
+
+        int32_t summs = 0;
+        for (int j = 0; j < QK_K / 32; j++) {
+            uint8_t sc_val, m_val;
+            q4_k_get_scale_min(j, sc, &sc_val, &m_val);
+            int32_t gsum = (int32_t)y[i].bsums[j * 2] + (int32_t)y[i].bsums[j * 2 + 1];
+            summs += m_val * gsum;
+        }
+
+        int isum = 0;
+        for (int j = 0; j < QK_K / 32; j++) {
+            uint8_t sc_val, m_val;
+            q4_k_get_scale_min(j, sc, &sc_val, &m_val);
+
+            const int byte_off = (j >> 1) * 32;
+            const int shift = (j & 1) * 4;
+
+            /* Load 32 q8 values for this group */
+            const int8x16x2_t q8v = vld1q_s8_x2(q8 + j * 32);
+
+            /* Unpack 32 q4 values from 32 bytes at qs[byte_off] with shift */
+            uint8_t q4_u[32];
+            if (shift == 0) {
+                for (int l = 0; l < 32; l++) q4_u[l] = qs[byte_off + l] & 0xF;
+            } else {
+                for (int l = 0; l < 32; l++) q4_u[l] = qs[byte_off + l] >> 4;
+            }
+
+            const int8x16_t q4a = vreinterpretq_s8_u8(vld1q_u8(q4_u));
+            const int8x16_t q4b = vreinterpretq_s8_u8(vld1q_u8(q4_u + 16));
+
+            isum += vaddvq_s32(vdotq_s32(zero, q4a, q8v.val[0])) * sc_val;
+            isum += vaddvq_s32(vdotq_s32(zero, q4b, q8v.val[1])) * sc_val;
+        }
+
+        sumf += d * (float)isum + dm * (float)summs;
+    }
+
+    *s = sumf;
+#else
+    float sumf = 0.0f;
+
+    for (int i = 0; i < nb; i++) {
+        const float d  = y[i].d * f16_to_f32(x[i].d);
+        const float dm = -y[i].d * f16_to_f32(x[i].dmin);
+
+        const uint8_t *qs = x[i].qs;
+        const uint8_t *sc = x[i].scales;
+        const int8_t  *q8 = y[i].qs;
+
+        int summs = 0;
+        for (int j = 0; j < QK_K / 32; j++) {
+            uint8_t sc_val, m_val;
+            q4_k_get_scale_min(j, sc, &sc_val, &m_val);
+            int32_t gsum = (int32_t)y[i].bsums[j * 2] + (int32_t)y[i].bsums[j * 2 + 1];
+            summs += m_val * gsum;
+        }
+
+        int isum = 0;
+        for (int j = 0; j < QK_K / 32; j++) {
+            uint8_t sc_val, m_val;
+            q4_k_get_scale_min(j, sc, &sc_val, &m_val);
+
+            const int byte_off = (j >> 1) * 32;
+            const int shift = (j & 1) * 4;
+
+            for (int l = 0; l < 32; l++) {
+                isum += ((qs[byte_off + l] >> shift) & 0xF) * (int)q8[j * 32 + l] * sc_val;
+            }
+        }
+
+        sumf += d * (float)isum + dm * (float)summs;
+    }
+
+    *s = sumf;
+#endif
+}
+
 static DS4_MAYBE_UNUSED void ds4_vec_dot_iq2_xxs_q8_K(int n, float *s, const block_iq2_xxs *x, const block_q8_K *y) {
     const int nb = n / QK_K;
 
@@ -4801,6 +4903,354 @@ static void matvec_q2_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint6
     }
 }
 
+/* =========================================================================
+ * Q4_K routed expert matrix-vector products.
+ * ========================================================================= */
+
+typedef struct {
+    float *mid;
+    const uint8_t *gate_base[DS4_MAX_EXPERT_USED];
+    const uint8_t *up_base[DS4_MAX_EXPERT_USED];
+    const block_q8_K *xq;
+    float expert_weight[DS4_MAX_EXPERT_USED];
+    float clamp;
+    uint64_t in_dim;
+    uint64_t out_dim;
+    uint64_t gate_row_bytes[DS4_MAX_EXPERT_USED];
+    uint64_t up_row_bytes[DS4_MAX_EXPERT_USED];
+    int n_expert;
+} matvec_q4_k_mid_ctx;
+
+static void matvec_q4_k_mid_worker(void *vctx, uint64_t row0, uint64_t row1) {
+    matvec_q4_k_mid_ctx *ctx = vctx;
+
+    for (uint64_t idx = row0; idx < row1; idx++) {
+        const int slot = (int)(idx / ctx->out_dim);
+        const uint64_t row = idx - (uint64_t)slot * ctx->out_dim;
+        float gate = 0.0f;
+        float up = 0.0f;
+
+        const block_q4_K *gate_row = (const block_q4_K *)(ctx->gate_base[slot] + row * ctx->gate_row_bytes[slot]);
+        ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &gate, gate_row, ctx->xq);
+
+        const block_q4_K *up_row = (const block_q4_K *)(ctx->up_base[slot] + row * ctx->up_row_bytes[slot]);
+        ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &up, up_row, ctx->xq);
+
+        if (ctx->clamp > 1.0e-6f) {
+            if (gate > ctx->clamp) gate = ctx->clamp;
+            if (up > ctx->clamp) up = ctx->clamp;
+            if (up < -ctx->clamp) up = -ctx->clamp;
+        }
+        ctx->mid[idx] = silu(gate) * up * ctx->expert_weight[slot];
+    }
+}
+
+static void matvec_q4_k_experts_mid_prequant(
+        float            *mid,
+        const ds4_model  *m,
+        const ds4_tensor *gate_w,
+        const ds4_tensor *up_w,
+        const block_q8_K *xq,
+        const int        *selected,
+        const float      *expert_weight,
+        int               n_expert,
+        float             clamp) {
+    if (gate_w->type != DS4_TENSOR_Q4_K || up_w->type != DS4_TENSOR_Q4_K)
+        ds4_die("expected Q4_K expert tensors");
+    if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count");
+
+    uint64_t in_dim0 = 0;
+    uint64_t out_dim0 = 0;
+    matvec_q4_k_mid_ctx ctx = {
+        .mid = mid,
+        .xq = xq,
+        .clamp = clamp,
+        .n_expert = n_expert,
+    };
+
+    for (int i = 0; i < n_expert; i++) {
+        uint64_t gate_in_dim, gate_out_dim;
+        uint64_t up_in_dim, up_out_dim;
+        ctx.gate_base[i] = tensor_expert_bytes(m, gate_w, (uint32_t)selected[i],
+                                               &gate_in_dim, &gate_out_dim, &ctx.gate_row_bytes[i]);
+        ctx.up_base[i] = tensor_expert_bytes(m, up_w, (uint32_t)selected[i],
+                                             &up_in_dim, &up_out_dim, &ctx.up_row_bytes[i]);
+        if (gate_in_dim != up_in_dim || gate_out_dim != up_out_dim) {
+            ds4_die("paired Q4_K expert tensors do not match");
+        }
+        if (i == 0) {
+            in_dim0 = gate_in_dim;
+            out_dim0 = gate_out_dim;
+        } else if (gate_in_dim != in_dim0 || gate_out_dim != out_dim0) {
+            ds4_die("Q4_K expert tensors do not share a layout");
+        }
+        ctx.expert_weight[i] = expert_weight[i];
+    }
+    if (in_dim0 % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned");
+
+    ctx.in_dim = in_dim0;
+    ctx.out_dim = out_dim0;
+    ds4_parallel_for((uint64_t)n_expert * out_dim0, matvec_q4_k_mid_worker, &ctx);
+}
+
+typedef struct {
+    float *out;
+    const uint8_t *base[DS4_MAX_EXPERT_USED];
+    const block_q8_K *xq[DS4_MAX_EXPERT_USED];
+    uint64_t in_dim;
+    uint64_t row_bytes[DS4_MAX_EXPERT_USED];
+    int n_expert;
+} matvec_q4_k_accum_ctx;
+
+static void matvec_q4_k_accum_worker(void *vctx, uint64_t row0, uint64_t row1) {
+    matvec_q4_k_accum_ctx *ctx = vctx;
+
+    for (uint64_t row = row0; row < row1; row++) {
+        float acc = 0.0f;
+        for (int i = 0; i < ctx->n_expert; i++) {
+            float v = 0.0f;
+            const block_q4_K *br = (const block_q4_K *)(ctx->base[i] + row * ctx->row_bytes[i]);
+            ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &v, br, ctx->xq[i]);
+            acc += v;
+        }
+        ctx->out[row] = acc;
+    }
+}
+
+static void matvec_q4_k_experts_accum_prequant(
+        float            *out,
+        const ds4_model  *m,
+        const ds4_tensor *w,
+        const block_q8_K *xq,
+        const int        *selected,
+        int               n_expert) {
+    if (w->type != DS4_TENSOR_Q4_K) ds4_die("expected a Q4_K expert tensor");
+    if (n_expert < 1 || (uint32_t)n_expert > DS4_N_EXPERT_USED) ds4_die("unexpected routed expert count");
+
+    uint64_t in_dim0 = 0;
+    uint64_t out_dim0 = 0;
+    const uint8_t *base[DS4_MAX_EXPERT_USED];
+    uint64_t row_bytes[DS4_MAX_EXPERT_USED];
+
+    for (int i = 0; i < n_expert; i++) {
+        uint64_t in_dim, out_dim;
+        base[i] = tensor_expert_bytes(m, w, (uint32_t)selected[i], &in_dim, &out_dim, &row_bytes[i]);
+        if (i == 0) {
+            in_dim0 = in_dim;
+            out_dim0 = out_dim;
+        } else if (in_dim != in_dim0 || out_dim != out_dim0) {
+            ds4_die("Q4_K expert tensors do not share a layout");
+        }
+    }
+    if (in_dim0 % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned");
+
+    const uint64_t n_blocks = in_dim0 / QK_K;
+    matvec_q4_k_accum_ctx ctx = {
+        .out = out,
+        .in_dim = in_dim0,
+        .n_expert = n_expert,
+    };
+    for (int i = 0; i < n_expert; i++) {
+        ctx.base[i] = base[i];
+        ctx.row_bytes[i] = row_bytes[i];
+        ctx.xq[i] = xq + (uint64_t)i * n_blocks;
+    }
+
+    ds4_parallel_for(out_dim0, matvec_q4_k_accum_worker, &ctx);
+}
+
+/* Q4_K batch mid worker: same structure as IQ2_XXS batch but uses Q4_K dot. */
+typedef struct {
+    float *mid;
+    const uint8_t *gate_base[DS4_MAX_EXPERT];
+    const uint8_t *up_base[DS4_MAX_EXPERT];
+    const block_q8_K *xq;
+    const ds4_expert_pair *pairs;
+    const uint32_t *pair_ids;
+    const uint32_t *expert_offset;
+    const uint32_t *active_expert;
+    const float *pair_weight;
+    float clamp;
+    uint64_t in_dim;
+    uint64_t out_dim;
+    uint64_t gate_row_bytes[DS4_MAX_EXPERT];
+    uint64_t up_row_bytes[DS4_MAX_EXPERT];
+    uint64_t xq_blocks;
+} matvec_q4_k_batch_mid_ctx;
+
+static void matvec_q4_k_batch_mid_worker(void *vctx, uint64_t task0, uint64_t task1) {
+    matvec_q4_k_batch_mid_ctx *ctx = vctx;
+
+    for (uint64_t task = task0; task < task1; task++) {
+        const uint32_t active_idx = (uint32_t)(task / ctx->out_dim);
+        const uint64_t row = task - (uint64_t)active_idx * ctx->out_dim;
+        const uint32_t expert = ctx->active_expert[active_idx];
+        const uint32_t begin = ctx->expert_offset[expert];
+        const uint32_t end = ctx->expert_offset[expert + 1];
+
+        const block_q4_K *gate_row = (const block_q4_K *)(ctx->gate_base[expert] + row * ctx->gate_row_bytes[expert]);
+        const block_q4_K *up_row = (const block_q4_K *)(ctx->up_base[expert] + row * ctx->up_row_bytes[expert]);
+
+        for (uint32_t i = begin; i < end; i++) {
+            const uint32_t pair_id = ctx->pair_ids[i];
+            const ds4_expert_pair pair = ctx->pairs[pair_id];
+            const block_q8_K *xq = ctx->xq + (uint64_t)pair.token * ctx->xq_blocks;
+            float gate = 0.0f;
+            float up = 0.0f;
+
+            ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &gate, gate_row, xq);
+            ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &up, up_row, xq);
+
+            if (ctx->clamp > 1.0e-6f) {
+                if (gate > ctx->clamp) gate = ctx->clamp;
+                if (up > ctx->clamp) up = ctx->clamp;
+                if (up < -ctx->clamp) up = -ctx->clamp;
+            }
+
+            ctx->mid[(uint64_t)pair_id * ctx->out_dim + row] = silu(gate) * up * ctx->pair_weight[pair_id];
+        }
+    }
+}
+
+/* Q4_K batch down accum worker: same structure as Q2_K batch but uses Q4_K dot. */
+typedef struct {
+    float *moe;
+    const uint8_t *base[DS4_MAX_EXPERT];
+    const block_q8_K *midq;
+    const ds4_expert_pair *pairs;
+    const uint32_t *pair_ids;
+    const uint32_t *expert_offset;
+    const uint32_t *active_expert;
+    uint32_t n_active;
+    uint32_t n_tok;
+    uint64_t in_dim;
+    uint64_t out_dim;
+    uint64_t row_bytes[DS4_MAX_EXPERT];
+    uint64_t midq_blocks;
+} matvec_q4_k_batch_accum_rows_ctx;
+
+static void matvec_q4_k_batch_accum_rows_worker(void *vctx, uint64_t row0, uint64_t row1) {
+    matvec_q4_k_batch_accum_rows_ctx *ctx = vctx;
+
+    for (uint64_t row = row0; row < row1; row++) {
+        for (uint32_t t = 0; t < ctx->n_tok; t++) {
+            ctx->moe[(uint64_t)t * ctx->out_dim + row] = 0.0f;
+        }
+
+        for (uint32_t ai = 0; ai < ctx->n_active; ai++) {
+            const uint32_t expert = ctx->active_expert[ai];
+            const uint32_t begin = ctx->expert_offset[expert];
+            const uint32_t end = ctx->expert_offset[expert + 1];
+            const block_q4_K *br = (const block_q4_K *)(ctx->base[expert] + row * ctx->row_bytes[expert]);
+
+            for (uint32_t i = begin; i < end; i++) {
+                const uint32_t pair_id = ctx->pair_ids[i];
+                const ds4_expert_pair pair = ctx->pairs[pair_id];
+                const block_q8_K *xq = ctx->midq + (uint64_t)pair_id * ctx->midq_blocks;
+                float v = 0.0f;
+
+                ds4_vec_dot_q4_K_q8_K((int)ctx->in_dim, &v, br, xq);
+                ctx->moe[(uint64_t)pair.token * ctx->out_dim + row] += v;
+            }
+        }
+    }
+}
+
+/* Dispatch: call the right gate/up mid builder based on tensor type. */
+static void matvec_experts_mid_prequant(
+        float            *mid,
+        const ds4_model  *m,
+        const ds4_tensor *gate_w,
+        const ds4_tensor *up_w,
+        const block_q8_K *xq,
+        const int        *selected,
+        const float      *expert_weight,
+        int               n_expert,
+        float             clamp) {
+    if (gate_w->type == DS4_TENSOR_IQ2_XXS) {
+        matvec_iq2_xxs_experts_mid_prequant(mid, m, gate_w, up_w, xq,
+                                            selected, expert_weight, n_expert, clamp);
+    } else if (gate_w->type == DS4_TENSOR_Q4_K) {
+        matvec_q4_k_experts_mid_prequant(mid, m, gate_w, up_w, xq,
+                                         selected, expert_weight, n_expert, clamp);
+    } else {
+        ds4_die("unsupported gate/up expert tensor type");
+    }
+}
+
+/* Dispatch: call the right down-projection accumulator based on tensor type. */
+static void matvec_experts_down_accum_prequant(
+        float            *out,
+        const ds4_model  *m,
+        const ds4_tensor *w,
+        const block_q8_K *xq,
+        const int        *selected,
+        int               n_expert) {
+    if (w->type == DS4_TENSOR_Q2_K) {
+        matvec_q2_k_experts_accum_prequant(out, m, w, xq, selected, n_expert);
+    } else if (w->type == DS4_TENSOR_Q4_K) {
+        matvec_q4_k_experts_accum_prequant(out, m, w, xq, selected, n_expert);
+    } else {
+        ds4_die("unsupported down expert tensor type");
+    }
+}
+
+/* Dispatch: single-expert gate/up pair for tracing. */
+static void matvec_expert_pair_prequant(
+        float            *out0,
+        float            *out1,
+        const ds4_model  *m,
+        const ds4_tensor *w0,
+        const ds4_tensor *w1,
+        const block_q8_K *xq,
+        uint32_t          expert) {
+    if (w0->type == DS4_TENSOR_IQ2_XXS) {
+        matvec_iq2_xxs_expert_pair_prequant(out0, out1, m, w0, w1, xq, expert);
+    } else if (w0->type == DS4_TENSOR_Q4_K) {
+        uint64_t in_dim0, out_dim0, rb0;
+        uint64_t in_dim1, out_dim1, rb1;
+        const uint8_t *base0 = tensor_expert_bytes(m, w0, expert, &in_dim0, &out_dim0, &rb0);
+        const uint8_t *base1 = tensor_expert_bytes(m, w1, expert, &in_dim1, &out_dim1, &rb1);
+        if (in_dim0 != in_dim1 || out_dim0 != out_dim1) ds4_die("paired Q4_K expert tensors do not match");
+
+        for (uint64_t row = 0; row < out_dim0; row++) {
+            const block_q4_K *gr = (const block_q4_K *)(base0 + row * rb0);
+            ds4_vec_dot_q4_K_q8_K((int)in_dim0, &out0[row], gr, xq);
+            const block_q4_K *ur = (const block_q4_K *)(base1 + row * rb1);
+            ds4_vec_dot_q4_K_q8_K((int)in_dim0, &out1[row], ur, xq);
+        }
+    } else {
+        ds4_die("unsupported gate/up expert tensor type");
+    }
+}
+
+/* Dispatch: single-expert down projection for tracing. */
+static void matvec_expert_down(
+        float            *out,
+        const ds4_model  *m,
+        const ds4_tensor *w,
+        const float      *x,
+        uint32_t          expert) {
+    if (w->type == DS4_TENSOR_Q2_K) {
+        matvec_q2_k_expert(out, m, w, x, expert);
+    } else if (w->type == DS4_TENSOR_Q4_K) {
+        uint64_t in_dim, out_dim, row_bytes;
+        const uint8_t *base = tensor_expert_bytes(m, w, expert, &in_dim, &out_dim, &row_bytes);
+        if (in_dim % QK_K != 0) ds4_die("Q4_K expert row is not QK_K aligned");
+
+        block_q8_K *xq = xmalloc((size_t)(in_dim / QK_K) * sizeof(xq[0]));
+        ds4_quantize_row_q8_K(x, xq, (int64_t)in_dim);
+
+        for (uint64_t row = 0; row < out_dim; row++) {
+            const block_q4_K *br = (const block_q4_K *)(base + row * row_bytes);
+            ds4_vec_dot_q4_K_q8_K((int)in_dim, &out[row], br, xq);
+        }
+        free(xq);
+    } else {
+        ds4_die("unsupported down expert tensor type");
+    }
+}
+
 typedef struct {
     float *moe;
     const float *down_pair;
@@ -5957,8 +6407,8 @@ static void layer_routed_moe_one(
     float *down = trace ? xmalloc((size_t)DS4_N_EMBD * sizeof(down[0])) : NULL;
     const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0];
     const uint64_t down_in_dim = layer->ffn_down_exps->dim[0];
-    if (expert_in_dim % QK_K != 0) ds4_die("IQ2_XXS expert input is not QK_K aligned");
-    if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) ds4_die("Q2_K expert input has an unexpected layout");
+    if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned");
+    if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) ds4_die("routed expert down input has an unexpected layout");
     block_q8_K *xq = xmalloc((size_t)(expert_in_dim / QK_K) * sizeof(xq[0]));
     block_q8_K *midq = trace ? NULL : xmalloc((size_t)DS4_N_EXPERT_USED * (down_in_dim / QK_K) * sizeof(midq[0]));
 
@@ -5973,29 +6423,29 @@ static void layer_routed_moe_one(
     }
 
     if (!trace) {
-        matvec_iq2_xxs_experts_mid_prequant(mid_all, model,
-                                            layer->ffn_gate_exps,
-                                            layer->ffn_up_exps,
-                                            xq,
-                                            selected,
-                                            expert_weight,
-                                            DS4_N_EXPERT_USED,
-                                            clamp);
+        matvec_experts_mid_prequant(mid_all, model,
+                                    layer->ffn_gate_exps,
+                                    layer->ffn_up_exps,
+                                    xq,
+                                    selected,
+                                    expert_weight,
+                                    DS4_N_EXPERT_USED,
+                                    clamp);
         for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) {
             ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim,
                                   midq + (uint64_t)i * (down_in_dim / QK_K),
                                   (int64_t)down_in_dim);
         }
-        matvec_q2_k_experts_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED);
+        matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED);
     } else {
         for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) {
             const uint32_t expert = (uint32_t)selected[i];
 
-            matvec_iq2_xxs_expert_pair_prequant(gate, up, model,
-                                                 layer->ffn_gate_exps,
-                                                 layer->ffn_up_exps,
-                                                 xq,
-                                                 expert);
+            matvec_expert_pair_prequant(gate, up, model,
+                                        layer->ffn_gate_exps,
+                                        layer->ffn_up_exps,
+                                        xq,
+                                        expert);
             char name[64];
             snprintf(name, sizeof(name), "blk.%u expert %u gate", il, expert);
             print_vec_stats(name, gate, DS4_N_FF_EXP);
@@ -6019,7 +6469,7 @@ static void layer_routed_moe_one(
             snprintf(name, sizeof(name), "blk.%u expert %u mid", il, expert);
             print_vec_stats(name, mid, DS4_N_FF_EXP);
 
-            matvec_q2_k_expert(down, model, layer->ffn_down_exps, mid, expert);
+            matvec_expert_down(down, model, layer->ffn_down_exps, mid, expert);
             snprintf(name, sizeof(name), "blk.%u expert %u down", il, expert);
             print_vec_stats(name, down, DS4_N_EMBD);
             for (uint32_t j = 0; j < DS4_N_EMBD; j++) out[j] += down[j];
@@ -6053,8 +6503,8 @@ static void layer_routed_moe_one_prealloc(
     const uint64_t expert_in_dim = layer->ffn_gate_exps->dim[0];
     const uint64_t down_in_dim = layer->ffn_down_exps->dim[0];
 
-    if (expert_in_dim % QK_K != 0) ds4_die("IQ2_XXS expert input is not QK_K aligned");
-    if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) ds4_die("Q2_K expert input has an unexpected layout");
+    if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned");
+    if (down_in_dim != DS4_N_FF_EXP || down_in_dim % QK_K != 0) ds4_die("routed expert down input has an unexpected layout");
 
     memset(out, 0, (size_t)DS4_N_EMBD * sizeof(out[0]));
     ds4_quantize_row_q8_K(x, xq, (int64_t)expert_in_dim);
@@ -6066,21 +6516,21 @@ static void layer_routed_moe_one_prealloc(
         layer_topk_selected_experts(selected, expert_weight, model, layer, x);
     }
 
-    matvec_iq2_xxs_experts_mid_prequant(mid_all, model,
-                                        layer->ffn_gate_exps,
-                                        layer->ffn_up_exps,
-                                        xq,
-                                        selected,
-                                        expert_weight,
-                                        DS4_N_EXPERT_USED,
-                                        clamp);
+    matvec_experts_mid_prequant(mid_all, model,
+                                layer->ffn_gate_exps,
+                                layer->ffn_up_exps,
+                                xq,
+                                selected,
+                                expert_weight,
+                                DS4_N_EXPERT_USED,
+                                clamp);
 
     for (uint32_t i = 0; i < DS4_N_EXPERT_USED; i++) {
         ds4_quantize_row_q8_K(mid_all + (uint64_t)i * down_in_dim,
                               midq + (uint64_t)i * (down_in_dim / QK_K),
                               (int64_t)down_in_dim);
     }
-    matvec_q2_k_experts_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED);
+    matvec_experts_down_accum_prequant(out, model, layer->ffn_down_exps, midq, selected, DS4_N_EXPERT_USED);
 
     (void)il;
 }
@@ -6100,8 +6550,8 @@ static void layer_routed_moe_batch(
     const uint64_t expert_out_dim = layer->ffn_gate_exps->dim[1];
     const uint64_t down_in_dim = layer->ffn_down_exps->dim[0];
     const uint64_t down_out_dim = layer->ffn_down_exps->dim[1];
-    if (expert_in_dim % QK_K != 0) ds4_die("IQ2_XXS expert input is not QK_K aligned");
-    if (down_in_dim % QK_K != 0) ds4_die("Q2_K expert input is not QK_K aligned");
+    if (expert_in_dim % QK_K != 0) ds4_die("routed expert input is not QK_K aligned");
+    if (down_in_dim % QK_K != 0) ds4_die("routed expert down input is not QK_K aligned");
     if (expert_out_dim != down_in_dim || down_out_dim != DS4_N_EMBD) {
         ds4_die("routed expert tensor layout is unexpected");
     }
@@ -6156,35 +6606,72 @@ static void layer_routed_moe_batch(
 
     float *mid = xmalloc((size_t)total_pairs * expert_out_dim * sizeof(mid[0]));
 
-    matvec_iq2_xxs_batch_mid_ctx mid_ctx = {
-        .mid = mid,
-        .xq = xq,
-        .pairs = pairs,
-        .pair_ids = pair_ids,
-        .expert_offset = counts,
-        .active_expert = active_expert,
-        .pair_weight = pair_weight,
-        .clamp = clamp,
-        .in_dim = expert_in_dim,
-        .out_dim = expert_out_dim,
-        .xq_blocks = xq_blocks,
-    };
+    const uint32_t gate_type = layer->ffn_gate_exps->type;
 
-    for (uint32_t ai = 0; ai < n_active; ai++) {
-        const uint32_t e = active_expert[ai];
-        uint64_t gate_in_dim, gate_out_dim;
-        uint64_t up_in_dim, up_out_dim;
-        mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e,
-                                                   &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]);
-        mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e,
-                                                 &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]);
-        if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim ||
-            gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) {
-            ds4_die("IQ2_XXS batch expert tensor layout mismatch");
+    /* Build mid vectors: dispatch based on gate/up tensor type. */
+    if (gate_type == DS4_TENSOR_IQ2_XXS) {
+        matvec_iq2_xxs_batch_mid_ctx mid_ctx = {
+            .mid = mid,
+            .xq = xq,
+            .pairs = pairs,
+            .pair_ids = pair_ids,
+            .expert_offset = counts,
+            .active_expert = active_expert,
+            .pair_weight = pair_weight,
+            .clamp = clamp,
+            .in_dim = expert_in_dim,
+            .out_dim = expert_out_dim,
+            .xq_blocks = xq_blocks,
+        };
+
+        for (uint32_t ai = 0; ai < n_active; ai++) {
+            const uint32_t e = active_expert[ai];
+            uint64_t gate_in_dim, gate_out_dim;
+            uint64_t up_in_dim, up_out_dim;
+            mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e,
+                                                       &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]);
+            mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e,
+                                                     &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]);
+            if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim ||
+                gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) {
+                ds4_die("batch expert tensor layout mismatch");
+            }
         }
-    }
 
-    ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_iq2_xxs_batch_mid_worker, &mid_ctx);
+        ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_iq2_xxs_batch_mid_worker, &mid_ctx);
+    } else if (gate_type == DS4_TENSOR_Q4_K) {
+        matvec_q4_k_batch_mid_ctx mid_ctx = {
+            .mid = mid,
+            .xq = xq,
+            .pairs = pairs,
+            .pair_ids = pair_ids,
+            .expert_offset = counts,
+            .active_expert = active_expert,
+            .pair_weight = pair_weight,
+            .clamp = clamp,
+            .in_dim = expert_in_dim,
+            .out_dim = expert_out_dim,
+            .xq_blocks = xq_blocks,
+        };
+
+        for (uint32_t ai = 0; ai < n_active; ai++) {
+            const uint32_t e = active_expert[ai];
+            uint64_t gate_in_dim, gate_out_dim;
+            uint64_t up_in_dim, up_out_dim;
+            mid_ctx.gate_base[e] = tensor_expert_bytes(model, layer->ffn_gate_exps, e,
+                                                       &gate_in_dim, &gate_out_dim, &mid_ctx.gate_row_bytes[e]);
+            mid_ctx.up_base[e] = tensor_expert_bytes(model, layer->ffn_up_exps, e,
+                                                     &up_in_dim, &up_out_dim, &mid_ctx.up_row_bytes[e]);
+            if (gate_in_dim != expert_in_dim || up_in_dim != expert_in_dim ||
+                gate_out_dim != expert_out_dim || up_out_dim != expert_out_dim) {
+                ds4_die("batch expert tensor layout mismatch");
+            }
+        }
+
+        ds4_parallel_for((uint64_t)n_active * expert_out_dim, matvec_q4_k_batch_mid_worker, &mid_ctx);
+    } else {
+        ds4_die("unsupported gate/up expert tensor type for batch");
+    }
 
     const uint64_t midq_blocks = down_in_dim / QK_K;
     block_q8_K *midq = xmalloc((size_t)total_pairs * midq_blocks * sizeof(midq[0]));
@@ -6197,31 +6684,64 @@ static void layer_routed_moe_batch(
     ds4_parallel_for(total_pairs, quantize_mid_pairs_worker, &quant_ctx);
     free(mid);
 
-    matvec_q2_k_batch_accum_rows_ctx down_ctx = {
-        .moe = moe,
-        .midq = midq,
-        .pairs = pairs,
-        .pair_ids = pair_ids,
-        .expert_offset = counts,
-        .active_expert = active_expert,
-        .n_active = n_active,
-        .n_tok = n_tok,
-        .in_dim = down_in_dim,
-        .out_dim = down_out_dim,
-        .midq_blocks = midq_blocks,
-    };
+    /* Down projection: dispatch based on down tensor type. */
+    const uint32_t down_type = layer->ffn_down_exps->type;
 
-    for (uint32_t ai = 0; ai < n_active; ai++) {
-        const uint32_t e = active_expert[ai];
-        uint64_t in_dim, out_dim;
-        down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e,
-                                               &in_dim, &out_dim, &down_ctx.row_bytes[e]);
-        if (in_dim != down_in_dim || out_dim != down_out_dim) {
-            ds4_die("Q2_K batch expert tensor layout mismatch");
+    if (down_type == DS4_TENSOR_Q2_K) {
+        matvec_q2_k_batch_accum_rows_ctx down_ctx = {
+            .moe = moe,
+            .midq = midq,
+            .pairs = pairs,
+            .pair_ids = pair_ids,
+            .expert_offset = counts,
+            .active_expert = active_expert,
+            .n_active = n_active,
+            .n_tok = n_tok,
+            .in_dim = down_in_dim,
+            .out_dim = down_out_dim,
+            .midq_blocks = midq_blocks,
+        };
+
+        for (uint32_t ai = 0; ai < n_active; ai++) {
+            const uint32_t e = active_expert[ai];
+            uint64_t in_dim, out_dim;
+            down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e,
+                                                   &in_dim, &out_dim, &down_ctx.row_bytes[e]);
+            if (in_dim != down_in_dim || out_dim != down_out_dim) {
+                ds4_die("batch expert tensor layout mismatch");
+            }
         }
-    }
 
-    ds4_parallel_for(down_out_dim, matvec_q2_k_batch_accum_rows_worker, &down_ctx);
+        ds4_parallel_for(down_out_dim, matvec_q2_k_batch_accum_rows_worker, &down_ctx);
+    } else if (down_type == DS4_TENSOR_Q4_K) {
+        matvec_q4_k_batch_accum_rows_ctx down_ctx = {
+            .moe = moe,
+            .midq = midq,
+            .pairs = pairs,
+            .pair_ids = pair_ids,
+            .expert_offset = counts,
+            .active_expert = active_expert,
+            .n_active = n_active,
+            .n_tok = n_tok,
+            .in_dim = down_in_dim,
+            .out_dim = down_out_dim,
+            .midq_blocks = midq_blocks,
+        };
+
+        for (uint32_t ai = 0; ai < n_active; ai++) {
+            const uint32_t e = active_expert[ai];
+            uint64_t in_dim, out_dim;
+            down_ctx.base[e] = tensor_expert_bytes(model, layer->ffn_down_exps, e,
+                                                   &in_dim, &out_dim, &down_ctx.row_bytes[e]);
+            if (in_dim != down_in_dim || out_dim != down_out_dim) {
+                ds4_die("batch expert tensor layout mismatch");
+            }
+        }
+
+        ds4_parallel_for(down_out_dim, matvec_q4_k_batch_accum_rows_worker, &down_ctx);
+    } else {
+        ds4_die("unsupported down expert tensor type for batch");
+    }
 
     free(midq);
     free(pair_ids);

tests/test_q4k_dot.c

@@ -0,0 +1,249 @@
+/*
+ * Unit test for Q4_K block layout, scale extraction, and dot product.
+ * Build:  cc -O2 -Wall -DDS4_NO_GPU -DDS4_Q4K_DOT_TEST_MAIN -I. -o tests/test_q4k_dot tests/test_q4k_dot.c -lm -pthread
+ * Run:    ./tests/test_q4k_dot
+ */
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <math.h>
+#include <stdint.h>
+
+#define QK_K 256
+
+typedef struct {
+    uint16_t d;
+    uint16_t dmin;
+    uint8_t  scales[12];
+    uint8_t  qs[QK_K / 2];
+} block_q4_K;
+
+typedef struct {
+    float   d;
+    int8_t  qs[QK_K];
+    int16_t bsums[QK_K / 16];
+} block_q8_K;
+
+static inline float f16_to_f32(uint16_t h) {
+    uint32_t sign = (h & 0x8000) << 16;
+    int32_t exp = (h >> 10) & 0x1F;
+    uint32_t frac = h & 0x3FF;
+    if (exp == 0) return *(float *)&sign;
+    if (exp == 31) return *(float *)(uint32_t[]){ sign | 0x7F800000 | (frac << 13) };
+    uint32_t f = sign | ((exp + 127 - 15) << 23) | (frac << 13);
+    float out;
+    memcpy(&out, &f, sizeof(out));
+    return out;
+}
+
+static inline void q4_k_get_scale_min(int j, const uint8_t *q, uint8_t *sc, uint8_t *m) {
+    if (j < 4) {
+        *sc = q[j] & 63;
+        *m  = q[j + 4] & 63;
+    } else {
+        *sc = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4);
+        *m  = (q[j + 4] >> 4)  | ((q[j - 0] >> 6) << 4);
+    }
+}
+
+/* The corrected dot product, matching ds4.c after the fix. */
+static void vec_dot_q4_K_q8_K(int n, float *s, const block_q4_K *x, const block_q8_K *y) {
+    const int nb = n / QK_K;
+    float sumf = 0.0f;
+
+    for (int i = 0; i < nb; i++) {
+        const float d  = y[i].d * f16_to_f32(x[i].d);
+        const float dm = -y[i].d * f16_to_f32(x[i].dmin);
+
+        const uint8_t *qs = x[i].qs;
+        const uint8_t *sc = x[i].scales;
+        const int8_t  *q8 = y[i].qs;
+
+        int summs = 0;
+        for (int j = 0; j < QK_K / 32; j++) {
+            uint8_t sc_val, m_val;
+            q4_k_get_scale_min(j, sc, &sc_val, &m_val);
+            int32_t gsum = (int32_t)y[i].bsums[j * 2] + (int32_t)y[i].bsums[j * 2 + 1];
+            summs += m_val * gsum;
+        }
+
+        int isum = 0;
+        for (int j = 0; j < QK_K / 32; j++) {
+            uint8_t sc_val, m_val;
+            q4_k_get_scale_min(j, sc, &sc_val, &m_val);
+
+            const int byte_off = (j >> 1) * 32;
+            const int shift = (j & 1) * 4;
+
+            for (int l = 0; l < 32; l++) {
+                isum += ((qs[byte_off + l] >> shift) & 0xF) * (int)q8[j * 32 + l] * sc_val;
+            }
+        }
+
+        sumf += d * (float)isum + dm * (float)summs;
+    }
+
+    *s = sumf;
+}
+
+/* Reference: fully dequantize Q4_K to float, then dot with Q8_K's dequantized values. */
+static float ref_dot(const block_q4_K *bx, const block_q8_K *by) {
+    float x[QK_K];
+    const float d  = f16_to_f32(bx->d);
+    const float dm = f16_to_f32(bx->dmin);
+
+    for (int j = 0; j < QK_K / 32; j++) {
+        uint8_t sc_val, m_val;
+        q4_k_get_scale_min(j, bx->scales, &sc_val, &m_val);
+        const int byte_off = (j >> 1) * 32;
+        const int shift = (j & 1) * 4;
+        for (int l = 0; l < 32; l++) {
+            int q = (bx->qs[byte_off + l] >> shift) & 0xF;
+            x[j * 32 + l] = d * sc_val * (float)q - dm * m_val;
+        }
+    }
+
+    float sum = 0.0f;
+    for (int i = 0; i < QK_K; i++) sum += x[i] * by->d * (float)by->qs[i];
+    return sum;
+}
+
+static void fill_q4_K(block_q4_K *bx, uint32_t seed) {
+    uint8_t *p = (uint8_t *)bx;
+    uint32_t s = seed;
+    for (size_t i = 0; i < sizeof(block_q4_K); i++) {
+        s = s * 1103515245u + 12345u;
+        p[i] = (uint8_t)(s >> 16);
+    }
+    /* Ensure d/dmin are valid finite f16 */
+    bx->d    &= 0x7BFF;
+    bx->dmin &= 0x7BFF;
+}
+
+static void fill_q8_K(block_q8_K *by, uint32_t seed) {
+    uint32_t s = seed;
+    by->d = ((s & 0xFFFF) / 65536.0f) * 2.0f + 0.01f;
+    for (int i = 0; i < QK_K; i++) {
+        s = s * 1103515245u + 12345u;
+        by->qs[i] = (int8_t)((s >> 16) & 0xFF);
+    }
+    for (int j = 0; j < QK_K / 16; j++) {
+        int32_t sum = 0;
+        for (int l = 0; l < 16; l++) sum += (int32_t)by->qs[j * 16 + l];
+        by->bsums[j] = (int16_t)sum;
+    }
+}
+
+static int test_block_sizes(void) {
+    int ok = (sizeof(block_q4_K) == 144 && sizeof(block_q8_K) == 292);
+    printf("  block sizes: Q4_K=%zu (expect 144), Q8_K=%zu (expect 292): %s\n",
+           sizeof(block_q4_K), sizeof(block_q8_K), ok ? "PASS" : "FAIL");
+    return ok ? 0 : 1;
+}
+
+static int test_scale_extraction(void) {
+    uint8_t scales[12] = {0};
+    uint8_t expected_sc[8] = {5, 12, 30, 60, 7, 15, 31, 63};
+    uint8_t expected_m[8]  = {3, 8, 20, 45, 1, 10, 25, 50};
+
+    for (int j = 0; j < 4; j++) {
+        scales[j] = expected_sc[j];
+        scales[j + 4] = expected_m[j];
+    }
+    for (int j = 4; j < 8; j++) {
+        scales[j + 4] = (expected_sc[j] & 0xF) | ((expected_m[j] & 0xF) << 4);
+        scales[j - 4] |= ((expected_sc[j] >> 4) << 6);
+        scales[j - 0] |= ((expected_m[j] >> 4) << 6);
+    }
+
+    int ok = 1;
+    for (int j = 0; j < 8; j++) {
+        uint8_t sc, m;
+        q4_k_get_scale_min(j, scales, &sc, &m);
+        if (sc != expected_sc[j] || m != expected_m[j]) {
+            printf("    j=%d: expected sc=%d m=%d, got sc=%d m=%d\n",
+                   j, expected_sc[j], expected_m[j], sc, m);
+            ok = 0;
+        }
+    }
+
+    printf("  scale extraction: %s\n", ok ? "PASS" : "FAIL");
+    return ok ? 0 : 1;
+}
+
+static int test_dot_reference(void) {
+    int ok = 1;
+    for (uint32_t seed = 1; seed <= 50; seed++) {
+        block_q4_K bx;
+        block_q8_K by;
+        fill_q4_K(&bx, seed);
+        fill_q8_K(&by, seed * 7 + 13);
+
+        float result = 0.0f;
+        vec_dot_q4_K_q8_K(QK_K, &result, &bx, &by);
+        float expected = ref_dot(&bx, &by);
+
+        float err = fabsf(result - expected);
+        float rel = fabsf(expected) > 1e-3f ? err / fabsf(expected) : err;
+        if (rel > 0.01f) {
+            printf("    seed=%u: result=%.6f expected=%.6f rel_err=%.6f\n",
+                   seed, result, expected, rel);
+            ok = 0;
+        }
+    }
+
+    printf("  dot vs reference (50 random blocks): %s\n", ok ? "PASS" : "FAIL");
+    return ok ? 0 : 1;
+}
+
+/* Test with a hand-crafted known block. */
+static int test_dot_known(void) {
+    block_q4_K bx;
+    block_q8_K by;
+    memset(&bx, 0, sizeof(bx));
+    memset(&by, 0, sizeof(by));
+
+    /* d=1.0 f16=0x3C00, dmin=0, all scales=1 (sc=1, m=0), all qs=0x88 (value 8) */
+    bx.d = 0x3C00;
+    bx.dmin = 0;
+    /* Groups 0-3: sc=1 m=0 */
+    for (int j = 0; j < 4; j++) {
+        bx.scales[j] = 1;
+        bx.scales[j + 4] = 0;
+    }
+    /* Groups 4-7: sc=1 m=0, high bits=0 */
+    for (int j = 4; j < 8; j++) {
+        bx.scales[j + 4] = (1 & 0xF) | (0 << 4);
+        /* high bits already 0 */
+    }
+    /* All qs = 0x88: low nibble=8, high nibble=8 */
+    memset(bx.qs, 0x88, sizeof(bx.qs));
+
+    /* Q8: all +1, d=1.0 */
+    by.d = 1.0f;
+    for (int i = 0; i < QK_K; i++) by.qs[i] = 1;
+    for (int j = 0; j < QK_K / 16; j++) by.bsums[j] = 16;
+
+    /* Expected: d=1, dm=0, isum = sum over 8 groups * 32 values * 8 * 1 * sc(=1) = 8*32*8 = 2048 */
+    float result = 0.0f;
+    vec_dot_q4_K_q8_K(QK_K, &result, &bx, &by);
+    float expected = 2048.0f;
+
+    int ok = (fabsf(result - expected) < 0.5f);
+    printf("  dot known: result=%.1f expected=%.1f: %s\n", result, expected, ok ? "PASS" : "FAIL");
+    return ok ? 0 : 1;
+}
+
+int main(void) {
+    int failures = 0;
+
+    printf("Q4_K unit tests:\n");
+    failures += test_block_sizes();
+    failures += test_scale_extraction();
+    failures += test_dot_known();
+    failures += test_dot_reference();
+
+    printf("\n%d/%d tests passed\n", 4 - failures, 4);
+    return failures ? 1 : 0;
+}