DFlash2: opt-in Q4_0 drafter pack (dtype q4_0, tools/dflash2_pack.py --q4_0): nibble-packed tiled weights, gemm4_kernel on the existing dp4a activations path, df2_q4_tile at load, CPU parity test dflash2_q4_parity (40 of 40 cases within 2.4e-6 of 19)

CMakeLists.txt

@@ -849,6 +849,9 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
   add_executable(dflash2_parity EXCLUDE_FROM_ALL src/core/dflash2_parity.cpp)
   target_include_directories(dflash2_parity PRIVATE ${_strata_gpu_include_directories})
   target_link_libraries(dflash2_parity PRIVATE strata_engine strata_kernels ${_strata_gpu_runtime_target})
+  add_executable(dflash2_q4_parity EXCLUDE_FROM_ALL src/core/dflash2_q4_parity.cpp)
+  target_include_directories(dflash2_q4_parity PRIVATE ${_strata_gpu_include_directories})
+  target_link_libraries(dflash2_q4_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target})
   add_executable(dflash2_gemm_bench EXCLUDE_FROM_ALL src/core/dflash2_gemm_bench.cpp)
   target_include_directories(dflash2_gemm_bench PRIVATE ${_strata_gpu_include_directories})
   target_link_libraries(dflash2_gemm_bench PRIVATE strata_kernels ${_strata_gpu_runtime_target})

include/strata/core/dflash2.hpp

@@ -101,6 +101,7 @@ private:
     std::vector<Tensor> tensors_;
     uint8_t* blob_ = nullptr;
     uint64_t vram_ = 0;
+    bool q4_ = false;   // the pack's matrices are Q4_0 (tools/dflash2_pack.py --q4_0)
     int N_ = 2560, L_ = 5, NH_ = 32, NKV_ = 8, HD_ = 128, I_ = 7680, window_ = 2048, group_ = 16, rank_ = 256, mask_token_ = 0;
     float theta_ = 10000.0f, eps_ = 1e-6f;
     std::vector<int> taps_;

include/strata/kernels/dflash2.hpp

