Skip to content

feat(gemma4): serve image input on the Gemma-4 releases - #467

Merged
jason-fxz merged 1 commit into
mainfrom
feat/gemma4-mm-squash
Sep 14, 2026
Merged

jason-fxz merged 1 commit into
mainfrom
feat/gemma4-mm-squash

Conversation

@jason-fxz

Copy link
Copy Markdown
Collaborator

Image input for the Gemma-4 releases: the gemma4 26B-A4B MoE and 31B dense through their ViT tower, and the 12B gemma4_unified release through its linear patch embedder, on the multimodal path from #454.

What Gemma-4 needs beyond the Qwen path

  • Processor. <|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-tokens picks the largest budget within it and refuses one below 70 at start-up; --image-min-tokens has no effect.
  • Tower. SigLIP-style ViT, 3x3 average pooling and the multimodal embedder on the text stack's projection layers, block weights host-streamed by default; the reader skips the tower tensors when no encoder is built.
  • Attention. Per the checkpoint's 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.
  • Chunking. On such a model a prefill chunk that would end inside an image span ends at the span start, decided after every other cap; a request whose image does not fit the budget left by other requests in the pass waits for a pass of its own. Only an image larger than a whole pass is split, with a one-time warning. Causal families keep the plain chunking.
  • 12B 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. Adds freetoken.layers.LayerNorm next to the RMSNorm family.

Flags

No new flags. docs/cli.md describes the multimodal flags in the engine's own terms; docs/models.md carries the per-family image-input table.

Verification

  • Tower and 12B embedder against the HF reference on real weights: mean cosine above 0.99 (tower, bf16 error within 1.2x of the reference's own noise) and 0.99999 (embedder).
  • Triton extend kernels with block_ends match 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.
  • Greedy image smokes describe the synthetic image correctly on Gemma-4-26B-A4B-NVFP4 (MoE offload, both encoder placements, --image-max-tokens 100), Gemma-4-31B-IT-NVFP4 and gemma-4-12B-it; --text-model-only unchanged. At --max-extend-tokens 200 the 26B logs the cut of its image once and still describes it.
  • Tokenizer tests on the real 26B and 12B checkpoints; full non-slow suite on the squashed tree 1800 passed (plus the four known environment failures on this box).

Not in this PR

  • Video and audio input, including the 12B release's embed_audio.

- 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
@jason-fxz jason-fxz added feature New feature or request multimodal labels Sep 14, 2026
@jason-fxz jason-fxz mentioned this pull request Sep 14, 2026
22 tasks
@jason-fxz
jason-fxz merged commit 84d236c into main 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature New feature or request multimodal

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant