feat(deepseek_v4): serve image input on the V4-Flash-Vision release by void-mckenzie · Pull Request #502 · FlashML-org/FreeToken

void-mckenzie · 2026-09-17T06:52:31Z

Image input for deepseek-ai/DeepSeek-V4-Flash-Vision-Exp (DeepseekV4ForCausalLM), the last open item under Multimodal in #79, on the multimodal path from #454, #467, #479, #480 and #481.

What

Processor. <｜deepseek_image｜> (id 129264) expands into the block the reference's prepare_vl_inputs builds: lead pads, IMAGE_START, the N-layout rows, the parity pads, IMAGE_END, one content pad id for the whole span instead of the reference's vocab_size + type ids. Lead pads align the block's first row token to a 4-token compression stride, so its length depends on where it lands: prompt_replacement receives that offset (a defaulted argument, existing families ignore it). --image-max-tokens lowers the checkpoint's per-image cap (vision_max_n_token, 384).
Tower. A 32-block native-resolution ViT with bidirectional attention and 2-D rope (fp32 tables), then the 3x3/stride-3 aligner into the 4096 text width. Blocks stream from pinned host banks under the default --mm-encoder-weights host: 0.59 GiB free after CUDA graph capture against 0.75 GiB text-only, so the tower is ~0.16 GiB resident.
Attention. An image span's tokens stay mutually visible in the sparse prefill path: the window candidates follow the span's own bounds, capped at vision_max_n_token, as the reference's get_image_visible / get_window_topk_idxs_visible pair does. The spans arrive as batch.mm_spans, derived in _gather_multimodal from the per-token block ends plan_mm_batch already reports. Chunks keep images whole through keep_images_whole / mm_chunk_end; only a block larger than the prefill budget is still cut, with the usual warning.
Router. bias_vl shifts expert selection on image rows only, on the hash layers and the score-routed ones alike; the routing weights stay the unbiased scores.
Reader. vision.*, aligner.* and image_* (tower, aligner, the four sentinel vectors) load when an encoder is built and are skipped under --text-model-only, which also serves this checkpoint text-only (both APIs then refuse images). get_mm_processor reads vocab_size off the flat config this checkpoint ships. docs/models.md gains the checkpoint and its row of the image-input table.

Verification

Layout and preprocessing against the checkpoint's own inference/ code: build_image_block bit-equal over 25 grid/offset combinations, load_image bit-equal (patches and grids) over 4 image sizes, and the tokenizer expansion matches prepare_vl_inputs on the real tokenizer (same token count, same block window, bit-equal patches; the last two gated on FREETOKEN_DSV4_VISION_CKPT / FREETOKEN_DSV4_VL_MODEL).
deepseek-ai/DeepSeek-V4-Flash-Vision-Exp on an RTX 5090 (32 GB), driver 580.95.05, started with ft serve --model /path/to/Flash_Vision_Exp --gpu 0 --num-tokens 311296 --moe-cache-size 750; /v1/stats reports input_modalities: ["text", "image"], image requests describe the image correctly on /v1/chat/completions (data URL) and /v1/messages (base64 block), two images in one request are described in order, decode is 20.4-21.0 tok/s at running-req 1.
Prefix caching through an image span, same 600-token prefix, temperature 0, #cached-token from the prefill log: 0 -> 512 (image A, first sight) -> 512 (text again) -> 640 (image A again) -> 512 (different image, same slot). The 640 is the span matched by content; the last request is the control.
pytest tests/ -m "not slow": 1768 passed, 208 skipped, 3 failed. The three are quant-config tests that fetch their checkpoints from the Hub and fail identically on a clean v0.1.3.

Not in this PR

An FTW conversion of this checkpoint: the reader side is in place (iter_vision_weights, the new prefixes), but the 163 GB conversion was not run, so that path is untested.

void-mckenzie · 2026-09-17T06:59:20Z

@jason-fxz since youve been working through the vision families, were you already on deepseek v4 flash vision exp or is this useful? its working on both apis here, details are in the description. happy to test more configs (num-tokens, moe cache size, longer prompts) if you want anything specific checked before you review.

