feat(gemma4): serve image input on the Gemma-4 releases - #467
Merged
Merged
Conversation
- Gemma4ForConditionalGeneration registers a vision encoder and the Gemma4MMProcessor: <|image|> expands to boi, the soft-token pad ids and eoi; 1-D rope; --image-max-tokens picks the largest accepted soft-token budget (70/140/280/560/1120) within it and a maximum below 70 is refused at start-up; --image-min-tokens has no effect because every image is scaled to its budget - the SigLIP tower streams its encoder layers from pinned host banks under --mm-encoder-weights host, rejects clipped-linear checkpoints instead of loading them without the clamps, and supports grouped kv heads; the reader skips the tower tensors when no encoder is built - SWAAttentionGroupConfig.bidirectional_mm_blocks and AttentionSpec.bidirectional_mm_blocks name the layers whose image spans attend in both directions; the scheduler hands the batch, per token, the end of the image span holding it (Batch.mm_block_ends) and the Triton extend kernels let those rows also attend to the span's later keys within the sliding window; the paged fallback raises and the other backends never take the spec - on such a model a prefill chunk that would end inside an image span ends at the span start instead, decided after the token budget, the sliding-window pool cap and the page and snapshot alignments; a request whose image does not fit the budget other requests left in the pass waits for a pass of its own; an image larger than a whole pass is split, with a one-time warning naming its size; causal families keep the plain chunking - Gemma4UnifiedForConditionalGeneration serves the 12B gemma4_unified release: Gemma4UnifiedVisionEmbedder (LayerNorm, dense projection, LayerNorm, factorized 2-D position embedding, LayerNorm) feeds the shared multimodal embedder, one 48x48 super-patch per soft token, resident under either --mm-encoder-weights; the config parser reads the unified vision section into UnifiedVisionConfig, the reader maps model.vision_embedder.* as vision weights and Gemma4UnifiedMMProcessor keeps the Gemma processor call and its budgets with 6912-wide rows - freetoken.layers.LayerNorm is the bias-carrying LayerNorm on torch's fused kernel, next to the RMSNorm family - docs/cli.md describes the multimodal flags in the engine's own terms and docs/models.md carries the per-family reading of the token budget and the processor kwargs - the Gemma vision, embedder and tokenizer tests that read a real checkpoint carry needs_weights next to their env-var skip
This was referenced Sep 14, 2026
gdevenyi
added a commit
to gdevenyi/FreeToken
that referenced
this pull request
Sep 15, 2026
…lashML-org#479) image input Conflicts: triton extend attention takes both the quantized-KV args (FlashML-org#408) and block_ends; models/config.py keeps the fp8 env helpers beside main's vision prefixes; docs drop the text-only note. Assisted-by: Claude Opus 5
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Image input for the Gemma-4 releases: the
gemma426B-A4B MoE and 31B dense through their ViT tower, and the 12Bgemma4_unifiedrelease through its linear patch embedder, on the multimodal path from #454.What Gemma-4 needs beyond the Qwen path
<|image|>expands to<start_of_image>, the soft-token pad ids and<end_of_image>; 1-D rope. The count is one of the checkpoint's fixed budgets (70 / 140 / 280 / 560 / 1120, default 280):--image-max-tokenspicks the largest budget within it and refuses one below 70 at start-up;--image-min-tokenshas no effect.use_bidirectional_attention: "vision", the sliding-window layers attend in both directions inside an image span; the full-attention layers stay causal. The scheduler hands the batch the end of each token's image span (Batch.mm_block_ends), the Triton extend kernels honor it within the window, the other backends never receive it.gemma4_unified. No tower: a linear patch embedder (LayerNorm, dense 6912 -> 3840, LayerNorm, factorized 2-D position embedding, LayerNorm) feeds the same multimodal embedder, one 48x48 super-patch per soft token, about 95 MiB resident. Same processor call and budgets, same bidirectional spans. Addsfreetoken.layers.LayerNormnext to the RMSNorm family.Flags
No new flags.
docs/cli.mddescribes the multimodal flags in the engine's own terms;docs/models.mdcarries the per-family image-input table.Verification
block_endsmatch a reference attention under the same mask, plain and split-KV; scheduler tests cover the span pull-back, page alignment, the shared-pass deferral and the sliding-window cap.--image-max-tokens 100), Gemma-4-31B-IT-NVFP4 and gemma-4-12B-it;--text-model-onlyunchanged. At--max-extend-tokens 200the 26B logs the cut of its image once and still describes it.Not in this PR
embed_audio.