dsv41: CPU lane takes Strata's AVX2 i-quant rows (switchable), ds4_gg_bench, pack_gguf repeatable --dense-override

CMakeLists.txt

@@ -787,6 +787,14 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
       target_link_libraries(strata-dsv41 PRIVATE strata_kernels_cpu)
       target_compile_definitions(strata-dsv41 PRIVATE STRATA_DSV41_TQ_CPU=1)
     endif()
+    if(STRATA_NATIVE_EXPERTS)
+      # the CPU lane's ggml-type rows: ggml-cpu's dots against Strata's AVX2 i-quant rows (bits and GB/s); no GPU, no model
+      add_executable(ds4_gg_bench tests/dsv41/gg_dot_bench.cpp src/dsv41/ds4_gg_cpu.cpp)
+      target_include_directories(ds4_gg_bench PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CURRENT_SOURCE_DIR}/src/dsv41
+                                 ${CMAKE_CURRENT_SOURCE_DIR}/third_party/ggml)
+      target_link_libraries(ds4_gg_bench PRIVATE strata_kernels_cpu Threads::Threads)
+      target_compile_definitions(ds4_gg_bench PRIVATE STRATA_DSV41_TQ_CPU=1)
+    endif()
     if(STRATA_NATIVE_EXPERTS)
       # the trellis experts (device kernels ds4_tq.cu, CPU lane ds4_tq_cpu.cpp) against a double reference; needs a GPU, no model
       set_source_files_properties(tests/dsv41/tq_expert_test.cpp PROPERTIES LANGUAGE HIP)

src/dsv41/ds4_gg_cpu.cpp

@@ -11,6 +11,7 @@
 #if defined(STRATA_DSV41_TQ_CPU)
 #include "ggml-cpu.h"
 #include "ggml.h"
+#include "strata/kernels/cpu/iq_avx2.hpp"
 
 #include <mutex>
 #include <vector>
@@ -72,7 +73,28 @@ void gg_cpu_quant(int act, const float* x, int k, void* dst) {
     t->from_float_ref(x, dst, k);
 }
 