void-mckenzie · 2026-09-30T19:04:56Z

@jason-fxz @andy-yang-1 Still open to feedback, if you have other ideas or directions you'd like me to explore, please let me know. happy to contribute.

FlashML-org#460 and FlashML-org#502 (DeepSeek V4.1 / V4-Flash vision) were rejected only because the models do not fit this machine; that says nothing about the patches. Deferred until after the queue. FlashML-org#535 (ROCm CI) deferred with FlashML-org#217. Assisted-by: Claude Opus 5.5

jason-fxz

@void-mckenzie
Thanks for the PR. DeepSeek-V4.1 (text + vision) landed on main in #603, so please rebase onto the current main. Conflicts: tokenizer/tokenize.py (see the inline comment), docs/models.md (the new DeepSeek-V4 row is joined onto the GLM-5.3-Flash row on one line) and tests/README.md. With the inline comments addressed, other models are unaffected; please rerun tests/models/deepseek_v41 and tests/mm after the rebase.

jason-fxz · 2026-10-07T23:34:55Z

    @abstractmethod
-    def prompt_replacement(self, item: MMItem) -> PromptReplacement:
-        """The token sequence that stands in for this item's placeholder in the prompt."""
+    def prompt_replacement(self, item: MMItem, start: int = 0) -> PromptReplacement:


