Skip to content

Commit 09ada61

Browse files
committed
[lenny] feat: record-granularity bank striping across N devices (exp180)
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
1 parent 7f1fbba commit 09ada61

3 files changed

Lines changed: 154 additions & 11 deletions

File tree

src/model.c

Lines changed: 57 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1067,7 +1067,10 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap,
10671067
memset(m, 0, sizeof *m);
10681068
pthread_mutex_init(&m->fetch_mu, NULL);
10691069
m->trunk_fd = -1;
1070-
for (int L = 0; L < WASTE_MAX_LAYERS; L++) m->bank[L].fd = -1;
1070+
for (int L = 0; L < WASTE_MAX_LAYERS; L++) {
1071+
for (int s = 0; s < WASTE_MAX_SHARDS; s++) m->bank[L].fd[s] = -1;
1072+
m->bank[L].n_shards = 1;
1073+
}
10711074
m->want_vision = opt->want_vision;
10721075
m->want_direct = opt->direct_io;
10731076
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,
13441347
js_free(&d); free(src); return -2;
13451348
}
13461349
m->bank[L].rec_bytes = m->bank[L].n_experts ? bytes / m->bank[L].n_experts : 0;
1347-
m->bank[L].fd = bank_open(path, m->bank[L].rec_bytes, m->want_direct,
1348-
&m->direct_io);
1349-
if (m->bank[L].fd < 0) { js_free(&d); free(src); return -1; }
1350+
/* Unstriped open is the N=1 case of the striped one: one fd on the
1351+
* bank's own file. A WASTE_BANK_SHARDS manifest (comma-separated
1352+
* directories) reopens the bank as N shard files named after the
1353+
* bank's basename in each directory. Every shard must exist and be
1354+
* exactly the size its share of experts demands — a short or long
1355+
* shard fails the load rather than serving a wrong record. */
1356+
{
1357+
const char *sh = getenv("WASTE_BANK_SHARDS");
1358+
char dirs[1024];
1359+
int n_sh = 1;
1360+
if (sh && *sh) {
1361+
snprintf(dirs, sizeof dirs, "%s", sh);
1362+
n_sh = 1;
1363+
for (const char *p = dirs; *p; p++) if (*p == ',') n_sh++;
1364+
if (n_sh > WASTE_MAX_SHARDS) { js_free(&d); free(src); return -2; }
1365+
}
1366+
const char *base = strrchr(fn, '/'); base = base ? base + 1 : fn;
1367+
for (int s = 0; s < n_sh; s++) {
1368+
char spath[1152];
1369+
if (n_sh == 1)
1370+
snprintf(spath, sizeof spath, "%s/%s", dir, fn);
1371+
else {
1372+
char *dstart = dirs;
1373+
for (int k = 0; k < s && dstart; k++) {
1374+
dstart = strchr(dstart, ',');
1375+
if (dstart) dstart++;
1376+
}
1377+
if (!dstart || !*dstart) { js_free(&d); free(src); return -2; }
1378+
char *comma = strchr(dstart, ',');
1379+
size_t dlen = comma ? (size_t)(comma - dstart) : strlen(dstart);
1380+
snprintf(spath, sizeof spath, "%.*s/%s", (int)dlen, dstart, base);
1381+
}
1382+
int sfd = bank_open(spath, m->bank[L].rec_bytes, m->want_direct,
1383+
&m->direct_io);
1384+
if (sfd < 0) { js_free(&d); free(src); return -1; }
1385+
m->bank[L].fd[s] = sfd;
1386+
m->bank[L].n_shards = n_sh;
1387+
}
1388+
}
1389+
if (m->bank[L].fd[0] < 0) { js_free(&d); free(src); return -1; }
13501390
}
13511391
js_free(&d);
13521392
free(src);
@@ -1539,7 +1579,8 @@ void waste_model_free(waste_model *m)
15391579
free(m->codebooksT);
15401580
for (int L = 0; L < 128; L++) {
15411581
free(m->S[L]); free(m->conv[L]); free(m->latcache[L]);
1542-
if (m->bank[L].fd >= 0) close(m->bank[L].fd);
1582+
for (int s = 0; s < WASTE_MAX_SHARDS; s++)
1583+
if (m->bank[L].fd[s] >= 0) close(m->bank[L].fd[s]);
15431584
}
15441585
free(m->x); free(m->h); free(m->tmp); free(m->att); free(m->logits);
15451586
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)
17771818
layer >= WASTE_MAX_LAYERS;
17781819
waste_bank *b = bad_layer ? NULL : &m->bank[layer];
17791820
if (bad_layer || expert < 0 || expert >= b->n_experts ||
1780-
b->fd < 0 || b->rec_bytes <= 0)
1821+
b->fd[0] < 0 || b->rec_bytes <= 0)
17811822
return bank_fail(m, REC_E_HEADER, layer, expert);
17821823

1783-
/* pread is positional, so the reader threads share the bank's fd
1784-
* without a seek to race over. */
1785-
const int64_t got = waste_pread(b->fd, dst, (size_t)b->rec_bytes,
1786-
(int64_t)expert * (int64_t)b->rec_bytes);
1824+
/* Striped banks: expert e lives on shard e % n_shards at offset
1825+
* (e / n_shards) * rec_bytes. Round-robin keeps the top-k experts of a
1826+
* single token spread across devices, which is the point — the per-token
1827+
* demand is k experts of ONE layer, so layer-granularity placement would
1828+
* leave every read of a token on one drive. Unstriped is N=1 and the
1829+
* division is exact: shard 0, offset e * rec_bytes. pread is positional,
1830+
* so the reader threads share each shard fd without a seek to race over. */
1831+
const int shard = expert % b->n_shards;
1832+
const int64_t off = ((int64_t)(expert / b->n_shards)) * (int64_t)b->rec_bytes;
1833+
const int64_t got = waste_pread(b->fd[shard], dst, (size_t)b->rec_bytes, off);
17871834
rec_status st = got == (int64_t)b->rec_bytes ? REC_OK : REC_E_READ;
17881835
if (st == REC_OK) st = record_check(m, layer, expert, dst);
17891836
if (st != REC_OK) return bank_fail(m, st, layer, expert);

