fix(gemv): instantiate mmvq_q4k_rows for K=6656 (Muse Glimmer)

kernels/csrc/cuda/gemm/gemv.cu

@@ -1337,6 +1337,14 @@ template __global__ void si_mmvq_q4k_rows_exact_kernel<float, 8, 6, SI_Q4K_OROWS
     const si_block_q8_1*, const unsigned char*, float*, int, int);
 template __global__ void si_mmvq_q4k_rows_exact_kernel<float, 16, 6, SI_Q4K_OROWS>(
     const si_block_q8_1*, const unsigned char*, float*, int, int);
+template __global__ void si_mmvq_q4k_rows_exact_kernel<__nv_bfloat16, 26, 8, SI_Q4K_OROWS>(
+    const si_block_q8_1*, const unsigned char*, __nv_bfloat16*, int, int);
+template __global__ void si_mmvq_q4k_rows_exact_kernel<__nv_bfloat16, 26, 6, SI_Q4K_OROWS>(
+    const si_block_q8_1*, const unsigned char*, __nv_bfloat16*, int, int);
+template __global__ void si_mmvq_q4k_rows_exact_kernel<float, 26, 8, SI_Q4K_OROWS>(
+    const si_block_q8_1*, const unsigned char*, float*, int, int);
+template __global__ void si_mmvq_q4k_rows_exact_kernel<float, 26, 6, SI_Q4K_OROWS>(
+    const si_block_q8_1*, const unsigned char*, float*, int, int);
 #endif
 // One block per row index: warps 0-3 -> qkv[row], warps 4-7 -> z[row], keeping vy hot
 // in L2 across both when row < min(n_qkv, n_z). Grid = max(n_qkv, n_z).
