Flash Next shared rounds priced from forwards timed at load; every stream's drafts in one command buffer

[sethforprivacy]

Posted by Claude (Anthropic's AI assistant) on behalf of @sethforprivacy, as part of our local-inference tuning work on a Mac Studio M5 Ultra.

As asked on ashhart#606: the measured round costs, alone against feat/flashnext-continuous-batching. It was rebased onto 04c8571 (1.0.8 with #7 and #8), with the review fixes as a second commit. It needs neither of the other two PRs (sized sessions, the lone driver). The receipts below are from 77ebc3f, before the rebase.

What it does

• batch_costs.zig times, at load, a shared forward at 1..16 rows and a chained head step, on the last session before any stream holds it. On this Mac: 10.3 / 12.1 / 15.1 / 21.0 / 32.2 ms at 1 / 2 / 4 / 8 / 16 rows, and 0.43 ms a chained step.
• The timing needs a session of batch_costs.ROWS (209) rows. On a smaller one (a --context under 177, since a session holds the window plus 32 rows) the host serves the shared rounds unpriced. FZ_BATCH_UNPRICED=1 skips the timing altogether, and the allocator, draft budgets and lone depth are then the ones from before this PR.
• Backend.facts hands them to the lane core: shared_costs for the allocator, and a lone stream's window_costs and mtp_step_ms for its depth (FZ_LONE_UNPRICED=1 keeps the heuristic).
• The lane core prices each draft past a stream's first at that step (Model.draft_step_ms, allocatePriced with a unit test) and leaves those steps out of the round's overhead estimate.
• draftTogether runs every stream's absorb and chained drafts in one command buffer, with one wait (FZ_BATCH_MTP_SERIAL=1 keeps the per-stream path).

Receipts on this branch. Mac Studio M5 Ultra, 256 GB, macOS 27.0, TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP @ 45032a1. --parallel 4 --context 16384 --prompt-cache-gib 0, greedy, the base and this branch served in turn.

	77ebc3f	This branch
One request alone, 512 tokens: code / prose / edit (tok/s)	187 / 125 / 179	215 / 129 / 205
Greedy ladder, ~200-token prompts, 400 tokens: 1 / 2 / 4 streams (aggregate tok/s)	114 / 85 / 140	123 / 114 / 161
Concurrent = solo, 3 streams	3/3 (175 tok/s)	3/3 (207 tok/s)

• Replies: every decode reply's sha is the same on both builds (code bb176f18, prose a49a7c9f, edit 3d5a75d5).
• Ladder first tokens: about 2.3 s at 2 and 4 streams on both builds. That is macOS re-wiring idle GPU memory, which Flash Next keeps its buffers wired between requests: a residency set, and a keepalive tick the GPU runs ashhart/TensorFold#732 addresses separately.

On our combined build (this branch's change plus sampling, images, sized sessions and the lone driver; add3da3 on sethforprivacy:flashnext-lane-parity), with 16 sessions:

• Sampled ladder at 0.7: 150 / 195 / 236 / 255 tok/s at 2 / 4 / 8 / 16 streams. Before this and the sized sessions, the same build managed 73 / 102 / 116 at 2 / 4 / 8 with three whole-window sessions.
• Greedy concurrent = solo: 3/3 and 8/8.
• Seeded concurrent = solo: three seeded sampled pairs (0.7) run two at a time equal their solo replies.

zig fmt --check passes on the changed files. zig build test at dd86083 passes 1522/1530 (8 skipped), including the new pricing and timing-boundary tests. One step fails: core.checkpoint_metal's split-checkpoint test, which fails the same way on 04c8571 here.

🤖 Generated with Claude Code

[cwschroeder]

Thanks, @sethforprivacy. This looks right to me and I'd like to merge it after a few small changes:

1. batch_costs.zig:36-43 always prefills 192 tokens and then runs 16-row windows without checking the session's capacity. With a small --context (e.g. 128, last session 160 tokens) calibration writes past the session's KV. Please bound the scratch prompt by s.m.cap or skip calibration when it cannot fit, with a boundary test.
2. native/flashnext_batch_host.zig:47 always calibrates, so FZ_BATCH_MTP_SERIAL=1 and FZ_LONE_UNPRICED=1 restore the serving paths but not the startup path. One switch that skips calibration entirely (and keeps the old allocator) would give a full escape hatch if calibration ever fails on a machine.
3. The pricing test in core/lanes/allocate.zig:124-134 catches pricing that is ignored entirely, but still passes if the first draft is charged too. Could you add asymmetric streams, the first-draft exemption and the zero-cost case?
4. Style: the multi-line doc comments (batch_costs.zig:1-3, session.zig:145-146) should be single lines per CONTRIBUTING.

Minor, not blocking: allocate.zig:30-43 still picks candidates by probability alone, so with unequal marginal costs it can miss a cheaper split (e.g. [1,1] beats [2,0] for one example we worked through). That's efficiency only, not acceptance; fine as a follow-up.

I have merged #8 and #7 and updated the branch to current main, so #9 now needs a rebase (a small conflict in backend.zig).

[sethforprivacy]

Posted by Claude (Anthropic's AI assistant) on behalf of @sethforprivacy, as part of our local-inference tuning work on a Mac Studio M5 Ultra.

@cwschroeder thanks for the reviews. First, the rebase you asked for: onto the branch's new tip 04c8571 (1.0.8 plus #7 and #8): ae9f4dc → 3f24b00. Three conflicts, all where #7's sampling meets this:

• Config.init sets .draft_step_ms beside main's .k.
• Rule's budget calls allocatePriced with Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits #7's overheadMs(whos.len, sampled). Engine still leaves the priced draft steps out of the overhead it observes, now per sampled/greedy list.
• backend.zig keeps both Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits #7's sampling / drawRow and batch_costs.

The two don't touch each other's paths. Sampled tokens are drawn from the target forward's logits, and batch_costs times its forwards through batch_forward.encode directly, before any stream exists. Drafts are greedy head chains either way.

Checks on the M5 Ultra (macOS 27.0, Zig 0.17.0, -Dcpu=apple_m1):

• zig fmt --check passes on the changed files, and zig build native builds.
• zig build test: 1515/1523 (this PR's two new tests included) pass, 8 skipped. One step fails: core.checkpoint_metal's split-checkpoint test, which fails the same way on 04c8571 itself here (1513/1521).

I haven't rerun it live, so the served receipts in the description are from the pre-rebase head. Each branch still sits directly on the batching branch, so they can be reviewed and merged in any order.

Your review is next, as separate commits on top: calibration bounded by s.m.cap (or skipped when it can't fit) with a boundary test; one switch that skips calibration and keeps the old allocator; asymmetric streams, the first-draft exemption and the zero-cost case in the pricing test; and single-line doc comments. I'll leave the cheaper-split allocator change for a follow-up, as you suggest.

[sethforprivacy]

Posted by Claude (Anthropic's AI assistant) on behalf of @sethforprivacy, as part of our local-inference tuning work on a Mac Studio M5 Ultra.

@cwschroeder your review points are in dd86083, on top of the rebase:

1. Capacity: the timing needs batch_costs.ROWS (192 scratch tokens plus a 16-row window past them, 209). timeCosts checks fits(s.m.cap) first, and on a smaller session the host serves the shared rounds unpriced, as before this PR. That's a --context under 177, since a session holds the window plus 32 rows, so your --context 128 case. time itself refuses with SessionTooSmall. A boundary test covers ROWS - 1, ROWS and 160.
2. Escape hatch: FZ_BATCH_UNPRICED=1 skips the timing entirely. With no costs, the allocator, draft budgets and lone depth are the ones from before this PR.
3. Pricing tests: three more, each worked out by hand with a flat 10 ms cost:
   • the first-draft exemption: a 100 ms step still takes a stream's lone 0.5-likely draft, and would drop it if the first draft were charged;
   • asymmetric streams: {0.9, 0.85, 0.8} and {0.9, 0.2, 0.1} price to [3, 1] at a 1 ms step, against [3, 3] unpriced;
   • a zero step, which prices exactly as allocate does.
4. Style: the doc comments in batch_costs.zig, session.zig and backend.zig are one line each.

I'll leave the cheaper-split allocator change for a follow-up, as you suggested.

One more fix: batch_costs.zig joins the core test root (tensorfold.zig). The file is only imported privately by backend.zig, so its test wasn't being compiled. The core test binary went from 282 to 283 tests.

M5 Ultra, Zig 0.17.0: zig fmt --check passes, zig build native builds, and zig build test passes 1522/1530 (8 skipped). The one failing step is the core.checkpoint_metal split-checkpoint test, which fails the same way on 04c8571 here. Description updated.

[cwschroeder]

@sethforprivacy thanks, dd86083 covers all four points. I checked the bound by hand (the widest timing window needs 208 rows, one under ROWS, and the chained drafts stay inside it) and the three pricing cases: [3, 1] against [3, 3], the free first draft, and the zero step. On an M5 Max the head builds, zig fmt --check and test-core-lane pass, and zig build test fails only the known core.checkpoint_metal step. Merging.

Every pair of #9, #10 and #11 conflicts where both add code at the same spot (backend.zig and flashnext_batch_host.zig, plus session.zig for #9 and #10). I'll take #11 next, then #10, so each of them needs a rebase onto the new tip after a merge. If that's easier for you, I can resolve those conflicts myself when merging.