From bc710d5b172b49d497a38dcaa4cf56118ef68d1d Mon Sep 17 00:00:00 2001 From: Michael Feth Date: Wed, 26 Aug 2026 20:27:42 -0400 Subject: [PATCH] [lenny] feat: record-granularity bank striping across N devices (exp180) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit WASTE_BANK_SHARDS=dirA,dirB,... opens each layer's bank as N shard files, expert e on shard e%N at offset (e/N)*rec_bytes. Round-robin because the per-token demand is k experts of ONE layer — layer-granularity placement would leave every read of a token on one drive. Unstriped is the N=1 case of the same code path; env unset = exact prior behavior. - waste_bank: fd[16] + n_shards; all 6 deref sites migrated - load: manifest-driven shard open, fail-closed on short shard sets - bank_fetch: shard+offset resolution, byte-exact by construction - tools/split_banks.py: split + byte-for-byte verify modes - evidence: tiny.waste striped across 2 dirs, logits byte-identical (sha de62689a...), make check 48/0, ASan+UBSan striped run clean 0 findings --- src/model.c | 67 +++++++++++++++++++++++++++----- src/model.h | 8 +++- tools/split_banks.py | 91 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 155 insertions(+), 11 deletions(-) create mode 100644 tools/split_banks.py diff --git a/src/model.c b/src/model.c index 231247c..0835218 100644 --- a/src/model.c +++ b/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); diff --git a/src/model.h b/src/model.h index 03547ba..ed0ed57 100644 --- a/src/model.h +++ b/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; diff --git a/tools/split_banks.py b/tools/split_banks.py new file mode 100644 index 0000000..391606d --- /dev/null +++ b/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 --dirs /mnt/a,/mnt/b [--mode split] + python3 tools/split_banks.py --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()