A container the suite can reach at index_bits 6, and a VQ4P arm in run.sh

tests/run.sh

@@ -813,6 +813,103 @@ else
     sk "engine checks" "no container at $MODEL"
 fi
 
+# --------------------------------------------------------------- VQ4P ----
+head_ "VQ4P engine (index_bits 6)"
+
+# vq_rows_p6 — the packed-index apply PR #41's AVX-512 kernel joins — had
+# no container `make check` could reach: the synthetic one was always
+# index_bits 8, and a real --index-bits 6 conversion costs hours plus the
+# source weights. So every green run until this arm covered VQ3R and
+# nothing covered VQ4P, on any platform. Same discipline as the engine
+# block above: no oracle, but the engine is compared against itself. The
+# p6 accumulate itself is integer until the per-block fold, but what
+# reaches the logits also went through the other dispatched kernels, so
+# agreement between backends lands in the fp-noise branch — that is a real
+# verdict, not a regression, and the checks below report it as such.
+# Scope: this container is 4 layers / 3 MoE, so the arm bounds gross
+# kernel errors; it does not and cannot bound the depth-amplified
+# discontinuity mode (a real index_bits 6 container shows max diff
+# ~0.58 at 27 layers with the kernel working correctly). Do not read a
+# green run here as more than that. The container is a few MB and
+# builds in milliseconds, so the arm runs on every host and in CI.
+VQ4P="$TMP/tiny6.waste"
+P6_IDS=3,7,11,5,9,13,2,17,4,8,19,23,6,29,12,31
+if python3 tools/make_test_container.py --index-bits 6 "$VQ4P" \
+        >/dev/null 2>&1; then
+    # Same struct check as above, on records whose payload is three packed
+    # bytes per row instead of four whole ones — the layout has to be the
+    # one test_container reads regardless of the packing.
+    banks=0; recs=0; bad=0
+    for bank in "$VQ4P"/experts-L*.bin; do
+        [ -f "$bank" ] || continue
+        banks=$((banks + 1))
+        out=$(./test_container "$bank" 2 2>/dev/null) || { bad=1; continue; }
+        n=$(printf '%s' "$out" | sed -n 's/^\([0-9]*\) records read, \([0-9]*\) problems$/\1 \2/p')
+        set -- $n
+        [ "${1:-0}" -gt 0 ] || bad=1
+        [ "${2:-1}" -eq 0 ] || bad=1
+        recs=$((recs + ${1:-0}))
+    done
+    if [ "$banks" -gt 0 ] && [ "$bad" = 0 ]; then
+        ok "VQ4P records read through the C structs ($recs records over $banks banks)"
+    else
+        no "VQ4P record layout"
+    fi
+
+    ./test_forward "$VQ4P" "$P6_IDS" "$TMP/p6_seq.bin" 0 >/dev/null 2>&1
+    WASTE_CHUNK=1 ./test_forward "$VQ4P" "$P6_IDS" "$TMP/p6_chunk.bin" 0 \
+        >/dev/null 2>&1
+    if python3 - "$TMP/p6_seq.bin" "$TMP/p6_chunk.bin" <<'PY'
+import struct, sys
+def L(p):
+    b = open(p, "rb").read()
+    return struct.unpack(f"<{len(b)//4}f", b)
+a, b = L(sys.argv[1]), L(sys.argv[2])
+d = max(abs(x - y) for x, y in zip(a, b))
+sys.exit(0 if d < 1e-3 and a.index(max(a)) == b.index(max(b)) else 1)
+PY
+    then ok "VQ4P chunked prefill == token-at-a-time"
+    else no "VQ4P chunked prefill diverges"
+    fi
+
+    WASTE_BACKEND=cpu ./test_forward "$VQ4P" "$P6_IDS" "$TMP/p6_cpu.bin" 0 \
+        >/dev/null 2>&1
+    if cmp -s "$TMP/p6_seq.bin" "$TMP/p6_cpu.bin"; then
+        ok "VQ4P SIMD backend bit-identical to the CPU baseline"
+    else
+        # same tolerance fallback as the VQ3R arm above. 1e-5, not 1e-3:
+        # the mutation table in the PR review puts the intact kernel at
+        # 9.54e-07 and every single-edit break at 0.002-0.003, so 1e-5
+        # leaves 10x headroom over noise and 300x under the break signal.
+        if python3 - "$TMP/p6_seq.bin" "$TMP/p6_cpu.bin" <<'PY'
+import struct, sys
+def L(p):
+    b = open(p, "rb").read()
+    return struct.unpack(f"<{len(b)//4}f", b)
+a, b = L(sys.argv[1]), L(sys.argv[2])
+sys.exit(0 if max(abs(x - y) for x, y in zip(a, b)) < 1e-5 else 1)
+PY
+        then ok "VQ4P SIMD backend matches the CPU baseline (within fp noise)"
+        else no "VQ4P SIMD backend diverges from the CPU baseline"
+        fi
+    fi
+
+    WASTE_CACHE_MB=512 ./test_forward "$VQ4P" "$P6_IDS" "$TMP/p6_cache.bin" 0 \
+        >/dev/null 2>&1
+    # same() + the 0/1/* case, not `cmp -s`: this is #42's own rule, and the
+    # same call site by name as the VQ3R cache check above it. An unguarded
+    # `cmp -s` on a PATH without diffutils (fresh MSYS2 UCRT64) is the one
+    # false FAIL on an otherwise clean board.
+    same "$TMP/p6_seq.bin" "$TMP/p6_cache.bin"
+    case $? in
+        0) ok "VQ4P expert cache is bit-identical to no cache" ;;
+        1) no "VQ4P expert cache changes results" ;;
+        *) sk "VQ4P expert cache is bit-identical to no cache" "$NO_CMP" ;;
+    esac
+else
+    sk "VQ4P engine" "cannot build a synthetic index_bits 6 container"
+fi
+
 # --------------------------------------------------------------- rotary ----
 head_ "rotary (MLA on a model that is not NoPE)"
 

