feat(qwen4_exp): image input (vision tower + mRoPE) for Qwen3.8-Flash-Next

[gdevenyi]

What this adds

Image input for qwen4_exp (Qwen3.8-Flash-Next) through /v1/chat/completions. FreeToken serves this model text-only today: the loader drops the 333 model.visual.* tensors and render_messages refuses every non-text part (#321, #348). Opt-in with FREETOKEN_LOAD_VISION=1 (the tower is 0.86 GiB of bf16 per rank). Video is not covered.

Stacked on #385 (the tensor-parallel PR): this branch contains that commit because the loader's _shard and the attention layer's TP-aware forward are the seams it hooks into. Review that one first; the vision commit is the last of the three here (the middle one is #385's loader fix).

How it works

• Vision tower (models/qwen4_exp/vision.py). The HF Qwen4ExpVisionModel (27 blocks) runs inside a BaseOP whose tensors travel as visual.* in the model state dict, so the engine loads them with the dense weights: the expert-cache planner counts them, --dummy-weight fills them, and every TP rank holds the same copy. The op is built on the meta device like the rest of the model; load_state_dict assigns the loaded tensors in place (no second GPU copy) and rebuilds the non-persistent rotary buffer on the device.
• mRoPE (models/qwen4_exp/mrope.py). Qwen3.8 ropes image tokens on three axes (T/H/W, interleaved sections 11/11/10) and text after an image continues from max(position) + 1. rope_index ports HF get_rope_index; mrope_cos_sin builds the interleaved cos|sin rows. A prefill batch with image tokens gets a per-token table [T, rotary_dim] in the same layout as RotaryEmbedding._cos_sin_cache, and the existing flashinfer / triton rope kernels index it with row numbers (Batch.rope_positions); decode reads the normal cache at position + delta. The QSA indexer ropes each compressed key at its group's first token, so the table carries index_ratio - 1 lead rows per request for groups that straddle a chunk boundary. No new kernels. Text-only batches alias positions and run exactly the kernels they ran before.
• Request path. image_url parts (inline data: URLs only, 16 MiB cap; remote URLs are refused rather than fetched on the client's behalf) are decoded in the API server and kept as {"type": "image"} parts for the chat template. The tokenizer worker runs the checkpoint's image processor (FREETOKEN_IMAGE_MAX_PIXELS, default 1280*28*28, about 1,000 tokens per image) and expands each <|image_pad|> to its soft-token count; the scheduler encodes the images on its rank and computes the rope positions before admission. Image prompts chunk like text prompts: each chunk scatters only the soft-token rows whose placeholders fall inside it (the old path raised NotImplementedError inside the scheduler). The wire encoder now carries N-D tensors (it asserted 1-D).

Other adapters (Anthropic, Responses) keep refusing image parts with the existing error.

Status

Verified on the real checkpoint on CPU (tower load through the state-dict contract, rope positions and cos|sin rows equal to HF, tokenizer path, weight iteration, wire round trip, 60 tests in tests/models/qwen4_exp plus the scheduler / tokenizer / server suites) and end to end on 2 x RTX 6000 Ada at TP=2: see the results comment below (image answers correct for one and two images and for a text follow-up; the loaded tower costs 2.6 points of expert residency and no decode or TTFT). The three GPU test failures on this box are pre-existing on main (see the comments). Draft only because of the open-PR cap; ready for review from my side.

• Prefix cache for image prompts (ed0982a). Image placeholders share one token id, so multimodal requests used to be kept out of the shared prefix cache and every turn of a conversation holding an image re-prefilled the whole context (measured before the change: 20 turns of 109k-122k tokens, #cached-token: 0 each, ~35 s per turn at TP=2). The tokenizer worker now emits cache_ids next to the expanded input_ids: the same tokens with each image's placeholder run replaced by ids derived from a blake2b hash of the image bytes (>= 2**30, above any vocabulary, hash + offset within the run). The cache manager keys match/insert on cache_ids when present, else on input_ids; the model still reads input_ids, and the per-chunk placeholder window already skips cached placeholders. Verified on the TP=2 server: turn 2 after an image hits (cached=320 of a 330-token prompt), an identical request hits fully, the same text with a different image matches nothing past the text (answer describes the new image), and a 26k-token prompt with the image after a 20k-token preamble prefills in 4 chunks then hits cached=26304 on turn 2. tests/scheduler/test_mm_cache_key.py, test_mrope.py::test_image_cache_ids.

Testing

• tests/models/qwen4_exp/test_mrope.py (CPU; the HF comparisons skip without a transformers that ships qwen4_exp).
• tests/models/qwen4_exp/test_mrope_gpu.py (GPU).

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt

[gdevenyi]

GPU results, 2 x RTX 6000 Ada at TP=2 (#385), RadixArk NVFP4, --num-tokens 262144 --memory-ratio 0.94 --max-running-requests 16, this branch with FREETOKEN_LOAD_VISION=1:

Image probe (rendered 640x480 PNGs as data: URLs, temperature 0, thinking off):

case	prompt tokens	answer
one image: "HELLO 42" over a blue circle, "what text / what colour"	332	HELLO 42
text follow-up in the same conversation: "spell it backwards"	357	24 OLLEH
two images in one message (second: "GREEN TEA" over a green circle)	635	first: "HELLO 42", solid blue circle; second: "GREEN TEA", solid green circle
text-only request	21	a normal one-sentence answer

The follow-up exercises decode after an image (rope at position + delta) and a second prefill that carries the image history; the two-image case exercises several images per prompt in order.

Cost of the loaded tower, same flags, same tree with vision off vs on:

	single-stream tok/s	8 concurrent tok/s	expert residency	TTFT (1k)	expert slots
vision off	88.9	325.5	94.5%	0.84 s	23,229
vision on	89.9	325.8	91.9%	0.83 s	22,594

So the 0.86 GiB tower per rank costs 2.6 points of expert residency and nothing else here. Greedy 256-token continuations of three text prompts: the code prompt and the ~1k-token prompt are word-for-word identical between the two runs (and to production without this branch); the short essay prompt diverges after 84 words, which is the same short-prompt run-to-run noise seen between two passes of the same TP=2 server (bf16 atomics in the expert kernels), not a rope difference: the 1k prompt is the one that exercises the QSA blocks.

tests/models/qwen4_exp on CPU: 60 passed. The GPU pass of the same package in the chain reported 3 failures whose names the harness did not keep; I am rerunning it with full output and will attribute them.

[gdevenyi]

GPU test pass of tests/models/qwen4_exp on one RTX 6000 Ada, same box: this branch 107 passed / 3 failed, #385 alone 97 / 3, plain main (af71ba4) 94 / 3. The three failures are the same on all three trees: test_qsa_backend.py::test_chunked_prefill_matches_one_shot[unaligned|page-boundary|boundary+1], whose torch.equal between chunked and one-shot prefill is off by bf16 noise here (torch 2.11.0+cu130, flashinfer 0.6.18, triton 3.6.0, sm_89). Pre-existing, unrelated to this PR. The new test_mrope_gpu.py (table path == cache path through the QSA layer, chunked continuation and decode) passes.

[gdevenyi]

Pushed 28fd56d: chunked prefill for image prompts. The first version refused image prompts longer than one prefill chunk (--max-extend-tokens, 8192 by default); a coding-agent client with a ~109k-token context hit that at once. prefill.py now chunks them like text and the scheduler scatters, per chunk, only the soft-token rows whose placeholders fall inside that chunk (_mm_embeds_window, unit test in tests/scheduler/test_mm_window.py); the mRoPE table already windowed per chunk. CPU packages on the box: 713 passed, 100 skipped. The GPU end-to-end run with a 20k-token preamble before the image is queued for the next restart window; I will post it here.

[gdevenyi]

GPU end-to-end for the chunked case (TP=2, production flags, 2026-09-04 18:45): a 23,504-token prompt (20k-token text preamble, then the image, then the question) prefilled in three chunks (8192 / 8192 / 7120) and answered Text: HELLO 42 / Circle colour: blue in 6.2 s; the one-, two-image and text-only cases are unchanged. The client whose 109k-token context was refused before goes through the same path.

[MT-z]

Ran this branch at TP=1 on one RTX 4090, 24 GiB, in a 61 GiB box. It
works, including the chunked path you added last night. The only thing standing between it
and a machine this size is that the model needs a second PR to fit in host RAM at all.

Both image cases answered correctly at temperature 0, FREETOKEN_LOAD_VISION=1:

case	prompt tokens	prefill chunks	answer
one 448x448 PNG, "list every shape with its colour"	263	1	Red triangle, blue rectangle, green circle
~20k words of filler, then the same image, then the question	22,439	4096 x 5 + 1959	Red triangle, blue rectangle, green circle.

The second is the case 28fd56d fixes. The same request against 70d9e8d came back
image prompts must fit in one prefill chunk: 22439 tokens > 4096 (--max-extend-tokens).
Here --max-extend-tokens is 4096, so the image lands in the sixth chunk rather than the
first, and the shapes still come out right.

Getting there needs #337 as well, on a box this size. Without the disk tier the expert
bank build reaches the cgroup ceiling and is OOM-killed in about half a minute, at 46 GiB
and again at 50 GiB; the checkpoint carries 73.3 GiB of expert tensors, which is simply
more than this machine has. With --moe-disk-tier on --expert-ram-experts 224 (224 of 512
per layer resident, the rest fetched from the checkpoint) the same tree peaks at 43 GiB and
serves. So the two PRs are a pair here rather than alternatives -- which also puts the
requirement at a 24 GB card and about 64 GB of system RAM. If anyone else has that, the
flags at the bottom are the whole configuration.

They merge almost cleanly, so the recipe is short. One conflict, models/nvfp4_banks.py,
two blocks: #337 adds rows_per_layer to bound how many expert rows the loader
materializes, #385 replaces the loose bank variables with the TP-aware _Placer, and the
two are orthogonal -- keeping both lines resolves it. The function bodies auto-merge into
the right hybrid on their own (place.put(...) for the writes, rows_per_layer for the
tracker and the skip), so there is nothing else to hand-resolve.

Tests on the merged tree, against plain main (af71ba4) on the same box:

                merged            main
tests/moe       127 passed        (disk-tier tests are #337's)
tests/models    4 failed 165 p    4 failed 151 p
tests/kernels   3 failed 214 p    3 failed 214 p
tests/scheduler 88 passed         88 passed
tests/engine    108 passed        108 passed

Same seven failures either way -- the test_qsa_backend.py::test_chunked_prefill_matches_one_shot
family you already attributed, plus four in tests/models that are on main too.

Measured on an RTX 4090 (24 GiB, sm_89) / i9-14900KF / 61 GiB box, driver 595.84, CUDA
13.3, torch 2.11.0+cu130, triton 3.6.0, sgl_kernel 0.4.5, freetoken 0.1.2, model
RadixArk/Qwen3.8-Flash-Next-NVFP4, TP=1. The tree is main af71ba4 + #337 @ a6bd5c0 +
this branch @ ed0982a. Server flags:

FREETOKEN_LOAD_VISION=1 ft serve --model-path RadixArk/Qwen3.8-Flash-Next-NVFP4 \
  --moe-disk-tier on --expert-ram-experts 224 --disable-moe-prefill-overlap \
  --cuda-graph-max-bs 0 --moe-cache-auto --max-running-requests 2 \
  --memory-ratio 0.90 --kv-reserve-tokens 32768 --max-prefill-length 4096

Run inside a systemd scope with MemoryMax=46G and MemorySwapMax=0, which is where the
OOM figures above come from.

The only edit anywhere was that one merge conflict in #337's nvfp4_banks.py; nothing in
this branch itself was changed. This box stays available if there is anything you want run
at TP=1 on a single 24 GiB card.

Assisted-by: Claude Opus 5