[BobClawblaw]

What

Image and video input for Flash Next on CUDA on two ranks. Today the engine serves images on one GPU with --parallel 2 or more and refuses --vision at --tp 2. Both ranks must feed identical hidden states into their halves of every layer; text does that by embedding the same ids, images only if both ranks hold the same features. So rank 0 runs the tower once per request and sends the result to rank 1 instead of rank 1 recomputing it:

• engine.py: the gate becomes "--parallel 2 or more", on one GPU or two ranks. The tower loads on rank 0 only (capacity_geometry and weight_transform already took the rank, so rank 1 reserves 128 MiB and no weights); rank 1 prints that it attaches what rank 0 encodes. --vision joins the settings both ranks must share at startup (_same_settings).
• multi.py: admit: on two ranks, rank 0 encodes the request's images before the admission message, so an image the tower refuses fails on rank 0 alone and no message reaches rank 1. A request admitted again after waiting for memory keeps the features it already has.
• multi_tp.py: the admission message and its digest carry the image rows, the decode offset and the feature shape. After the plan agreed and both ranks' memory checks passed, _share_vision sends rank 0's bf16 feature rows and int32 rotary positions to rank 1 point to point (comm.exchange), inside the admission lockstep and before any round collective; both ranks reach it or neither does. shape() includes whether a slot holds image positions. An older leader's nine-field message still parses (no image).
• multi_plan.py: an image prompt reuses no kept prefix, keeps no prompt states and never takes the lone-stream graph slot, as on one GPU (Shadow.images).
• image_rows.py: begin attaches features another rank encoded; a prepared prompt still needs the tower.

Nothing changes in the prompt passes, the rotary kernel or the rounds: image_rows.embed and attn_multi read each rank's own copy of identical data. Video frames are image rows to the engine, so video_url parts work the same way when PyAV is installed. Image prompts on two ranks keep the one-GPU limits: no prefix reuse, no grammars or logprobs. Size: at most the request's visual-token budget x 2560 x 2 bytes (80 MB at 16,384 tokens) plus 12 bytes a prompt token, a few milliseconds on a RoCE link.

Rank 0 gives up the tower (0.84 GiB) and its 4 GiB encode workspace (TENSORFOLD_VISION_WORKSPACE_MIB) from stream-cache room: 53.4 GiB free for caches on rank 0 against 58.0 on rank 1 at --parallel 16 --kv-dtype int8 on this pair, two streams at the full window.

I read #286: the Python engine is frozen and new work goes to the Zig engine. Flash Next is not on CUDA there yet, so this is offered as the reference and the measurements for when it is, like #468 (independent of it: this branch is main plus this one commit).

Receipt

Environment. TensorFold cb2ebf0 (v0.6.6 = main) plus this branch, inside nvcr.io/nvidia/pytorch:26.07-py3 (PyTorch 2.13.0a0, CUDA 13.3 forward-compat on kernel driver 580.178.04, transformers 5.18.0, PyAV 19.0.1 added for video). Two NVIDIA DGX Spark (GB10, 128 GB) at --tp 2 --parallel 16 --vision, --kv-dtype int8, --mtp-drafts 15 --mtp-confidence 0.70, context 262144, NCCL over one RoCE HCA. Checkpoint TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP at 2b170fa6309d5d1ee380b35636075fac7945f286 (it carries the 27-layer vision tower). "Before" is the same image and host session started without --vision (main refuses the flag at --tp 2, so the before cells are text only).

Correctness. Greedy, thinking off, three probes with known answers:

input	prompt tokens	reply
drawn PNG, a red circle left and a blue square right: "which shapes, what colour"	150	"The image contains a red circle and a blue square."
3840 x 2160 photo (a lake under clouds with two posts): "describe in one sentence"	4,101	"A serene, mirror-like lake reflects a soft, cloudy sky, with two weathered, rust-stained posts rising from the water's surface, their forms echoed perfectly below."
4 s 320 x 240 video (a green frame turning red, a white square moving right): "what happens"	373	"The video starts with a green screen, which then changes to a red screen."

Exactness, by tensorfold.token_sha, 160-token replies, four image prompts (the PNG and three questions on the photo):

cell	result
drafted against the same request with "draft": false	4 of 4 equal
four image requests sent together against each sent alone	4 of 4 equal
a text request sent beside the four image requests against alone	equal

Startup lines, rank 0 and rank 1:

[tensorfold] vision: image and video input, a 0.84 GiB tower with 4.00 GiB of workspace reserved; rank 1 receives each request's image features from this rank
[tensorfold] vision: image input on two ranks; this rank attaches the features rank 0 encodes

Decode speed, text, same host session, --vision off then on (tools/bench_openai.py --reps 2 --temperatures 0):

prompt	--vision off	--vision on
fibonacci-raw	92.4 tok/s	89.3 tok/s
gpu-chat-no-think	95.2	94.4

Frozen ruler (bench_decode.py, prose and structured, one stream, 3-run median):

phase	--vision off	--vision on
prose (ends at EOS near 100 tokens)	66.2 tok/s, TTFT 0.058 s	65.2 tok/s, TTFT 0.057 s
structured (200 tokens)	231.3	232.0

Within this pair's run-to-run band (the two fibonacci reps were 90.4 / 88.8 against 93.3 / 92.5 earlier in the day on the same image without --vision; the other three cells move under 1.5% either way). Text rounds take no vision code path: attn_multi sees no image positions and image_rows.embed finds no features.

