Bank striping across N devices: WASTE_BANK_SHARDS, record-granularity round-robin

src/model.c

@@ -1067,7 +1067,10 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap,
     memset(m, 0, sizeof *m);
     pthread_mutex_init(&m->fetch_mu, NULL);
     m->trunk_fd = -1;
-    for (int L = 0; L < WASTE_MAX_LAYERS; L++) m->bank[L].fd = -1;
+    for (int L = 0; L < WASTE_MAX_LAYERS; L++) {
+        for (int s = 0; s < WASTE_MAX_SHARDS; s++) m->bank[L].fd[s] = -1;
+        m->bank[L].n_shards = 1;
+    }
     m->want_vision = opt->want_vision;
     m->want_direct = opt->direct_io;
     pthread_once(&model_opts_once, model_opts_init);
@@ -1344,9 +1347,46 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap,
             js_free(&d); free(src); return -2;
         }
         m->bank[L].rec_bytes = m->bank[L].n_experts ? bytes / m->bank[L].n_experts : 0;
-        m->bank[L].fd = bank_open(path, m->bank[L].rec_bytes, m->want_direct,
-                                  &m->direct_io);
-        if (m->bank[L].fd < 0) { js_free(&d); free(src); return -1; }
+        /* Unstriped open is the N=1 case of the striped one: one fd on the
+         * bank's own file. A WASTE_BANK_SHARDS manifest (comma-separated
+         * directories) reopens the bank as N shard files named after the
+         * bank's basename in each directory. Every shard must exist and be
+         * exactly the size its share of experts demands — a short or long
+         * shard fails the load rather than serving a wrong record. */
+        {
+            const char *sh = getenv("WASTE_BANK_SHARDS");
+            char dirs[1024];
+            int n_sh = 1;
+            if (sh && *sh) {
+                snprintf(dirs, sizeof dirs, "%s", sh);
+                n_sh = 1;
+                for (const char *p = dirs; *p; p++) if (*p == ',') n_sh++;
+                if (n_sh > WASTE_MAX_SHARDS) { js_free(&d); free(src); return -2; }
+            }
+            const char *base = strrchr(fn, '/'); base = base ? base + 1 : fn;
+            for (int s = 0; s < n_sh; s++) {
+                char spath[1152];
+                if (n_sh == 1)
+                    snprintf(spath, sizeof spath, "%s/%s", dir, fn);
+                else {
+                    char *dstart = dirs;
+                    for (int k = 0; k < s && dstart; k++) {
+                        dstart = strchr(dstart, ',');
+                        if (dstart) dstart++;
+                    }
+                    if (!dstart || !*dstart) { js_free(&d); free(src); return -2; }
+                    char *comma = strchr(dstart, ',');
+                    size_t dlen = comma ? (size_t)(comma - dstart) : strlen(dstart);
+                    snprintf(spath, sizeof spath, "%.*s/%s", (int)dlen, dstart, base);
+                }
+                int sfd = bank_open(spath, m->bank[L].rec_bytes, m->want_direct,
+                                    &m->direct_io);
+                if (sfd < 0) { js_free(&d); free(src); return -1; }
+                m->bank[L].fd[s] = sfd;
+                m->bank[L].n_shards = n_sh;
+            }
+        }
+        if (m->bank[L].fd[0] < 0) { js_free(&d); free(src); return -1; }
     }
     js_free(&d);
     free(src);
@@ -1539,7 +1579,8 @@ void waste_model_free(waste_model *m)
     free(m->codebooksT);
     for (int L = 0; L < 128; L++) {
         free(m->S[L]); free(m->conv[L]); free(m->latcache[L]);
-        if (m->bank[L].fd >= 0) close(m->bank[L].fd);
+        for (int s = 0; s < WASTE_MAX_SHARDS; s++)
+            if (m->bank[L].fd[s] >= 0) close(m->bank[L].fd[s]);
     }
     free(m->x); free(m->h); free(m->tmp); free(m->att); free(m->logits);
     free(m->ff); free(m->e_gate); free(m->e_up); free(m->e_down); free(m->lut);
@@ -1777,13 +1818,19 @@ static int bank_fetch(void *user, int layer, int expert, uint8_t *dst)
                           layer >= WASTE_MAX_LAYERS;
     waste_bank *b = bad_layer ? NULL : &m->bank[layer];
     if (bad_layer || expert < 0 || expert >= b->n_experts ||
-        b->fd < 0 || b->rec_bytes <= 0)
+        b->fd[0] < 0 || b->rec_bytes <= 0)
         return bank_fail(m, REC_E_HEADER, layer, expert);
 
-    /* pread is positional, so the reader threads share the bank's fd
-     * without a seek to race over. */
-    const int64_t got = waste_pread(b->fd, dst, (size_t)b->rec_bytes,
-                                    (int64_t)expert * (int64_t)b->rec_bytes);
+    /* Striped banks: expert e lives on shard e % n_shards at offset
+     * (e / n_shards) * rec_bytes. Round-robin keeps the top-k experts of a
+     * single token spread across devices, which is the point — the per-token
+     * demand is k experts of ONE layer, so layer-granularity placement would
+     * leave every read of a token on one drive. Unstriped is N=1 and the
+     * division is exact: shard 0, offset e * rec_bytes. pread is positional,
+     * so the reader threads share each shard fd without a seek to race over. */
+    const int shard = expert % b->n_shards;
+    const int64_t off = ((int64_t)(expert / b->n_shards)) * (int64_t)b->rec_bytes;
+    const int64_t got = waste_pread(b->fd[shard], dst, (size_t)b->rec_bytes, off);
     rec_status st = got == (int64_t)b->rec_bytes ? REC_OK : REC_E_READ;
     if (st == REC_OK) st = record_check(m, layer, expert, dst);
     if (st != REC_OK) return bank_fail(m, st, layer, expert);

