perf(qwen36): Q4_K shared-expert FFN kernels + Q8_0/Q5_K→Q4_K requant of the MoE FFN stage (+2.8–3.2%)

.gitignore

@@ -8,6 +8,9 @@ __pycache__/
 .prebuilt/
 .llamacpp/
 /models/
+/models35/
+/models36/
+*.gguf
 dist/
 .claude
 .env.eval

kernels/csrc/cuda/moe/expert_ffn_q4k.cu

@@ -588,6 +588,63 @@ __global__ void shared_down_q8_mmvq_kernel(
     }
 }
 
+// ---- Qwen3.6 shared expert, Q4_K gate/up/down (requant of the UD Q8_0 weights) -------
+// Byte-for-byte the same fused shared-expert math as the Q8_0 kernels above, but the
+// per-token weight read is Q4_K super-blocks through the faithful si_vec_dot_q4_K dp4a
+// path the routed experts and dense FFN already use — ~47% fewer weight bytes/token on
+// the always-on shared expert (40 layers). One 4-warp block per gate/up output row
+// (F rows, K=H -> NB=H/256 super-blocks, 16 lanes/super-block); one warp per down output
+// row (H rows, K=F -> F/256 super-blocks, two half-warps stride alternate super-blocks).
+// The Q8_1 activation buffer (hn / h) is unchanged.
+template <int H, int F>
+__global__ void shared_gate_up_q4k_mmvq_kernel(
+    const si_block_q8_1* __restrict__ vy, const unsigned char* __restrict__ gate_q,
+    const unsigned char* __restrict__ up_q, const float* __restrict__ dw,
+    float* __restrict__ h_scratch) {
+    constexpr int NW = 4, WS = 32, NB = H >> 8;
+    const int f = blockIdx.x, lane = threadIdx.x & 31, warp = threadIdx.x >> 5, tid = threadIdx.x;
+    const int kbx0 = tid >> 4, kqs = 2 * (tid & 15);
+    const si_block_q4_K* g_row = (const si_block_q4_K*)(gate_q + (size_t)f * NB * 144);
+    const si_block_q4_K* u_row = (const si_block_q4_K*)(up_q   + (size_t)f * NB * 144);
+    float tg = 0.f, tu = 0.f;
+    for (int kbx = kbx0; kbx < NB; kbx += 8) {
+        tg += si_vec_dot_q4_K(g_row + kbx, vy + (size_t)kbx * 8, kqs);
+        tu += si_vec_dot_q4_K(u_row + kbx, vy + (size_t)kbx * 8, kqs);
+    }
+    __shared__ float sg[NW - 1][WS], su[NW - 1][WS];
+    if (warp > 0) { sg[warp - 1][lane] = tg; su[warp - 1][lane] = tu; }
+    __syncthreads();
+    if (warp > 0) return;
+    #pragma unroll
+    for (int l = 0; l < NW - 1; l++) { tg += sg[l][lane]; tu += su[l][lane]; }
+    #pragma unroll
+    for (int m = 16; m > 0; m >>= 1) { tg += __shfl_xor_sync(0xffffffff, tg, m); tu += __shfl_xor_sync(0xffffffff, tu, m); }
+    if (lane == 0) {
+        float w = dw ? __ldg(dw) : 1.f;
+        h_scratch[f] = w * q4kf_silu(tg) * tu;
+    }
+}
+
+template <int H, int F, bool ACCUM>
+__global__ void shared_down_q4k_mmvq_kernel(
+    const si_block_q8_1* __restrict__ hq8, const unsigned char* __restrict__ down_q,
+    __nv_bfloat16* __restrict__ out) {
+    const int h = blockIdx.x * WPB + (threadIdx.x >> 5), lane = threadIdx.x & 31;
+    if (h >= H) return;
+    constexpr int NB = F >> 8;
+    const si_block_q4_K* d_row = (const si_block_q4_K*)(down_q + (size_t)h * NB * 144);
+    const int sub = lane >> 4, kqs = (lane & 15) << 1;
+    float acc = 0.f;
+    #pragma unroll
+    for (int kbx = sub; kbx < NB; kbx += 2)
+        acc += si_vec_dot_q4_K(d_row + kbx, hq8 + (size_t)kbx * 8, kqs);
+    acc = q4kf_wsum(acc);
+    if (lane == 0) {
+        if constexpr (ACCUM) acc += __bfloat162float(out[h]);
+        out[h] = __float2bfloat16(acc);
+    }
+}
+
 template <int H, int F, int TOPK>
 __global__ void gate_up_mmvq2_qwen_kernel(
     const si_block_q8_1* __restrict__ vy, const unsigned char* __restrict__ gate_q,
@@ -1540,6 +1597,41 @@ void launch_shared_expert_q8_mmvq(
             reinterpret_cast<__nv_bfloat16*>(output));
     }
 }
