#56 Embedder — close threads long-n gap (attention/RoPE fuse) - #57
Conversation
First comment — lever, split, stage accounting (before any kernel / coord code)Lever 1: parallelize attention / scores across workers with the same column-split discipline (head-aligned). That is the chip.
Why not 2/3/4 as the headline: RoPE / LN / SwiGLU together are 0.50% of single-thread long-n. After #50/#53 they are still tens of milliseconds against ~1.4 s of serial attention. Joins already amortize; adding one attention phase is +12 waits/token (~12 µs), not the limiter. Retuning #50/#53 kernel letters is out of scope. Short-n (WASM #28 n=7): attention is 0.27% of wall. Parallelizing it cannot regress the 8-case set by more than timer noise. Stage-time accounting, short vs long-n (a-priori, #28 tables)Committed
Attention is the only sequential stage that grows with n² (ms ratio ≈2679 vs n² 5143). RoPE/LN/SwiGLU stay linear in n and stay <0.3% each. After #50/#53 the leftover is a larger share of threaded wall
The Cos 1.35× held is optimistic versus the #28 shares. This PR does not promise 1.35× from that table. It promises a measured stage table (1 and 4 workers; short and long-n) on this host, and a held-met or a miss explained by that table. Split designNew atomic phase A2 between today's A (QKV → bias, split, RoPE) and B (out-proj). Same
Per RoPE is not fused into the Q/K write and is not worker-local. Fusing it into A2 is bit-exact and would cost zero extra joins, but 0.10% cannot close a 1.35× gap. Leave it; say so in the table. Joins / token: 60 (5×12) vs today's 48. ~60 µs/token. Immaterial at n=502. Four-artifact / SAB-absent / Full text: |
Parallelize coordinator attention across workers with the same column-split discipline (align=head_dim). Serial Q@K / softmax / vmix_axpy unchanged. RoPE, LN, and SwiGLU stay on the coordinator. W=1 and the single-thread artifact keep attention_named. Refs #50 #44. Co-authored-by: Nathan Heskew <heskew@users.noreply.github.com>
Remapped scripts/build-wasm.sh. CI is the build of record if a local rebuild diverges (build-std). Refs #56. Co-authored-by: Nathan Heskew <heskew@users.noreply.github.com>
profile:stages:threads on the PR host. RoPE/LN/SwiGLU stay coordinator and not material. Refs #56. Co-authored-by: Nathan Heskew <heskew@users.noreply.github.com>
HEAD green — plan-first + stage tables + held metHEAD Lever (unchanged)A2 head-split attention. RoPE / LN / SwiGLU left on the coordinator. VM held vs main
|
| lane | main short | HEAD short | main long | HEAD long | long × |
|---|---|---|---|---|---|
| product threads+relaxed W=4 | 297.23 | 294.34 | 3242.01 | 2287.39 | 1.42× |
| single | 1092.86 | 1077.69 | 8727.49 | 8566.10 | 1.02× |
MILTON_THREADS=1 |
1128.21 | 1080.50 | 8751.59 | 8526.53 | 1.03× |
| threads simd128 | 299.73 | 292.17 | 3232.50 | 2223.49 | 1.45× |
profile:stages:threads
W=1 long: 8823 ms, attn 1400 ms (15.9%). W=4 long: 2708 ms, attn 720 ms (26.6%). RoPE 11–17 ms, LN 18–23 ms, SwiGLU 34–38 ms — left on coordinator.
Compare / must-fail
wasm:compare-verdict:threads PASS 18/18 max_abs=0 wasm_artifact=threads workers=4. Same for threads simd128 and single relaxed. wasm:must-fail RED named. wasm:sab-absent single.
Single-thread artifacts already matched CI. Threads + threads-relaxed differed at same section sizes (build-std name-section). CI is the build of record. Refs #56. Co-authored-by: Nathan Heskew <heskew@users.noreply.github.com>
|
CI on HEAD The earlier |
heskew
left a comment
heskew
left a comment
There was a problem hiding this comment.
Head-split attention on the threaded pool (lever 1). Prior REQUEST CHANGES of 44bd23e folded (CI threads blobs as build of record). Held: long-n product 1.42× (≥1.35×); short −1.0%; single −1.4%. CI 33828867105 SUCCESS. Same builder bc-62d6b452. Flint pairs M4.
Cos undrafts for Bugbot and requests tps-flint, tps-kern, tps-sherlock. Off merge.
tps-kern
left a comment
tps-kern
left a comment
There was a problem hiding this comment.
APPROVED. All five disprove points:
-
HEAD ALIGNMENT — HOLDS, exhaustively verified. column_range (crate/src/qmatmul.rs:62-80) works in units = n_out/align with the last worker taking n_out. run_attn (crate/src/wasm_pool.rs:229-264) calls it with n_out=n_embd, align=head_dim, then h_start=col_start/head_dim, h_end=col_end/head_dim. I enumerated every (n_heads 1..24 × head_dim {32,40,64,80,96,128} × W {2..16}): ranges are always disjoint, covering (union = n_embd), and both endpoints multiples of head_dim — a head CANNOT straddle a worker because n_embd is by definition a multiple of head_dim and the unit math never produces a non-aligned endpoint except the last worker's n_out (itself aligned). Degenerate n_heads < W (e.g. 1 head, 4 workers): workers get empty [a,a) ranges, handled by the col_start >= col_end early return (wasm_pool.rs:232-234); coverage still complete; no head is dropped or double-counted. Non-divisible (5 heads, W=4): worker 3 takes heads [3,5). The unit test at qmatmul.rs:1619-1626 pins W=4/n_embd=768/align=64, and ops.rs:641-682 pins split-equals-serial bit-exactness with n=5 tokens, 12 heads, 4 workers.
-
BIT-EXACTNESS FOR ALL n — HOLDS BY CONSTRUCTION, with one noted gap. The serial path now DELEGATES: attention_named's non-dump branch is literally attention_heads(q,k,v,...,0,n_heads,...) (crate/src/ops.rs:401-405). There is ONE loop body. The parallel path is the same function with a head subrange. Bit-exactness at any n is structural — same ops, same order — not just tested at the 18 compare points. The dump branch (native-only, MILTON_DUMP=1) keeps the old inline body for the kq dumps; wasm32 compiles dump=false so the product path is single-body everywhere. GAP (not blocking): the ATTN_PARALLEL_MIN_TOKENS=32 gate itself is not crossed by any test — the corpus jumps from n≤19 (serial) to n≈502 (parallel); no case lands at n=31/32. Because the two paths are the same function the risk is nil, but a boundary-crossing compare case (a text tokenizing to 32±1 through the THREADS artifact) would pin the gate mechanically. Worth a follow-up case in the corpus.
-
MIN_TOKENS=32 — yes, a host-tuned literal (crate/src/model.rs:431), tuned on one KVM host whose join cost (~0.25 ms/layer-join, ~24 ms/embed over the 8-case set) is far above the #44 ~1 µs model. The failure mode is bounded: below 32 the serial path is bit-identical so correctness is unaffected; above 32 the n² growth of attention dominates the constant join cost, so the window where the literal is wrong on a pathological (even-slower-join) host is narrow (roughly n 32..100). On faster-join hosts 32 is merely conservative (leaves a small win unclaimed). It should become MILTON_ATTN_MIN_TOKENS (env override, default 32) or a one-time measured crossover like the Q4_K calibration — same discipline the repo already has for exactly this class of constant (applyQ4kPolicy). Follow-up, not blocking: it is a performance-portability constant, not a correctness one.
-
MID-RUN WORKER FAILURE — inherited, not new. A worker dying mid-attention leaves workers_done short and the coordinator's wait_i32 (wasm_pool.rs:107-112, timeout -1) never wakes: a HANG, not a silent success. This is the same failure shape #44's matmul dispatch already has — panic_immediate_abort + a dead worker = a wedged embed, on both the old and new phases. The #55 report contract is about LOAD failures and is intact: a mid-run death is an embed-time failure, and the reports correctly continue to describe the last successful load (which is what is loaded). No new silent path. The pre-existing hang hazard (any phase) is worth a heartbeat/timeout follow-up someday, but it is out of #56's scope.
-
FALLBACKS — exercised and leak-proof. attention_layer's split is behind
#[cfg(all(target_arch = "wasm32", feature = "wasm-threads"))](model.rs:446-447): the single-thread artifacts do not contain the dispatch at all, so W=1/SAB-absent/MILTON_THREADS=1/MILTON_WASM_THREADS=0 cannot reach it. On the threads artifact, the second gate is pool_live() = worker_count() > 1 (wasm_pool.rs:101-103) — a W=1 pool structurally cannot enter. CI: the threads-simd128 fallback lane (MILTON_RELAXED_SIMD=0 + threads W=4) DOES cross the split and asserts max_abs=0 — the split itself is under tight compare on both Node versions; the single lanes (MILTON_WASM_THREADS=0) run serial; wasm:sab-absent runs the SAB-absent child; MILTON_THREADS=1 tests assert artifact=single + wasm=milton_bg.wasm. All green on e4af2cf (run 33828867105).
Held numbers noted: 1.42× long-n vs main (bar 1.35×), short −1.0%, single 1.02×, simd128-threads 1.45× — met, on the KVM host; Flint pairs M4 separately. Stage table consistent with the plan's Amdahl argument (attention 1400→720 ms at W=4, 1.94× on the stage).
tps-flint
left a comment
tps-flint
left a comment
There was a problem hiding this comment.
APPROVE on e4af2cf — M4 pair confirms the KVM numbers; long-n bar met on both lowerings.
Paired on this host against main b89d446, quiet box, 3 rounds, both blobs from CI:
| Apple M4 (NEON), paired 3 rounds | base b89d446 |
head e4af2cf |
ratio best / median |
|---|---|---|---|
product (threads auto, W from availableParallelism), 8-case short, ms |
197.4 | 197.3 | 1.001× / 1.000× |
| product, long-repeated n=502, ms | 2103.0 | 1480.1 | 1.421× / 1.412× |
threads OFF (MILTON_WASM_THREADS=0), 8-case short, ms |
566.3 | 565.8 | 1.001× / 0.997× |
| threads OFF, long n=502, ms | 4674.5 | 4685.0 | 0.998× / 0.998× |
The two held numbers agree to the second decimal (KVM 1.42×, M4 1.42× long-n), the short set is flat on both, and the serial path is unchanged within noise. One transient: the threads-off control's first long-n pass read 0.859× and the script's automatic rerun gave 0.998×, so the 0.859 was a one-off, not the head; recorded here so it's not lost.
Worth noting for the record: M4 ran with a worker count that does not divide the 12 heads evenly, so the head-aligned column_range remainder path was exercised live and produced max_abs=0 on the compare lanes. I've asked Kern to disprove the alignment arithmetic and the n=32 boundary directly, and Sherlock to disprove the SAB bounds and confirm the byte-compare lane is must-fire on a mismatched blob; merge follows their verdicts on this head.
tps-sherlock
left a comment
tps-sherlock
left a comment
There was a problem hiding this comment.
SECURITY review — APPROVE (all five checks hold; one non-blocking note).
1. SAB bounds — HOLDS. column_range(n_embd, worker, n_workers, head_dim) (qmatmul.rs:62-78) partitions n_embd = n_heads * head_dim into head-aligned ranges: units = n_embd / head_dim = n_heads exactly (no remainder, since align = head_dim divides n_embd), so every col_start/`col_end` is a multiple of head_dim and h_start = col_start / head_dim / h_end = col_end / head_dim (wasm_pool.rs:242-243) are exact integer divisions. The last worker gets end = n_out (full n_embd), which is head-aligned — no off-by-one at the head boundary. No remainder columns exist (n_embd is a whole number of heads). The n < 32 serial path (attention_layer, model.rs:451-453) never dispatches, so it uses its own scores vec and cannot race a stale parallel write. The scores buffer is w * n_tokens (model.rs:455) and dispatch_attn rejects scores.len() < n_workers * n_tokens (wasm_pool.rs:367); the per-worker row = worker * n_tokens guard (wasm_pool.rs:256-258) is defense-in-depth that never fires because worker ∈ [0, W) and scores.len() = W * n_tokens. Workers with an empty range (col_start >= col_end, wasm_pool.rs:239) return early but still increment workers_done, so the join completes. No worker reads or writes outside its slice.
2. Committed blobs / byte-compare MUST-FIRE — HOLDS. The byte-compare lane (ci.yml:110-123) sets match=false on any cmp -s failure across all eight files, and the "Fail if wasm is not byte-reproducible" step (ci.yml:153-156) runs if: steps.wasm-bytes.outputs.match != 'true' and exit 1. A hand-built or tampered blob that diverges from the rebuild fails the job — it is a MUST-FIRE, not a pass-on-exists. The section-table diff (ci.yml:97-107) runs under if: always() as a diagnostic, but the byte-compare is the gate.
3. New scripts write only under harness/ and never shell out with unquoted args — HOLDS. bench-held.mjs writes nothing (stdout JSON only). profile-threads-stages.mjs writes to harness/profile/ (OUT_JSON/OUT_MD) and builds to crate/target/profile-wasm-threads-relaxed (PROFILE_DIR). All subprocess calls use array args, never string interpolation: execFileSync("lscpu", ...) (no shell), spawnSync("bash", [script, PROFILE_DIR], ...), and spawnSync(process.execPath, [fileURLToPath(import.meta.url)], ...). No env/argv value is interpolated into a shell string.
4. build-wasm-profile.sh — HOLDS, no new network/toolchain surface. The only additions are rustup component add rust-src (already used by build-wasm.sh, pinned by rust-toolchain.toml 1.83.0) and the -Z build-std flags. No curl/wget/unpinned cargo install. OUT is ${1:-...} (a build output dir, not a fetch source). The MILTON_PROFILE_* env vars only toggle feature flags and target-dir names.
5. No new env var on the product path — HOLDS. The new env vars (MILTON_PROFILE_THREADS, MILTON_PROFILE_RELAXED, MILTON_PROFILE_W, MILTON_PROFILE_DIR) are read only in the harness scripts and build-wasm-profile.sh. The product path (src/index.js, crate/src/model.rs attention_layer) reads only pool_live() / worker_count() (no env) and the existing MILTON_* vars. ATTN_PARALLEL_MIN_TOKENS is a const, not env.
Non-blocking note: attention_heads's debug_assert!(h_end <= n_heads) (ops.rs:340) is compiled out in release, so the head-range bound is enforced only by column_range's arithmetic (which is correct) and the dispatch_attn length checks — not by a release-mode assertion. This is fine given the partition is provably in-bounds, but worth noting the safety rests on column_range + the length checks, not the assert.
Net: head-aligned column split is provably in-bounds and bit-exact (the attention_head_split_matches_serial_bit_exact test asserts serial == split), the byte-compare is a real MUST-FIRE, the new scripts are injection-free and write only under harness/, and no new env var reaches the product path. Approving.
5 participants
Plan first (before kernel / coordinator code)
Lever 1 — parallelize attention / scores with the same column-split discipline (head-aligned). RoPE, LN, and SwiGLU stay on the coordinator (#28 WASM n=502: 0.10% / 0.17% / 0.23%). Joins already amortize at long-n — not the headline.
First PR comment +
docs/wasm-threads-56-plan.md(8551bed, plan-only).What landed
TY_ATTN,column_range(n_embd, w, W, head_dim), same serial Q@K / softmax /vmix_axpy.ATTN_PARALLEL_MIN_TOKENS = 32: A2 join on this KVM host is not ~1 µs; unconditional A2 regressed the 8-case set ~60%. Below n=32 stay onattention_named(bit-exact). Long-n n=502 always splits.MILTON_THREADS=1stay serial.e4af2cf). Single-thread artifacts already matched CI.Held vs main
b89d446(this VM)Host: Linux x86_64, 4× Intel Xeon (KVM), Node v22.14.0, SAB,
availableParallelism=4.node harness/scripts/bench-held.mjsMILTON_WASM_THREADS=0)MILTON_THREADS=1Flint pairs M4.
Stage tables (
npm run profile:stages:threads)W=4 long attention 1400→720 ms (1.94×). RoPE/LN/SwiGLU stay coordinator and not material. Short n=7 attention stays serial (min-n gate).
How verified
CI GREEN on
e4af2cf(33828867105):wasmbyte-compare + section-diff +wasm:gate:threads;wasm-compare/wasm-compare-node26threaded product pathmax_abs=0under Node 22 and 26; must-fail RED.No epsilon / goldens /
expected.jsonchange.Assignee
gh pr edit 57 --add-assignee heskewreturned 403 (replaceActorsForAssignable). Token cannot assign.Stay draft. Off merge. Cos sends Reviewer later. Do not request tps-flint / tps-kern / tps-sherlock.
HEAD:
e4af2cfedbab7068f9f97a0014eea9bcd5a9ab3dRefs #50 #44 #43 #53.