tests/test_container.c

@@ -105,7 +105,10 @@ int main(int argc, char **argv)
         }
         if (off % WASTE_ALIGN) { printf("  not 4 KiB aligned!\n"); bad++; }
         if (h.lowrank_id != 0) { printf("  lowrank_id != 0 (v0 violation)\n"); bad++; }
-        if (h.fmt != WQ_VQ3R && h.fmt != WQ_VQ2R) { printf("  unexpected fmt %u\n", h.fmt); bad++; }
+        /* VQ4P shares VQ3R's record size and layout — only the index
+         * packing inside the payload differs — so it is accepted here for
+         * the same reason: the struct reads are the same either way. */
+        if (h.fmt != WQ_VQ3R && h.fmt != WQ_VQ2R && h.fmt != WQ_VQ4P) { printf("  unexpected fmt %u\n", h.fmt); bad++; }
 
         /* the whole expert in ONE read — the point of the layout */
         const size_t bytes = (size_t)h.rec_4k_blocks * WASTE_ALIGN;

tools/make_test_container.py

@@ -33,8 +33,13 @@
 MAGIC_EXPERT = 0x50584557        # 'WEXP'
 MAGIC_CODEBOOK = 0x4B424357      # 'WCBK'
 ALIGN = 4096
-FMT_F32, FMT_Q8G, FMT_Q4G, FMT_VQ3R = 0, 2, 3, 4
+FMT_F32, FMT_Q8G, FMT_Q4G, FMT_VQ3R, FMT_VQ4P = 0, 2, 3, 4, 8
 VEC_DIM, CB_ENTRIES, STAGES, IDX_BLOCK = 8, 256, 3, 64
+# --index-bits 6 switches all three: the engine accepts 6 only as 4 stages
+# of 64 entries (src/model.c), which is also the only combination
+# convert.py writes. PACKED and INDEX_BITS are set in main() after the
+# arguments are read; the default leaves every byte below identical.
+PACKED, INDEX_BITS = False, 8
 GROUP = 128
 KINDS = ("gate", "up", "down")
 
@@ -186,6 +191,33 @@ def block_indices(idx, M, N):
     return bytes(out)
 
 
+def block_indices_packed6(idx, M, N):
+    """block_indices, then four 6-bit stages squeezed into three bytes per
+    row — the VQ4P layout. Same [M/B][pos][row_in_block] blocking; only the
+    trailing per-row run changes, from four whole bytes to three packed
+    ones, which keeps a VQ4P record the same size as VQ3R's. Little-endian
+    bit order, LSB of stage 0 at bit 0, byte-for-byte the packing
+    tools/convert.py's block_indices_packed writes, so the engine's
+    P6_J0..P6_J3 unpack recovers the stages in order."""
+    nvr = N // VEC_DIM
+    pad = (-M) % IDX_BLOCK
+    nb = (M + pad) // IDX_BLOCK
+    out = bytearray(nb * nvr * IDX_BLOCK * 3)
+    for b in range(nb):
+        for v in range(nvr):
+            for r in range(IDX_BLOCK):
+                row = b * IDX_BLOCK + r
+                if row >= M:
+                    continue                      # padding stays zero
+                src = row * nvr + v
+                s0, s1, s2, s3 = (idx[s][src] for s in range(4))
+                dst = ((b * nvr + v) * IDX_BLOCK + r) * 3
+                out[dst]     = (s0 | (s1 << 6)) & 0xFF
+                out[dst + 1] = ((s1 >> 2) | (s2 << 4)) & 0xFF
+                out[dst + 2] = ((s2 >> 4) | (s3 << 2)) & 0xFF
+    return bytes(out)
+
+
 def write_expert(f, layer, eid, cb_base, shapes, rng):
     hdr_size = 48
     off, offsets, body = hdr_size, [], bytearray()
