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

kernels/csrc/cuda/moe/expert_ffn_q4k.cu

@@ -588,6 +588,65 @@ __global__ void shared_down_q8_mmvq_kernel(
     }
 }
 
+// ---- Qwen3.6 shared expert (Q4_K gate/up/down, F=512) -------------------------
+// Byte-for-byte the same shared-expert control flow as the Q8_0 kernels above
+// (FNQ Q8_1(hn) -> gate/up SiLU -> quant_h Q8_1 -> down), only the weight dot
+// swaps to the faithful int8 dp4a Q4_K path (si_vec_dot_q4_K). Fed by a load-time
+// Q8_0->Q4_K requant of ffn_{gate,up,down}_shexp, so decode reads ~47% fewer weight
+// bytes on the always-on shared FFN. gate/up rows are Q4_K [F,H] (H/256 super-blocks
+// per row); one 4-warp block per output f with the vdr=2 (kbx,kqs) indexing.
+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;
+    const int 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;
+    }
+}
+
+// down: Q4_K [H,F] (F/256 super-blocks per row). one warp per output row h, striding
+// the vdr=2 positions of the F/256 super-blocks against the pre-quantized Q8_1(h).
+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, WORK = NB * 16;
+    const si_block_q4_K* drow = (const si_block_q4_K*)(down_q + (size_t)h * NB * 144);
+    float acc = 0.f;
+    for (int wi = lane; wi < WORK; wi += 32) {
+        const int kbx = wi >> 4, kqs = (wi & 15) << 1;
+        acc += si_vec_dot_q4_K(drow + 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 +1599,41 @@ void launch_shared_expert_q8_mmvq(
             reinterpret_cast<__nv_bfloat16*>(output));
     }
 }
+
+// Qwen3.6 UD shared expert: Q4_K weights + int8 dp4a MMVQ. Same pipeline as
+// launch_shared_expert_q8_mmvq (reuse FNQ Q8_1(hn) for gate/up, overwrite h_q8_buf
+// with Q8_1(h) for down); fed by the load-time Q8_0->Q4_K shared-expert requant.
+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,14 @@ void launch_shared_expert_q8_mmvq(
     int hidden, int ffn, cudaStream_t stream = nullptr,
     bool accum = false);
 
+// Qwen3.6 UD shared expert: Q4_K gate/up/down via int8 dp4a MMVQ (same pipeline as
+// the Q8_0 variant). Fed by a load-time Q8_0->Q4_K requant of ffn_{gate,up,down}_shexp,
+// so decode reads ~47% fewer weight bytes on the always-on shared FFN.
+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

@@ -728,7 +728,8 @@ int Qwen35Model::forward_token(int token_id, int position) {
             kernels::launch_add_rmsnorm2(s.x, s.ao, w.post_attn_norm, s.h, s.hn, 1, H, c.rms_eps, st);
 
         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;
+                       && (w.shared_gate_qtype == 8 || 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);
@@ -776,11 +777,18 @@ int Qwen35Model::forward_token(int token_id, int position) {
                 // races MoE (shared finishes first, MoE overwrites routed). Always write s.shared;
                 // fold happens after both complete. SPARKINFER_SHEXP_ACCUM=1 only applies on the
                 // non-pipelined path where MoE has already landed in routed.
-                kernels::launch_shared_expert_q8_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);
+                if (w.shared_gate_qtype == 12)   // Q8_0->Q4_K requant (Qwen3.6): new Q4_K shared kernels
+                    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_shared_expert_q8_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);
@@ -892,12 +900,20 @@ int Qwen35Model::forward_token(int token_id, int position) {
                     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(
-                        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,
-                        sx_accum ? s.routed : s.shared, s.mf_h, s.aq81, H, c.moe_ffn, st,
-                        sx_accum);
+                    if (w.shared_gate_qtype == 12)   // Q8_0->Q4_K requant (Qwen3.6): new Q4_K shared kernels
+                        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,
+                            sx_accum ? s.routed : s.shared, s.mf_h, s.aq81, H, c.moe_ffn, st,
+                            sx_accum);
+                    else
+                        kernels::launch_shared_expert_q8_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,
+                            sx_accum ? s.routed : s.shared, s.mf_h, s.aq81, H, c.moe_ffn, st,
+                            sx_accum);
                     if (sx_accum) {
                         if (fnq)
                             kernels::launch_add_rmsnorm2_q8(s.h, s.routed, nextnorm, s.x, s.xn, s.aq81, H, c.rms_eps, st);
@@ -1157,8 +1173,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 != 13 && qtype != 8) || !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;
@@ -1185,6 +1201,30 @@ bool Qwen35Model::load_gguf(const std::string& path) {
         if (req < 0) { const char* e = getenv("SPARKINFER_DOWN_REQUANT_Q4K"); req = (e && e[0] == '0') ? 0 : 1; }
         return dev_quant_requant_q4k(name, qtype, req != 0);
     };
+    // Qwen3.6-35B-A3B UD ships the always-on shared expert (ffn_{gate,up,down}_shexp)
+    // as Q8_0 (34 B/block) on a read that fires every layer, every token. Requantize it
+    // to Q4_K at load (~47% fewer weight bytes) so it decodes through the Q4_K shared-expert
+    // MMVQ kernels. On by default for the Qwen3.6 fingerprint; SPARKINFER_SHEXP_REQUANT_Q4K=0
+    // restores native Q8_0. Strict no-op on non-matching models.
+    const bool q36_shexp_requant = is_qwen35_or_qwen36_hybrid_moe(g) && [] {
+        const char* e = getenv("SPARKINFER_SHEXP_REQUANT_Q4K");
+        return !(e && e[0] == '0');
+    }();
+    auto dev_quant_shexp = [&](const std::string& name, int& qtype) -> const void* {
+        return dev_quant_requant_q4k(name, qtype, q36_shexp_requant);
+    };
+    // Routed-expert down (ffn_down_exps) stays at Q5_K in the UD Q4_K_M while gate/up are
+    // already Q4_K. That down read sits in the same decode-critical FFN stage, so requantize
+    // Q5_K->Q4_K at load and route it through the pre-existing Q4_K routed-down MMVQ dispatch
+    // (down_type == 12). On by default for the Qwen3.6 fingerprint; SPARKINFER_ROUTED_DOWN_REQUANT_Q4K=0
+    // restores native Q5_K reads.
+    const bool q36_routed_down_requant = is_qwen35_or_qwen36_hybrid_moe(g) && [] {
+        const char* e = getenv("SPARKINFER_ROUTED_DOWN_REQUANT_Q4K");
+        return !(e && e[0] == '0');
+    }();
+    auto dev_quant_routed_down = [&](const std::string& name, int& qtype) -> const void* {
+        return dev_quant_requant_q4k(name, qtype, q36_routed_down_requant);
+    };
     // dense weight -> bf16 (optionally transpose [out,in] -> [in,out])
     auto dense = [&](const std::string& name, bool transpose) -> const void* {
         const GGUFTensor* t = g.tensor(name);
@@ -1448,7 +1488,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_routed_down(b + "ffn_down_exps.weight", w.down_qtype);  // Q5_K->Q4_K (Qwen3.6)
             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 +1499,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_shexp(b + "ffn_gate_shexp.weight", w.shared_gate_qtype);
+                w.shared_up_q   = dev_quant_shexp(b + "ffn_up_shexp.weight",   w.shared_up_qtype);
+                w.shared_down_q = dev_quant_shexp(b + "ffn_down_shexp.weight", w.shared_down_qtype);
             }
             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);