perf(bonsai): decode and small packed batches read the FFN from a ternary copy through a dp4a GEMV (1.35x decode @128) by kaivaryn · Pull Request #1154 · gittensor-ai-lab/sparkinfer

kaivaryn · 2026-09-24T20:00:58Z

Summary

Main decodes Ternary-Bonsai-2 from weights folded into Q4_K (15.2 GB/token), and the 64 dense FFNs are ~9.6 GB of that. Single-row decode now reads the FFN from its stored ternary blocks (a 3.7 GB copy, released when a new session needs the VRAM) through a dp4a GEMV with its int8 activation staged in shared memory and produced by the Hadamard rotation, and packed steps of up to 8 rows read the same copy with each row bit-identical to its single-row decode; the dp4a scratch is allocated with that copy, so no other model's VRAM changes. SPARKINFER_BONSAI_DECODE_SHADOW=0 restores main.

Proof of speedup

Tested on RTX 5090 (sm_120)

Target model(s)

Muse Glimmer
Qwen3.8-27B (ModelOpt NVFP4 / DSpark)
Ternary-Bonsai-2-27B (PTQ1_0 GGUF)
Shared / all models — the change is in code every model uses and should help more than one

Decode tok/s (end-to-end, from bench/scripts/bench.sh — fill if this PR targets decode):

	decode tok/s
before (main)	98.2
after (this PR)	132.8

Prefill pp tok/s (fill if this PR targets prefill; use --ctx 4096, 32768, 65536, or
131072 and copy the prefill pp line — report your best context. Keep the row labels as-is):

	prefill pp tok/s
before prefill (main)
after prefill (this PR)