+// Strata's AVX2 multi-token rows for the i-quants (iq256_rows: each 32-value chunk decoded once, every token applies it;
+// the integer sums are ggml's, only the float additions come in another order).  DS4_GG_IQ256=0: ggml-cpu's row dots.
+static int g_iq256 = -1;   // -1: from the environment and the CPU, 0 off, 1 on
+void gg_cpu_set_iq256(int mode) { g_iq256 = mode; }
+static bool use_iq256(int type) {
+    if (g_iq256 < 0) {
+        const char* v = std::getenv("DS4_GG_IQ256");
+        bool avx2 = false;
+#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__))
+        __builtin_cpu_init();
+        avx2 = __builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma") && __builtin_cpu_supports("f16c");
+#endif
+        g_iq256 = avx2 && (v == nullptr || std::atoi(v) != 0) ? 1 : 0;
+    }
+    return g_iq256 == 1 && strata::kernels::cpu::iq256_supported(type);
+}
+
 static void dots(int type, const uint8_t* w, size_t rb, int n, const void* const* act, int nt, float* const* out, int r0, int r1) {
+    if (nt > 0 && use_iq256(type)) {
+        strata::kernels::cpu::iq256_rows(type, w, rb, n, act, nt, out, r0, r1);
+        return;
+    }
     const ggml_type_traits_cpu* tc = ggml_get_type_traits_cpu((ggml_type) type);
     if (dbg()) std::fprintf(stderr, "gg dots type %d n %d rows %d..%d nt %d\n", type, n, r0, r1, nt);
     for (int r = r0; r < r1; ++r)

src/dsv41/ds4_gguf_cpu.hpp

@@ -52,6 +52,8 @@ bool gg_cpu_fmt(int tg, int td, int n_embd, int n_ff, IqCpu& f, std::string& err
 void gg_cpu_quant(int act, const float* x, int k, void* dst);
 void gg_cpu_gate_up(const IqCpu& f, const uint8_t* gate, const uint8_t* up, const void* const* act, int nt, float* const* ff,
                     int r0, int r1, float limit);
+/// -1: i-quant rows from Strata's AVX2 kernels when the CPU has AVX2 (DS4_GG_IQ256=0 turns it off), 0 ggml-cpu, 1 forced (benches).
+void gg_cpu_set_iq256(int mode);
 void gg_cpu_down(const IqCpu& f, const uint8_t* down, const void* const* hq, int nt, float* const* out, int r0, int r1);
 
 /// The CPU lane's MXFP4 rows in the device's arithmetic (ds4_gguf.cu's mx_dot): the activations rounded to int8 per 32

tests/dsv41/gg_dot_bench.cpp

@@ -0,0 +1,142 @@
+// Strata. The DeepSeek-V4.1 CPU lane's ggml-type rows (src/dsv41/ds4_gg_cpu.cpp): ggml-cpu's row dot against Strata's AVX2
+// multi-token i-quant rows, on random blocks of IQ2_XXS / IQ3_XXS.
+//   ds4_gg_bench [--type 16|18] [--threads N] [--mb M] [--nt 1..8] [--reps R]
+// Checks first (every row, every token): the two paths' results agree to a few float ulps (the integer sums are the same;
+// only the order of the float additions differs) and, with --exact, bit for bit.  Then the time of a pass over M MB of
+// weights (default 3072, far past L3) on N threads, as GB/s of weights read, beside a plain streaming read of the same
+// bytes (the memory ceiling the rows are measured against).
+#include "ds4_gguf_cpu.hpp"
+
+#include "ggml-cpu.h"
+#include "ggml.h"
+
+#include <algorithm>
+#include <atomic>
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <thread>
+#include <vector>
+
+using namespace strata::ds4::gguf;
+using clk = std::chrono::steady_clock;
+
+static double secs(clk::time_point a) { return std::chrono::duration<double>(clk::now() - a).count(); }
+
+int main(int argc, char** argv) {
+    int type = 18, threads = 32, nt = 1, reps = 3;
+    long mb = 3072;
+    bool exact = false;
+    for (int i = 1; i < argc; ++i) {
+        std::string a = argv[i];
+        auto val = [&] { return i + 1 < argc ? std::atol(argv[++i]) : 0L; };
+        if (a == "--type") type = (int) val();
+        else if (a == "--threads") threads = (int) val();
+        else if (a == "--mb") mb = val();
+        else if (a == "--nt") nt = (int) val();
+        else if (a == "--reps") reps = (int) val();
+        else if (a == "--exact") exact = true;
+    }
+    ggml_cpu_init();
+    const int n = 5120;   // a gate / up row
+    IqCpu f;
+    std::string err;
+    if (!gg_cpu_fmt(type, type, n, 2304, f, err)) {
+        std::fprintf(stderr, "gg_cpu_fmt: %s\n", err.c_str());
+        return 2;
+    }
+    const size_t rb = f.gu_row;
+    const long rows = (long) ((size_t) mb * 1000000 / rb) / 8 * 8;
+    std::vector<uint8_t> w((size_t) rows * rb);
+    uint64_t st = 0x9E3779B97F4A7C15ull;
+    auto next = [&st] { st ^= st << 13; st ^= st >> 7; st ^= st << 17; return st; };
+    for (size_t o = 0; o + 8 <= w.size(); o += 8) { const uint64_t v = next(); std::memcpy(&w[o], &v, 8); }
+    const size_t bb = ggml_type_size((ggml_type) type);
+    for (size_t o = 0; o + bb <= w.size(); o += bb) {   // the blocks' F16 scale: a small normal value
+        const ggml_fp16_t d = ggml_fp32_to_fp16(0.002f + 0.001f * (float) (next() % 16));
+        std::memcpy(&w[o], &d, sizeof d);
+    }
+    std::vector<std::vector<uint8_t>> a((size_t) nt, std::vector<uint8_t>(f.act_bytes));
+    std::vector<const void*> ap((size_t) nt);
+    for (int t = 0; t < nt; ++t) {
+        std::vector<float> x((size_t) n);
+        for (auto& v : x) v = ((float) (next() % 2001) - 1000.0f) * 1e-3f;
+        gg_cpu_quant(f.gu_act, x.data(), n, a[(size_t) t].data());
+        ap[(size_t) t] = a[(size_t) t].data();
+    }
+    std::vector<std::vector<float>> o0((size_t) nt, std::vector<float>((size_t) rows)), o1 = o0;
+    std::vector<float*> p0, p1;
+    for (int t = 0; t < nt; ++t) { p0.push_back(o0[(size_t) t].data()); p1.push_back(o1[(size_t) t].data()); }
+
+    // a pass over all rows on `threads` threads, mode 0 ggml-cpu, 1 Strata's rows; returns the best time of reps
+    auto pass = [&](int mode, float* const* out) {
+        gg_cpu_set_iq256(mode);
+        double best = 1e30;
+        for (int r = 0; r < reps + 1; ++r) {
+            const auto t0 = clk::now();
+            std::vector<std::thread> th;
+            std::atomic<long> cursor{0};
+            const long chunk = 64;
+            for (int k = 0; k < threads; ++k)
+                th.emplace_back([&] {
+                    for (;;) {
+                        const long r0 = cursor.fetch_add(chunk);
+                        if (r0 >= rows) break;
+                        // one "expert part" is 2304 rows; the engine hands a thread a few rows of it at a time
+                        gg_cpu_down(f, w.data(), ap.data(), nt, out, (int) r0, (int) std::min(rows, r0 + chunk));
+                    }
+                });
+            for (auto& t : th) t.join();
+            if (r > 0) best = std::min(best, secs(t0));
+        }
+        return best;
+    };
+    // gg_cpu_down uses f.td / d_row; here td == tg and the activation has n_embd values: f.d_row would be for n_ff.
+    // Use the gate rows through a second IqCpu whose down geometry is the gate's.
+    f.td = f.tg; f.d_row = f.gu_row; f.n_ff = n; f.d_act = f.gu_act; f.h_bytes = f.act_bytes;
+
+    const double t_gg = pass(0, p0.data());
+    const double t_own = pass(1, p1.data());
+    long bad = 0, differ = 0;
+    double worst = 0;
+    for (int t = 0; t < nt; ++t)
+        for (long r = 0; r < rows; ++r) {
+            const float x = o0[(size_t) t][(size_t) r], y = o1[(size_t) t][(size_t) r];
+            if (std::memcmp(&x, &y, 4) != 0) {
+                ++differ;
+                const double rel = std::fabs((double) x - y) / (std::fabs((double) x) + 1e-3);
+                worst = std::max(worst, rel);
+                if (rel > 1e-5) ++bad;
+            }
+        }
+    // memory ceiling: a plain read of the same bytes on the same threads
+    double t_rd = 1e30;
+    for (int r = 0; r < reps + 1; ++r) {
+        const auto t0 = clk::now();
+        std::vector<std::thread> th;
+        std::vector<uint64_t> sink((size_t) threads);
+        for (int k = 0; k < threads; ++k)
+            th.emplace_back([&, k] {
+                const size_t lo = w.size() / (size_t) threads * (size_t) k, hi = k + 1 == threads ? w.size() : w.size() / (size_t) threads * (size_t) (k + 1);
+                uint64_t s = 0, v;
+                for (size_t o = lo; o + 8 <= hi; o += 8) { std::memcpy(&v, &w[o], 8); s += v; }
+                sink[(size_t) k] = s;
+            });
+        for (auto& t : th) t.join();
+        if (r > 0) t_rd = std::min(t_rd, secs(t0));
+    }
+    const double gb = (double) w.size() / 1e9;
+    std::printf("type %d nt %d threads %d: %.2f GB of weights (%ld rows)\n", type, nt, threads, gb, rows);
+    std::printf("  ggml-cpu rows %.3f s = %.1f GB/s | Strata AVX2 rows %.3f s = %.1f GB/s (x%.2f) | plain read %.1f GB/s\n", t_gg, gb / t_gg,
+                t_own, gb / t_own, t_gg / t_own, gb / t_rd);
+    std::printf("  results: %ld of %ld differ in bits, worst relative %.2e, %ld beyond 1e-5\n", differ, (long) nt * rows, worst, bad);
+    if (bad != 0 || (exact && differ != 0)) {
+        std::printf("FAIL\n");
+        return 1;
+    }
+    std::printf("OK\n");
+    return 0;
+}

tools/dsv41/pack_gguf.py

@@ -333,7 +333,7 @@ def main() -> int:
                                      "checkpoint's vision.* / aligner.* (default: the first mmproj*.gguf in <model>/..)")
     ap.add_argument("--tq-experts", nargs="+", help="folders tq_convert.py wrote (tq.json + tq-L<layer>.bin): the routed experts of the layers "
                                          "listed there are read from it as trellis types (TQK5X ..), the others keep their MXFP4")
-    ap.add_argument("--dense-override", help="a folder dense_requant.py wrote (dense.json + dense.bin): its tensors replace the same-named "
+    ap.add_argument("--dense-override", action="append", help="a folder dense_requant.py wrote (dense.json + dense.bin): its tensors replace the same-named "
                                              "dense matrices of the GGUF (K-quants of the attention / shared experts)")
     ap.add_argument("--override-experts", nargs="+", help="requantized GGUF folders, each DIR or DIR:roles (requant_gguf.py output): their routed experts replace the main "
                                                "model's for --override-roles (gate,up,down) in --override-layers (default all): mixed-bit packs")
@@ -492,8 +492,8 @@ def main() -> int:
         print("experts overridden from %s: %d tensors (roles %s, layers %s)" % (od, len(new), drl or a.override_roles, a.override_layers or "all"))
     # dense matrices requantized by dense_requant.py replace the GGUF's BF16 ones
     n_dense = 0
-    if a.dense_override:
-        dd = pathlib.Path(a.dense_override).expanduser().resolve()
+    for dsrc in (a.dense_override or []):
+        dd = pathlib.Path(dsrc).expanduser().resolve()
         dj = json.loads((dd / "dense.json").read_text())
         try:
             relf = os.path.relpath(dd / "dense.bin", model).replace(os.sep, "/")
@@ -503,12 +503,14 @@ def main() -> int:
             files.append(relf)
         dfi = files.index(relf)
         pos = {t[0]: i for i, t in enumerate(tensors)}
+        n_here = 0
         for name, (tname, fname, off, nb, shape) in dj.items():
             if name not in pos:
                 raise SystemExit("dense override: %s is not in the pack" % name)
             tensors[pos[name]] = (name, tname, dfi, off, nb, shape)
-            n_dense += 1
-        print("dense override from %s: %d tensors" % (dd, n_dense))
+            n_here += 1
+        n_dense += n_here
+        print("dense override from %s: %d tensors (a later folder wins)" % (dd, n_here))
     # the DSpark drafter from unsloth's GGUF (architecture dflash): blk.N.* -> mtp.N.*, the head tensors under the
     # checkpoint's names (fc -> main_proj, enc.output_norm -> main_norm, output_norm -> norm of the last block,
     # conf_proj -> confidence_head.proj, markov_w1 / w2 -> markov_head.embed / head)