+
+// Qwen3.6 UD shared expert with the gate/up/down weights requantized Q8_0->Q4_K at load.
+// Identical control flow to launch_shared_expert_q8_mmvq (reuse FNQ Q8_1(hn), gate/up ->
+// silu*up -> quant_h Q8_1 -> down), only the weight dot changes to the Q4_K dp4a path.
+void launch_shared_expert_q4k_mmvq(
+    const void* input, const void* input_q8,
+    const void* gate_q, const void* up_q, const void* down_q,
+    const float* dw, void* output, float* h_scratch, void* h_q8_buf,
+    int hidden, int ffn, cudaStream_t stream, bool accum = false) {
+    const si_block_q8_1* vy;
+    si_block_q8_1* qbuf = reinterpret_cast<si_block_q8_1*>(h_q8_buf);
+    if (input_q8) {
+        vy = reinterpret_cast<const si_block_q8_1*>(input_q8);
+    } else {
+        si_quant_bf16_q8_1<<<(hidden >> 5), 32, 0, stream>>>(
+            reinterpret_cast<const __nv_bfloat16*>(input), qbuf, hidden);
+        vy = qbuf;
+    }
+    shared_gate_up_q4k_mmvq_kernel<2048, 512><<<ffn, 4 * 32, 0, stream>>>(
+        vy, reinterpret_cast<const unsigned char*>(gate_q),
+        reinterpret_cast<const unsigned char*>(up_q), dw, h_scratch);
+    const int nqb = ffn >> 5;
+    quant_h_q8_1_kernel<<<(nqb + 7) / 8, 8 * 32, 0, stream>>>(
+        h_scratch, qbuf, nqb, 0);
+    dim3 dn((hidden + WPB - 1) / WPB);
+    if (accum) {
+        shared_down_q4k_mmvq_kernel<2048, 512, true><<<dn, WPB * 32, 0, stream>>>(
+            qbuf, reinterpret_cast<const unsigned char*>(down_q),
+            reinterpret_cast<__nv_bfloat16*>(output));
+    } else {
+        shared_down_q4k_mmvq_kernel<2048, 512, false><<<dn, WPB * 32, 0, stream>>>(
+            qbuf, reinterpret_cast<const unsigned char*>(down_q),
+            reinterpret_cast<__nv_bfloat16*>(output));
+    }
+}
 #endif
 
 } // namespace kernels

kernels/include/sparkinfer/kernels/moe.h

@@ -86,4 +86,12 @@ void launch_shared_expert_q8_mmvq(
     int hidden, int ffn, cudaStream_t stream = nullptr,
     bool accum = false);
 
+// Same as launch_shared_expert_q8_mmvq but for gate/up/down requantized Q8_0->Q4_K at load.
+void launch_shared_expert_q4k_mmvq(
+    const void* input, const void* input_q8,
+    const void* gate_q, const void* up_q, const void* down_q,
+    const float* dw, void* output, float* h_scratch, void* h_q8_buf,
+    int hidden, int ffn, cudaStream_t stream = nullptr,
+    bool accum = false);
+
 }} // namespace sparkinfer::kernels

runtime/src/models/qwen35.cpp