@@ -52,7 +52,8 @@ void df2_attention(const float* q, const float* kring, const float* vring, const
 void df2_gemm_variant(int v);
 /// Bench only: attention variants (bit 0 skips Q K^T, bit 1 skips P V).
 void df2_attention_variant(int v);
-struct Df2Mat { const void* w; float* y; int n_out; };
+/// q4: the weights are the Q4_0 nibble plane (n_out * n_in / 2 bytes, see gemm4_kernel in dflash2.cu) then the fp16 scale plane (df2_q4_tile); all matrices of a launch share the format.
+struct Df2Mat { const void* w; float* y; int n_out; bool q4 = false; };
 void df2_gemm(const Df2Mat* mats, int n, const void* xq, const float* xd, int n_in, int ncols, void* stream);
 /// Planar quantization of rows x[ncols][K].
 void df2_quant_rows(const float* x, int K, int ncols, void* xq, float* xd, void* stream);
@@ -64,6 +65,8 @@ void df2_conv_quant(const float* x, const float* dyn, int dyn_stride, const floa
 void df2_swiglu_quant(const float* gate, const float* up, int rows, int I, void* xq, float* xd, void* stream);
 /// GGUF Q8_0 blocks (34 bytes) of a matrix with rows of n_in -> the tiled qs plane and the tiled d plane (count = rows * n_in / 32 blocks).
 void df2_q8_tile(const void* blocks, void* qs, void* dpl, int n_in, int64_t count, void* stream);
+/// ggml Q4_0 blocks (18 bytes) of a matrix with rows of n_in -> the tiled nibble plane (n_out * n_in / 2 bytes) and the tiled d plane (Df2Mat::q4).
+void df2_q4_tile(const void* blocks, void* qs, void* dpl, int n_in, int64_t count, void* stream);
 
 /// h += conv(x) (the finish convolution with the residual), then the RMS norm of the new h (weights gamma) -> xn (optional) and the planar quantization.
 /// rows <= 8, H <= 3072 (H a multiple of 32).

src/core/dflash2.cpp

@@ -187,6 +187,11 @@ bool Dflash2Drafter::load(const std::string& path, int device, std::string& err)
             if (cudaMemcpy(tmp, stage.data(), t.bytes, cudaMemcpyHostToDevice) != cudaSuccess) { err = "dflash2: weight upload failed"; return false; }
             kernels::df2_q8_tile(tmp, blob_ + t.off, blob_ + t.off + (uint64_t) t.rows * (uint64_t) t.cols, (int) t.cols, t.rows * t.cols / 32, nullptr);
             if (cudaDeviceSynchronize() != cudaSuccess) { err = "dflash2: weight conversion failed"; return false; }
+        } else if (t.dtype == "q4_0") {   // GGML Q4_0 blocks -> the tiled nibble plane at the tensor's offset, then the d plane (same total bytes)
+            if (cudaMemcpy(tmp, stage.data(), t.bytes, cudaMemcpyHostToDevice) != cudaSuccess) { err = "dflash2: weight upload failed"; return false; }
+            kernels::df2_q4_tile(tmp, blob_ + t.off, blob_ + t.off + (uint64_t) t.rows * (uint64_t) t.cols / 2, (int) t.cols, t.rows * t.cols / 32, nullptr);
+            if (cudaDeviceSynchronize() != cudaSuccess) { err = "dflash2: weight conversion failed"; return false; }
+            q4_ = true;
         } else if (cudaMemcpy(blob_ + t.off, stage.data(), t.bytes, cudaMemcpyHostToDevice) != cudaSuccess) {
             err = "dflash2: weight upload failed";
             return false;
@@ -328,8 +333,8 @@ bool Dflash2Drafter::bind(const WeightTable& wt, const NativeHead* head, const s
         head_rows = (int64_t) vocab.size();
     }
     if (!bind_raw(embed, hw, head->type(), head_rows, vocab, err)) return false;
-    std::fprintf(stderr, "strata dflash2: %lld-token draft head (%.1f MiB), drafter %.0f MiB of VRAM\n", (long long) n_dvocab_,
-                 (double) (n_dvocab_ * row_bytes) / 1048576.0, (double) vram_ / 1048576.0);
+    std::fprintf(stderr, "strata dflash2: %lld-token draft head (%.1f MiB), drafter %.0f MiB of VRAM (%s weights)\n", (long long) n_dvocab_,
+                 (double) (n_dvocab_ * row_bytes) / 1048576.0, (double) vram_ / 1048576.0, q4_ ? "Q4_0" : "Q8_0");
     return true;
 }
 
@@ -339,7 +344,7 @@ void Dflash2Drafter::gemm(std::initializer_list<const char*> names, std::initial
     auto o = outs.begin();
     for (const char* name : names) {
         const Tensor& t = find(name);
-        mats[n] = kernels::Df2Mat{blob_ + t.off, *o++, (int) t.rows};
+        mats[n] = kernels::Df2Mat{blob_ + t.off, *o++, (int) t.rows, t.dtype == "q4_0"};
         if (n == 0) n_in = (int) t.cols;
         ++n;
     }

src/core/dflash2_q4_parity.cpp

@@ -0,0 +1,164 @@
+// src/core/dflash2_q4_parity.cpp - the drafter's Q4_0 GEMM (kernels/dflash2.hpp df2_gemm with Df2Mat::q4, after df2_q4_tile) against a CPU reference.
+//
+//     dflash2_q4_parity
+//
+// Random float matrices are quantized to ggml Q4_0 blocks on the CPU (the reference of tools/dflash2_pack.py quant_q4_0), tiled on the device, and multiplied by
+// random planar-int8 activations (df2_quant_rows).  The reference reads the same blocks and the same activation codes back and sums per block in integers,
+// so the only difference to the kernel is the float order of the sum over blocks.  Shapes: the drafter's (n_in 2560, 4096, 7680, 12800), one, two and three
+// matrices to a launch, 1..8 columns.  Exit 0 when every output is within 2e-5 of the largest value (relative), 1 otherwise.
+#include "strata/kernels/dflash2.hpp"
+
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <random>
+#include <vector>
+
+namespace {
+
+uint16_t f2h(float f) {   // IEEE half, round to nearest even (the d of a block is stored this way)
+    uint32_t x;
+    std::memcpy(&x, &f, 4);
+    const uint32_t sign = (x >> 16) & 0x8000u;
+    int32_t e = (int32_t) ((x >> 23) & 0xff) - 127 + 15;
+    uint32_t m = x & 0x7fffffu;
+    if (e <= 0) {
+        if (e < -10) return (uint16_t) sign;
+        m |= 0x800000u;
+        const int sh = 14 - e;
+        uint32_t r = m >> sh;
+        const uint32_t rem = m & ((1u << sh) - 1), half = 1u << (sh - 1);
+        if (rem > half || (rem == half && (r & 1))) ++r;
+        return (uint16_t) (sign | r);
+    }
+    if (e >= 31) return (uint16_t) (sign | 0x7c00u);
+    uint32_t r = ((uint32_t) e << 10) | (m >> 13);
+    const uint32_t rem = m & 0x1fffu;
+    if (rem > 0x1000u || (rem == 0x1000u && (r & 1))) ++r;
+    return (uint16_t) (sign | r);
+}
+
+float h2f(uint16_t h) {
+    const uint32_t sign = (uint32_t) (h & 0x8000u) << 16;
+    const uint32_t e = (h >> 10) & 0x1f, m = h & 0x3ff;
+    uint32_t x;
+    if (e == 0) {
+        if (m == 0) x = sign;
+        else {
+            int s = 0;
+            uint32_t mm = m;
+            while (!(mm & 0x400)) { mm <<= 1; ++s; }
+            x = sign | ((uint32_t) (127 - 15 - s + 1) << 23) | ((mm & 0x3ff) << 13);
+        }
+    } else if (e == 31) x = sign | 0x7f800000u | (m << 13);
+    else x = sign | ((e - 15 + 127) << 23) | (m << 13);
+    float f;
+    std::memcpy(&f, &x, 4);
+    return f;
+}
+
+// ggml quantize_row_q4_0_ref: 18 bytes a block
+void quant_q4_0(const float* x, int n, uint8_t* out) {
+    for (int b = 0; b < n / 32; ++b) {
+        float amax = 0, mx = 0;
+        for (int i = 0; i < 32; ++i)
+            if (std::fabs(x[b * 32 + i]) > amax) { amax = std::fabs(x[b * 32 + i]); mx = x[b * 32 + i]; }
+        const float d = mx / -8.0f, id = d != 0.0f ? 1.0f / d : 0.0f;
+        const uint16_t dh = f2h(d);
+        out[b * 18] = (uint8_t) (dh & 0xff);
+        out[b * 18 + 1] = (uint8_t) (dh >> 8);
+        for (int i = 0; i < 16; ++i) {
+            const int q0 = std::min(15, (int) (x[b * 32 + i] * id + 8.5f)), q1 = std::min(15, (int) (x[b * 32 + 16 + i] * id + 8.5f));
+            out[b * 18 + 2 + i] = (uint8_t) (q0 | (q1 << 4));
+        }
+    }
+}
+
+struct Case { int n_in; std::vector<int> n_out; };
+
+}  // namespace
+
+int main() {
+    std::mt19937 rng(20261031);
+    std::normal_distribution<float> nd(0.0f, 0.1f);
+    const Case cases[] = {{2560, {64}}, {2560, {64, 32}}, {2560, {4096, 1024, 1024}}, {4096, {2560}}, {7680, {2560}}, {12800, {2560}}, {2560, {7680, 7680}}, {2560, {16}}};
+    int bad = 0, total = 0;
+    for (const Case& cs : cases) {
+        const int n_in = cs.n_in, nb = n_in / 32, nm = (int) cs.n_out.size();
+        std::vector<std::vector<uint8_t>> blocks(nm);
+        std::vector<uint8_t*> dev_w(nm);
+        std::vector<float*> dev_y(nm);
+        for (int m = 0; m < nm; ++m) {
+            const int n_out = cs.n_out[m];
+            std::vector<float> w((size_t) n_out * n_in);
+            for (float& v : w) v = nd(rng);
+            blocks[m].resize((size_t) n_out * nb * 18);
+            for (int r = 0; r < n_out; ++r) quant_q4_0(&w[(size_t) r * n_in], n_in, &blocks[m][(size_t) r * nb * 18]);
+            uint8_t* tmp;
+            cudaMalloc((void**) &tmp, blocks[m].size());
+            cudaMemcpy(tmp, blocks[m].data(), blocks[m].size(), cudaMemcpyHostToDevice);
+            cudaMalloc((void**) &dev_w[m], blocks[m].size());
+            strata::kernels::df2_q4_tile(tmp, dev_w[m], dev_w[m] + (size_t) n_out * n_in / 2, n_in, (int64_t) n_out * nb, nullptr);
+            cudaDeviceSynchronize();
+            cudaFree(tmp);
+            cudaMalloc((void**) &dev_y[m], (size_t) 8 * n_out * 4);
+        }
+        std::vector<float> x((size_t) 8 * n_in);
+        for (float& v : x) v = nd(rng) * 4.0f;
+        float* dx;
+        int8_t* dxq;
+        float* dxd;
+        cudaMalloc((void**) &dx, x.size() * 4);
+        cudaMalloc((void**) &dxq, (size_t) 8 * n_in + 4096);
+        cudaMalloc((void**) &dxd, (size_t) 8 * nb * 4 + 64);
+        cudaMemcpy(dx, x.data(), x.size() * 4, cudaMemcpyHostToDevice);
+        strata::kernels::df2_quant_rows(dx, n_in, 8, dxq, dxd, nullptr);
+        cudaDeviceSynchronize();
+        std::vector<int8_t> xq((size_t) 8 * n_in);
+        std::vector<float> xd((size_t) 8 * nb);
+        cudaMemcpy(xq.data(), dxq, xq.size(), cudaMemcpyDeviceToHost);
+        cudaMemcpy(xd.data(), dxd, xd.size() * 4, cudaMemcpyDeviceToHost);
+        for (int ncols : {1, 2, 3, 5, 8}) {
+            strata::kernels::Df2Mat mats[3];
+            for (int m = 0; m < nm; ++m) mats[m] = strata::kernels::Df2Mat{dev_w[m], dev_y[m], cs.n_out[m], true};
+            strata::kernels::df2_gemm(mats, nm, dxq, dxd, n_in, ncols, nullptr);
+            cudaDeviceSynchronize();
+            double worst = 0, big = 0;
+            for (int m = 0; m < nm; ++m) {
+                const int n_out = cs.n_out[m];
+                std::vector<float> y((size_t) 8 * n_out);
+                cudaMemcpy(y.data(), dev_y[m], y.size() * 4, cudaMemcpyDeviceToHost);
+                for (int c = 0; c < ncols; ++c)
+                    for (int r = 0; r < n_out; ++r) {
+                        double ref = 0;
+                        for (int b = 0; b < nb; ++b) {
+                            const uint8_t* blk = &blocks[m][((size_t) r * nb + b) * 18];
+                            const uint16_t dh = (uint16_t) blk[0] | ((uint16_t) blk[1] << 8);
+                            long isum = 0;
+                            for (int i = 0; i < 16; ++i) {
+                                isum += (long) ((blk[2 + i] & 15) - 8) * xq[(size_t) c * n_in + b * 32 + i];
+                                isum += (long) ((blk[2 + i] >> 4) - 8) * xq[(size_t) c * n_in + b * 32 + 16 + i];
+                            }
+                            ref += (double) h2f(dh) * (double) xd[(size_t) c * nb + b] * (double) isum;
+                        }
+                        const double got = y[(size_t) c * n_out + r];
+                        worst = std::max(worst, std::fabs(got - ref));
+                        big = std::max(big, std::fabs(ref));
+                    }
+            }
+            const bool ok = worst <= 2e-5 * std::max(big, 1e-9);
+            ++total;
+            if (!ok) ++bad;
+            std::printf("n_in %5d mats %d rows %5d cols %d: max |kernel - cpu| %.3g of %.3g (%.2g)  %s\n", n_in, nm, cs.n_out[0], ncols, worst, big, worst / std::max(big, 1e-9), ok ? "ok" : "MISMATCH");
+        }
+        for (int m = 0; m < nm; ++m) { cudaFree(dev_w[m]); cudaFree(dev_y[m]); }
+        cudaFree(dx); cudaFree(dxq); cudaFree(dxd);
+    }
+    std::printf("%s: %d of %d cases match the CPU reference\n", bad == 0 ? "PASS" : "FAIL", total - bad, total);
+    return bad == 0 ? 0 : 1;
+}

src/kernels/cuda/dflash2.cu

@@ -344,6 +344,75 @@ __global__ void __launch_bounds__(WARPS * 32) gemm8_kernel(Df2Segs segs, const i
         }
 }
 
+// ---- Q4_0 variant of the same GEMM (opt-in drafter pack, dtype "q4_0"; a weight is 4 bits, the activations stay int8)
+//
+// Same tiling as the Q8 plane: a group of 8 rows is a unit and a step is 512 values of the reduction axis.  A row's 512 values are 256 bytes in the
+// step (2 KB for the group), lane = block * 2 + half takes 8 bytes = 16 values as two words.  Word 0 holds values 0-3 in its four low nibbles and
+// 4-7 in its four high nibbles, word 1 holds 8-11 and 12-15, so (w & 0x0F0F0F0F) and ((w >> 4) & 0x0F0F0F0F) are the int8 words dp4a wants, as
+// unsigned q = 0..15 (value = (q - 8) * d, ggml Q4_0).  sum((q - 8) x) = sum(q x) - 8 sum(x): the sum of the lane's 16 activations is one more dp4a
+// per word and column, shared by the 8 rows.  The fp16 block scales follow the nibble plane (n_out * n_in / 2 bytes) in the Q8 plane's order.
+template <int WARPS, int UNR = 2>
+__global__ void __launch_bounds__(WARPS * 32) gemm4_kernel(Df2Segs segs, const int8_t* __restrict__ xq, const float* __restrict__ xd, int n_in, int ncols) {
+    constexpr int R = 8;
+    const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
+    const int row_g = (blockIdx.x * WARPS + warp) * R;
+    int si = 0;
+    while (si + 1 < segs.n && row_g >= segs.s[si + 1].row0) ++si;
+    const Df2Seg sg = segs.s[si];
+    const int row0 = row_g - sg.row0;
+    const int nb = n_in / 32, steps = nb / 16;
+    const uint2* wq = reinterpret_cast<const uint2*>(sg.q + (size_t) row0 * (n_in / 2)) + lane;
+    const __half* wd = sg.d + (size_t) row0 * nb + (lane >> 1);
+    const int4* xv = reinterpret_cast<const int4*>(xq) + lane;
+    float acc[R][8];
+#pragma unroll
+    for (int r = 0; r < R; ++r)
+#pragma unroll
+        for (int c = 0; c < 8; ++c) acc[r][c] = 0.0f;
+#pragma unroll UNR
+    for (int st = 0; st < steps; ++st) {
+        int4 w[R];
+        float dw[R];
+#pragma unroll
+        for (int r = 0; r < R; ++r) {
+            const uint2 raw = wq[(size_t) st * 256 + r * 32];   // 8 rows x 32 lanes x 8 bytes = 2 KB a step
+            w[r] = make_int4((int) (raw.x & 0x0F0F0F0Fu), (int) ((raw.x >> 4) & 0x0F0F0F0Fu), (int) (raw.y & 0x0F0F0F0Fu), (int) ((raw.y >> 4) & 0x0F0F0F0Fu));
+            dw[r] = __half2float(wd[(size_t) st * 128 + r * 16]);
+        }
+        float xs[8];
+#pragma unroll
+        for (int c = 0; c < 8; ++c) xs[c] = xd[(size_t) c * nb + st * 16 + (lane >> 1)];
+#pragma unroll
+        for (int c = 0; c < 8; ++c) {
+            const int4 x = xv[(size_t) c * (n_in / 16) + st * 32];
+            int xsum = STRATA_DP4A(0x01010101, x.x, 0);
+            xsum = STRATA_DP4A(0x01010101, x.y, xsum);
+            xsum = STRATA_DP4A(0x01010101, x.z, xsum);
+            xsum = STRATA_DP4A(0x01010101, x.w, xsum);
+#pragma unroll
+            for (int r = 0; r < R; ++r) {
+                int isum = STRATA_DP4A(w[r].x, x.x, 0);
+                isum = STRATA_DP4A(w[r].y, x.y, isum);
+                isum = STRATA_DP4A(w[r].z, x.z, isum);
+                isum = STRATA_DP4A(w[r].w, x.w, isum);
+                acc[r][c] = fmaf(dw[r] * xs[c], (float) (isum - 8 * xsum), acc[r][c]);
+            }
+        }
+    }
+#pragma unroll
+    for (int r = 0; r < R; ++r)
+#pragma unroll
+        for (int c = 0; c < 8; ++c) {
+            float v = acc[r][c];
+            v += __shfl_xor_sync(0xffffffffu, v, 16);
+            v += __shfl_xor_sync(0xffffffffu, v, 8);
+            v += __shfl_xor_sync(0xffffffffu, v, 4);
+            v += __shfl_xor_sync(0xffffffffu, v, 2);
+            v += __shfl_xor_sync(0xffffffffu, v, 1);
+            if (lane == 0 && c < ncols) sg.y[(size_t) c * sg.n_out + row0 + r] = v;
+        }
+}
+
 // quantize one 32-block of one row to the planar activation layout (a warp: lane = element)
 __device__ __forceinline__ void quant_block(float v, int blk, int col, int K, int8_t* __restrict__ xq, float* __restrict__ xd) {
     const int lane = threadIdx.x & 31;
@@ -519,6 +588,35 @@ __global__ void q8_tile_kernel(const uint8_t* __restrict__ blocks, int8_t* __res
     for (int i = 0; i < 32; ++i) o[i] = (int8_t) src[2 + i];
 }
 
+// ggml Q4_0 blocks (18 bytes: fp16 d, 16 bytes with value i in the low nibble of byte i and value i + 16 in the high nibble) of a [n_out][n_in] matrix -> the
+// nibble plane (n_out * n_in / 2 bytes, the layout above) and the tiled d plane; one thread per block
+__global__ void q4_tile_kernel(const uint8_t* __restrict__ blocks, uint8_t* __restrict__ qs, __half* __restrict__ dpl, int n_in, int64_t count) {
+    const int64_t b = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    if (b >= count) return;
+    const int nb = n_in / 32;
+    const int64_t row = b / nb;
+    const int blk = (int) (b % nb);
+    const int64_t g = row / 8;
+    const int r = (int) (row % 8), st = blk / 16, bi = blk % 16;
+    const uint8_t* src = blocks + b * 18;
+    const uint16_t dh = (uint16_t) src[0] | ((uint16_t) src[1] << 8);
+    dpl[(g * 8 * nb) + (int64_t) st * 128 + r * 16 + bi] = __ushort_as_half(dh);
+    for (int half = 0; half < 2; ++half) {   // values half * 16 .. +15 of the block
+        uint32_t wd[2];
+        for (int w = 0; w < 2; ++w) {   // word w holds values w * 8 .. w * 8 + 7 of the half: byte p = (value w*8+p, value w*8+4+p) as (low, high) nibble
+            uint32_t v = 0;
+            for (int p = 0; p < 4; ++p) {
+                const uint8_t b_lo = src[2 + w * 8 + p], b_hi = src[2 + w * 8 + 4 + p];
+                const uint32_t q_lo = half == 0 ? (b_lo & 15u) : (b_lo >> 4), q_hi = half == 0 ? (b_hi & 15u) : (b_hi >> 4);
+                v |= (q_lo | (q_hi << 4)) << (p * 8);
+            }
+            wd[w] = v;
+        }
+        uint8_t* o = qs + g * 8 * (n_in / 2) + (int64_t) st * 2048 + r * 256 + (bi * 2 + half) * 8;
+        *reinterpret_cast<uint2*>(o) = make_uint2(wd[0], wd[1]);
+    }
+}
+
 // ---- top-16
 constexpr int TK_CHUNK_MAX = 4096;
 __device__ __forceinline__ bool better(float va, int ia, float vb, int ib) { return va > vb || (va == vb && ia < ib); }
@@ -730,13 +828,20 @@ void df2_gemm(const Df2Mat* mats, int n, const void* xq, const float* xd, int n_
     int rows = 0;
     for (int i = 0; i < n; ++i) {
         if (mats[i].n_out % 16 != 0) throw std::invalid_argument("df2_gemm: n_out must be a multiple of 16");
+        if (mats[i].q4 != mats[0].q4) throw std::invalid_argument("df2_gemm: one launch, one weight format");
         const int8_t* q = static_cast<const int8_t*>(mats[i].w);
-        segs.s[i] = Df2Seg{q, reinterpret_cast<const __half*>(q + (size_t) mats[i].n_out * n_in), mats[i].y, mats[i].n_out, rows};
+        segs.s[i] = Df2Seg{q, reinterpret_cast<const __half*>(q + (size_t) mats[i].n_out * n_in / (mats[i].q4 ? 2 : 1)), mats[i].y, mats[i].n_out, rows};
         rows += mats[i].n_out;
     }
     segs.n = n;
     const int8_t* xqp = static_cast<const int8_t*>(xq);
     cudaStream_t st = (cudaStream_t) stream;
+    if (mats[0].q4) {   // the Q4_0 pack: 4 warps a block when the rows allow, else 2
+        if (rows % 32 == 0) gemm4_kernel<4><<<rows / 32, 128, 0, st>>>(segs, xqp, xd, n_in, ncols);
+        else gemm4_kernel<2><<<rows / 16, 64, 0, st>>>(segs, xqp, xd, n_in, ncols);
+        check("gemm4");
+        return;
+    }
     const int variant = g_gemm_variant != 0 ? g_gemm_variant : (rows % 32 == 0 ? 1 : 2);
     switch (variant) {
         case 3: gemm8_kernel<2, false><<<rows / 16, 64, 0, st>>>(segs, xqp, xd, n_in, ncols); break;
@@ -752,6 +857,12 @@ void df2_gemm(const Df2Mat* mats, int n, const void* xq, const float* xd, int n_
     check("gemm8");
 }
 
+void df2_q4_tile(const void* blocks, void* qs, void* dpl, int n_in, int64_t count, void* stream) {
+    q4_tile_kernel<<<(unsigned) ((count + 255) / 256), 256, 0, (cudaStream_t) stream>>>(static_cast<const uint8_t*>(blocks), static_cast<uint8_t*>(qs),
+                                                                                         static_cast<__half*>(dpl), n_in, count);
+    check("q4_tile");
+}
+
 void df2_quant_rows(const float* x, int K, int ncols, void* xq, float* xd, void* stream) {
     quant_rows_kernel<<<dim3((K / 32 + 7) / 8, ncols), 256, 0, (cudaStream_t) stream>>>(x, K, static_cast<int8_t*>(xq), xd);
     check("quant_rows");

tools/dflash2_pack.py

@@ -1,13 +1,14 @@
 #!/usr/bin/env python3
 """Pack the incoai DFlash2 drafter for Qwen3.8-Flash-Next (safetensors) into the engine's runtime file.
 
-    python tools/dflash2_pack.py model.safetensors config.json out.pack [--bf16]
+    python tools/dflash2_pack.py model.safetensors config.json out.pack [--q4_0]
 
 Format (little endian): "DF2P", u32 version (1), u32 index_len, a text index (padded with spaces so the data starts on a 4096 multiple), then
 the tensors, each 256-byte aligned at its offset (from the end of the index).  Index lines: "cfg <key> <value>", "t <name> <dtype> <rows>
 <cols> <offset> <bytes>" (a matrix [out][in] is rows x cols; a vector or codebook row set is the leading dims folded into rows), "end".
 dtypes: "q8_0" (ggml Q8_0 blocks, 34 bytes per 32 values, rows of the weight matrix [out][in] one after the other: what the engine's
-multi-column MMVQ reads), "f32" (norms, the two conv base kernels), "bf16" (the selector's two codebooks: only gathered rows are read).
+multi-column MMVQ reads), "f32" (norms, the two conv base kernels), "q4_0" (--q4_0: ggml Q4_0 blocks, 18 bytes per 32 values, the matrices only; the engine tiles them like Q8_0, see gemm4_kernel in
+src/kernels/cuda/dflash2.cu), "bf16" (the selector's two codebooks: only gathered rows are read).
 The matrices' quantization matches ggml's quantize_row_q8_0 (d = amax / 127, q = round(x / d), d stored as fp16); the research harness
 (tools/dflash2/q8test.py) measured no acceptance loss from it.
 """
@@ -80,7 +81,24 @@ def quant_q8_0(w):
     return out.tobytes()
 
 
-def pack(st_path, cfg_path, out_path, say=print):
+def quant_q4_0(w):
+    """w: float32 [rows, cols], cols % 32 == 0 -> bytes of rows * cols/32 ggml Q4_0 blocks (fp16 d, 16 bytes: value i in the low nibble of byte i, value i + 16 in the high
+    nibble).  ggml's quantize_row_q4_0_ref: d = (the value of largest magnitude) / -8, q = min(15, int(x / d + 8.5)); the scale is stored as fp16, q uses the float d."""
+    rows, cols = w.shape
+    assert cols % 32 == 0
+    b = w.reshape(rows, cols // 32, 32).astype(np.float32)
+    idx = np.abs(b).argmax(axis=-1)
+    mx = np.take_along_axis(b, idx[..., None], axis=-1)[..., 0]
+    d = mx / np.float32(-8.0)
+    inv = np.where(d != 0, 1.0 / np.where(d == 0, 1, d), 0.0).astype(np.float32)
+    q = np.minimum(15, (b * inv[..., None] + np.float32(8.5)).astype(np.int32)).astype(np.uint8)
+    out = np.empty((rows, cols // 32, 18), np.uint8)
+    out[..., :2] = d.astype(np.float16).view(np.uint8).reshape(rows, cols // 32, 2)
+    out[..., 2:] = q[..., :16] | (q[..., 16:] << 4)
+    return out.tobytes()
+
+
+def pack(st_path, cfg_path, out_path, say=print, q4=False):
     """Write the engine's runtime file for the drafter in `st_path` (+ its config.json). Returns the number of bytes written."""
     with open(cfg_path, encoding="utf-8") as fh:
         cfg = json.load(fh)
@@ -96,7 +114,10 @@ def pack(st_path, cfg_path, out_path, say=print):
     def mat(name, new=None):
         w = load_tensor(f, name)
         assert w.ndim == 2
-        add(new or name, "q8_0", w.shape, quant_q8_0(w))
+        if q4:
+            add(new or name, "q4_0", w.shape, quant_q4_0(w))
+        else:
+            add(new or name, "q8_0", w.shape, quant_q8_0(w))
 
     def vec(name, new=None):
         w = load_tensor(f, name)
@@ -155,9 +176,10 @@ def pack(st_path, cfg_path, out_path, say=print):
 
 
 def main():
-    if len(sys.argv) < 4:
+    args = [a for a in sys.argv[1:] if a != "--q4_0"]
+    if len(args) < 3:
         sys.exit(__doc__)
-    pack(*sys.argv[1:4])
+    pack(*args[:3], q4="--q4_0" in sys.argv[1:])
 
 
 if __name__ == "__main__":

tools/test_dflash2_pack.py

@@ -123,6 +123,35 @@ class PackWithoutTorch(unittest.TestCase):
         widened = (bf16(self.arrays["fc.weight"]).astype(np.uint32) << 16).view(np.float32)
         self.assertEqual(got, new.quant_q8_0(widened))
 
+    def test_q4_0_blocks_are_ggml_q4_0_and_only_the_matrices_change(self):
+        q8, q4 = self.dir / "q8b.pack", self.dir / "q4.pack"
+        self.pack(q8)
+        new.pack(str(self.dir / "model.safetensors"), str(self.dir / "config.json"), str(q4), say=lambda *_: None, q4=True)
+        rows = {}
+        for path in (q8, q4):
+            data = path.read_bytes()
+            index_len = struct.unpack("<I", data[8:12])[0]
+            rows[path] = (data, 12 + index_len, {l.split()[1]: l.split() for l in data[12:12 + index_len].decode().splitlines() if l.startswith("t ")})
+        d8, b8, t8 = rows[q8]
+        d4, b4, t4 = rows[q4]
+        self.assertEqual(set(t8), set(t4))
+        for name, r in t4.items():
+            dtype4, nbytes4 = r[2], int(r[6])
+            if t8[name][2] == "q8_0":
+                self.assertEqual(dtype4, "q4_0", name)
+                self.assertEqual(nbytes4, int(r[3]) * int(r[4]) // 32 * 18, name)
+            else:   # norms, conv kernels and the codebooks are byte for byte the same
+                self.assertEqual(d4[b4 + int(r[5]): b4 + int(r[5]) + nbytes4], d8[b8 + int(t8[name][5]): b8 + int(t8[name][5]) + int(t8[name][6])], name)
+        # one matrix: dequantize the blocks and check the error is within one step of each block scale (the +8 end is clamped to 7)
+        w = (bf16(self.arrays["fc.weight"]).astype(np.uint32) << 16).view(np.float32)
+        r = t4["fc.weight"]
+        blocks = np.frombuffer(d4[b4 + int(r[5]): b4 + int(r[5]) + int(r[6])], np.uint8).reshape(w.shape[0], w.shape[1] // 32, 18)
+        d = blocks[..., :2].copy().view(np.float16)[..., 0].astype(np.float32)
+        q = np.concatenate([blocks[..., 2:] & 15, blocks[..., 2:] >> 4], axis=-1).astype(np.float32) - 8
+        deq = (q * d[..., None]).reshape(w.shape)
+        step = np.abs(d)[..., None].repeat(32, -1).reshape(w.shape)
+        self.assertTrue(np.all(np.abs(deq - w) <= 1.01 * step + 1e-6))
+
     def test_it_matches_the_torch_version_when_torch_is_here(self):
         try:
             import torch  # noqa: F401