Please keep MMProcessor.prompt_replacement(item) and apply() as on main and revert the start parameter in every processor and the test fakes. DSV4MMProcessor and DeepseekV41MMProcessor (#603) are separate processors: only the V4 block depends on its position, so handle the offset inside DSV4MMProcessor by overriding its own apply(). After the rebase, the shared signature change makes every V4.1 image request raise TypeError.

Done in 9a3a04a. MMProcessor.prompt_replacement(item) / apply() are back to main's signatures, and start is removed from every other processor and the test fakes. Only DSV4MMProcessor keeps it, called from its own apply() override which computes the block offset itself. The V4.1 TypeErrors are gone: tests/models/deepseek_v41 passes 43/43.

jason-fxz · 2026-10-07T23:34:55Z



-__all__ = ["cut_image_spans", "mm_chunk_end", "mm_rows_after", "plan_mm_batch", "plan_mm_chunk"]
+def mm_spans(block_ends: List[int], reqs) -> dict[int, List[Tuple[int, int]]]:


Please drop mm_spans, Batch.mm_spans and the scheduler line: it runs on every model's image batch and only DeepSeek-V4 reads it. The spans are already on the batch: each request in batch.reqs has its table_idx and mm_items, and item.offsets are absolute and cover the whole block, so DeepseekV4ForCausalLM.forward can build spans_by_ti itself.

The span should also start at IMAGE_START rather than at the first lead pad. The reference get_image_visible counts within [IMAGE_START, IMAGE_END], so the lead pads keep the causal window. With the whole block as the span, each lead pad sees the entire block and rows 128+ tokens in also see the pads: after a 201-token prefix, a 20x16 grid (2 lead pads) gets different window candidates than the reference on 218 rows. For example:

spans_by_ti = { req.table_idx: [(lo + COMPRESS_PAD_TO - 1 - lo % COMPRESS_PAD_TO, hi) for item in req.mm_items for lo, hi in item.offsets] for req in batch.reqs if req.mm_items } or None

A test against the reference pair (get_image_visible + get_window_topk_idxs_visible) would catch this.

Done in 16002fa. mm_spans, Batch.mm_spans and the scheduler line are dropped (scheduler files identical to main). DeepseekV4ForCausalLM.forward now builds spans_by_ti itself from req.table_idx / item.offsets, with the span starting at IMAGE_START (lo + COMPRESS_PAD_TO - 1 - lo % COMPRESS_PAD_TO, the lead-pad count from build_image_block). Added a parity test against get_image_visible + get_window_topk_idxs_visible: your 201-prefix / 20x16 case diverges on exactly 218 rows with the old lead-pad start and matches on all rows with the fix.

jason-fxz · 2026-10-07T23:34:55Z

 # State-dict key prefixes of the vision stack; load_weight drops them when the engine serves text-only.
-VISION_KEY_PREFIXES = ("vision_tower.", "embed_vision.", "vision_embedder.", "visual.")
+# ``image_*`` covers deepseek_v4's sentinel vectors, which carry no tower prefix of their own.
+VISION_KEY_PREFIXES = ("vision_tower.", "embed_vision.", "vision_embedder.", "visual.", "vision.", "aligner.", "image_")


Please leave VISION_KEY_PREFIXES as on main: it filters the weights of every engine that builds no vision tower (all text models, and VLMs under --text-model-only), so image_ would silently drop any future text parameter with that name. Mounting the tower, the aligner and the four sentinels under an existing prefix (for example a visual container on the wrapper) and renaming the checkpoint names in deepseek_v4/weight.py keeps the text-only and FTW filters working without touching the list. On main, is_checkpoint_tower_name in scripts/ftw_hotfix.py lists image_start / image_end / image_newline; please add image_pad there.

Done in a962fe2. VISION_KEY_PREFIXES is untouched. The tower, aligner and four sentinels now mount under visual.* (the same container pattern as V4.1), with deepseek_v4/weight.py renaming the checkpoint names on load; text-only and FTW filters work unchanged. image_pad added to is_checkpoint_tower_name in scripts/ftw_hotfix.py.

jason-fxz · 2026-10-07T23:34:55Z

    return module


+def _reinject_dsv4_image_data(messages: list[dict], images: list[bytes] | None) -> list[dict]:


Main already does this: render_prompt passes msg.images through _render to _apply_dsv4_chat_encoder, which puts the bytes (or the prerender freetoken_ref) on each image part and raises if the parts outnumber the images, and the encoder lookup includes encoding_dsv4.py. Please drop _reinject_dsv4_image_data, _NO_IMAGE_DATA and the _render_msg split when rebasing; this hunk conflicts with main anyway.

Done in the rebase (6e9a6cf) , tokenize.py is unchnged from main; _reinject_dsv4_image_data, _NO_IMAGE_DATA and the _render_msg split are dropped repo-wide.

jason-fxz · 2026-10-07T23:34:55Z

        and concatenates.
+
+        ``spans`` (VL requests) are the request's image-block ranges ``[(start, end)]`` inside
+        ``[start_pos, start_pos + n)`` -- chunking never splits one. Inside a span, every token


The docstring says chunking never splits a block, but keep_images_whole / mm_chunk_end only engage when an attention group sets bidirectional_mm_blocks (scheduler.py:94), which DSV4AttentionGroupConfig does not have. So a chunk can still end inside an image block: the block becomes two spans, one per chunk, the rows before the cut never see the rest of the block, and the cut warning stays off too. Adding bidirectional_mm_blocks: bool = False to DSV4AttentionGroupConfig and setting it in deepseek_v4/config.py when the checkpoint has a vision tower turns on the path Gemma 4 uses, with no scheduler change. A test whose chunk budget ends inside a block would cover it.

Done in 9c40ee8. DSV4AttentionGroupConfig gains bidirectional_mm_blocks: bool = False, set in deepseek_v4/config.py when the checkpoint has a vision tower; no scheduler change; it goes through the existing gemma4 path. Docstring updated to state the real guarantee. New test with the chunk budget ending inside a block: vision config keeps the block whole (no cut warning), text-only keeps the old behavior.

jason-fxz · 2026-10-07T23:34:55Z

@@ -155,6 +159,8 @@ class ModelSpec:
        packed_modules_mapping=_EXPERTS_W123_PACKED,
        # the head, the KV compressors and the indexer's scorer ship bf16; the fp8 config has no modules_to_not_convert
        unquantized_modules=("head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj"),


This checkpoint's fp8 config has no modules_to_not_convert, so vision.* and aligner.* resolve to fp8_block although they ship bf16. Please add both to unquantized_modules; tests/models/test_quant_config.py::test_scheme_for_agrees_with_the_stored_tensors fails for this checkpoint without them.

Done in a962fe2. vision.* and aligner.* added to unquantized_modules. test_scheme_for_agrees_with_the_stored_tensors runs against the real checkpoint: 0 mismatches with the patterns, 98 (all vision.* / aligner.*) without.

The block layout depends on where the image lands: lead pads align its first row token to a 4-token compression stride, so the item hash carries the layout and two alignments never share a radix key. The block's tokens all carry one content pad id (the shared pad scheme); the sparse attention reads its spans from batch.mm_spans, which the scheduler recovers from the flat block ends its existing plan already reports.

…t in its apply

…age_start

…ights unquantized

void-mckenzie · 2026-10-08T17:32:40Z

Hi @jason-fxz , thanks for the response. I have now rebased onto current main and all six comments addressed (I have per-comment replies, so that you can crosscheck against the actual commits). tests/mm 23/23 and tests/models/deepseek_v41 43/43 after the rebase. The gated tests were also run against the real checkpoint (deepseek-ai/DeepSeek-V4-Flash-Vision-Exp, 48-shard HF release): the reference-parity vision test and test_scheme_for_agrees_with_the_stored_tensors both pass on it, and image input was served end-to-end on both API surfaces with a synthetic probe verified correct. Let me know if you need any more changes, happy to contribute more :)

