Flash Next: a lone request on the engine's GPU-side rounds, handed to the shared rounds when a second arrives

[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 lone driver and its handoff, alone against feat/flashnext-continuous-batching. It was rebased onto 04c8571 (1.0.8 with #7 and #8), with the review fix as a second commit. It needs neither of the other two PRs. The receipts below are from 77ebc3f, before the rebase. With the sized sessions, a handed-over stream would also move to the smallest free session that holds it, keeping the engine's long session free. That step would follow whichever of the two lands second.

What it does

• A lone driver: the batch host gives LaneHost one (api.Lone, as the 27B's host does). A lone request runs the shared sessions' own prompt pass into the engine's session (session 0, whose buffers the engine's rounds use), so the states the prompt cache keeps are the shared path's.
• Decode: the request then decodes on the engine's GPU-side rounds, generateFrom's decode split out as Engine.decode, with two rounds in flight.
• Handoff: when LaneHost's yield hook asks (a request waits, or a cancel), the loop stops after a landed round.
  • The round ahead has already started the way rounds do: the kept DeltaNet states, the n-gram tail, the head's absorb of the last window.
  • The session's host fields are set to match: position, n-gram history, head position, pooled counts clamped.
  • The lane core then adopts the stream, and its next window starts at the cache's end with nothing held.
• On by default, so unset changes execution: a lone greedy request takes the lone driver. FLASHNEXT_NO_LONE=1 serves every request through the shared rounds.
• Scope: sampled requests stay on the shared rounds since Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits #7 (api.Lone.sampled stays false), as do requests with a thinking budget or the loop guard (lane_host_lone.fits). A sampled reply would need one more step at handoff, the round start run by hand since nothing runs ahead of it, which our combined build has.

Receipts on this branch. Mac Studio M5 Ultra, 256 GB, macOS 27.0, TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP @ 45032a1, greedy, --context 32768 --parallel 2 (base and branch alike).

One request alone, 512 tokens: code / prose / edit (tok/s)
77ebc3f, shared rounds	187 / 125 / 179
This branch, the lone driver	213 / 136 / 196
1.0.7's single-reply driver, for reference (our earlier run)	215 / 136 / 197

• Same replies: the decode replies' shas equal the base's (code bb176f18, prose a49a7c9f, edit 3d5a75d5).
• Handoffs mid-reply, greedy (handoff_exact.py): each first request is sent alone, and the second follows once the first has streamed ~400 characters. All three pairs equal their solo replies, both requests.
• A long handoff: the 18,893-token count was handed over at 3,041 tokens and equals the base's reply on the shared rounds (sha ca7e464a7a501db9).
• Concurrent = solo: 3/3.
• Handoffs in the window: 5, at 77 to 3,041 tokens.

On our combined build (add3da3 on sethforprivacy:flashnext-lane-parity):

• Mid-reply handoffs: seven equal their solo replies: three greedy, three sampled at 0.7 with a seed, and the 18,893-token reply, which then also moved from a 16K session to a 32K one.
• Concurrent = solo: 3/3 and 8/8, greedy, with solo on the lone driver.

zig fmt --check passes on the changed files. zig build test at 09c7262 passes 1514/1522 (8 skipped), including a test that the prompt is read after the first token's commit moves the context. 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. One fix before merge:

• families/flashnext/backend.zig:194,201,231: prompt aliases stream.context. Committing the first token can reallocate that array, and decode then reads the stale slice into GPU history (engine.zig:687). Please take stream.prompt() again after the commit (or reserve capacity first), and add a test with a prompt length that forces the reallocation.

When this lands after #7: sampled requests should stay on the shared rounds (.sampled = false), since the lone handoff does not sample. After #10: keep session 0 for lone decode and smallest fit for shared admission. The lone driver is on by default, so the description should say that unset changes execution.

[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): bff378b → a3085be. The one real change is in engine.zig. 1.0.8's thinking budget (think.Watch) runs through generateFrom's decode loop, where this PR splits out Engine.decode, so I redid the split on the new loop:

• decode takes the reply's *think.Watch, and generateFrom passes its own after watch.first(pick). The forced close spans, learn gating and kv_len handoff count are otherwise as before.
• The lone driver passes an empty Watch. lane_host_lone.fits already keeps think-budget, loop-guarded and sampled requests in the lane core (api.Lone.sampled stays false), so with Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits #7 the lone driver still only takes greedy requests. The commit message now says so.
• fill drops its old GreedyOnly refusal, as Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits #7 did in prefill. A sampled request's first token comes from Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits #7's drawRow in both paths.

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: 1513/1521 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: prompt taken again from stream.prompt() after the first token's commit, with a test whose prompt length forces the reallocation. Session 0 stays the lone driver's, and the description will say the driver is on by default. On #7: as above, the rebased branch keeps .sampled = false, so sampled requests stay on the shared rounds.

[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 fixed in 09c7262, on top of the rebase:

• The stale prompt: Backend.commitFirst commits the first token and returns stream.prompt() from the context as it is after that commit, and that's what decode gets. decode reads the prompt's contents only while setting up, when it copies them into the rounds' history (engine.zig), before any round commits. So later commits that move the context leave it alone.
• The test: it fills the context to capacity and refuses growth in place (FailingAllocator with resize_fail_index = 0), so the commit must move the buffer. It then checks that the returned prompt is the moved context's, with the prompt's tokens. backend.zig joins the core test root so the test runs; the core test binary went from 282 to 283.
• After Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits #7: .sampled stays false, so sampled requests stay on the shared rounds, as do think-budget and loop-guarded ones.
• After Flash Next shared sessions sized to the traffic: FLASHNEXT_SLOT_CONTEXTS, and a stream that outgrows its session moves #10: the lone driver keeps session 0, and shared admission keeps smallest fit.
• Description: it now says the driver is on by default, so unset changes execution.

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

[cwschroeder]

@sethforprivacy thanks, 09c7262 fixes it: decode reads prompt only while it copies it into the rounds' history, before any round commits, and the test forces the move. One nit: commitFirst landed between lone's doc comment and lone, so that comment now documents commitFirst. Could you move it when you rebase onto the tip with #9? I'll merge this one next.

[cwschroeder]

Rechecked 09c72626, including an independent Astra review. The stale prompt fix is correct: commitFirst takes the prompt from the context after the commit, and decode copies that history before subsequent commits. The regression test forces relocation and checks both pointer ownership and token contents.

Local checks on an M5 Max with Zig 0.17.0 and -Dcpu=apple_m1: zig fmt --check passes; zig build native -j4 passes 45/45 steps; zig build test -j4 reports 1514/1522 tests passed, 8 skipped. All 283 core tests pass, including the new relocation test, but the test step still exits nonzero because the split-checkpoint test logs the previously reported text.weight error. This does not replace a served model receipt.

No new correctness blocker found in this head. The misplaced doc comment identified above remains. Please include that small move in the rebase after #9; the combined head will need another check, especially session 0 and admission when #10 is also included. #10 now has a separate follow-up review with two remaining issues.

[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 up to date with the tip after #9 (3cb70f0), now at b1a4204. I merged the tip into the branch rather than rebasing, so there's no force-push and the reviewed commits keep their hashes.

• Merge (59df001): backend.zig keeps the think import beside batch_costs, and the lone handoff counter beside Flash Next shared rounds priced from forwards timed at load; every stream's drafts in one command buffer #9's serial_drafts and costs. The startup line names the drafts mode and the lone driver. The test root takes both backend.zig and batch_costs.zig.
• Your nit (0f679e2): lone's doc comment is back above lone. I also cut this PR's other multi-line doc comments (Handoff, decode, handOver) to one line each, as on Flash Next shared rounds priced from forwards timed at load; every stream's drafts in one command buffer #9.
• One the merge introduced (b1a4204): Flash Next shared rounds priced from forwards timed at load; every stream's drafts in one command buffer #9's facts() names its priced-lone flag lone, which shadows Backend.lone once both are in, so the native build failed (local constant shadows declaration of 'lone'). It's now lone_priced, with no change in behaviour.

M5 Ultra, Zig 0.17.0, -Dcpu=apple_m1: zig fmt --check passes, zig build native builds, and zig build test passes 1523/1531 (8 skipped), the tip's 1522 plus the relocation test. The one failing step is the core.checkpoint_metal split-checkpoint test.

This head doesn't include #10, so the combined check you mention (session 0 and admission with #10) still waits for #10.