@@ -194,7 +226,8 @@ def write_expert(f, layer, eid, cb_base, shapes, rng):
         idx = [[rng.randrange(CB_ENTRIES) for _ in range(nvec)]
                for _ in range(STAGES)]
         offsets.append(off)
-        b = block_indices(idx, M, N)
+        b = block_indices_packed6(idx, M, N) if PACKED \
+            else block_indices(idx, M, N)
         body += b
         off += len(b)
     corr_off = off
@@ -206,7 +239,8 @@ def write_expert(f, layer, eid, cb_base, shapes, rng):
     total = hdr_size + len(body)
     blocks = (total + ALIGN - 1) // ALIGN
     hdr = struct.pack("<IHHBBHHHIIIIIIII",
-                      MAGIC_EXPERT, layer, eid, FMT_VQ3R, 0, cb_base, 0, 0,
+                      MAGIC_EXPERT, layer, eid,
+                      FMT_VQ4P if PACKED else FMT_VQ3R, 0, cb_base, 0, 0,
                       blocks, offsets[0], offsets[1], offsets[2], corr_off,
                       zlib.crc32(bytes(body)) & 0xFFFFFFFF, 0, 0)
     assert len(hdr) == hdr_size
@@ -281,6 +315,13 @@ def main():
                     help="also write a small tiktoken-style vocabulary with "
                          "K3's XTML specials, so the container can be driven "
                          "with text instead of raw ids")
+    ap.add_argument("--index-bits", type=int, default=8, choices=(6, 8),
+                    help="6 writes a VQ4P container: 4 stages of 64 entries, "
+                         "indices packed 4x6 bits into 3 bytes, the only "
+                         "combination the engine accepts at index_bits 6 and "
+                         "the one convert.py writes. Default is 8, and at 8 "
+                         "every byte of the container is exactly what this "
+                         "script wrote before the flag existed.")
     ap.add_argument("--prefix", default="", metavar="PFX",
                     help="put the text tensors under a tensor_prefix, e.g. "
                          "language_model., and add one tensor outside it — "
@@ -314,6 +355,13 @@ def main():
                          "at 1.0. Unequal mscales put a ratio on cos/sin that "
                          "the engine does not apply, so it refuses instead")
     args = ap.parse_args()
+    if args.index_bits == 6:
+        # The engine validates index_bits 6 only as 4 stages of 64 entries
+        # (src/model.c), so there is no other combination to write. Both
+        # constants are read by the writers below; reassigning them here is
+        # the whole of the switch.
+        global STAGES, CB_ENTRIES, PACKED, INDEX_BITS
+        STAGES, CB_ENTRIES, PACKED, INDEX_BITS = 4, 64, True, 6
     rng = random.Random(args.seed)
     os.makedirs(args.out, exist_ok=True)
 
@@ -438,14 +486,27 @@ def main():
                               "bytes": total, "codebook_base": cb_base}
             cb_base += len(KINDS) * STAGES
 
+    # index_bits is absent from every container written before VQ4P and the
+    # engine reads the absence as 8 (src/model.c), so the default container
+    # keeps omitting the key — adding it would change manifest bytes that
+    # the rotary fixture hashes. Same for fmt and bits_per_weight, which at
+    # index_bits 8 stay exactly what this script wrote before the flag
+    # existed.
+    eq = {"fmt": "VQ3R", "stages": STAGES, "vec_dim": VEC_DIM,
+          "entries": CB_ENTRIES, "index_block": IDX_BLOCK,
+          "bits_per_weight": STAGES}
+    if PACKED:
+        eq["fmt"] = "VQ4P"
+        eq["index_bits"] = INDEX_BITS
+        # stages*index_bits/vec_dim = 3.00; 4*6/8, same as convert.py
+        eq["bits_per_weight"] = STAGES * INDEX_BITS / VEC_DIM
+
     manifest = {
         "format_version": 0,
         "arch": cfg["model_type"],
         "tensor_prefix": args.prefix,
         "config": cfg,
-        "expert_quant": {"fmt": "VQ3R", "stages": STAGES, "vec_dim": VEC_DIM,
-                         "entries": CB_ENTRIES, "index_block": IDX_BLOCK,
-                         "bits_per_weight": STAGES},
+        "expert_quant": eq,
         "layers": layers,
         "trunk": t.index,
     }