@@ -729,6 +729,9 @@ int Qwen35Model::forward_token(int token_id, int position) {
 
         const bool qmoe = w.shared_gate_q && w.shared_up_q && w.shared_down_q
                        && w.shared_gate_qtype == 8 && c.hidden == 2048 && c.moe_ffn == 512;
+        // Same shape, gate/up/down requantized Q8_0->Q4_K at load (Qwen3.6 UD shared expert).
+        const bool qmoe_q4k = w.shared_gate_q && w.shared_up_q && w.shared_down_q
+                       && w.shared_gate_qtype == 12 && c.hidden == 2048 && c.moe_ffn == 512;
         const bool shexp_pipelined = (c.n_shared > 0) && s.gguf && s.use_shexp_pipe;
         if (shexp_pipelined) {
             cudaEventRecord(s.ev_pipe_fork, st);
@@ -781,6 +784,12 @@ int Qwen35Model::forward_token(int token_id, int position) {
                     w.shared_gate_q, w.shared_up_q, w.shared_down_q,
                     w.shared_gate_inp ? s.d_shared_w : nullptr,
                     s.shared, s.sx_h, s.sx_q8, H, c.moe_ffn, s.stream_k, false);
+            } else if (qmoe_q4k) {
+                kernels::launch_shared_expert_q4k_mmvq(
+                    s.hn, fnq ? s.aq81 : nullptr,
+                    w.shared_gate_q, w.shared_up_q, w.shared_down_q,
+                    w.shared_gate_inp ? s.d_shared_w : nullptr,
+                    s.shared, s.sx_h, s.sx_q8, H, c.moe_ffn, s.stream_k, false);
             } else {
                 kernels::launch_gemv(s.hn, w.shared_gate, s.sh_gate, c.moe_ffn, H, s.stream_k);
                 kernels::launch_gemv(s.hn, w.shared_up,   s.sh_up,   c.moe_ffn, H, s.stream_v);
@@ -886,13 +895,15 @@ int Qwen35Model::forward_token(int token_id, int position) {
                 }
             }
             if (s.gguf) {
-                if (qmoe) {
+                if (qmoe || qmoe_q4k) {
                     // MoE already wrote routed; safe to accum shared down into it (opt-in).
                     static int shexp_accum = -1;
                     if (shexp_accum < 0) { const char* e = getenv("SPARKINFER_SHEXP_ACCUM");
                         shexp_accum = (e && e[0] == '1') ? 1 : 0; }
                     const bool sx_accum = shexp_accum != 0;
-                    kernels::launch_shared_expert_q8_mmvq(
+                    auto shexp_launch = qmoe_q4k ? kernels::launch_shared_expert_q4k_mmvq
+                                                 : kernels::launch_shared_expert_q8_mmvq;
+                    shexp_launch(
                         s.hn, fnq ? s.aq81 : nullptr,
                         w.shared_gate_q, w.shared_up_q, w.shared_down_q,
                         w.shared_gate_inp ? s.d_shared_w : nullptr,
@@ -1157,8 +1168,8 @@ bool Qwen35Model::load_gguf(const std::string& path) {
     // requant; qtype flips to 12 on success.
     auto dev_quant_requant_q4k = [&](const std::string& name, int& qtype, bool req) -> const void* {
         const void* q6 = dev_quant(name, qtype);
-        if (!req || (qtype != 14 && qtype != 8) || !q6) return q6;
-        const int src_type = qtype;            // 14 (Q6_K) or 8 (Q8_0) -> Q4_K
+        if (!req || (qtype != 14 && qtype != 8 && qtype != 13) || !q6) return q6;
+        const int src_type = qtype;            // 14 (Q6_K), 13 (Q5_K) or 8 (Q8_0) -> Q4_K
         const GGUFTensor* t = g.tensor(name);
         const long nv = t->n_values;
         if (nv % 256 != 0) return q6;
@@ -1237,6 +1248,20 @@ bool Qwen35Model::load_gguf(const std::string& path) {
     // (~+3.3% decode at short context, gate-passing). On by default for the Qwen3.6
     // fingerprint; a no-op on the dense Qwythos path (which uses its own qkv default).
     const bool q36_ud_requant_default = is_qwen35_or_qwen36_hybrid_moe(g);
+    // Qwen3.6-35B-A3B UD ships the always-on shared-expert FFN (ffn_{gate,up,down}_shexp)
+    // as Q8_0. It runs on every token across all 40 layers (~13% of decode), so requantizing
+    // it Q8_0->Q4_K at load (~47% fewer weight bytes) is a broad short-context decode win.
+    // Gated to the exact shared-expert shape the Q4_K shared kernels specialize (<2048,512>);
+    // default on for the Qwen3.6 fingerprint, no-op elsewhere. SPARKINFER_SHEXP_REQUANT_Q4K=0
+    // restores the native Q8_0 shared-expert reads.
+    const bool req_shexp_q4 = env_enabled("SPARKINFER_SHEXP_REQUANT_Q4K",
+        q36_ud_requant_default && H == 2048 && c.moe_ffn == 512);
+    // Qwen3.6 routed experts keep their down projection at Q5_K (the gate/up are already Q4_K).
+    // Only the top-8 experts' down runs per token, but that read sits in the same decode-critical
+    // FFN stage as the shared expert, so requantizing ffn_down_exps Q5_K->Q4_K at load stacks a
+    // further byte reduction on the same stage. It reuses the existing Q4_K routed-down MMVQ
+    // dispatch (no new kernel). Default on for the Qwen3.6 fingerprint; =0 keeps native Q5_K.
+    const bool req_rdown_q4 = env_enabled("SPARKINFER_ROUTED_DOWN_REQUANT_Q4K", q36_ud_requant_default);
     const char* attn_env = getenv("SPARKINFER_ATTN_REQUANT_Q4K");
     const std::string attn_requant_mode =
         attn_env ? std::string(attn_env)
@@ -1448,7 +1473,7 @@ bool Qwen35Model::load_gguf(const std::string& path) {
             }
             w.gate_q = dev_quant(b + "ffn_gate_exps.weight", w.gate_qtype);   // kept quantized
             w.up_q   = dev_quant(b + "ffn_up_exps.weight",   w.up_qtype);
-            w.down_q = dev_quant(b + "ffn_down_exps.weight", w.down_qtype);
+            w.down_q = dev_quant_requant_q4k(b + "ffn_down_exps.weight", w.down_qtype, req_rdown_q4);
             if (s.cfg.n_shared > 0) {
             if (!expect_dims(b + "ffn_gate_shexp.weight", {H, c.moe_ffn}) ||
                 !expect_dims(b + "ffn_up_shexp.weight", {H, c.moe_ffn}) ||
@@ -1459,12 +1484,12 @@ bool Qwen35Model::load_gguf(const std::string& path) {
             const bool qmoe = []{ const char* a = getenv("SPARKINFER_QMOE");
                                    return !(a && a[0] == '0'); }();
             if (qmoe) {
-                w.shared_gate_q = dev_quant(b + "ffn_gate_shexp.weight", w.shared_gate_qtype);
-                w.shared_up_q   = dev_quant(b + "ffn_up_shexp.weight",   w.shared_up_qtype);
-                w.shared_down_q = dev_quant(b + "ffn_down_shexp.weight", w.shared_down_qtype);
+                w.shared_gate_q = dev_quant_requant_q4k(b + "ffn_gate_shexp.weight", w.shared_gate_qtype, req_shexp_q4);
+                w.shared_up_q   = dev_quant_requant_q4k(b + "ffn_up_shexp.weight",   w.shared_up_qtype,   req_shexp_q4);
+                w.shared_down_q = dev_quant_requant_q4k(b + "ffn_down_shexp.weight", w.shared_down_qtype, req_shexp_q4);
             }
             if (!qmoe || !w.shared_gate_q || !w.shared_up_q || !w.shared_down_q ||
-                w.shared_gate_qtype != 8) {
+                (w.shared_gate_qtype != 8 && w.shared_gate_qtype != 12)) {
                 w.shared_gate = dense(b + "ffn_gate_shexp.weight", false);
                 w.shared_up   = dense(b + "ffn_up_shexp.weight", false);
                 w.shared_down = dense(b + "ffn_down_shexp.weight", false);