[sethforprivacy]

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

What this fixes. On Flash Next's shared sessions, an image chat that adds an image in a later turn reads its whole history again: 0 cached tokens, and every image back through the tower. #986 names this as a property of the shifted frame ("a turn that adds an image reads its prompt again"). On #565 you put the rule as "a prefix has to sit where every longer prompt puts it". This PR restores that on the shared sessions. The single-reply host is unchanged.

Why the shifted frame can't resume. vis_prompt.frame moves every row below the last image by the images' total shift, so the rows past the last image rotate at their slot. RoPE only sees differences, so replies are right. But the shift counts every image in the prompt, so a new image moves the whole prefix. That's why the frame is part of the cache key (LaneHost.frame_key).

What changes.

• The standard frame (vis_prompt.standardFrame). An image stream on the shared sessions takes get_rope_index's own positions. A prefix's rows don't depend on later images, so its kept state resumes as text does, and the shared host's keys carry no frame. Rows past the last image rotate at their slot plus VisPrompt.delta (minus the shift): in the prompt pass through positions(), and in the shared rounds by adding the delta to the window's rows in pos8 (Backend.verify). The single-reply host still builds the shifted frame (flashnext_images.build(..., standard)).
• The indexer. Pooled index keys rotated at 4 b would sit off by delta past the last image, and off the M-RoPE positions inside it. So an image stream's blocks pool at their first row's three positions in a new kernel, fz_idx_pool_pos: q4_idx_pool's arithmetic, axes chosen as fn_attn_prep_mrope chooses them. A text block has all three alike, so it keeps q4_idx_pool's bits at that position.
  • The positions live per session (Session.blocks, Model.bpos), are cleared by Model.reset, and are carried by Session.moveTo.
  • Select.encodeAt and catchLayerAt take them; encode and catchLayer are unchanged for their other callers.
• What stays. The engine's own GPU-side rounds keep the shifted frame. So image requests stay off the lone driver: api.Lone.images is a new field, default true, false for the Flash Next shared host. A lone image request runs on the shared rounds. The MTP head keeps slot positions, since it only drafts. Text requests run the code they ran before.
• The kernel contract.
  • fz_idx_pool_pos has its own source, prefill/fn_pool_pos.metal, compiled as its own library, so fn_select.metal and every pinned source are unchanged.
  • The contract gains a sixth rotary role, positioned: every slot (stride 1), on the M-RoPE axes, as mrope.
  • Receipts now hold 12 rows and 5 pinned sources. The sweep covers 302,008,320 cases.
  • rotary_interval.pooledAt and vision_bounds.pooledAt take a block's three positions; pooled is the splat.
  • The pool probe binds each block's positions before the output: even rows a text block, odd rows an image block's three, rotary-long rows the long fixture positions.
  • The runner, the limits records, the cache check (96 rows), and the sweep and receipt host tests count the new role.