@@ -2471,7 +2479,7 @@ bool launch_mmvq_q4k_rows(const void* q81, const void* W, void* y,
     // the token loop, and DFlash speculation could never engage on Qwen3.8 at all. 20 (5120) and
     // 24 (6144) are added for exactly that.
     if (M < 1 || M > 8 || N < 1) return false;
-    if (K != 2048 && K != 4096 && K != 5120 && K != 6144) return false;
+    if (K != 2048 && K != 4096 && K != 5120 && K != 6144 && K != 6656) return false;
     // Dispatch the tightest instantiated row width: MMAX bounds tmp[]/partial[] and the
     // number of predicated row bodies, so a 6-row block should not pay an 8-row footprint.
     const auto* q = reinterpret_cast<const si_block_q8_1*>(q81);
@@ -2486,7 +2494,8 @@ bool launch_mmvq_q4k_rows(const void* q81, const void* W, void* y,
     if      (K == 2048) SI_Q4K_ROWS_DISPATCH(8);
     else if (K == 4096) SI_Q4K_ROWS_DISPATCH(16);
     else if (K == 5120) SI_Q4K_ROWS_DISPATCH(20);
-    else                SI_Q4K_ROWS_DISPATCH(24);
+    else if (K == 6144) SI_Q4K_ROWS_DISPATCH(24);
+    else                SI_Q4K_ROWS_DISPATCH(26);
     #undef SI_Q4K_ROWS_DISPATCH
     return true;
 }
@@ -2541,7 +2550,7 @@ bool launch_mmvq_rows_f32(int qtype, const void* q81, const void* W, float* y,
     // Same instantiated-width limit as launch_mmvq_q4k_rows above, and the same consequence:
     // this is the LM-head path, so K=5120 (Qwen3.8's hidden) refused here made the verify decline
     // AFTER every layer had already succeeded -- "unsupported LM head type=12 H=5120".
-    if (qtype == 12 && (K == 2048 || K == 4096 || K == 5120 || K == 6144)) {
+    if (qtype == 12 && (K == 2048 || K == 4096 || K == 5120 || K == 6144 || K == 6656)) {
         const int grid = (N + SI_Q4K_OROWS - 1) / SI_Q4K_OROWS;
         #define SI_Q4K_ROWS_F32_DISPATCH(KB) \
             do { \
@@ -2551,7 +2560,8 @@ bool launch_mmvq_rows_f32(int qtype, const void* q81, const void* W, float* y,
         if      (K == 2048) SI_Q4K_ROWS_F32_DISPATCH(8);
         else if (K == 4096) SI_Q4K_ROWS_F32_DISPATCH(16);
         else if (K == 5120) SI_Q4K_ROWS_F32_DISPATCH(20);
-        else                SI_Q4K_ROWS_F32_DISPATCH(24);
+        else if (K == 6144) SI_Q4K_ROWS_F32_DISPATCH(24);
+        else                SI_Q4K_ROWS_F32_DISPATCH(26);
         #undef SI_Q4K_ROWS_F32_DISPATCH
         return true;
     }

runtime/CMakeLists.txt

@@ -97,6 +97,11 @@ if(BUILD_EXAMPLES)
     add_executable(nvfp4_gemm_check examples/nvfp4_gemm_check.cpp)
     target_include_directories(nvfp4_gemm_check PRIVATE include)
     target_link_libraries(nvfp4_gemm_check PRIVATE sparkinfer_runtime CUDA::cudart)
+    # Model-free check that the batched-row Q4_K path matches serial decode calls at K=6656
+    # (Muse Glimmer's hidden size), which the batched-row launcher did not previously support.
+    add_executable(mmvq_q4k_rows_6656_check examples/mmvq_q4k_rows_6656_check.cpp)
+    target_include_directories(mmvq_q4k_rows_6656_check PRIVATE include)
+    target_link_libraries(mmvq_q4k_rows_6656_check PRIVATE sparkinfer_runtime CUDA::cudart)
 
     add_executable(qwen3_gguf_prefill_check examples/qwen3_gguf_prefill_check.cpp)
     target_include_directories(qwen3_gguf_prefill_check PRIVATE include)

runtime/examples/mmvq_q4k_rows_6656_check.cpp

@@ -0,0 +1,96 @@
+#include "sparkinfer/kernels/gemm.h"
+#include "sparkinfer/kernels/quant.h"
+
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+namespace {
+uint16_t bf16(float x) {
+    uint32_t u;
+    std::memcpy(&u, &x, sizeof(u));
+    return static_cast<uint16_t>(u >> 16);
+}
+template <class T> T* device_copy(const std::vector<T>& h) {
+    T* d = nullptr;
+    cudaMalloc(&d, h.size() * sizeof(T));
+    cudaMemcpy(d, h.data(), h.size() * sizeof(T), cudaMemcpyHostToDevice);
+    return d;
+}
+template <class T> bool equal_device(const T* a, const T* b, size_t n, const char* what) {
+    std::vector<T> ha(n), hb(n);
+    cudaMemcpy(ha.data(), a, n * sizeof(T), cudaMemcpyDeviceToHost);
+    cudaMemcpy(hb.data(), b, n * sizeof(T), cudaMemcpyDeviceToHost);
+    if (ha == hb) return true;
+    size_t i = 0;
+    while (i < n && ha[i] == hb[i]) i++;
+    std::printf("[FAIL] %s differs at %zu (%d vs %d)\n", what, i, (int)ha[i], (int)hb[i]);
+    return false;
+}
+}
+
+int main() {
+    int ndev = 0;
+    if (cudaGetDeviceCount(&ndev) != cudaSuccess || ndev == 0) {
+        std::printf("[SKIP] no CUDA device\n");
+        return 0;
+    }
+    constexpr int N = 4, MH = 6656, MN = 8192, SB = MH / 256; // SB=26, Muse Glimmer's hidden size
+    auto make_bf16 = [](size_t n, int salt, float scale) {
+        std::vector<uint16_t> h(n);
+        for (size_t i = 0; i < n; i++) {
+            const int v = (int)((i * 1103515245u + 12345u + salt) % 2001u) - 1000;
+            h[i] = bf16(scale * v / 1000.f);
+        }
+        return h;
+    };
+    auto hact = make_bf16((size_t)N * MH, 11, 0.2f);
+    uint16_t* dact = device_copy(hact);
+    const size_t q81_row = sparkinfer::kernels::llama_q8_1_bytes(MH);
+    void* dq81 = nullptr;
+    cudaMalloc(&dq81, (size_t)N * q81_row);
+    sparkinfer::kernels::launch_quantize_q8_1_rows(dact, dq81, MH, N, MH);
+
+    std::vector<unsigned char> hw((size_t)MN * SB * 144);
+    for (int row = 0; row < MN; row++) for (int sb = 0; sb < SB; sb++) {
+        unsigned char* p = hw.data() + ((size_t)row * SB + sb) * 144;
+        p[0] = 0x1f; p[1] = 0x21; p[2] = 0x1f; p[3] = 0x21;
+        for (int i = 4; i < 144; i++) p[i] = (unsigned char)((row * 17 + sb * 29 + i * 13) & 255);
+    }
+    unsigned char* dw = device_copy(hw);
+    uint16_t *ym = nullptr, *ys = nullptr;
+    cudaMalloc(&ym, (size_t)N * MN * 2); cudaMalloc(&ys, (size_t)N * MN * 2);
+
+    bool ok = sparkinfer::kernels::launch_mmvq_q4k_rows(dq81, dw, ym, N, MN, MH);
+    if (!ok) { std::printf("[FAIL] launch_mmvq_q4k_rows returned false at K=%d\n", MH); return 1; }
+    for (int m = 0; m < N; m++)
+        sparkinfer::kernels::launch_mmvq_q4k((const char*)dq81 + (size_t)m * q81_row,
+                                              dw, ys + (size_t)m * MN, MN, MH);
+    cudaDeviceSynchronize();
+    ok = equal_device(ym, ys, (size_t)N * MN, "K=6656 Q4_K exact rows") && ok;
+
+    // Timing: batched rows vs. serial single-row calls
+    cudaEvent_t t0, t1, t2;
+    cudaEventCreate(&t0); cudaEventCreate(&t1); cudaEventCreate(&t2);
+    cudaEventRecord(t0);
+    for (int rep = 0; rep < 200; rep++)
+        sparkinfer::kernels::launch_mmvq_q4k_rows(dq81, dw, ym, N, MN, MH);
+    cudaEventRecord(t1);
+    for (int rep = 0; rep < 200; rep++)
+        for (int m = 0; m < N; m++)
+            sparkinfer::kernels::launch_mmvq_q4k((const char*)dq81 + (size_t)m * q81_row,
+                                                  dw, ys + (size_t)m * MN, MN, MH);
+    cudaEventRecord(t2);
+    cudaEventSynchronize(t2);
+    float ms_batch = 0, ms_serial = 0;
+    cudaEventElapsedTime(&ms_batch, t0, t1);
+    cudaEventElapsedTime(&ms_serial, t1, t2);
+    std::printf("K=6656 rows: batch %.4f ms, %d serial calls %.4f ms, %.2fx\n",
+                ms_batch / 200, N, ms_serial / 200, ms_serial / ms_batch);
+
+    std::printf(ok ? "[PASS] K=6656 batched-row Q4_K matches serial reference\n"
+                    : "[FAIL] K=6656 mismatch\n");
+    return ok ? 0 : 1;
+}