Gemma-4 PLE/MoE VLM support

ARCHITECTURE.md

@@ -29,10 +29,10 @@ which models are supported on Spyre.
 | DeepSeek-Coder 1.3B | llama | 128 | 64 | Yes | Yes | Yes | Yes |
 | Yi 1.5 6B | llama | 128 | 64 | Yes | Yes | Yes | Yes |
 | Granite Vision 4.1 4B (text backbone) | granite (text) | 64→128 | 64 | Yes (padded) | Yes | Yes | Yes |
-| Gemma 4 12B (bf16) | gemma4\_unified | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
+| Gemma 4 12B | gemma4\_unified | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
 | Gemma 4 26B-A4B (MoE) | gemma4 (MoE, `enable_moe_block`) | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
-| Gemma 4 E2B (bf16) | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
-| Gemma 4 E4B (bf16) | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
+| Gemma 4 E2B | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
+| Gemma 4 E4B | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
 | Gemma 4 31B | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
 | Gemma 4 31B Instruct | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes |
 | Gemma 3 1B | gemma3\_text | 256 | 128 | Yes | Yes | Yes | Yes |
@@ -61,18 +61,20 @@ chat-template tokenization (PR#385).
 Full multimodal VLMs: a vision tower (image → patch features) plus a causal text
 decoder (features + prompt → generated text). Both towers compile and run on
 Spyre; the projector / patch-embed / feature-merge ops that don't lower run on
-CPU (see Multimodal VLM Path below).
+CPU where integer/control-flow-heavy frontend or pooling operations do not lower
+(see Multimodal VLM Path below).
 
 | Model | model\_type | Towers | Stick Aligned | CPU Accurate | Spyre Compiles | Spyre Runs |
 |-------|-----------|--------|-----------|--------------|-------------|-----------|
 | Granite Vision 4.1 4B | granite4\_vision | SigLIP vision + Granite text | Yes (padded) | Yes | Yes | Yes |
 | Mistral-Small-3.1-24B-Instruct-2503 | mistral3 | Pixtral + Mistral text | Yes (padded) | Yes | Yes | Yes |
-| Ministral-3-14B-Instruct-2512 (bf16) | mistral3 | Pixtral + Ministral3 text | Yes (padded) | Yes | Yes | Yes |
-| Gemma 4 12B IT (bf16) | gemma4\_unified | Encoder-free vision embedder + Gemma 4 text | Yes | Yes | Yes | Yes |
-| Gemma 4 12B Base (bf16) | gemma4\_unified | Encoder-free vision embedder + Gemma 4 text | Yes | Yes | Yes | Yes |
+| Ministral-3-14B-Instruct-2512 | mistral3 | Pixtral + Ministral3 text | Yes (padded) | Yes | Yes | Yes |
+| Gemma 4 12B IT | gemma4\_unified | Encoder-free vision embedder + Gemma 4 text | Yes | Yes | Yes | Yes |
+| Gemma 4 E2B | gemma4 | Gemma 4 vision tower + PLE/KV-sharing text | Yes | Yes | Yes | Yes |
+| Gemma 4 26B-A4B | gemma4 | Gemma 4 vision tower + MoE text | Yes | Not run | Yes | Yes |
 
 **CPU Accurate** = adapter `generate` matches stock `model.generate` token-for-token on CPU (`test_vlm_e2e_cpu.py`).
-**Spyre Runs** = `test_vlm_e2e_spyre.py` drives the adapter teacher-forced on stock's tokens and asserts per-step logit cosine ≥ 0.999 vs the CPU reference over prefill + decode steps (top-1 agreement is reported, not asserted — an open-ended caption hits near-ties where the fp16-substrate winner is numerically arbitrary; see Multimodal VLM Path). granite-vision-4.1 holds cosine ≥ 0.99991 at every step and produces a correct, coherent caption. Gemma 4 12B IT runs in **bf16** (like the rest of the Gemma family, it overflows its residual stream in fp16) and holds cosine ≥ 0.99856 at every step with 5/5 top-1 agreement, producing a caption byte-identical to stock. Gemma 4 12B Base routes via the same adapter; its causal-LM path (text-decoder only) holds 5/5 top-1 on Spyre.
+**Spyre Runs** = `test_vlm_e2e_spyre.py` drives the adapter teacher-forced on stock's tokens and asserts per-step logit cosine ≥ 0.99 vs the CPU reference over prefill + decode steps (top-1 agreement is reported, not asserted — an open-ended caption hits near-ties where the bf16-substrate winner is numerically arbitrary; see Multimodal VLM Path). granite-vision-4.1 holds cosine ≥ 0.99991 at every step and produces a correct, coherent caption. Gemma 4 12B IT runs in **bf16** (like the rest of the Gemma family, it overflows its residual stream in fp16) and holds cosine ≥ 0.99939 at every step with 5/5 top-1 agreement. Gemma 4 E2B holds cosine ≥ 0.99502 with 4/5 top-1 agreement after preserving stock PLE and mask semantics and running its 16-layer vision transformer on Spyre. Gemma 4 26B-A4B holds cosine ≥ 0.99956 with 4/5 top-1 agreement while running its 27-layer vision transformer on Spyre.
 
 ### Embedding
 
@@ -159,19 +161,20 @@ and four decode steps. Gemma 4 E2B and E4B both pass that check.
 > adapter or verify a checkpoint, update *only* this file (and the badge
 > counts in README.md, noted below).
 
-**Coverage:** 33 adapters · 55 verified checkpoints · 10K+ compatible models.
-The 55 verified rows are 34 generative + 13 embedding + 2 seq-classification +
-2 token-classification + 4 vision-language (see the Verified Checkpoints tables
+**Coverage:** 33 adapters · 57 verified checkpoints · 10K+ compatible models.
+The 57 verified rows are 34 generative + 13 embedding + 2 seq-classification +
+2 token-classification + 6 vision-language (see the Verified Checkpoints tables
 above). `hf_siglip_vision` and `hf_pixtral_vision` are bare vision-tower components
 used by VLM adapters and are not included in the adapter count. The three DSpark
 speculative-decoding drafter adapters are included in the adapter count.
 Granite Vision 4.1 is verified both as a text backbone (generative) and as a full VLM.
 `hf_mistral3_vision_mm` covers both the ``mistral`` text-backbone variant
 (Mistral-Small-3.1/3.2) and the ``ministral3`` variant (Ministral-3-14B-Instruct-2512,
 blocked-FP8, bf16) — both share the Pixtral vision tower; CPU and Spyre verified
-(token-exact on CPU, 5/5 top-1 match on Spyre). `hf_gemma4_mm` is the first
-**encoder-free** VLM: no vision tower, just a projection of processor-merged pixel
-patches into the text embedding space (runs bf16; CPU and Spyre verified).
+(token-exact on CPU, 5/5 top-1 match on Spyre). `hf_gemma4_mm` supports both the encoder-free unified checkpoints and the full
+Gemma 4 multimodal architecture. The latter currently runs its stock vision tower
+on CPU, then feeds the projected features into the Spyre PLE/KV-sharing or MoE
+text decoder (bf16; CPU reference and Spyre verified).
 
 Each adapter handles a HuggingFace `model_type`. Once verified with
 one checkpoint, all size variants and fine-tunes of that architecture
@@ -193,7 +196,7 @@ pattern, norms, and weight layout.
 | hf\_smollm3.py | smollm3 | 1 | — |
 | hf\_lfm2.py | lfm2 | 1 | LFM2 700M/1.2B and dense LFM2 fine-tunes with hybrid convolution/attention layers |
 | hf\_gemma4.py | gemma4\_unified / gemma4 (dense + PLE/KV-share) | 4 | Gemma 4 31B (dense). Not 26B-A4B (MoE). |
-| hf\_gemma4\_mm.py | gemma4\_unified (multimodal) | 2 | Gemma 4 31B (dense unified VLM). Not E2B/E4B (PLE) or 26B-A4B (MoE). |
+| hf\_gemma4\_mm.py | gemma4\_unified / gemma4 (multimodal) | 4 | Encoder-free dense unified VLMs plus full-vision PLE/KV-share and MoE variants. Combined MoE+PLE/KV-share remains unsupported. |
 | hf\_gemma4\_moe.py | gemma4 (MoE, `enable_moe_block`) | 1 | Gemma 4 26B-A4B (128 experts, top-8 routing). Persistent prefill + gathered decode, 5/5 token match. |
 | hf\_gemma3.py | gemma3\_text / gemma3 (dense) | 2 | Gemma 3 4B/12B/27B (text decoder of the multimodal checkpoints); EmbeddingGemma (bidirectional embedder). Not Gemma 3n (PLE). |
 | hf\_gemma2.py | gemma2 | 1 | Gemma 2 2B and Gemma 2 fine-tunes. |
@@ -736,23 +739,47 @@ Multimodal-specific Spyre adaptations (beyond those shared with Granite VLM):
 - **`Mistral3PatchMerger`** (`nn.functional.unfold` + `merging_layer`) runs on
   CPU inside the projector — `unfold` doesn't lower on Spyre.
 
-**Encoder-free VLM — Gemma 4** (`hf_gemma4_mm.py`): the first VLM with **no
-vision tower**. Gemma 4 is unified (`model_type=gemma4_unified`); vision is a pure
-projection of processor-merged raw pixel patches (the image processor already
-merges `pooling_kernel_size²` 16×16 patches into each 48×48 merged patch, 1:1 to a
-soft token) into the LM embedding space — no attention, no RoPE, no KV cache. The
-adapter loads the full VLM via `AutoModelForImageTextToText`, projects the patches,
-scatters them into the `<image>` token slots, and runs the `hf_gemma4` text decoder
-unchanged (both towers live under the one model, so a single `prepare_for_spyre`
-covers them). Gemma4-specific Spyre adaptations:
-
-- **Vision embedder compiled on Spyre.** The attention-free projection core
-  (`LN₁ → Dense → LN₂ → +pos_embs → pos_norm → RMSNorm → Linear`) is
-  `torch.compile`d and runs on Spyre. Only the integer-XY positional-embedding
-  gather (with `-1` padding validity masking) and the final padding-patch strip run
-  on **CPU** — those integer-gather / boolean-index ops don't lower (same doctrine
-  as the SigLIP CPU patch-embed). The CPU-built per-patch positional-embedding
-  tensor is passed into the compiled core as a device argument.
+**Gemma 4 VLM** (`hf_gemma4_mm.py`) supports two vision architectures behind one
+multimodal frontend. `Gemma4UnifiedConfig` is encoder-free: it projects
+processor-merged raw pixel patches directly into the LM embedding space. The
+`Gemma4Config` checkpoints have a full Gemma 4 vision transformer. Its dense
+transformer blocks run on Spyre; integer position lookup, spatial pooling, and the
+text-space projector stay on CPU before features enter the shared Spyre
+image-scatter and decoder path. Text preparation dispatches by the nested config:
+dense and PLE/KV-sharing checkpoints reuse `hf_gemma4`, while
+`enable_moe_block=True` reuses `hf_gemma4_moe`. Gemma4-specific Spyre adaptations:
+
+- **Encoder-free vision embedder compiled on Spyre.** For `gemma4_unified`, the
+  attention-free projection core (`LN₁ → Dense → LN₂ → +pos_embs → pos_norm →
+  RMSNorm → Linear`) is `torch.compile`d and runs on Spyre. Only the integer-XY
+  positional gather and padding strip run on CPU.
+- **Full vision transformer blocks on Spyre** (`hf_gemma4_vision.py`). The CPU
+  patch projection and learned XY-position lookup produce a fixed 2520-patch
+  sequence; it is right-padded to 2560 so the attention score matmul has no ragged
+  final stick. Added and processor-padding key columns are masked once, the
+  sandwich-RMSNorm/GELU-gated encoder blocks run on Spyre, and the output is
+  cropped back to 2520 before the stock CPU spatial pooler. The text-space
+  `embed_vision` projector remains on CPU.
+- **Two-axis RoPE and head padding.** E2B's 64-wide and 26B-A4B's 72-wide heads
+  are padded to 128. Q/K channels and learned norm weights are rearranged so a
+  single matrix-form RoPE applies the stock independent X/Y rotations without
+  head-dimension slicing. Padded Q/K/V RMSNorm compensates for the added zero
+  lanes, preserving the native-width denominator; attention keeps stock's
+  explicit scale of 1.0. E2B's finite linear clipping bounds are preserved.
+- **Full-tower MLP padding.** The 26B-A4B vision tower's private intermediate
+  width is padded 4304→4352 by zero-extending gate/up outputs and down-projection
+  inputs. E2B's 3072-wide MLP is already stick-aligned. The 26B-only fp32 output
+  standardization remains on CPU with the pooler.
+- **PLE preserves stock multimodal semantics.** The token-identity component uses
+  image-placeholder IDs replaced with the pad token, while the contextual component
+  projects embeddings with multimodal positions replaced by the raw, unscaled pad
+  embedding weight.
+  Decode passes each generated
+  token ID through the generic VLM hook so PLE can be recomputed per step. Existing
+  producer-cache mapping handles the trailing KV-sharing layers.
+- **MoE text composition.** `enable_moe_block=True` selects the existing persistent
+  prefill and gathered-decode MoE blocks. Checkpoints combining MoE with PLE or KV
+  sharing remain explicitly unsupported.
 - **Vision LayerNorms un-fused + fp32 reduction** (`patch_layernorm`). The three
   vision `nn.LayerNorm`s NaN on Spyre's fused lowering on near-constant (small but
   nonzero variance) rows — in **both** bf16 and fp16, so it's a genuine lowering
