Flash Next shared sessions sized to the traffic: FLASHNEXT_SLOT_CONTEXTS, and a stream that outgrows its session moves

[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 sized, growing sessions, alone against feat/flashnext-continuous-batching. It was rebased onto 04c8571 (1.0.8 with #7 and #8), with review fixes on top, and later merged with the tip after #9. It needs neither of the other two PRs (round costs, the lone driver). The receipts below are from 77ebc3f, before the rebase and the fixes.

What it does

• FLASHNEXT_SLOT_CONTEXTS lists the sessions past the engine's own, as token counts with CONTEXTxN for N alike. On our 256 GB lane, 65536,32768x3,16384x11 is 16 streams in 11.6 GiB, where three whole-window sessions took 15.4 GiB.
• Unset, nothing changes: each stream takes the first free session as before, with no admit hook and no growth checks.
• A new stream takes the smallest free session that holds its prompt and its reply's next FLASHNEXT_SLOT_GROW tokens (default 8192).
• With it set, the batch host gives main's Info.admit (merged with 1.0.8) Backend.fits, so a request waits for a session that holds it rather than failing.
• Before each round (slots.room), a stream whose session can't hold its reply to the end (prompt, max_new and a window) and is running out of room moves to the smallest free larger session. Session.moveTo copies its cache rows, DeltaNet states, n-gram tail, head rows and pooled blocks on the GPU at a round's end, so the bits are the same. With none free, the stream pauses (Stream.paused) on a round that still fits, until one frees.
• A stream whose session holds its reply never moves or pauses. The largest session always does (admission checks it), so a stream there runs to its end however many drafts are rejected. A stream only waits for a larger session that exists.
• A paused stream waits only on LaneMemory. Any other failed move ends it with an error the host reports, and frees its session (slots.resumePaused, Backend.failed).
• Lane core: since a release can now end another stream, Engine.sweep releases finished live streams until none is left, from step and discard, so none outlives its job.

Receipts on this branch. Mac Studio M5 Ultra, 256 GB, macOS 27.0, TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP @ 45032a1, greedy.

• Eight sessions: FLASHNEXT_SLOT_CONTEXTS=16384,4096x6 FLASHNEXT_SLOT_GROW=2048 --context 32768 --parallel 8 opens 32800, 16416 and 6 × 4128 tokens under the 70 % share; the base at --context 16384 fits four.
  • Concurrent = solo: 3/3 and 8/8 (231 tok/s aggregate at 8, against 132 solo).
  • Greedy ladder at 1 / 2 / 4 / 8 streams: 127 / 84 / 143 / 206 tok/s (the base's four sessions: 114 / 85 / 140).
  • Decode alone equals the base: same shas, same speed.
• A growing reply (count to 4000, 18,893 tokens) started in a 4K session, moved to the 16K one at 4,092 tokens, then to the engine's session at 16,382. It equals the base's unmoved reply character for character (sha ca7e464a7a501db9).
• Two growing at once (FLASHNEXT_SLOT_CONTEXTS=4096x3 --parallel 4): the two replies outgrew 4K together. One moved to the engine's session; the other logged waits for a session larger than 4128 tokens to free, then moved once the first finished. Both equal the base's reply. (My harness marked that step failed because of its job control; both reply files show identical.)

On our combined build (add3da3 on sethforprivacy:flashnext-lane-parity, 16 sessions with sampling and images):

• One reply outgrowing: an 18,893-token reply moved from 16K to 32K and equals a run that never moved.
• Two outgrowing at once: one moved and the other waited for it; both equal it.
• Greedy concurrent = solo: 3/3 and 8/8.
• Seeded sampled: replies equal their solo replies.

zig fmt --check passes on the changed files. zig build test at 069ef4d passes 1531/1539 (8 skipped), including slots.zig's tests and the lane core's sweep tests. One step fails: core.checkpoint_metal's split-checkpoint test, which fails the same way on the tip here. zig build tf-flashnext-batch builds.

🤖 Generated with Claude Code

[cwschroeder]

Thanks, @sethforprivacy. I can't merge this one yet:

1. families/flashnext/backend.zig:144-152: makeRoom pauses based on the window width plus headroom, not on whether the rest of the reply fits. In the largest session, a partly accepted window near the context limit can leave tokens outstanding after paused = true, and move needs a strictly larger session (:166), so the stream never resumes. Growth should depend on the remaining work and never wait for a larger session that does not exist. Please add near-limit tests with rejected drafts, also with the option unset.
2. native/flashnext_batch_host.zig:35-40 bounds the repetition counts after allocation, and tokens + MARGIN can overflow before the clamp. Please validate the slot count before appending and use checked arithmetic.
3. With FLASHNEXT_SLOT_CONTEXTS unset the behaviour still changes: smallest-fit replaces first-free allocation, and the growth checks run. Either keep the old path when the option is unset or say so in the description so the receipts cover it.

The move itself (waiting for settled state, copying caches, recurrent and head state) looks right to me.

[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): b9ed2ae → 5551aba. Two conflicts:

• LaneHost.admitOne: main's lane_host_admit.pick (merged with 1.0.8) already does what this PR's hook did: it holds a request whose caches don't fit beside the running streams. It also lets a later request that fits go ahead of a waiting head. So I dropped this PR's copy. The host passes Backend.fits as main's Info.admit, which has the same shape (engine_api.zig came out identical).
• flashnext_batch_host.zig: Flash Next shared sessions: FLASHNEXT_LOAD_LIMIT_GB lets the slots take the server past 70% of RAM #8's ceiling() (FLASHNEXT_LOAD_LIMIT_GB) sits beside slotContexts() / slotGrow(). The sized sessions are budgeted from Flash Next shared sessions: FLASHNEXT_LOAD_LIMIT_GB lets the slots take the server past 70% of RAM #8's cap (room = cap - ready - margin - reserved), so FLASHNEXT_SLOT_CONTEXTS stays under the same limit as the prompt cache's budgetUnder(cap, …). greedy_only stays dropped, as Flash Next shared sessions sample: a request's temperature draws each token on the host from the head's logits #7 left it.

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, as separate commits on top:

• makeRoom will grow on the reply's remaining work and never wait for a larger session that doesn't exist, with near-limit tests (rejected drafts, and the option unset).
• The slot count will be validated before appending, with checked arithmetic.
• With FLASHNEXT_SLOT_CONTEXTS unset, the old first-free path and no growth checks.

[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 ba31b91, on top of the rebase. The decisions moved into a small slots.zig, so they can be tested without a GPU.

1. Growth from the remaining work:
   • slots.room decides before each round from the reply's end (prompt + max_new + a window).
   • A stream whose session holds that end is never moved or paused. Admission already requires the largest session to hold the whole reply, so a stream there runs to its end however many drafts are rejected near the limit.
   • A stream that does need more moves while a larger session is free. Otherwise it pauses on a round that still fits.
   • It only ever waits for a larger session that exists (.full if none could hold the reply). Waits only point upward in size, and the largest session's stream never pauses, so a paused stream resumes once that session frees.
   • Tests: near-limit rounds where every draft but one is rejected, in a session that holds the reply (never pauses). Outgrowing a smaller session: it moves while a larger one is free, else pauses on a fitting round. And no wait when nothing larger exists.
2. Parsing: slots.parse checks each repetition count against the sessions still allowed before appending, and adds the margin with checked arithmetic. An overflowing count, an overflowing repeat, or too many sessions in all is refused before anything is allocated. Tests cover each.
3. Unset: with FLASHNEXT_SLOT_CONTEXTS unset, the host takes the old path: the first free session (ContextFull if it can't hold the reply), no admit hook and no growth checks. A test pins first-free against smallest-fit, and the description says so.

slots.zig is in the core test root, and its four tests run (the core test binary went from 282 to 286). M5 Ultra, Zig 0.17.0: zig fmt --check passes, zig build native builds, and zig build test passes 1517/1525 (8 skipped). The one failing step is the core.checkpoint_metal split-checkpoint test, as on 04c8571 here.

I haven't rerun the served receipts in the description since the fixes. They're from 77ebc3f, and I'll rerun the growing-reply checks in our next lane window.

[cwschroeder]

@sethforprivacy thanks, ba31b91 covers all three. With FLASHNEXT_SLOT_CONTEXTS unset, the session pick, the refusal and the admission now match the branch before this PR, and with it set slots.room only waits for a larger session that exists. A stream only ever waits for a larger session than its own, and a stream in the largest never waits, so a paused stream always resumes. I'll merge it after #11, once it's rebased; the rerun of the growing-reply checks can follow in the description.

[cwschroeder]

Follow-up to my merge-readiness comment above: the three requested fixes are present in ba31b918, but two remaining issues need fixing before merge.

1. The model parity harness still calls Backend.init(gpa, eng, count, capacity) at zig/tests/flashnext_batch.zig:147, while this PR changes it to Backend.init(gpa, eng, capacities). Please update that caller and build zig build tf-flashnext-batch -Dcpu=apple_m1. On Zig 0.17.0, the unmodified harness first stops at the two older generateFrom calls that lack the thinking-budget argument. In a temporary diagnostic copy I supplied that argument to both calls; the compiler then confirmed line 147: "expected 3 argument(s), found 4". The generateFrom mismatch already exists on 04c8571; the Backend.init mismatch is introduced by this PR. No model was loaded or run.

2. resumePaused logs and swallows every migration error, including GpuFailed and StateNotSettled, leaving the stream paused. If the last running larger stream has just released its session, there may be no further release to retry the move. The lane core counts paused streams as active but skips them in step, so the request can remain unfinished while the host spins. Only LaneMemory should keep waiting; other migration errors need a terminal path visible to the host. Please add a fault-injection test where the last larger session frees and the move fails. This is a source-review finding, not a reproduced GPU failure.

The successful-growth cases, bounded parsing and unset-option path look correct. Astra independently reviewed the correction commits. Growing-session model parity still needs a receipt on the corrected/rebased head.

[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 both points are in 069ef4d, after merging the tip with #9 (198b3c5, no force-push).

1. The harness. zig/tests/flashnext_batch.zig now passes Backend.init the sessions' capacities. Like init, it refuses more than MAXR sessions. Its two generateFrom calls take the thinking budget (.{}) that 1.0.8 added, the break you found already on 04c8571. zig build tf-flashnext-batch -Dcpu=apple_m1 builds. I haven't run it against a model.

2. A paused stream's failed move.

• The resume rule: slots.resumePaused keeps a stream waiting only on LaneMemory. Any other failure (GpuFailed, StateNotSettled, ...) goes to Backend.failed, which ends the stream (slots.endPaused: finished, reason error, a problem the host reports as a failed request) and frees its session at once. The fault-injection test runs your scenario: the last larger session frees, one paused stream's move fails on the GPU and ends, one finds nothing free and waits, one moves.
• A lane-core change, so you can judge it separately: a release can now end another stream, and the host assumes a finished stream has already left the engine's live list. failAll, cancels and drop all skip discard for finished streams, then free the job. So a stream ended during another's release could outlive its job in e.live.
  • Engine.sweep releases finished live streams, again while a pass ends another. It's called from step's prune, replacing the old single pass, and from discard.
  • Two tests on the fake target, which gains a release_hook: a stream ended by another's release in a step, and by a discard, leaves the live list. Against the old engine.zig both tests fail; with sweep they pass.

M5 Ultra, Zig 0.17.0: zig fmt --check passes, zig build native builds, and zig build test passes 1531/1539 (8 skipped). The one failing step is the core.checkpoint_metal split-checkpoint test.

Still open: the growing-session model parity receipt on this head. I'll run it in our next lane window and add it to the description. #10 and #11 touch the same spots in backend.zig and the host, so once #11 lands I'll merge the tip into this branch again, unless you'd rather resolve it while merging.