jason-fxz

LGTM, thanks!

…L-org#639 871c600 carried only the renderer's upstream-state refresh (FlashML-org#502 and FlashML-org#134 merged, FlashML-org#133 closed); this adds the FlashML-org#305 note it was meant to. Assisted-by: Claude Opus 5.5

…lashML-org#502) * feat(deepseek_v4): serve image input on the V4-Flash-Vision release The block layout depends on where the image lands: lead pads align its first row token to a 4-token compression stride, so the item hash carries the layout and two alignments never share a radix key. The block's tokens all carry one content pad id (the shared pad scheme); the sparse attention reads its spans from batch.mm_spans, which the scheduler recovers from the flat block ends its existing plan already reports. * fix(mm): keep the shared mmprocessor signature and own the dsv4 offset in its apply * fix(models): build dsv4 image spans in the model and start them at image_start * fix(models): mount the dsv4 vision tower under visual and keep its weights unquantized * fix(models): keep dsv4 image blocks whole across chunked prefill

jason-fxz reviewed Oct 7, 2026

View reviewed changes

void-mckenzie added 5 commits October 8, 2026 08:39

fix(mm): keep the shared mmprocessor signature and own the dsv4 offse…

9a3a04a

…t in its apply

fix(models): build dsv4 image spans in the model and start them at im…

16002fa

…age_start

fix(models): mount the dsv4 vision tower under visual and keep its we…

a962fe2

…ights unquantized

fix(models): keep dsv4 image blocks whole across chunked prefill

9c40ee8

void-mckenzie force-pushed the dsv4-vision-upstream branch from d3ac96c to 9c40ee8 Compare October 8, 2026 17:23

jason-fxz approved these changes Oct 8, 2026

View reviewed changes

jason-fxz merged commit 23fc922 into FlashML-org:main Oct 8, 2026

jason-fxz mentioned this pull request Oct 8, 2026

FreeToken Roadmap (2026) #79

Open

17 of 26 tasks

github-actions Bot mentioned this pull request Oct 9, 2026

FreeToken ROCm Watch — notification channel Castoff995/FreeToken-rocm-test#1

Draft



		__all__ = ["cut_image_spans", "mm_chunk_end", "mm_rows_after", "plan_mm_batch", "plan_mm_chunk"]
		def mm_spans(block_ends: List[int], reqs) -> dict[int, List[Tuple[int, int]]]:

		return module


		def _reinject_dsv4_image_data(messages: list[dict], images: list[bytes] \| None) -> list[dict]:

Conversation

void-mckenzie commented Sep 17, 2026

What

Verification

Not in this PR

Uh oh!

void-mckenzie commented Sep 17, 2026

Uh oh!

void-mckenzie commented Sep 30, 2026

Uh oh!

jason-fxz left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

void-mckenzie commented Oct 8, 2026

Uh oh!

jason-fxz left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jason-fxz left a comment •

edited

Loading