@@ -761,11 +788,12 @@ covers them). Gemma4-specific Spyre adaptations:
   HF RMSNorm uses), keeping the affine multiply in bf16. Without it the VLM logits are
   all-NaN (see docs/gemma4_mm_vision_layernorm_spyre.md).
 - **Bidirectional vision attention at prefill.** `use_bidirectional_attention ==
-  "vision"`: within one image the soft-tokens attend bidirectionally. Stock OR-s a
-  blockwise band (same image group ⇒ allowed) into the causal mask for **both**
-  full and sliding layers, so at prefill the adapter builds `full = OR(causal,
-  blockwise)` and `sliding = AND(sliding_window, OR(causal, blockwise))`. Decode
-  steps are pure text (one new causal token), so no blockwise band after prefill.
+  "vision"`: full-vision `Gemma4Config` checkpoints keep full-attention layers
+  causal and use `AND(sliding_window, OR(causal, blockwise))` for sliding layers.
+  Encoder-free `Gemma4UnifiedConfig` checkpoints use `OR(causal, blockwise)` for
+  full layers and `OR(AND(sliding_window, causal), blockwise)` for sliding layers,
+  matching the stock Unified forward path. Decode steps are pure text (one new
+  causal token), so no blockwise band is needed after prefill.
 - **Runs in bf16.** Like the rest of the Gemma family (Gemma 3 / EmbeddingGemma),
   Gemma 4 overflows its residual stream in fp16 (`inf` → NaN end-to-end), so it
   runs in bf16. This is a separate concern from the vision LayerNorm defect above,

README.md

@@ -13,7 +13,7 @@ from `transformers`.
 
 ## Supported Models
 
-**34 adapters · 55 verified checkpoints · 10K+ compatible models**
+**34 adapters · 10K+ compatible models**
 
 Coverage spans **generative** (causal-LM), **embedding** (sentence-transformers),
 **sequence classification** (sentiment / text categorisation),
@@ -22,7 +22,7 @@ Coverage spans **generative** (causal-LM), **embedding** (sentence-transformers)
 Llama / Qwen / Granite / Mistral / Phi / Gemma / OLMo / GPT decoders to BERT /
 XLM-RoBERTa / MPNet / ModernBERT encoders, the Granite Vision 4.1 (SigLIP tower +
 Granite text), Mistral3 Vision (Pixtral tower + Mistral text), and Gemma 4
-(encoder-free) multimodal VLMs, plus the DSpark block-propose drafters for
+(encoder-free and full-vision) multimodal VLMs, plus the DSpark block-propose drafters for
 Qwen 3 / Granite / Gemma 4.
 
 Each adapter covers all size variants and fine-tuned checkpoints sharing the same
@@ -247,8 +247,9 @@ A multimodal checkpoint's config is registered under both auto classes:
 `AutoSpyreModelForCausalLM` selects the text-only adapter (vision tower
 discarded), while `AutoSpyreModelForImageTextToText` selects the combined
 multimodal adapter. This works for Granite Vision (`Granite4VisionConfig`),
-Mistral3 Vision (`Mistral3Config`), and Gemma 4 (`Gemma4UnifiedConfig`, an
-encoder-free VLM — no vision tower; see [ARCHITECTURE.md](ARCHITECTURE.md#multimodal-vlm-path-vision-tower--text-decoder)).
+Mistral3 Vision (`Mistral3Config`), and Gemma 4 (`Gemma4UnifiedConfig` for the
+encoder-free variant, or `Gemma4Config` for the full vision-tower variants; see
+[ARCHITECTURE.md](ARCHITECTURE.md#multimodal-vlm-path-vision-tower--text-decoder)).
 
 ## Repo Structure
 

hf_adapters/auto_spyre_model.py

@@ -204,6 +204,7 @@
 IMAGE_TEXT_TO_TEXT_CONFIG_TO_ADAPTER_MODULE_MAPPING: dict[
     type[PretrainedConfig], ModuleType
 ] = {
+    Gemma4Config: hf_gemma4_mm,
     Gemma4UnifiedConfig: hf_gemma4_mm,
     Granite4VisionConfig: hf_granite_vision_mm,
     Mistral3Config: hf_mistral3_vision_mm,
@@ -846,6 +847,7 @@ def _run_vlm_text_forward(
         key_caches,
         value_caches,
         cache_index,
+        input_ids=input_ids,
     )