src/model.h

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ typedef struct {
5050
int kv_lora, q_lora, qk_nope, qk_rope, v_head;
5151
int kda_heads, kda_dim, conv_k;
5252
#define WASTE_MAX_LAYERS 128
53+
#define WASTE_MAX_SHARDS 16
5354

5455
/* Vector positions sharing one fp32 scale in the int8 LUT (WQ_VQ4P).
5556
* Bounds the int16 accumulator: 4 stages x 32 positions x 127 = 16256, so
@@ -103,7 +104,12 @@ typedef struct {
103104
} waste_config;
104105

105106
typedef struct {
106-
int fd; /* positional reads, no page cache */
107+
/* Positional reads, no page cache. With N>1 shards the bank's experts are
108+
* round-robin split across devices: expert e lives on shard e % n_shards
109+
* at offset (e / n_shards) * rec_bytes. fd[0] is the unstriped fd when
110+
* n_shards == 1, so the unstriped path is the same code with N=1. */
111+
int fd[WASTE_MAX_SHARDS];
112+
int n_shards;
107113
int64_t rec_bytes;
108114
int n_experts, cb_base;
109115
} waste_bank;

tools/split_banks.py

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
#!/usr/bin/env python3
2+
"""split_banks.py — split a WASTE container's expert banks into N shard sets.
3+
4+
Round-robin placement matching src/model.c bank_fetch: expert e lives on
5+
shard e % N at slot e // N. Byte-exact by construction — the same record
6+
bytes land in the same logical order; only device placement changes.
7+
8+
Two modes:
9+
--mode split write N shard directories, each holding shard files
10+
--mode verify check an existing shard set against the source bank
11+
(reads both, compares every record byte-for-byte)
12+
13+
Usage:
14+
python3 tools/split_banks.py <container.waste> --dirs /mnt/a,/mnt/b [--mode split]
15+
python3 tools/split_banks.py <container.waste> --dirs /mnt/a,/mnt/b --mode verify
16+
17+
The container's manifest.json lists per-layer bank files under "layers".
18+
This tool reads the manifest, never the trunk. Shards are plain files named
19+
after the bank's basename, placed one per directory. The engine opens them
20+
via WASTE_BANK_SHARDS="/mnt/a,/mnt/b".
21+
"""
22+
import argparse
23+
import json
24+
import os
25+
import sys
26+
27+
def bank_manifest(container):
28+
mf = os.path.join(container, "manifest.json")
29+
if not os.path.exists(mf):
30+
# some containers nest it
31+
for cand in ("waste.json", "index.json"):
32+
p = os.path.join(container, cand)
33+
if os.path.exists(p):
34+
mf = p
35+
break
36+
else:
37+
sys.exit(f"no manifest.json under {container}")
38+
with open(mf) as f:
39+
return json.load(f), mf
40+
41+
def main():
42+
ap = argparse.ArgumentParser()
43+
ap.add_argument("container")
44+
ap.add_argument("--dirs", required=True, help="comma-separated shard dirs")
45+
ap.add_argument("--mode", choices=("split", "verify"), default="split")
46+
a = ap.parse_args()
47+
dirs = [d for d in a.dirs.split(",") if d]
48+
if len(dirs) < 2:
49+
sys.exit("need at least 2 dirs")
50+
man, mf = bank_manifest(a.container)
51+
root = os.path.dirname(mf)
52+
layers = man.get("layers", {})
53+
n_done = 0
54+
for key, ent in sorted(layers.items(), key=lambda kv: int(kv[0])):
55+
fn = ent["file"]
56+
src = os.path.join(root, fn)
57+
n_exp = int(ent["experts"])
58+
if not os.path.exists(src):
59+
sys.exit(f"bank missing: {src}")
60+
rec = os.path.getsize(src) // n_exp
61+
assert os.path.getsize(src) == rec * n_exp, f"bank not divisible: {src}"
62+
if a.mode == "split":
63+
outs = []
64+
for d in dirs:
65+
os.makedirs(d, exist_ok=True)
66+
outs.append(open(os.path.join(d, os.path.basename(fn)), "wb"))
67+
with open(src, "rb") as f:
68+
for e in range(n_exp):
69+
blob = f.read(rec)
70+
assert len(blob) == rec
71+
outs[e % len(dirs)].write(blob)
72+
for o in outs:
73+
o.close()
74+
print(f"layer {key}: {n_exp} experts x {rec}B -> {len(dirs)} shards")
75+
else:
76+
fhs = [open(os.path.join(d, os.path.basename(fn)), "rb") for d in dirs]
77+
with open(src, "rb") as f:
78+
for e in range(n_exp):
79+
want = f.read(rec)
80+
got = fhs[e % len(dirs)].read(rec)
81+
if want != got:
82+
sys.exit(f"MISMATCH layer {key} expert {e}")
83+
for fh in fhs:
84+
fh.close()
85+
print(f"layer {key}: VERIFY OK ({n_exp} records byte-identical)")
86+
n_done += 1
87+
print(f"{a.mode} complete: {n_done} layers, {len(dirs)} shard dirs")
88+
89+
if __name__ == "__main__":
90+
main()

0 commit comments

Comments
 (0)