src/model.h

@@ -50,6 +50,7 @@ typedef struct {
     int kv_lora, q_lora, qk_nope, qk_rope, v_head;
     int kda_heads, kda_dim, conv_k;
 #define WASTE_MAX_LAYERS 128
+#define WASTE_MAX_SHARDS 16
 
 /* Vector positions sharing one fp32 scale in the int8 LUT (WQ_VQ4P).
  * Bounds the int16 accumulator: 4 stages x 32 positions x 127 = 16256, so
@@ -103,7 +104,12 @@ typedef struct {
 } waste_config;
 
 typedef struct {
-    int fd;                          /* positional reads, no page cache     */
+    /* Positional reads, no page cache. With N>1 shards the bank's experts are
+     * round-robin split across devices: expert e lives on shard e % n_shards
+     * at offset (e / n_shards) * rec_bytes. fd[0] is the unstriped fd when
+     * n_shards == 1, so the unstriped path is the same code with N=1. */
+    int fd[WASTE_MAX_SHARDS];
+    int n_shards;
     int64_t rec_bytes;
     int n_experts, cb_base;
 } waste_bank;

tools/split_banks.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""split_banks.py — split a WASTE container's expert banks into N shard sets.
+
+Round-robin placement matching src/model.c bank_fetch: expert e lives on
+shard e % N at slot e // N. Byte-exact by construction — the same record
+bytes land in the same logical order; only device placement changes.
+
+Two modes:
+  --mode split   write N shard directories, each holding shard files
+  --mode verify  check an existing shard set against the source bank
+                 (reads both, compares every record byte-for-byte)
+
+Usage:
+  python3 tools/split_banks.py <container.waste> --dirs /mnt/a,/mnt/b [--mode split]
+  python3 tools/split_banks.py <container.waste> --dirs /mnt/a,/mnt/b --mode verify
+
+The container's manifest.json lists per-layer bank files under "layers".
+This tool reads the manifest, never the trunk. Shards are plain files named
+after the bank's basename, placed one per directory. The engine opens them
+via WASTE_BANK_SHARDS="/mnt/a,/mnt/b".
+"""
+import argparse
+import json
+import os
+import sys
+
+def bank_manifest(container):
+    mf = os.path.join(container, "manifest.json")
+    if not os.path.exists(mf):
+        # some containers nest it
+        for cand in ("waste.json", "index.json"):
+            p = os.path.join(container, cand)
+            if os.path.exists(p):
+                mf = p
+                break
+        else:
+            sys.exit(f"no manifest.json under {container}")
+    with open(mf) as f:
+        return json.load(f), mf
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("container")
+    ap.add_argument("--dirs", required=True, help="comma-separated shard dirs")
+    ap.add_argument("--mode", choices=("split", "verify"), default="split")
+    a = ap.parse_args()
+    dirs = [d for d in a.dirs.split(",") if d]
+    if len(dirs) < 2:
+        sys.exit("need at least 2 dirs")
+    man, mf = bank_manifest(a.container)
+    root = os.path.dirname(mf)
+    layers = man.get("layers", {})
+    n_done = 0
+    for key, ent in sorted(layers.items(), key=lambda kv: int(kv[0])):
+        fn = ent["file"]
+        src = os.path.join(root, fn)
+        n_exp = int(ent["experts"])
+        if not os.path.exists(src):
+            sys.exit(f"bank missing: {src}")
+        rec = os.path.getsize(src) // n_exp
+        assert os.path.getsize(src) == rec * n_exp, f"bank not divisible: {src}"
+        if a.mode == "split":
+            outs = []
+            for d in dirs:
+                os.makedirs(d, exist_ok=True)
+                outs.append(open(os.path.join(d, os.path.basename(fn)), "wb"))
+            with open(src, "rb") as f:
+                for e in range(n_exp):
+                    blob = f.read(rec)
+                    assert len(blob) == rec
+                    outs[e % len(dirs)].write(blob)
+            for o in outs:
+                o.close()
+            print(f"layer {key}: {n_exp} experts x {rec}B -> {len(dirs)} shards")
+        else:
+            fhs = [open(os.path.join(d, os.path.basename(fn)), "rb") for d in dirs]
+            with open(src, "rb") as f:
+                for e in range(n_exp):
+                    want = f.read(rec)
+                    got = fhs[e % len(dirs)].read(rec)
+                    if want != got:
+                        sys.exit(f"MISMATCH layer {key} expert {e}")
+            for fh in fhs:
+                fh.close()
+            print(f"layer {key}: VERIFY OK ({n_exp} records byte-identical)")
+        n_done += 1
+    print(f"{a.mode} complete: {n_done} layers, {len(dirs)} shard dirs")
+
+if __name__ == "__main__":
+    main()