Ternary-Bonsai-2-27B-PTQ1_0.gguf, RTX 5090, bot build (default archs, BUILD_SERVER=ON), main a4cf6f4, interleaved
decode ctx 128/512/4k/16k/32k: 98.2/97.6/96.1/92.7/87.9 -> 132.8/131.9/129.2/123.6/115.0 (+35.2..+30.9%)
cb c2/c4/c8/c16/c32: 162.6/253.0/489.4/765.4/1032.6 -> 203.2/330.5/544.3/766.5/1030.7; prefill 1.00-1.04x
SHADOW=0 (same binary): decode 98.1/97.5/96.0/92.6/87.7, reproduces main
accuracy vs main: top-1 0.961, KL 0.0179, PPL 3.675 vs 3.685; prefill_check P=128/1024 inside the bot's bars
eval/bonsai_regression.py: OK (score, generate, serve folded + native); batched rows vs alone: 20/20 bit-identical
qwen3.6 @32k: 471.5/470.9 vs 471.5/471.0; free VRAM after its load 9644 MB on both (#1149 read 9634)
rebased onto 31bf490 (#1153 touches only prefill_gemm_i8.cu); the PR's own diff is byte-identical to the measured one

@128

…nary copy through a dp4a GEMV (1.35x decode @128) Main decodes Ternary-Bonsai-2 from weights folded into Q4_K, 15.2 GB per token, of which the 64 dense FFNs are ~9.6 GB. The load now also keeps the FFN's stored ternary blocks (3.7 GB, freed the first time a new session cannot allocate its state, which a c32 batch triggers) and single-row decode reads them through a dp4a GEMV: the CTA stages the int8 activation in bank-padded shared memory once, and the Hadamard rotation emits that int8 copy for gate and up to share. A packed step of up to 8 rows reads the same copy through a batched kernel with each row's lane walk, per-block fold and shuffle tree unchanged, so every row decodes bit-identically batched or alone. The head and attention/GDN projections stay folded (KL 0.018 vs main). The dp4a scratch is allocated with the shadow, not as static device memory, so no other model's VRAM changes. SPARKINFER_BONSAI_DECODE_SHADOW=0 restores main.

kaivaryn · 2026-09-24T21:12:51Z

Hi! I noticed this got the hold label right after I rebased it onto the latest main. Is there anything you'd like me to change or double-check before it goes through eval? Happy to adjust. Thanks!

skyrocket2026 · 2026-09-24T21:31:34Z

sparkinfer bonsai auto-eval — `eval-bonsai:XL`

metric	value
label	`eval-bonsai:XL`
model	Ternary-Bonsai-2-27B PTQ1_0 GGUF, default (folded) loader
scored at	decode + prefill @ 128/512/4k/16k/32k · concurrent decode @ c2/c4/c8/c16/c32 — 15 axes, each also a regression floor; the label is the best
tier from	`bonsai-decode@128` (+35.4%)
accuracy vs main (teacher-forced)	✅ top-1 0.9515 (bar ≥0.93) · KL 0.02092 (bar ≤0.03) · PPL ×0.9915 of main (bar ≤1.02) over 1177 tokens
prefill path vs main	✅ batched prefill within main's spread at prefix 128/1024
bonsai_regression.py	✅ tensors · score · generate · serve
qwen3.6 guard	✅ no regression (decode + prefill @ 32k)
modelopt qwen3.8 guard	✅ no regression (decode + prefill @ 32k)
unsloth qwen3.8 guard	✅ no regression (decode + prefill @ 32k)
muse glimmer guard	✅ no regression (decode + prefill @ 32k)
PPL PR / main	3.6654 / 3.6969
Polaris receipt	collected, not signed (no key configured)
commit	`5a823be2d`, measured merged onto `main` `f90d6b5` (this round's baseline)

ctx	phase	main	PR	delta
128	decode	98.96	133.94	+35.4%
128	prefill	4824.02	4917.39	+1.9%
512	decode	98.30	133.05	+35.4%
512	prefill	4812.48	4885.02	+1.5%
4k	decode	96.83	130.24	+34.5%
4k	prefill	8921.93	9012.73	+1.0%
16k	decode	93.43	124.40	+33.1%
16k	prefill	8925.66	8970.67	+0.5%
32k	decode	88.49	115.85	+30.9%
32k	prefill	8452.68	8465.95	+0.2%

Concurrent decode — aggregate tok/s with N requests in flight, median of 3 runs

concurrency	main	PR	delta
c2	168.00	208.20	+23.9%
c4	299.00	361.20	+20.8%
c8	525.70	554.80	+5.5%
c16	817.20	817.40	+0.0%
c32	1109.70	1107.20	-0.2%

Prefill path — batched prefill against the token loop, 64 teacher-forced positions, mean of 3 runs per side

prefix	main top-1	PR top-1	top-1 bar	main KL	PR KL	KL bar
128	0.984	0.995	≥0.884	0.0052	0.0062	≤0.0552
1024	0.948	0.953	≥0.848	0.0057	0.0064	≤0.0557

ok

_{Measured on the pinned RTX 5090 against a same-box origin/main from the same round. Any axis regressing below 98% of main is a REJECT; otherwise the label is the best measured delta. none only means no Ternary-Bonsai-2-27B speedup was measured, which is expected for a change aimed at another model. The round's best passing speedup is auto-merged as bonsai-merge-first (only at the exact commit scored); a separate comment says so. It does not close PRs.}

…first (#1158) reconcile_bonsai_merge_labels skipped only needs-rebase PRs when picking the round's winner. A PR under hold -- or any other AUTOMERGE_BLOCK label -- could still win merge-first, have its merge refused, and push the next-best PR to needs-rebase for a merge that never happened. On 2026-09-24 #1154 was held for review with the round's best score (XL, +35.4%), which would have done that to every other Bonsai PR for as long as the hold lasted. Blocked PRs are now left out of the ranking, and a stale merge-first on one is removed.

kaivaryn · 2026-09-24T21:56:43Z

Thanks for running it! It came back XL with all the gates and guards green, but it looks like it didn't get merged: bonsai-merge-first was added and then removed a minute later, and hold is still on. One thing that might matter: the eval measured 5a823be, while the head is now c0521b9, which is just that same commit rebased onto f90d6b5 with an identical diff. Is there anything I should do on my side to get it merged, or does it just need the hold lifted? Thanks!

skyrocket2026 · 2026-09-24T23:08:23Z

Closed: long-context serving regression — `eval-bonsai:REJECT`

The automated round scored this XL: decode +35.4% at every context, with every gate passing. The decode gain is real. But the bot runs qwen3_gguf_bench, which sizes its KV cache for at most 32k tokens. The server sizes it for the whole --ctx, and at Ternary-Bonsai-2-27B's native context the decode shadow leaves too little memory for prefill.

Measured on the eval box's RTX 5090: sparkinfer_server --ctx 262144, main 7586d04 against this PR (c0521b9) merged onto that same commit, with identical prompts.

	`main`	this PR
VRAM after startup	26.2 GB	30.0 GB
short decode	95.8 tok/s	129.0 tok/s
60k-token prompt	8.1 s	479.6 s
120k-token prompt	19.6 s	1,056.6 s
200k-token prompt	40.7 s	timed out at 1,200 s
short decode afterwards	97.2 tok/s	0.2 tok/s

The server log for this PR shows [prefill] scratch alloc failed (ctx=16384, chunk=4096, held=1731 MB, free=63/32109 MB) -> fallback three times. The shadow was never released.

Why:

After the 262,144-token KV pool (8.0 GiB, int8), the shadow's +3.85 GB leaves only about 60–100 MB free.
A long prompt's prefill can't allocate its scratch, so it falls back to a far slower path.
release_bonsai_shadow only runs when open_session can't allocate a session's state, never when prefill scratch fails. So the server doesn't recover.

To resubmit:

Don't build the shadow, or release it, when the VRAM left after the KV pool can't hold the prefill working set. That's known at startup from --ctx.
Release it on any allocation failure it causes, not only a new session's.
Ideally, rebuild it once the memory is free again, so one busy moment doesn't lose the speedup for the rest of the process's life.
In the PR, include the numbers above for --ctx 262144: 60k, 120k and 200k-token prompts, then a short decode.

hold stays on this PR so it isn't re-scored and merged unchanged if reopened. A maintainer will remove it once a resubmission addresses the above.

…t only a new session (#1160) #1154 shipped a ternary-copy decode shadow for Ternary-Bonsai-2 (+35.4% decode) but only ever released its ~3.7 GB when a new session's state couldn't be allocated. At long --ctx the KV pool already leaves little headroom, so the shadow crowded out batched prefill's own scratch arena: every long prompt fell back to the token loop (479s/1056s/timeout at 60k/120k/200k tokens where main took 8.1s/19.6s/40.7s), and the server never recovered because nothing released the shadow for that failure. Batched prefill reads the folded Q4_K weights regardless of the shadow, so on a scratch-alloc failure it's pure margin to give back: prefill_batched_run now reports (via a new scratch_oom_out flag on Qwen35PrefillCtx) when a failure was VRAM-shaped rather than an unsupported config, and prefill_batched releases the shadow and retries once when that happens. prefill_batched_chunked's windowed path goes through prefill_batched per window, so a long prompt recovers window by window instead of dropping to the token loop for its whole length. Also skip building the shadow at load when the numbers already look this tight: the KV pool is sized before load_gguf runs (ModelEngine::load), so free VRAM at that point already reflects it. The threshold is a scaled estimate (the rejection's own ctx=16384/chunk=4096/1731 MB data point, scaled to this run's worst-case single pass, +50% margin) rather than a byte-accurate model of the scratch arena -- the retry above is the real safety net regardless of how tight the estimate is. Rebuilding the shadow once memory frees back up (so one busy moment doesn't cost the speedup for the rest of the process's life) is left as a follow-up: it needs the shadow's source tensors kept reachable past load_gguf, which is a bigger change than this fix. Co-authored-by: inference2026 <299740695+inference2026@users.noreply.github.com>

* fix(eval): the eval bots record, gate and close correctly Review of the Ternary-Bonsai bot against the Muse Glimmer and Qwen3.8 bots, and of all three together. Every fix below comes from a failure seen on 2026-09-24/25. Commit recorded (all three): - A verdict names the tip the box fetched and built (arb.measured_commit), not the head listed at the start of the round. #1167 was force-pushed mid-round; its verdict, scores entry and log named a commit that was never measured, the next round re-measured it, and auto-merge refused it. merged_checkout_script prints the full PR_TIP; the Muse bot prints REMOTE_SHA. The Polaris attestation checks out that tip instead of fetching the ref a third time. Merging (all three): - Only a PR auto-merge would accept can be made merge-first (auto_merge_ok_*, require_merge_first=False), so a refused winner no longer holds merge-first while every other speedup PR is pushed to needs-rebase. The Muse and Qwen3.8 bots now also exclude every AUTOMERGE_BLOCK label there (the #1154 hold bug), merge only the commit they scored with --match-head-commit, and refuse a PR another bot REJECTed. Bonsai regression gate: - bonsai_regression.py's serve check fails only on 2 of up to 3 trials, each on a fresh server after the GPU drains; a build equal to main failed a single trial in 2 of 8 runs, so one failure forced a REJECT about one time in five. A trial whose two alone-baselines differ is inconclusive, never a failure. A check that raises is a named failure of that check. - The bot gates each check separately: a serve failure on main no longer switches off the tensors, score and generate checks for every PR in the round. Box faults are not the PR's (all three): - A compiler killed for memory, a full disk or the overlayfs EFAULT rebuilds once at -j4, then is RETRYABLE infra. The Muse and Qwen3.8 bots now post nothing for a box fault, an unmeasured guard (Muse's ModelOpt, unsloth and Qwen3.6 guards, Qwen3.8's Qwen3.6 guard), an exception or a lost main dump; they used to write REJECT, and Muse closed the PR. Muse's llama.cpp reference build failing is infra; its PR fetch is RETRYABLE. The main baseline is wrapped against transport exceptions. - A failed build's verdict shows the compiler's own error lines (arb.BUILD_FAILURE_SH, first_build_error, failure_excerpt), not 80 lines of `ptxas info`: #1163's sm_89 errors never reached its comment. Bonsai measurement: - A crashed concurrent-decode run is one attempt, not the end of the width; main must measure every width; a width only the PR fails is a REJECT on the second round on the same commit (strikes file), and a GPU that never drained is infra. - A crashed prefill-check run is retried up to three times (not a timeout). Closing: - Each bot's stale close touches only PRs routed to its own model and never any bot's merge-first (arb.stale_close_skip_reason); the Qwen3.8 bot had closed #1157, a Ternary-Bonsai PR. The Bonsai bot gets its own. The stale-close comments no longer say a push revives a closed PR. - Bonsai closes like its siblings (decision 2026-09-26): a measured REJECT closes, a `none` closes only a PR declared for Ternary-Bonsai-2-27B alone. Selection: - Only comments from members and collaborators count as an existing verdict. - Open PRs are listed up to 300 (the 80 newest hid the oldest). - Bonsai honours hold and drafts in any run that posts; --no-post may still measure one by name. The Muse bot no longer evaluates a PR that edits its measuring harness, which it builds from the PR itself. Tests: test_pr_bonsai_bot.py 73 (was 42); new test_sibling_bot_fixes.py (19) runs every sibling check on both bots -- the Muse bot had none. All suites pass. * fix(eval): address the pre-merge review Found by an independent review of the previous commit, each verified in code or by running the rendered bash: - Qwen3.8: SCORE_FAILED (the PR's own score binary crashing) is an explicit exit, so it read as a silent kill and was retried every round with nothing posted. It is in _EXPLICIT_FAIL_MARKERS now, and reported. - All three bots pinned the merge to a SECOND head lookup, so a push landing between the gate and the merge was merged unscored. They pin to the scored commit. - A run killed at the 2 h ssh limit is a hang (main finished the same script): posted once with a labelled marker (arb.exception_result), not retried every hour. - A reconcile that could not read a PR from GitHub demoted the real merge-first holder; it now leaves every label alone (arb.PR_UNREADABLE). - bonsai_regression.py: a request that errors, or a server that exits, fails the serve trial -- a server dying on its first request used to PASS (its two "connection refused" strings matched). The script runs unbuffered, its exit code is reported, and a run killed without naming a check is judged over two rounds, not REJECTed and closed at once. The serve check is gated per path. - Bonsai: requiring main to measure every width could stall the bot for every PR if main itself broke at a width; such a width is dropped for the round, loudly, again. The strike key names the kind ("cb"), not the exact widths, which could reset the count every round. A box-shaped build fault recurring at one commit for three rounds is charged to the PR. Early REJECTs carry the built tip. - A guard that measured nothing on main now skips the round in all three bots, instead of every PR being measured in full only to be deferred. - Stale close read UTC times as local time (closing ~2 h early on the controller) and could close a PR opened from old commits before it was ever evaluated; it uses calendar.timegm and counts the PR's own opening. Tests: Bonsai 81, sibling 23 (+ a real _serve_trial against a fake server, the per-path gate, both strike paths, the unreadable-PR reconcile, the stale clock).

kaivaryn force-pushed the perf/bonsai-ffn-shadow-v2 branch from 5a823be to c0521b9 Compare September 24, 2026 20:57

skyrocket2026 added the hold Maintainer override: never auto-merge this PR label Sep 24, 2026

skyrocket2026 added eval-bonsai:XL sparkinfer Ternary-Bonsai-2-27B vs-main speed tier: XL eval:XL eval verdict parked by a 3-day noise ban — restores to eval:XL labels Sep 24, 2026

skyrocket2026 mentioned this pull request Sep 24, 2026

fix(eval): a held or otherwise blocked PR does not take bonsai-merge-first #1158

Merged

skyrocket2026 added bonsai-merge-first round winner: biggest verified Ternary-Bonsai-2-27B speedup — auto-merge candidate and removed bonsai-merge-first round winner: biggest verified Ternary-Bonsai-2-27B speedup — auto-merge candidate labels Sep 24, 2026

skyrocket2026 closed this Sep 24, 2026

inference2026 mentioned this pull request Sep 24, 2026

fix(bonsai): release the decode shadow's VRAM for batched prefill, not only a new session #1160

Merged

5 tasks

skyrocket2026 mentioned this pull request Sep 26, 2026

fix(eval): the eval bots record, gate and close correctly #1171

Merged

Navigation Menu

perf(bonsai): decode and small packed batches read the FFN from a ternary copy through a dp4a GEMV (1.35x decode @128) - #1154

perf(bonsai): decode and small packed batches read the FFN from a ternary copy through a dp4a GEMV (1.35x decode @128)#1154
kaivaryn wants to merge 1 commit into
gittensor-ai-lab:mainfrom
kaivaryn:perf/bonsai-ffn-shadow-v2

kaivaryn commented Sep 24, 2026

Uh oh!

kaivaryn commented Sep 24, 2026

Uh oh!

skyrocket2026 commented Sep 24, 2026

Uh oh!

kaivaryn commented Sep 24, 2026

Uh oh!

skyrocket2026 commented Sep 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

kaivaryn commented Sep 24, 2026

Summary

Proof of speedup

Uh oh!

kaivaryn commented Sep 24, 2026

Uh oh!

skyrocket2026 commented Sep 24, 2026

sparkinfer bonsai auto-eval — eval-bonsai:XL

Uh oh!

kaivaryn commented Sep 24, 2026

Uh oh!

skyrocket2026 commented Sep 24, 2026

Closed: long-context serving regression — eval-bonsai:REJECT

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

sparkinfer bonsai auto-eval — `eval-bonsai:XL`

Closed: long-context serving regression — `eval-bonsai:REJECT`