metal DFlash2 reads the Q4_0 selector codebooks (df2_select_q4; the pack loader keeps them raw rows) and prints 'selector codebooks: Q4_0 rows', the text strata_control/drafting.py looks for in an engine before it uses the converted pack; STRATA_DF2_HEAD=native runs the head over the target's whole head and ignores the draft vocabulary

metal/src/engine/dflash2.inc

@@ -202,8 +202,11 @@ bool Window::df2_load_pack(std::string& err) {
     auto up256 = [](uint64_t x) { return (x + 255) / 256 * 256; };
     uint64_t wq_bytes = 0, misc_bytes = 0;
     std::map<std::string, uint64_t> woff, soff, moff;
+    auto is_codebook = [](const std::string& n) { return n.find("codebook") != std::string::npos; };   // the selector's tables are gathered by row, never repacked
     for (const auto& [n, t] : ten) {
-        if (t.type == ref::Q4_0) {
+        if (t.type == ref::Q4_0 && is_codebook(n)) {
+            moff[n] = misc_bytes; misc_bytes += up256(t.bytes);                   // Q4_0 rows of 8 blocks (144 bytes), read as they are by df2_select_q4
+        } else if (t.type == ref::Q4_0) {
             const uint64_t K = t.shape[0], rows = t.shape[1];
             if (K % 32 != 0 || rows % 8 != 0) { err = "dflash2: " + n + " is not a whole number of 8-row, 32-column Q4_0 tiles"; return false; }
             woff[n] = wq_bytes; wq_bytes += up256(rows * K / 2);                  // fragment-order words: [row tile][block][lane] uint32
@@ -227,7 +230,7 @@ bool Window::df2_load_pack(std::string& err) {
         tmp.resize(t.bytes);
         fseeko(f, (off_t) (data_off + t.offset), SEEK_SET);
         if (std::fread(tmp.data(), 1, t.bytes, f) != t.bytes) { std::fclose(f); err = "dflash2: the pack is shorter than its index"; return false; }
-        if (t.type == ref::Q4_0) {
+        if (t.type == ref::Q4_0 && !is_codebook(n)) {
             // Q4_0 blocks (half d, 16 bytes: byte j = element j low nibble, element j + 16 high nibble) -> the matrix units' fragment order (df2_mma4s):
             // for a tile of 8 rows and one block, lane l (fragment row fm, columns fn, fn + 1) holds the codes of W[fm][8 s + fn + e] as nibble 2 s + e
             const uint64_t K = t.shape[0], rows = t.shape[1], nb = K / 32, nrt = rows / 8;
@@ -282,6 +285,11 @@ bool Window::df2_bind(const std::string& vocab_file, std::string& err) {
     if (d.bound) return true;
     if (!taps_on_) { err = "dflash2: the window keeps no taps"; return false; }
     if (!df2_load_pack(err)) return false;
+    {   // the text strata_control/drafting.py looks for in the program to know that this engine reads a pack whose selector tables are Q4_0 rows
+        const auto cb = d.ten.find("candidate_selector.predecessor_codebook");
+        if (cb != d.ten.end() && cb->second.type == ref::Q4_0)
+            std::fprintf(stderr, "strata metal: DFlash2 selector codebooks: Q4_0 rows (about 174 MiB less than bf16)\n");
+    }
     const double t0 = now_ms();
     auto A = [&](size_t bytes) { d.bytes += bytes; return cx_->alloc(bytes); };
     // ---- the draft vocabulary's head rows: the MTP's when it is the same list, else gathered here
@@ -308,6 +316,8 @@ bool Window::df2_bind(const std::string& vocab_file, std::string& err) {
     {
         const char* hm = std::getenv("STRATA_DF2_HEAD");
         d.head_q8 = !(hm && std::string(hm) == "native");
+        // STRATA_DF2_HEAD=native: the target's own head, whole (the community drafter's): no draft vocabulary, no gathered copy, the engine's GEMV on the target's rows
+        if (!d.head_q8) vocab.clear();
     }
     d.nd = vocab.empty() ? n_vocab_ : (int64_t) vocab.size();
     d.use_map = !vocab.empty();
@@ -655,7 +665,12 @@ void Window::df2_block_list(std::vector<mtl::Dispatch>& ds) {
     {
         struct SelArgs { uint32_t rank, rows; } sa{(uint32_t) d.rank, (uint32_t) ND};
         mtl::Dispatch sd;
-        sd.p = P(ldf2_, "df2_select", {}, "df2_select");
+        {   // the selector's tables: BF16 rows (the incoai pack) or Q4_0 rows (the community GGUF: 8 blocks of 18 bytes a row)
+            const gguf::Tensor& cbt = d.ten.at("candidate_selector.predecessor_codebook");
+            if (cbt.type != d.ten.at("candidate_selector.successor_codebook").type) { std::fprintf(stderr, "strata metal: DFlash2: the two selector tables differ in type\n"); }
+            const bool cb4 = cbt.type == ref::Q4_0;
+            sd.p = P(ldf2_, cb4 ? "df2_select_q4" : "df2_select", {}, cb4 ? "df2_select_q4" : "df2_select");
+        }
         sd.args = {mtl::Arg::B(d.hsel), mtl::Arg::B(d.topval), mtl::Arg::B(d.toptok),
                    mtl::Arg::B(d.pack, d.ten.at("candidate_selector.predecessor_codebook").offset),
                    mtl::Arg::B(d.pack, d.ten.at("candidate_selector.successor_codebook").offset), mtl::Arg::B(d.anchor, 0),

metal/src/engine/msl/dflash2.metal

@@ -349,6 +349,54 @@ kernel void df2_select(device const float* hp [[buffer(0)]], device const float*
     }
 }
 
+// The same walk over Q4_0 selector tables (the community GGUF's): a row is 8 ggml Q4_0 blocks of 18 bytes (half d, 16 bytes: byte j = element j low nibble,
+// element j + 16 high nibble); a value is (q - 8) * d, in fp32 as ggml dequantizes it.
+static inline float q4_0_val(device const uchar* row, uint e) {
+    device const uchar* b = row + (e >> 5) * 18;
+    const float d = float(*(device const half*)b);
+    const uint j = e & 31;
+    const uint q = j < 16 ? (b[2 + j] & 15u) : (b[2 + j - 16] >> 4);
+    return d * (float(q) - 8.0f);
+}
+
+kernel void df2_select_q4(device const float* hp [[buffer(0)]], device const float* unary [[buffer(1)]], device const int* tok [[buffer(2)]],
+                          device const uchar* pred_cb [[buffer(3)]], device const uchar* succ_cb [[buffer(4)]],
+                          device const int* anchor_tok [[buffer(5)]], device int* out_tok [[buffer(6)]], device float* out_prob [[buffer(7)]],
+                          constant SelArgs& a [[buffer(8)]], ushort tid [[thread_index_in_threadgroup]],
+                          ushort w [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
+    threadgroup float vp[256];
+    threadgroup float sc[TOPK];
+    threadgroup int prev;
+    const uint rowb = (a.rank / 32) * 18;   // 144 bytes a row
+    if (tid == 0) prev = *anchor_tok;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    for (uint i = 0; i < a.rows; ++i) {
+        vp[tid] = q4_0_val(pred_cb + uint(prev) * rowb, tid) * hp[i * a.rank + tid];
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for (int j = w * 2; j < w * 2 + 2; ++j) {
+            const int t = tok[i * TOPK + j];
+            device const uchar* row = succ_cb + uint(t) * rowb;
+            float acc = 0.0f;
+            for (int e = 0; e < 8; ++e) acc = fma(q4_0_val(row, lane * 8 + e), vp[lane * 8 + e], acc);
+            acc = simd_sum(acc);
+            if (lane == 0) sc[j] = unary[i * TOPK + j] + acc;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (tid == 0) {
+            int bj = 0;
+            float bv = sc[0];
+            for (int j = 1; j < TOPK; ++j)
+                if (sc[j] > bv) { bv = sc[j]; bj = j; }
+            float z = 0.0f;
+            for (int j = 0; j < TOPK; ++j) z += exp(sc[j] - bv);
+            out_tok[i] = tok[i * TOPK + bj];
+            out_prob[i] = 1.0f / z;
+            prev = out_tok[i];
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+}
+
 // ---- the tap: out[r][out_off + i] = (sum_c R[r][c][i] + bo[r][i] * sum_c 2 sigmoid(inj[r][c] / hc)) / hc.  Grid (n / 256, rows), 256 threads.
 struct TapArgs { uint n; uint hc; uint R_stride; uint bo_stride; uint inj_stride; uint out_stride; uint out_off; };