Image requests, one stream: the PNG prompt decodes at 84 tok/s (160 tokens, prompt pass 0.09 s); the 4,101-token photo prompt passes in 3.2 s where 4,124 text tokens take 1.7 s (the tower's share about 1.5 s) and decodes at 70 tok/s. Four 4k-token image requests arriving together finish in 12-14 s each against 2-6 s alone: their prompt passes run one after another with --decode-share 0, as text prompts of that size do.

Prompt speed, tools/prefill_cold.py (2k to 64k cold text prompts, 3 reps), --vision off then on in the same session:

prompt tokens	--vision off: TTFT, tok/s	--vision on: TTFT, tok/s
2,048	0.850 s, 2,411	0.853 s, 2,401
8,192	3.260 s, 2,513	3.259 s, 2,514
16,383	6.569 s, 2,494	6.503 s, 2,519
32,768	13.283 s, 2,467	13.240 s, 2,475
65,536	27.647 s, 2,371	27.517 s, 2,382

Tests. Inside the image on the CPU: tests/test_flashnext_vision_two_ranks_host.py (5 tests: the admission message and digest with and without images, both ranks share the tensors or neither does and they fall out of step; begin without a tower; image prompts off prefix reuse, prompt states and the graph slot). All five fail on main and pass here. The full suite in the same container: this branch 4095 passed / 127 failed / 11 errors / 175 skipped; pristine cb2ebf0 4090 / 127 / 11 / 175, the same 138 failing tests in both (Metal, MLX and network-bound tests on a Linux box without them). The vision, concurrent-decoder and two-rank planning tests (test_vision_cuda.py, test_cuda_batch_admit.py, test_cuda_growing_caches.py, test_flashnext_round_plan.py, test_flash_next_tp_readiness_host.py, test_flashnext_absolute_grow_host.py): 78 passed.

Second commit: image prompts on the serial engine too

The --parallel 2 gate on --vision goes. Its reason was the serial engine's CUDA graphs: a capture bakes in the rotary mode, and the captured graphs know no image positions. prefill() already took an image payload (it attaches the rows and refuses prefix resume for them); nothing called it. Now Engine.forward and mtp_forward run eagerly while the state holds image positions, text states keep their graphs; generate() encodes on rank 0 before the request is shared, _share/_unpack carry the image description (an older leader's message still parses), follow() receives the features on rank 1, _decode prefills an image prompt from its start and keeps no snapshot, _serial takes the payload for "draft": false. The two-rank transfer moved to vision_ranks.py (describe, exchange), shared by the concurrent decoder and the serial engine; both ranks validate the same description.

An image request on the serial engine decodes at the eager rate, text keeps the serial rate, and the image request's fresh prefill overwrites the engine's kept text states.

Receipt, serial engine (--tp 2, no --parallel, same checkpoint and settings as above; --vision off then on in the same host session):

cell	--vision off	--vision on
tools/bench_openai.py --reps 2 --temperatures 0, fibonacci-raw / gpu-chat-no-think	101.6 / 101.6 tok/s	100.7 / 99.2 tok/s
frozen prose / structured ruler, one stream, 3-run median	72.3 / 250.2	72.1 / 239.0
tools/prefill_cold.py, 2,048 / 8,192 / 16,383 / 32,768 / 65,536 cold prompt tokens	2,413 / 2,522 / 2,534 / 2,479 / 2,400 tok/s	2,420 / 2,545 / 2,550 / 2,509 / 2,404 tok/s

Image prompts on the serial engine, greedy: the three probes above give the same replies and the same reply token shas as on the concurrent decoder (a7588cf7f668, e99c86a30d35, 333c343e9195); the 4,101-token photo prompt passes in 1.64 s. Drafted against "draft": false: 4 of 4 equal. A 200-token text reply before and after an image request: same sha, 73.4 then 72.4 tok/s over 106 rounds, so the graphs come back after the eager image rounds; the image reply itself ran at 99 tok/s (its answer repeats the question's words). Both ranks' startup lines are as above; rank 1 attaches what rank 0 encodes.

Tests. Three more in tests/test_flashnext_vision_two_ranks_host.py: the serial request message with and without images, describe/exchange validation (and the point-to-point call on rank 0), the eager rounds and fresh prefill for an image state. They fail on main and pass here. Full suite in the same container: this branch 127 failed, 4114 passed, 175 skipped, 11 errors; pristine cb2ebf0 127 failed, 4090 passed, 175 skipped, 11 errors, the same 138 failing tests in both.

[ashhart]

Thank you! Image and video input on two ranks is merged into main with your two commits as they are (812cef3, 8884aef). Encoding once on rank 0 and sending the features keeps both halves of every layer on identical hidden states. Putting --vision in the settings both ranks must share, and letting a new rank 1 still read an older rank 0's requests, is careful work.

One fix of ours follows in 6e8f29e. The vision start line had a backslash inside an f-string expression. Python 3.11 rejects that, and 0.6.6 supports 3.11, so the engine module wouldn't import there. After that, in d19e30a, we folded the docstrings onto one line each, which is our house style.

The Python engine is frozen at 0.6.6, so this lives in main rather than a tagged release. The Zig engine has no image input yet, and when it does, your two-rank exchange is the design to follow.