• Two commits: 825288b7 (the standard frame) and 8e46fd2f (the kernel's own source and the contract role).

Environment. Mac Studio M5 Ultra, 80-core GPU, 256 GB, macOS 27.0 (26A428); compiler 382.5.3. Zig 0.17.0. Checkpoint TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP @ 45032a1. Binary sha256 cb200a6e2e06b5d8… (zig build native -Dcpu=apple_m1), with libjpeg-turbo 3.1.4.1 beside it. Linux: DGX Spark, Ubuntu 24.04 aarch64.

Served, on this head (--parallel 16, FLASHNEXT_SLOT_CONTEXTS=65536,32768x3,16384x11, FLASHNEXT_SLOT_GROW=2048, --context 262144, --prompt-cache-gib 16, --vision --vision-max-images 32 --vision-image-tokens 16384, greedy, thinking off):

Image chat, four turns: a screenshot, text, a second screenshot, text (a ~5K-token system prompt, a few hundred tokens of notes each turn)	Turn 1	Turn 2	Turn 3 (adds an image)	Turn 4
Cached tokens, main (da275a36, same frame code as 1.1.2; measured on 1.1.2)	0	5,551	0 (2.78 s)	6,735
Cached tokens, this PR	0	5,551	6,050 (1.12 s)	6,735
Reply = a --prompt-cache-gib 0 server's, this PR	yes	yes	yes	yes
Reply sha256 (first 16), this PR	cbf03fede93662e3	3250774adeac2682	d4f4bb22fa9534ad	c32472b3cfb221c2

• Our image gate, 6 cases (colour, shape, two images, text in an image, drafted = serial, an image stream beside text streams = solo): pass. JPEG (red, blue, green): pass.
• 8 text streams: concurrent = solo 8/8, at 144.9 / 228.9 tok/s. Solo sha256s (first 12): fe44df3683ae 90b46aec0d6f 1fa38f08cf7f ec3780ef850a ea1fd937fb0e 63f2d38634c5 8f78abf04ae1 63e2e6a7de87. These equal the same commit's frame change served on 1.1.2.
• The frame changes bits by rounding only. Turn 3's reply is 556 characters on the standard frame and 555 on the shifted one; each equals its own cache-off server.
• 0 errors in the server log.

Served, on the same change on 1.1.2 (dca5e388, the lane's build, which also carries our shared-session --learn spill). These ran before this rebase:

• A long image reply (the image's words, then counting to 4,000: 18,900 tokens) on 15 sessions of 16K moves to the 262K session at 16,377 tokens. It equals the same request served where it never moves (sha c14bebdf9c631e3b). Session.moveTo carries the block positions.
• qualify_mac 9/9; structured output 5/5 (greedy, sampled, thinking, guided_choice, 4 constrained beside 4 plain); the 262K edge; handoffs; Flash Next on Metal (1.1.0): a decoding stream stops for the whole of another stream's prefill #1032 replies = --parallel 1's; first token after idle 0.096 s.

Host checks.

	main da275a36	this PR
Mac: zig build test	not run on main (v1.1.2: 393/393 steps, 5,697 passed)	397/397 steps, 5,762 passed, 22 skipped
Mac: test-native	not run on main (v1.1.2: 52/52)	52/52 steps
Mac: test-kernel-contract, test-rotary-sweep-host, test-rotary-cache-host, test-flashnext-{vision-contracts,rotary-bound,bound-probes,precise}-host	pass	pass
Linux: zig build test	278/278 steps, 3,491 passed, 15 skipped	the same
Linux: test-cuda-host, test-kernel-contract (72/72), the six rotary/contract host steps	pass	pass
tf-rotary-sweep	234,895,360 cases	302,008,320; positioned's cos/sin worst 3.241/3.261 ulp (YaRN 2.987/3.313), the same as mrope's
kernel-contract --rotary-accuracy (M5 Ultra)	101 kernels pass gates 1-4; 14/14 controls red; 12 limit problems	103 pass (positioned and its YaRN x4: nonfinite, repeat, shapes within the FP64 bound at worst 1.000 of it, sums); 14/14 controls red; 14 limit problems

zig fmt --check is clean.

Not run, or known:

• Limits. The 12 limit problems on main are "no compiled thread limit recorded" on M1 and M2 for the Flash Next vision and rotary pipelines (core/contract/flashnext_pending.zig). fz_idx_pool_pos adds the same two. Our M5 limits pass. We have no M1-M4 Macs.
• Not run: CUDA devices (no CUDA code changes), video (no libtfvideo here), the single-reply host's image path (unchanged), and Gemma, GLM and Nemotron.
• Pre-existing: the same 9 Gemma 4 FAIL shapes lines show in kernel-contract on main and here.
• Not covered by this PR: a turn that adds images can still miss when the per-image budget binds. image_tokens / n_images (server/images.zig:169) resizes the earlier images once the share is smaller than an image, which changes their tokens. That's --vision-image-tokens is capped at 16,384 again: d5e69ee4's raise to 262,144 didn't survive #940's merge #1198.

Thanks to @BobClawblaw for Flash Next images on Metal (#986), which this builds on, and to @cwschroeder for the shared sessions and the contract review that shaped the rotary probes.

🤖 Generated with Claude Code

[ashhart]

Thanks, Seth. This is on main as e08019c: your three commits unchanged, then one maintainer commit and one small contract commit after them.

We ran it on an M3 Ultra with Flash Next MLX 4-bit. An image chat that adds a second image now resumes from 12,556 of its 12,563 cached tokens, where main resumes from none, and the resumed, drafted, changed-image and concurrent replies all match a fresh server exactly. The maintainer commit refuses malformed image spans before anything is allocated, and fixes the rotary cache check, which seeded absolute block starts while your positioned pooling reads three position axes. With that fix the 96-case cache check, the rotary sweep and contract gates 1 to 4 pass on both the M3 Ultra and an M5 Max. The contract commit records the M1 and M2 thread limits an offline compile reports for the seven Flash Next vision and rotary pipelines (1024 each), including your new positioned one.

• Synthema, Ash's AI Assistant.