diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 6e65a392..0ac9c3d2 100644 --- a/ARCHITECTURE.md +++ b/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,9 +161,9 @@ 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. @@ -169,9 +171,10 @@ Granite Vision 4.1 is verified both as a text backbone (generative) and as a ful `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 `` 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, diff --git a/README.md b/README.md index 0f3a7b2f..fb15c1c7 100644 --- a/README.md +++ b/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 diff --git a/hf_adapters/auto_spyre_model.py b/hf_adapters/auto_spyre_model.py index 436cf5a3..dc40826c 100644 --- a/hf_adapters/auto_spyre_model.py +++ b/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, ) diff --git a/hf_adapters/hf_gemma4.py b/hf_adapters/hf_gemma4.py index b8346ad8..6ce7510c 100644 --- a/hf_adapters/hf_gemma4.py +++ b/hf_adapters/hf_gemma4.py @@ -61,15 +61,15 @@ - **Per-Layer Embeddings (PLE).** E-variants inject a per-layer residual after the MLP: ``embed_tokens_per_layer`` + a projected/normed context term, gated - and added back per layer (``_compute_per_layer_inputs`` + ``_ple_tail``, - mirroring stock ``get_per_layer_inputs`` / ``project_per_layer_inputs`` and - the decoder tail). Gated off (``has_ple=False``) for the dense 12B/31B + and added back per layer (``_compute_per_layer_inputs`` plus the compiled + decoder tail), mirroring stock ``get_per_layer_inputs`` / + ``project_per_layer_inputs``. Gated off (``has_ple=False``) for the dense 12B/31B variants, which carry no PLE submodules. - **KV-sharing across layers.** The trailing ``num_kv_shared_layers`` layers reuse the KV cache of the nearest preceding non-shared layer of the same ``layer_type`` (stock ``store_full_length_kv`` semantics), so they run a lean Q-only block (no k/v proj, no cache write) against that producer's cache — see - ``_shared_producer_map`` and ``Gemma4SharedBlock``. + ``_shared_producer_map`` and the ``shared`` kind of ``Gemma4Block``. The MoE 26B-A4B variant is handled by the sibling ``hf_gemma4_moe`` adapter, which reuses this module's attention-side setup and forward driver. @@ -85,6 +85,8 @@ outputs = model.generate(**encoded, max_new_tokens=32) """ +from dataclasses import dataclass + import torch import torch.nn as nn import torch.nn.functional as F @@ -92,6 +94,7 @@ from hf_adapters.hf_common import ( InvFreqShim, PrecomputedRotaryEmbedding, + SpyreUnsupportedModelError, add_causal_sliding_window_band, apply_rope_matmul, get_backbone, @@ -155,22 +158,6 @@ def _compute_per_layer_inputs(model, inputs_embeds, input_ids): return (context + token_identity) * backbone.per_layer_input_scale -def _ple_tail(block, h, per_layer_input): - """Gemma 4 PLE per-layer residual injection (stock modeling_gemma4 tail). - - ``block`` is the registered decoder block carrying the per-layer PLE - submodules (``per_layer_input_gate`` / ``per_layer_projection`` / - ``post_per_layer_input_norm``), captured in ``__init__`` when ``has_ple``. - """ - residual = h - x = block.per_layer_input_gate(h) - x = F.gelu(x, approximate="tanh") - x = x * per_layer_input - x = block.per_layer_projection(x) - x = block.post_per_layer_input_norm(x) - return residual + x - - def _offset_zero_per_layer_input(per_layer_inputs, layer_index): """Copy one PLE layer slice into fresh, offset-zero storage. @@ -297,13 +284,6 @@ def __init__(self, attn, num_q_heads, num_kv_heads, head_dim, is_kv_eq_v): self.head_dim = head_dim self.is_kv_eq_v = is_kv_eq_v self.scaling = attn.scaling # 1.0 for Gemma 4 - self._use_compiled_rms_norm = False - - def _rms_norm(self, hidden_states, norm): - if self._use_compiled_rms_norm: - weight = norm.weight if norm.with_scale else None - return _compiled_gemma4_rms_norm(hidden_states, weight, norm.eps) - return norm(hidden_states) def forward( self, @@ -315,8 +295,6 @@ def forward( cache_index, ): bsz, seq_len, _ = hidden_states.shape - # Q/K/V projections viewed as [B, L, n_heads, head_dim]; norms are - # applied per-head (last dim = head_dim) before the transpose. q = self.q_proj(hidden_states).view( bsz, seq_len, self.num_q_heads, self.head_dim ) @@ -325,31 +303,19 @@ def forward( ) if self.is_kv_eq_v: - # V reuses the raw k_proj output (pre-k_norm, pre-RoPE) but still - # passes through v_norm: stock HF aliases value_states = key_states - # *before* k_norm/RoPE, then applies self.v_norm(value_states) - # unconditionally (modeling_gemma4 Gemma4TextAttention.forward). The - # norm exists on these layers even though v_proj is None. - v = self._rms_norm(k_lin, self.v_norm).transpose(1, 2) + v = _gemma4_rms_norm(k_lin, None, self.v_norm.eps).transpose(1, 2) else: v = self.v_proj(hidden_states).view( bsz, seq_len, self.num_kv_heads, self.head_dim ) - v = self._rms_norm(v, self.v_norm).transpose(1, 2) + v = _gemma4_rms_norm(v, None, self.v_norm.eps).transpose(1, 2) - q = self._rms_norm(q, self.q_norm).transpose(1, 2) - k = self._rms_norm(k_lin, self.k_norm).transpose(1, 2) - # Materialize the transpose returned by RoPE before the cache scatter. - # A view here can make index_copy_ consume the wrong physical layout. + q = _gemma4_rms_norm(q, self.q_norm.weight, self.q_norm.eps).transpose(1, 2) + k = _gemma4_rms_norm(k_lin, self.k_norm.weight, self.k_norm.eps).transpose(1, 2) q = apply_rope_matmul(q, selected_freqs).contiguous() k = apply_rope_matmul(k, selected_freqs).contiguous() - key_cache, value_cache = kv_cache_update( - k, - v, - key_cache, - value_cache, - cache_index, + k, v, key_cache, value_cache, cache_index ) attn_out = F.scaled_dot_product_attention( q, @@ -365,17 +331,42 @@ def forward( class Gemma4Block(nn.Module): - """Registered dense Gemma 4 decoder block used by the Spyre adapter.""" + """Registered writer or KV-sharing Gemma 4 decoder block.""" - def __init__(self, layer, num_q_heads, num_kv_heads, head_dim, is_kv_eq_v, has_ple): + def __init__( + self, + layer, + kind, + num_q_heads, + num_kv_heads, + head_dim, + is_kv_eq_v, + has_ple, + ): super().__init__() - self.self_attn = Gemma4Attention( - layer.self_attn, - num_q_heads, - num_kv_heads, - head_dim, - is_kv_eq_v, - ) + if kind not in ("writer", "shared"): + raise ValueError(f"Unsupported Gemma 4 block kind: {kind!r}") + self.kind = kind + + if kind == "writer": + self.self_attn = Gemma4Attention( + layer.self_attn, + num_q_heads, + num_kv_heads, + head_dim, + is_kv_eq_v, + ) + else: + # A shared block runs lean Q-only attention against its producer's KV + # cache: it has no K/V projections or norms and never updates a cache. + attn = layer.self_attn + self.q_proj = attn.q_proj + self.q_norm = attn.q_norm + self.o_proj = attn.o_proj + self.scaling = attn.scaling + self.num_q_heads = num_q_heads + self.head_dim = head_dim + self.mlp = layer.mlp self.input_layernorm = layer.input_layernorm self.post_attention_layernorm = layer.post_attention_layernorm @@ -394,154 +385,390 @@ def __init__(self, layer, num_q_heads, num_kv_heads, head_dim, is_kv_eq_v, has_p self.post_per_layer_input_norm = layer.post_per_layer_input_norm self.train(layer.training) - def forward( - self, - hidden_states, - selected_freqs, - attn_mask, - key_cache, - value_cache, - cache_index, - layer_scalar, - per_layer_input=None, - query_row_mask=None, - ): - residual = hidden_states - h = self.input_layernorm(hidden_states) - attn_out, key_cache, value_cache = self.self_attn( - h, - selected_freqs, - attn_mask, - key_cache, - value_cache, - cache_index, + # This module intentionally owns state without implementing forward(). The + # shared compiled executor receives its tensors explicitly so structurally + # identical layers reuse one graph. An eager forward could delegate to that + # same functional executor in the future without duplicating block semantics. + + +@dataclass(frozen=True) +class _Gemma4BlockSpec: + kind: str + layer_type: str + num_q_heads: int + num_kv_heads: int + head_dim: int + intermediate_size: int + activation: str + is_kv_eq_v: bool + attention_bias: bool + has_ple: bool + ple_dim: int + scaling: float + input_norm_eps: float + q_norm_eps: float + k_norm_eps: float + v_norm_eps: float + post_attention_norm_eps: float + pre_feedforward_norm_eps: float + post_feedforward_norm_eps: float + post_ple_norm_eps: float + + +def _gemma4_mlp_activation(hidden_states): + return F.gelu(hidden_states, approximate="tanh") + + +def _linear_weight(linear, name): + if linear.bias is not None: + raise ValueError(f"Gemma 4 {name} bias is not supported") + return linear.weight + + +def _block_kind_and_attention(block): + if not isinstance(block, Gemma4Block): + raise TypeError(f"Unsupported Gemma 4 block type: {type(block).__name__}") + kind = block.kind + if kind == "writer": + return kind, block.self_attn + if kind == "shared": + return kind, block + raise ValueError(f"Unsupported Gemma 4 block kind: {kind!r}") + + +def _block_state(block): + kind, attn = _block_kind_and_attention(block) + if kind == "writer": + state = [attn.q_proj.weight, attn.k_proj.weight] + if attn.q_proj.bias is not None: + state.extend([attn.q_proj.bias, attn.k_proj.bias]) + if not attn.is_kv_eq_v: + state.append(attn.v_proj.weight) + if attn.v_proj.bias is not None: + state.append(attn.v_proj.bias) + state.append(attn.o_proj.weight) + if attn.o_proj.bias is not None: + state.append(attn.o_proj.bias) + state.extend([attn.q_norm.weight, attn.k_norm.weight]) + else: + state = [attn.q_proj.weight] + if attn.q_proj.bias is not None: + state.append(attn.q_proj.bias) + state.append(attn.o_proj.weight) + if attn.o_proj.bias is not None: + state.append(attn.o_proj.bias) + state.append(attn.q_norm.weight) + + state.extend( + [ + _linear_weight(block.mlp.gate_proj, "gate_proj"), + _linear_weight(block.mlp.up_proj, "up_proj"), + _linear_weight(block.mlp.down_proj, "down_proj"), + block.input_layernorm.weight, + block.post_attention_layernorm.weight, + block.pre_feedforward_layernorm.weight, + block.post_feedforward_layernorm.weight, + ] + ) + if block.has_ple: + state.extend( + [ + _linear_weight(block.per_layer_input_gate, "per_layer_input_gate"), + _linear_weight(block.per_layer_projection, "per_layer_projection"), + block.post_per_layer_input_norm.weight, + ] ) - # Sandwich: norm the attention output BEFORE adding the residual. - h = residual + self.post_attention_layernorm(attn_out) + state.append(block.layer_scalar) + return tuple(state) + + +def _finish_block( + spec, + residual, + attn_out, + o_weight, + o_bias, + post_attn_norm_weight, + pre_ffn_norm_weight, + post_ffn_norm_weight, + gate_weight, + up_weight, + down_weight, + ple_gate_weight, + ple_projection_weight, + post_ple_norm_weight, + layer_scalar, + per_layer_input, + query_row_mask, +): + attn_out = F.linear(attn_out, o_weight, o_bias) + h = residual + _gemma4_rms_norm( + attn_out, post_attn_norm_weight, spec.post_attention_norm_eps + ) + residual = h + h = _gemma4_rms_norm(h, pre_ffn_norm_weight, spec.pre_feedforward_norm_eps) + h = F.linear( + _gemma4_mlp_activation(F.linear(h, gate_weight)) * F.linear(h, up_weight), + down_weight, + ) + h = residual + _gemma4_rms_norm( + h, post_ffn_norm_weight, spec.post_feedforward_norm_eps + ) + if spec.has_ple: residual = h - h = self.pre_feedforward_layernorm(h) - h = self.mlp(h) - h = self.post_feedforward_layernorm(h) + h = F.gelu(F.linear(h, ple_gate_weight), approximate="tanh") + h = h * per_layer_input + h = F.linear(h, ple_projection_weight) + h = _gemma4_rms_norm(h, post_ple_norm_weight, spec.post_ple_norm_eps) h = residual + h - if self.has_ple: - h = _ple_tail(self, h, per_layer_input) - h = h * layer_scalar - if query_row_mask is not None: - h = h * query_row_mask - return h, key_cache, value_cache - - -class Gemma4SharedBlock(nn.Module): - """KV-sharing Gemma 4 decoder block (E-variant trailing layers). + h = h * layer_scalar + if query_row_mask is not None: + h = h * query_row_mask + return h - Runs a lean Q-only attention against a *producer* layer's KV cache: no - k/v projection, no k_norm/v_norm, no RoPE-on-K, no cache update. The - producer's cache is passed in by the driver (``_run_blocks_over_embeds`` - selects ``key_caches[producer_of[i]]``), so this block returns only the - updated hidden state (never a cache tuple) and needs no ``cache_index`` — - it never writes. - """ - - def __init__(self, layer, num_q_heads, head_dim, has_ple): - super().__init__() - attn = layer.self_attn - self.q_proj = attn.q_proj - self.q_norm = attn.q_norm - self.o_proj = attn.o_proj - self.scaling = attn.scaling # 1.0 for Gemma 4 - self.num_q_heads = num_q_heads - self.head_dim = head_dim - self.mlp = layer.mlp - self.input_layernorm = layer.input_layernorm - self.post_attention_layernorm = layer.post_attention_layernorm - self.pre_feedforward_layernorm = layer.pre_feedforward_layernorm - self.post_feedforward_layernorm = layer.post_feedforward_layernorm - self.register_buffer( - "layer_scalar", - layer.layer_scalar, - persistent="layer_scalar" not in layer._non_persistent_buffers_set, - ) - self.has_ple = has_ple - if has_ple: - self.per_layer_input_gate = layer.per_layer_input_gate - self.per_layer_projection = layer.per_layer_projection - self.post_per_layer_input_norm = layer.post_per_layer_input_norm - self.train(layer.training) +def _make_block_forward(spec): def forward( - self, + state, hidden_states, selected_freqs, attn_mask, key_cache, value_cache, - layer_scalar, - per_layer_input=None, - query_row_mask=None, + cache_index, + per_layer_input, + query_row_mask, ): + pos = 0 + if spec.kind == "writer": + q_weight, k_weight = state[pos : pos + 2] + pos += 2 + if spec.attention_bias: + q_bias, k_bias = state[pos : pos + 2] + pos += 2 + else: + q_bias = k_bias = None + if spec.is_kv_eq_v: + v_weight = v_bias = None + else: + v_weight = state[pos] + pos += 1 + if spec.attention_bias: + v_bias = state[pos] + pos += 1 + else: + v_bias = None + else: + q_weight = state[pos] + pos += 1 + if spec.attention_bias: + q_bias = state[pos] + pos += 1 + else: + q_bias = None + + o_weight = state[pos] + pos += 1 + if spec.attention_bias: + o_bias = state[pos] + pos += 1 + else: + o_bias = None + if spec.kind == "writer": + q_norm_weight, k_norm_weight = state[pos : pos + 2] + pos += 2 + else: + q_norm_weight = state[pos] + pos += 1 + gate_weight, up_weight, down_weight = state[pos : pos + 3] + pos += 3 + input_norm_weight, post_attn_norm_weight = state[pos : pos + 2] + pos += 2 + pre_ffn_norm_weight, post_ffn_norm_weight = state[pos : pos + 2] + pos += 2 + ple_gate_weight = ple_projection_weight = post_ple_norm_weight = None + if spec.has_ple: + ple_gate_weight, ple_projection_weight, post_ple_norm_weight = state[ + pos : pos + 3 + ] + pos += 3 + layer_scalar = state[pos] + + bsz, seq_len, _ = hidden_states.shape residual = hidden_states - h = self.input_layernorm(hidden_states) - bsz, seq_len, _ = h.shape - q = self.q_proj(h).view(bsz, seq_len, self.num_q_heads, self.head_dim) - q = self.q_norm(q).transpose(1, 2) - q = apply_rope_matmul(q, selected_freqs) + h = _gemma4_rms_norm(hidden_states, input_norm_weight, spec.input_norm_eps) + q = F.linear(h, q_weight, q_bias).view( + bsz, seq_len, spec.num_q_heads, spec.head_dim + ) + q = _gemma4_rms_norm(q, q_norm_weight, spec.q_norm_eps).transpose(1, 2) + + if spec.kind == "writer": + k_lin = F.linear(h, k_weight, k_bias).view( + bsz, seq_len, spec.num_kv_heads, spec.head_dim + ) + if spec.is_kv_eq_v: + v = _gemma4_rms_norm(k_lin, None, spec.v_norm_eps).transpose(1, 2) + else: + v = F.linear(h, v_weight, v_bias).view( + bsz, seq_len, spec.num_kv_heads, spec.head_dim + ) + v = _gemma4_rms_norm(v, None, spec.v_norm_eps).transpose(1, 2) + q = apply_rope_matmul(q, selected_freqs).contiguous() + k = apply_rope_matmul( + _gemma4_rms_norm(k_lin, k_norm_weight, spec.k_norm_eps).transpose(1, 2), + selected_freqs, + ).contiguous() + key_cache, value_cache = kv_cache_update( + k, v, key_cache, value_cache, cache_index + ) + else: + q = apply_rope_matmul(q, selected_freqs) + attn_out = F.scaled_dot_product_attention( q, key_cache, value_cache, attn_mask=attn_mask, dropout_p=0.0, - scale=self.scaling, + scale=spec.scaling, enable_gqa=True, ) attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) - attn_out = self.o_proj(attn_out) - h = residual + self.post_attention_layernorm(attn_out) + h = _finish_block( + spec, + residual, + attn_out, + o_weight, + o_bias, + post_attn_norm_weight, + pre_ffn_norm_weight, + post_ffn_norm_weight, + gate_weight, + up_weight, + down_weight, + ple_gate_weight, + ple_projection_weight, + post_ple_norm_weight, + layer_scalar, + per_layer_input, + query_row_mask, + ) + return h, key_cache, value_cache - residual = h - h = self.pre_feedforward_layernorm(h) - h = self.mlp(h) - h = self.post_feedforward_layernorm(h) - h = residual + h - if self.has_ple: - h = _ple_tail(self, h, per_layer_input) - h = h * layer_scalar - if query_row_mask is not None: - h = h * query_row_mask - return h + return forward + + +def _block_spec(block, layer_type): + kind, attn = _block_kind_and_attention(block) + if kind == "writer": + num_kv_heads = attn.num_kv_heads + is_kv_eq_v = attn.is_kv_eq_v + k_norm_eps = attn.k_norm.eps + v_norm_eps = attn.v_norm.eps + else: + num_kv_heads = 0 + is_kv_eq_v = False + k_norm_eps = 0.0 + v_norm_eps = 0.0 + activation = block.mlp.config.hidden_activation + if activation != "gelu_pytorch_tanh": + raise SpyreUnsupportedModelError( + "Gemma 4 checkpoints must use hidden_activation='gelu_pytorch_tanh'; " + f"got {activation!r}" + ) + attention_bias = attn.q_proj.bias is not None + attention_projections = [attn.q_proj, attn.o_proj] + if kind == "writer": + attention_projections.append(attn.k_proj) + if not is_kv_eq_v: + attention_projections.append(attn.v_proj) + if any( + (projection.bias is not None) != attention_bias + for projection in attention_projections + ): + raise ValueError( + "Gemma 4 attention projections must use a consistent bias setting" + ) + ple_dim = block.per_layer_input_gate.out_features if block.has_ple else 0 + post_ple_eps = block.post_per_layer_input_norm.eps if block.has_ple else 0.0 + scaled_norms = [ + ("input_layernorm", block.input_layernorm), + ("q_norm", attn.q_norm), + ("post_attention_layernorm", block.post_attention_layernorm), + ("pre_feedforward_layernorm", block.pre_feedforward_layernorm), + ("post_feedforward_layernorm", block.post_feedforward_layernorm), + ] + if kind == "writer": + scaled_norms.append(("k_norm", attn.k_norm)) + if block.has_ple: + scaled_norms.append( + ("post_per_layer_input_norm", block.post_per_layer_input_norm) + ) + unscaled = [name for name, norm in scaled_norms if not norm.with_scale] + if unscaled: + raise SpyreUnsupportedModelError( + "Gemma 4 requires scaled RMSNorm for " + ", ".join(unscaled) + ) + if kind == "writer" and attn.v_norm.with_scale: + raise SpyreUnsupportedModelError( + "Gemma 4 requires an unscaled RMSNorm for v_norm" + ) + return _Gemma4BlockSpec( + kind=kind, + layer_type=layer_type, + num_q_heads=attn.num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=attn.head_dim, + intermediate_size=block.mlp.gate_proj.out_features, + activation=activation, + is_kv_eq_v=is_kv_eq_v, + attention_bias=attention_bias, + has_ple=block.has_ple, + ple_dim=ple_dim, + scaling=attn.scaling, + input_norm_eps=block.input_layernorm.eps, + q_norm_eps=attn.q_norm.eps, + k_norm_eps=k_norm_eps, + v_norm_eps=v_norm_eps, + post_attention_norm_eps=block.post_attention_layernorm.eps, + pre_feedforward_norm_eps=block.pre_feedforward_layernorm.eps, + post_feedforward_norm_eps=block.post_feedforward_layernorm.eps, + post_ple_norm_eps=post_ple_eps, + ) def prepare_gemma4_blocks( - layers, num_q_heads_per_layer, kv_shapes, is_kv_eq_v_per_layer, producer_of, has_ple + layers, + layer_types, + num_q_heads_per_layer, + kv_shapes, + is_kv_eq_v_per_layer, + producer_of, + has_ple, ): - """Replace Gemma 4 decoder layers with registered blocks and compile them. - - ``producer_of[i]`` is ``None`` for a normal (KV-writing) layer and an int - for a KV-sharing layer, in which case a lean ``Gemma4SharedBlock`` is built - instead of a full ``Gemma4Block``. - """ - blocks = [] + """Install registered blocks and share compiled executors by structure.""" + compiled_by_spec = {} + compiled_blocks = [] for i, layer in enumerate(list(layers)): - if producer_of[i] is None: - block = Gemma4Block( - layer, - num_q_heads_per_layer[i], - kv_shapes[i][0], - kv_shapes[i][1], - is_kv_eq_v_per_layer[i], - has_ple, - ) - else: - block = Gemma4SharedBlock( - layer, - num_q_heads_per_layer[i], - kv_shapes[i][1], - has_ple, - ) + kind = "writer" if producer_of[i] is None else "shared" + block = Gemma4Block( + layer, + kind, + num_q_heads_per_layer[i], + kv_shapes[i][0], + kv_shapes[i][1], + is_kv_eq_v_per_layer[i], + has_ple, + ) layers[i] = block - blocks.append(torch.compile(block, dynamic=False)) - return blocks + spec = _block_spec(block, layer_types[i]) + if spec not in compiled_by_spec: + compiled_by_spec[spec] = torch.compile( + _make_block_forward(spec), dynamic=False, fullgraph=True + ) + compiled_blocks.append(compiled_by_spec[spec]) + return compiled_blocks def _build_layer_masks( @@ -662,9 +889,7 @@ def _run_blocks_over_embeds( else None ) p = producer_of[i] - # Pass the per-layer scalar as a tensor read fresh from the registered, - # device-moved block — NOT as a Python float — so Dynamo guards on tensor - # metadata instead of recompiling for each distinct learned value. + block = backbone_layers[i] if is_moe: # The dedicated MoE blocks predate the dense/E optional inputs and # compile their attention and FFN regions internally. Preserve that @@ -680,25 +905,26 @@ def _run_blocks_over_embeds( ) elif p is None: h, key_caches[i], value_caches[i] = compiled_block( + _block_state(block), h, freqs[lt], masks[lt], key_caches[i], value_caches[i], cache_index, - backbone_layers[i].layer_scalar, pli, query_row_mask, ) else: # KV-sharing layer: read the producer's cache, write nothing. - h = compiled_block( + h, _, _ = compiled_block( + _block_state(block), h, freqs[lt], masks[lt], key_caches[p], value_caches[p], - backbone_layers[i].layer_scalar, + None, pli, query_row_mask, ) @@ -866,17 +1092,19 @@ def prepare_text_decoder_for_spyre(model): Runs the shared attention-side setup (``_setup_gemma4_text_decoder``: feature dispatch, RMSNorm patch, per-type RoPE, KV shapes, LM-head padding) - then compiles a dense ``Gemma4Block`` per decoder layer. The MoE - adapter (``hf_gemma4_moe``) calls the same seam with ``allow_moe=True`` and - compiles its own MoE blocks instead. + then installs registered dense blocks and compiles one parameter-explicit + executor per structural layer class. The MoE adapter (``hf_gemma4_moe``) + calls the same seam with ``allow_moe=True`` and compiles its own MoE blocks. """ backbone = _gemma4_backbone(model) + cfg = text_config(model.config) num_q_heads_per_layer, kv_shapes, is_kv_eq_v_per_layer = _setup_gemma4_text_decoder( model, allow_moe=False ) model._spyre_compiled_blocks = prepare_gemma4_blocks( backbone.layers, + cfg.layer_types, num_q_heads_per_layer, kv_shapes, is_kv_eq_v_per_layer, diff --git a/hf_adapters/hf_gemma4_mm.py b/hf_adapters/hf_gemma4_mm.py index 86092374..91109fa0 100644 --- a/hf_adapters/hf_gemma4_mm.py +++ b/hf_adapters/hf_gemma4_mm.py @@ -13,21 +13,17 @@ # limitations under the License. """ -Unified (encoder-free) HuggingFace adapter for Gemma 4 12B on Spyre — image→text. +HuggingFace adapter for Gemma 4 multimodal models on Spyre — image→text. -Supports both the base checkpoint (``google/gemma-4-12b``) and the instruction-tuned -variant (``google/gemma-4-12B-it``); both use ``model_type=gemma4_unified`` and -``Gemma4UnifiedForConditionalGeneration``. +Supports both encoder-free ``Gemma4UnifiedConfig`` checkpoints and full-vision +``Gemma4Config`` checkpoints. Full checkpoints run their transformer encoder +blocks on Spyre while keeping position lookup and spatial pooling on CPU. The +shared multimodal frontend scatters projected image features into ```` +token slots, then delegates to the existing Gemma 4 dense/PLE or MoE text +implementation selected by the nested text config. -Where ``hf_gemma4`` runs only the text decoder (``AutoSpyreModelForCausalLM``), -this module loads the full unified multimodal model -(``Gemma4UnifiedForConditionalGeneration``, ``model_type=gemma4_unified``) via -``AutoModelForImageTextToText`` and runs the image→text pipeline. It is the -adapter behind ``AutoSpyreModelForImageTextToText``. - -Gemma 4 is **encoder-free**: there is no vision tower. Vision is a pure -projection of raw (processor-merged) pixel patches into the LM embedding space, -scattered into the ```` token slots of the text embeddings: +For encoder-free checkpoints, vision is a pure projection of raw +(processor-merged) pixel patches into the LM embedding space: pixel_values [B, P, 48²·3] (processor already merged │ image_position_ids [B, P, 2] pooling_kernel_size² raw @@ -46,21 +42,12 @@ logits ──► final_logit_softcapping **Bidirectional vision attention.** ``text_config.use_bidirectional_attention == -"vision"``: within one image, the soft-tokens attend to each other -bidirectionally. Stock builds every layer-type mask via -``create_causal_mask(block_sequence_ids=...)``, which OR-s a "blockwise" overlay -(same image group ⇒ allowed, from ``mm_token_type_ids``) into the causal mask for -**both** full and sliding layers — NOT sliding-only. So at prefill we OR the -blockwise band into the base causal mask for both types: - - - full_attention = OR(causal, blockwise) - - sliding_attention = AND(sliding_window, OR(causal, blockwise)) - -Decode steps are pure text (one new causal token), so no blockwise band is -needed after prefill. (Verified against stock ``create_causal_mask``: the -``create_masks_for_vision_model`` docstring claiming globals stay causal is not -the path the forward takes — traced directly through -``create_masks_for_generate`` → ``create_causal_mask``.) +"vision"``: within one image, soft tokens attend bidirectionally. Full-vision +``Gemma4Config`` checkpoints keep full-attention layers causal and build sliding +masks as ``AND(sliding_window, OR(causal, blockwise))``. Encoder-free +``Gemma4UnifiedConfig`` checkpoints build full masks as ``OR(causal, blockwise)`` +and sliding masks as ``OR(AND(sliding_window, causal), blockwise)``. Decode steps +are pure text, so no blockwise band is needed after prefill. **Vision embedder on Spyre.** The compilable core (LN₁→Dense→LN₂→+posemb→ pos_norm→RMSNorm→Linear) is ``torch.compile``d and runs on Spyre. The @@ -81,7 +68,7 @@ import torch -from hf_adapters import hf_gemma4 +from hf_adapters import hf_gemma4, hf_gemma4_moe, hf_gemma4_vision from hf_adapters.hf_common import ( DEVICE, get_backbone, @@ -147,31 +134,44 @@ def prepare_for_spyre(model): vision-capable unified checkpoint and that audio is out of scope. """ cfg = text_config(model.config) - assert getattr(cfg, "use_bidirectional_attention", None) == "vision", ( - "hf_gemma4_mm expects a unified Gemma 4 with " - "use_bidirectional_attention='vision'; got " - f"{getattr(cfg, 'use_bidirectional_attention', None)!r}." + assert getattr(cfg, "use_bidirectional_attention", None) in (None, "vision"), ( + "hf_gemma4_mm supports causal or vision-bidirectional Gemma 4 text; got " + f"use_bidirectional_attention={getattr(cfg, 'use_bidirectional_attention', None)!r}." ) assert getattr(model.model, "embed_vision", None) is not None, ( "hf_gemma4_mm requires a vision embedder (model.model.embed_vision); " "this checkpoint has no vision_config." ) - # Shared text decoder (mirrors hf_gemma4.prepare_for_spyre). - hf_gemma4.prepare_text_decoder_for_spyre(model) - assert not model._spyre_has_ple, ( - "hf_gemma4_mm does not support PLE (E-variant) checkpoints; " - "the VLM embed path does not compute per_layer_inputs." - ) - - # Vision projection core, compiled for Spyre. The three vision LayerNorms - # (patch_ln1/patch_ln2/pos_norm) must be patched to the un-fused - # decomposition BEFORE compiling: the fused F.layer_norm lowering NaNs on - # near-constant patch rows (see patch_layernorm / the doc). Patch first, then - # compile so the core captures the patched forward. - embedder = _vision_embedder(model) - patch_layernorm(embedder.patch_ln1, embedder.patch_ln2, embedder.pos_norm) - model._spyre_vision_core = _make_compiled_vision_core(embedder) + # Reuse the text adapter selected by the nested decoder configuration. + if getattr(cfg, "enable_moe_block", False): + hf_gemma4_moe.prepare_text_decoder_for_spyre(model) + else: + hf_gemma4.prepare_text_decoder_for_spyre(model) + + if getattr(model.model, "vision_tower", None) is not None: + # Full checkpoints run the dense transformer body on Spyre. The position + # lookup, spatial pooler, and text-space projection remain on CPU. + hf_gemma4_vision.prepare_for_spyre(model) + cpu_submodules = list(getattr(model, "_spyre_cpu_submodules", [])) + cpu_submodules.extend( + [ + "model.vision_tower.patch_embedder", + "model.vision_tower.pooler", + "model.embed_vision", + ] + ) + if getattr(model.model, "audio_tower", None) is not None: + cpu_submodules.extend(["model.audio_tower", "model.embed_audio"]) + model._spyre_cpu_submodules = cpu_submodules + model._spyre_vision_core = None + else: + # Unified checkpoints use an attention-free projection core on Spyre. + # Patch the three LayerNorms before compiling so the graph captures the + # unfused fp32-reduction implementation. + embedder = _vision_embedder(model) + patch_layernorm(embedder.patch_ln1, embedder.patch_ln2, embedder.pos_norm) + model._spyre_vision_core = _make_compiled_vision_core(embedder) def _build_pos_embs(embedder, image_position_ids): @@ -192,21 +192,21 @@ def _build_pos_embs(embedder, image_position_ids): def _image_features(model, pixel_values, image_position_ids): - """Run the Spyre vision core and return stripped features [valid_patches, H]. - - CPU: build the positional embeddings and the padding mask. Spyre: the - LN/Dense/RMSNorm projection core. CPU: strip padding patches - (``image_position_ids == -1`` on both axes), matching stock - ``get_image_features``. - """ - embedder = _vision_embedder(model) - dtype = get_model_dtype(model) - - # anyres / multi-image: [B, T, P, ...] -> [B*T, P, ...] (stock flattens too) + """Run the checkpoint's vision path and return flattened text-space features.""" + # anyres / multi-image: [B, T, P, ...] -> [B*T, P, ...] if pixel_values.dim() == 4: pixel_values = pixel_values.flatten(0, 1) image_position_ids = image_position_ids.flatten(0, 1) + if getattr(model.model, "vision_tower", None) is not None: + features = hf_gemma4_vision.prefill_vision_tower( + model, pixel_values, image_position_ids + ) + return model.model.embed_vision(features.to("cpu")) + + embedder = _vision_embedder(model) + dtype = get_model_dtype(model) + pos_embs = _build_pos_embs(embedder, image_position_ids).to(dtype) features = model._spyre_vision_core( pixel_values.to(dtype).to(DEVICE), pos_embs.to(DEVICE) @@ -219,24 +219,29 @@ def _image_features(model, pixel_values, image_position_ids): def _embed_and_scatter(model, input_ids, image_features): - """Scaled word embeddings with image features scattered into slots. + """Build decoder and PLE-context embeddings for multimodal prefill. ``embed_tokens`` is ``Gemma4UnifiedTextScaledWordEmbedding`` (×√hidden runs as-is). Stock does ``inputs_embeds.masked_scatter(image_mask, features)``; Spyre can't ``masked_scatter``, so we zero the image-token slots (elementwise mul by a CPU-built keep factor) and add a CPU-built additive tensor holding the features at the image positions — bit-identical given the zeroed slots - (same doctrine as hf_granite_vision_mm._inject_deepstack). Asserts the - token/feature counts match (mirrors stock's shape check). + (same doctrine as hf_granite_vision_mm._inject_deepstack). The separately + returned PLE context keeps scaled text embeddings but replaces image slots + with the raw, unscaled pad embedding, matching stock Gemma 4. The returned + token ids apply the same image-to-pad substitution for PLE token identity. + Asserts the token/feature counts match (mirrors stock's shape check). """ backbone = get_backbone(model) image_token_id = model.config.image_token_id dtype = get_model_dtype(model) - ids = input_ids.to(backbone.embed_tokens.weight.device) + input_ids_cpu = input_ids.to("cpu").clone() + image_mask = input_ids_cpu == image_token_id # [B, L] bool + input_ids_cpu[image_mask] = text_config(model.config).pad_token_id + ids = input_ids_cpu.to(backbone.embed_tokens.weight.device) h = backbone.embed_tokens(ids) # scaled word embeddings, on embed device - image_mask = input_ids == image_token_id # [B, L] bool, CPU n_image_tokens = int(image_mask.sum()) hidden = h.shape[-1] feats = image_features.to("cpu", dtype) @@ -247,11 +252,18 @@ def _embed_and_scatter(model, input_ids, image_features): ) keep = (~image_mask).to(dtype).unsqueeze(-1).to(h.device) - h = h * keep + text_embeds = h * keep - additive = torch.zeros(h.shape[0], h.shape[1], hidden, dtype=dtype) - additive[image_mask] = feats.view(n_image_tokens, hidden) - return h + additive.to(h.device) + image_additive = torch.zeros(h.shape[0], h.shape[1], hidden, dtype=dtype) + image_additive[image_mask] = feats.view(n_image_tokens, hidden) + inputs_embeds = text_embeds + image_additive.to(h.device) + + ple_context_embeds = None + if model._spyre_has_ple: + raw_pad = backbone.embed_tokens.weight[text_config(model.config).pad_token_id] + pad_additive = image_mask.to(dtype).unsqueeze(-1).to(h.device) * raw_pad + ple_context_embeds = text_embeds + pad_additive + return inputs_embeds, ple_context_embeds, input_ids_cpu def _blockwise_band(mm_token_type_ids, padded_len, max_cache_len, dtype): @@ -298,22 +310,12 @@ def _blockwise_band(mm_token_type_ids, padded_len, max_cache_len, dtype): def _sliding_window_lower_band(mask, sliding_window): - """Add stock's sliding-window *lower bound* to an additive causal prefill mask. + """Apply stock's sliding-window lower bound to an additive prefill mask. Masks keys further back than ``sliding_window`` (``q - k >= window``) but — - unlike ``hf_common.add_causal_sliding_window_band`` — does NOT mask future - keys (``q - k < 0``). This is stock's ``sliding_window_overlay`` - (``kv_idx > q_idx - window``), an ``and_mask`` over the *causal* base: the - causal upper bound already lives in ``mask``, so the window only supplies the - backward cutoff. - - NOTE: this must be applied to the causal base *before* the blockwise vision - band is OR-ed on top — stock composes the sliding mask as - ``OR(AND(sliding_window_overlay, causal), blockwise)`` with the blockwise - overlay as the OUTERMOST op (``masking_utils.create_sliding_window_causal_mask``, - line 1189 then 1222), so the window must NOT gate the bidirectional image - pairs. Applying it after the OR would clip intra-image-block attention to the - window (wrong for a single image/video block longer than ``sliding_window``). + unlike ``hf_common.add_causal_sliding_window_band`` — does not independently + mask future keys. For multimodal prefill, call this after OR-ing the causal and + blockwise masks so the window gates bidirectional image edges too. Prefill only (``cache_index`` starts at cache slot 0), so a query row's cache coordinate is its row index ``q`` and the key column is the cache slot ``k``. @@ -332,46 +334,25 @@ def _sliding_window_lower_band(mask, sliding_window): return (mask.to("cpu") + band[None, None, :, :]).to(orig_device) -def _build_mm_masks(prefill_mask, blockwise_band, sliding_window): - """Per-layer-type masks for a multimodal prefill: {full, sliding}. - - Stock builds each mask type from the same blockwise vision overlay, but ORs - it in as the OUTERMOST op for **both** full and sliding layers (traced through - ``create_masks_for_generate`` → ``create_causal_mask`` / - ``create_sliding_window_causal_mask``, ``masking_utils.py`` L997 / L1222): - - - full_attention = OR(causal, blockwise) - - sliding_attention = OR(AND(sliding_window_lowerbound, causal), blockwise) - - Crucially the sliding window is AND-ed onto the *causal base only* and the - blockwise band is OR-ed on *after* — the window never gates the bidirectional - image pairs. (An earlier version AND-ed the window over ``OR(causal, - blockwise)``, which wrongly clipped intra-image-block attention to the window - for image/video blocks longer than ``sliding_window``.) - - ``prefill_mask`` is the additive causal base (``build_prefill_mask``: causal + - left-pad + unused-cache masking). ``blockwise_band`` is the additive 0/-inf - image band; its allowed cells are only within a same image group, and padded - columns are group ``-1`` so the band never re-admits a padded key. The OR is an - elementwise ``max`` of the two additive masks, the AND (window) an elementwise - add; both done on CPU to avoid the bf16 ``-inf + -inf`` NaN hazard. Prefill - only. ``prefill_mask``/``blockwise_band`` are ``[B, 1, padded_len, - max_cache_len]``. - """ +def _build_mm_masks(prefill_mask, blockwise_band, sliding_window, *, unified): + """Build the architecture-specific Gemma 4 multimodal attention masks.""" orig_device = prefill_mask.device prefill_cpu = prefill_mask.to("cpu") blockwise_cpu = blockwise_band.to("cpu") - # full_attention = OR(causal, blockwise) - full_mask = torch.maximum(prefill_cpu, blockwise_cpu).to(orig_device) + if unified: + full_mask = torch.maximum(prefill_cpu, blockwise_cpu) + sliding_causal = _sliding_window_lower_band(prefill_cpu, sliding_window) + sliding_mask = torch.maximum(sliding_causal.to("cpu"), blockwise_cpu) + else: + full_mask = prefill_cpu + blockwise_causal = torch.maximum(prefill_cpu, blockwise_cpu) + sliding_mask = _sliding_window_lower_band(blockwise_causal, sliding_window) - # sliding_attention = OR(AND(window, causal), blockwise): window gates only the - # causal base, then the image band is OR-ed back on top (ungated by the window). - windowed_causal = _sliding_window_lower_band(prefill_cpu, sliding_window) - sliding_mask = torch.maximum(windowed_causal.to("cpu"), blockwise_cpu).to( - orig_device - ) - return {"full_attention": full_mask, "sliding_attention": sliding_mask} + return { + "full_attention": full_mask.to(orig_device), + "sliding_attention": sliding_mask.to(orig_device), + } def _logits_from_embeds( @@ -383,6 +364,8 @@ def _logits_from_embeds( value_caches, cache_index, masks=None, + input_ids=None, + ple_context_embeds=None, ): """Text decoder over image-scattered embeds → logits (+ softcap). @@ -397,7 +380,29 @@ def _logits_from_embeds( ``[0, padded_len)`` at prefill, a single slot per decode step. Every computed position is written, so the shared walk's sliding-window ``block_base`` is simply ``cache_index[0]``: the cache column the (only) query row occupies. + ``ple_context_embeds`` optionally carries stock's pre-scatter multimodal PLE + context: scaled text embeddings with image positions replaced by the raw, + unscaled pad embedding. Decode omits it and uses ``inputs_embeds`` directly. """ + cfg = text_config(model.config) + query_row_mask = None + if inputs_embeds.shape[1] > 1 and not getattr(cfg, "enable_moe_block", False): + query_row_mask = hf_gemma4._query_row_mask(inputs_embeds, attn_mask) + inputs_embeds = inputs_embeds * query_row_mask + + if model._spyre_has_ple: + if input_ids is None: + raise ValueError( + "Gemma 4 PLE decoding requires input_ids alongside inputs_embeds." + ) + + per_layer_inputs = hf_gemma4._compute_per_layer_inputs( + model, + inputs_embeds if ple_context_embeds is None else ple_context_embeds, + input_ids, + ) + else: + per_layer_inputs = None h = hf_gemma4._run_blocks_over_embeds( model, inputs_embeds, @@ -407,6 +412,8 @@ def _logits_from_embeds( value_caches, cache_index, masks=masks, + per_layer_inputs=per_layer_inputs, + query_row_mask=query_row_mask, ) logits = model.lm_head(h) cap = text_config(model.config).final_logit_softcapping @@ -433,21 +440,30 @@ def _prefill_forward( """Shared multimodal prefill: padded ids + image → full-sequence logits. Builds scaled text embeddings with the image features scattered into the - ```` slots, then the per-layer-type masks with the bidirectional - vision overlay OR-ed into both full and sliding layers, and runs the decoder - once (writing the KV caches). ``mm_token_type_ids`` has already undergone + ```` slots, then the architecture-specific blockwise vision masks, + and runs the decoder once (writing the KV caches). + ``mm_token_type_ids`` has already undergone the same prompt compaction and block padding as ``input_ids``. """ dtype = get_model_dtype(model) cfg = text_config(model.config) image_features = _image_features(model, pixel_values, image_position_ids) - inputs_embeds = _embed_and_scatter(model, input_ids, image_features) + inputs_embeds, ple_context_embeds, ple_input_ids = _embed_and_scatter( + model, input_ids, image_features + ) - padded_len = input_ids.shape[1] - max_cache_len = attention_mask.shape[-1] - blockwise = _blockwise_band(mm_token_type_ids, padded_len, max_cache_len, dtype) - masks = _build_mm_masks(attention_mask, blockwise, cfg.sliding_window) - masks = {lt: m.to(DEVICE) for lt, m in masks.items()} + masks = None + if getattr(cfg, "use_bidirectional_attention", None) == "vision": + padded_len = input_ids.shape[1] + max_cache_len = attention_mask.shape[-1] + blockwise = _blockwise_band(mm_token_type_ids, padded_len, max_cache_len, dtype) + masks = _build_mm_masks( + attention_mask, + blockwise, + cfg.sliding_window, + unified=getattr(model.model, "vision_tower", None) is None, + ) + masks = {lt: m.to(DEVICE) for lt, m in masks.items()} return _logits_from_embeds( model, inputs_embeds.to(DEVICE), @@ -457,4 +473,8 @@ def _prefill_forward( value_caches, cache_index=cache_index, masks=masks, + input_ids=ple_input_ids.to(DEVICE), + ple_context_embeds=( + ple_context_embeds.to(DEVICE) if ple_context_embeds is not None else None + ), ) diff --git a/hf_adapters/hf_gemma4_moe.py b/hf_adapters/hf_gemma4_moe.py index 9522da4a..0c16ec49 100644 --- a/hf_adapters/hf_gemma4_moe.py +++ b/hf_adapters/hf_gemma4_moe.py @@ -33,7 +33,12 @@ _setup_gemma4_text_decoder, ) -__all__ = ["prepare_for_spyre", "_run_forward", "_run_backbone_forward"] +__all__ = [ + "prepare_for_spyre", + "prepare_text_decoder_for_spyre", + "_run_forward", + "_run_backbone_forward", +] _MOE_TILE = 32 # Decode gather requires tiles with at least two rows. @@ -444,8 +449,8 @@ def _prepare_experts(experts): experts.down_proj = _move_expert_weight(down) -def prepare_for_spyre(model): - """Prepare a Gemma 4 MoE causal LM for Spyre in place.""" +def prepare_text_decoder_for_spyre(model): + """Prepare only the Gemma 4 MoE text decoder for Spyre in place.""" from torch_spyre._C import get_elem_in_stick from torch_spyre.model_utils import dma_moe_per_expert_scale_to_spyre @@ -485,3 +490,8 @@ def prepare_for_spyre(model): blocks.append(block) model._spyre_compiled_blocks = blocks + + +def prepare_for_spyre(model): + """Prepare a Gemma 4 MoE causal LM for Spyre in place.""" + prepare_text_decoder_for_spyre(model) diff --git a/hf_adapters/hf_gemma4_vision.py b/hf_adapters/hf_gemma4_vision.py new file mode 100644 index 00000000..aacd4211 --- /dev/null +++ b/hf_adapters/hf_gemma4_vision.py @@ -0,0 +1,581 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Spyre execution for the transformer body of a full Gemma 4 vision tower.""" + +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from hf_adapters.hf_common import ( + BLOCK_SIZE, + DEVICE, + SpyreUnsupportedFeatureError, + SpyreUnsupportedModelError, + _pad_proj_input_simple, + _pad_proj_output_simple, + apply_rope_matmul, +) +from hf_adapters.hf_gemma4 import _gemma4_mlp_activation, _gemma4_rms_norm + + +def _pad_qk_linear(proj, num_heads, orig_head_dim, padded_head_dim): + """Pad and reorder two-axis RoPE channels into one matrix-RoPE layout.""" + linear = proj.linear + weight = linear.weight.detach().view(num_heads, orig_head_dim, -1) + new_weight = torch.zeros( + num_heads, padded_head_dim, weight.shape[-1], dtype=weight.dtype + ) + quarter = orig_head_dim // 4 + padded_half = padded_head_dim // 2 + new_weight[:, :quarter] = weight[:, :quarter] + new_weight[:, quarter : 2 * quarter] = weight[:, 2 * quarter : 3 * quarter] + new_weight[:, padded_half : padded_half + quarter] = weight[ + :, quarter : 2 * quarter + ] + new_weight[:, padded_half + quarter : padded_half + 2 * quarter] = weight[ + :, 3 * quarter : + ] + padded = nn.Linear( + linear.in_features, + num_heads * padded_head_dim, + bias=linear.bias is not None, + ) + padded.weight = nn.Parameter( + new_weight.reshape(num_heads * padded_head_dim, -1), requires_grad=False + ) + if linear.bias is not None: + bias = linear.bias.detach().view(num_heads, orig_head_dim) + new_bias = torch.zeros(num_heads, padded_head_dim, dtype=bias.dtype) + new_bias[:, :quarter] = bias[:, :quarter] + new_bias[:, quarter : 2 * quarter] = bias[:, 2 * quarter : 3 * quarter] + new_bias[:, padded_half : padded_half + quarter] = bias[ + :, quarter : 2 * quarter + ] + new_bias[:, padded_half + quarter : padded_half + 2 * quarter] = bias[ + :, 3 * quarter : + ] + padded.bias = nn.Parameter(new_bias.reshape(-1), requires_grad=False) + return padded + + +def _pad_norm_weight(norm, orig_head_dim, padded_head_dim): + weight = norm.weight.detach() + padded = torch.ones(padded_head_dim, dtype=weight.dtype) + quarter = orig_head_dim // 4 + padded_half = padded_head_dim // 2 + padded[:quarter] = weight[:quarter] + padded[quarter : 2 * quarter] = weight[2 * quarter : 3 * quarter] + padded[padded_half : padded_half + quarter] = weight[quarter : 2 * quarter] + padded[padded_half + quarter : padded_half + 2 * quarter] = weight[3 * quarter :] + return nn.Parameter(padded, requires_grad=False) + + +def _padded_rms_norm(hidden_states, weight, eps, orig_head_dim): + dtype = hidden_states.dtype + hidden_states = hidden_states.float() + variance = (hidden_states * hidden_states).mean(-1, keepdim=True) + variance = variance * (hidden_states.shape[-1] / orig_head_dim) + hidden_states = hidden_states * torch.rsqrt(variance + eps) + if weight is not None: + hidden_states = hidden_states * weight.float() + return hidden_states.to(dtype) + + +def _clamp(hidden_states, bounds): + if bounds is None: + return hidden_states + return torch.maximum(torch.minimum(hidden_states, bounds[1]), bounds[0]) + + +def _pad_mlp(layers, orig_intermediate, padded_intermediate): + for layer in layers: + mlp = layer.mlp + mlp.gate_proj.linear = _pad_proj_output_simple( + mlp.gate_proj.linear, 1, orig_intermediate, padded_intermediate + ) + mlp.up_proj.linear = _pad_proj_output_simple( + mlp.up_proj.linear, 1, orig_intermediate, padded_intermediate + ) + mlp.down_proj.linear = _pad_proj_input_simple( + mlp.down_proj.linear, 1, orig_intermediate, padded_intermediate + ) + + +@dataclass(frozen=True) +class _Gemma4VisionBlockSpec: + hidden_size: int + num_heads: int + orig_head_dim: int + padded_head_dim: int + intermediate_size: int + activation: str + scaling: float + input_norm_eps: float + q_norm_eps: float + k_norm_eps: float + v_norm_eps: float + post_attention_norm_eps: float + pre_feedforward_norm_eps: float + post_feedforward_norm_eps: float + projection_biases: tuple[bool, ...] + clipped_projections: tuple[bool, ...] + + +def _assert_output_clamp_preserves_zero(module, name): + if not module.use_clipped_linears: + return + preserves_zero = torch.all((module.output_min <= 0) & (module.output_max >= 0)) + if not bool(preserves_zero.item()): + raise ValueError( + f"{name} output clipping must include zero when attention heads are " + "padded; otherwise padded channels become nonzero before RMSNorm. " + f"Got output_min={module.output_min}, output_max={module.output_max}." + ) + + +def _prepare_clip_bounds(module): + if not module.use_clipped_linears: + return + for which in ("input", "output"): + for limit in ("min", "max"): + value = getattr(module, f"{which}_{limit}") + module.register_buffer( + f"_spyre_{which}_{limit}", + value.detach().clone(), + persistent=False, + ) + + +def _projection_state(module): + linear = module.linear + state = [linear.weight] + if linear.bias is not None: + state.append(linear.bias) + if module.use_clipped_linears: + state.extend( + [ + module._spyre_input_min, + module._spyre_input_max, + module._spyre_output_min, + module._spyre_output_max, + ] + ) + return tuple(state) + + +def _vision_block_state(layer): + attn = layer.self_attn + state = [] + for module in (attn.q_proj, attn.k_proj, attn.v_proj, attn.o_proj): + state.extend(_projection_state(module)) + state.extend([attn.q_norm.weight, attn.k_norm.weight]) + for module in (layer.mlp.gate_proj, layer.mlp.up_proj, layer.mlp.down_proj): + state.extend(_projection_state(module)) + state.extend( + [ + layer.input_layernorm.weight, + layer.post_attention_layernorm.weight, + layer.pre_feedforward_layernorm.weight, + layer.post_feedforward_layernorm.weight, + ] + ) + return tuple(state) + + +def _vision_block_spec(layer, num_heads, orig_head_dim, padded_head_dim): + attn = layer.self_attn + projections = ( + attn.q_proj, + attn.k_proj, + attn.v_proj, + attn.o_proj, + layer.mlp.gate_proj, + layer.mlp.up_proj, + layer.mlp.down_proj, + ) + activation = layer.mlp.config.hidden_activation + if activation != "gelu_pytorch_tanh": + raise SpyreUnsupportedModelError( + "Gemma 4 vision checkpoints must use " + "hidden_activation='gelu_pytorch_tanh'; " + f"got {activation!r}" + ) + if attn.v_norm.with_scale: + raise SpyreUnsupportedFeatureError( + "Scaled Gemma 4 vision V normalization is not supported on Spyre." + ) + return _Gemma4VisionBlockSpec( + hidden_size=layer.input_layernorm.weight.numel(), + num_heads=num_heads, + orig_head_dim=orig_head_dim, + padded_head_dim=padded_head_dim, + intermediate_size=layer.mlp.gate_proj.linear.out_features, + activation=activation, + scaling=float(attn.scaling), + input_norm_eps=layer.input_layernorm.eps, + q_norm_eps=attn.q_norm.eps, + k_norm_eps=attn.k_norm.eps, + v_norm_eps=attn.v_norm.eps, + post_attention_norm_eps=layer.post_attention_layernorm.eps, + pre_feedforward_norm_eps=layer.pre_feedforward_layernorm.eps, + post_feedforward_norm_eps=layer.post_feedforward_layernorm.eps, + projection_biases=tuple( + module.linear.bias is not None for module in projections + ), + clipped_projections=tuple(module.use_clipped_linears for module in projections), + ) + + +def _take_projection(state, pos, has_bias, is_clipped): + weight = state[pos] + pos += 1 + if has_bias: + bias = state[pos] + pos += 1 + else: + bias = None + if is_clipped: + input_bounds = (state[pos], state[pos + 1]) + output_bounds = (state[pos + 2], state[pos + 3]) + pos += 4 + else: + input_bounds = output_bounds = None + return weight, bias, input_bounds, output_bounds, pos + + +def _make_vision_forward(spec): + def block_forward(state, hidden_states, rope_matrices, attn_mask): + pos = 0 + ( + q_weight, + q_bias, + q_input_bounds, + q_output_bounds, + pos, + ) = _take_projection( + state, pos, spec.projection_biases[0], spec.clipped_projections[0] + ) + ( + k_weight, + k_bias, + k_input_bounds, + k_output_bounds, + pos, + ) = _take_projection( + state, pos, spec.projection_biases[1], spec.clipped_projections[1] + ) + ( + v_weight, + v_bias, + v_input_bounds, + v_output_bounds, + pos, + ) = _take_projection( + state, pos, spec.projection_biases[2], spec.clipped_projections[2] + ) + ( + o_weight, + o_bias, + o_input_bounds, + o_output_bounds, + pos, + ) = _take_projection( + state, pos, spec.projection_biases[3], spec.clipped_projections[3] + ) + q_norm_weight, k_norm_weight = state[pos : pos + 2] + pos += 2 + ( + gate_weight, + gate_bias, + gate_input_bounds, + gate_output_bounds, + pos, + ) = _take_projection( + state, pos, spec.projection_biases[4], spec.clipped_projections[4] + ) + ( + up_weight, + up_bias, + up_input_bounds, + up_output_bounds, + pos, + ) = _take_projection( + state, pos, spec.projection_biases[5], spec.clipped_projections[5] + ) + ( + down_weight, + down_bias, + down_input_bounds, + down_output_bounds, + pos, + ) = _take_projection( + state, pos, spec.projection_biases[6], spec.clipped_projections[6] + ) + input_norm_weight, post_attn_norm_weight = state[pos : pos + 2] + pos += 2 + pre_ffn_norm_weight, post_ffn_norm_weight = state[pos : pos + 2] + + bsz, seq_len, _ = hidden_states.shape + residual = hidden_states + hidden_states = _gemma4_rms_norm( + hidden_states, input_norm_weight, spec.input_norm_eps + ) + + query = F.linear(_clamp(hidden_states, q_input_bounds), q_weight, q_bias) + query = _clamp(query, q_output_bounds).view( + bsz, seq_len, spec.num_heads, spec.padded_head_dim + ) + query = _padded_rms_norm( + query, q_norm_weight, spec.q_norm_eps, spec.orig_head_dim + ).transpose(1, 2) + + key = F.linear(_clamp(hidden_states, k_input_bounds), k_weight, k_bias) + key = _clamp(key, k_output_bounds).view( + bsz, seq_len, spec.num_heads, spec.padded_head_dim + ) + key = _padded_rms_norm( + key, k_norm_weight, spec.k_norm_eps, spec.orig_head_dim + ).transpose(1, 2) + + value = F.linear(_clamp(hidden_states, v_input_bounds), v_weight, v_bias) + value = _clamp(value, v_output_bounds).view( + bsz, seq_len, spec.num_heads, spec.padded_head_dim + ) + value = _padded_rms_norm( + value, None, spec.v_norm_eps, spec.orig_head_dim + ).transpose(1, 2) + + query = apply_rope_matmul(query, rope_matrices).contiguous() + key = apply_rope_matmul(key, rope_matrices).contiguous() + attn_output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=False, + scale=spec.scaling, + ) + attn_output = _clamp(attn_output, o_input_bounds) + attn_output = attn_output.transpose(1, 2).reshape(bsz, seq_len, -1) + attn_output = F.linear(attn_output, o_weight, o_bias) + attn_output = _clamp(attn_output, o_output_bounds) + hidden_states = residual + _gemma4_rms_norm( + attn_output, post_attn_norm_weight, spec.post_attention_norm_eps + ) + + residual = hidden_states + hidden_states = _gemma4_rms_norm( + hidden_states, pre_ffn_norm_weight, spec.pre_feedforward_norm_eps + ) + gate = F.linear( + _clamp(hidden_states, gate_input_bounds), gate_weight, gate_bias + ) + gate = _clamp(gate, gate_output_bounds) + up = F.linear(_clamp(hidden_states, up_input_bounds), up_weight, up_bias) + up = _clamp(up, up_output_bounds) + hidden_states = _gemma4_mlp_activation(gate) * up + hidden_states = F.linear( + _clamp(hidden_states, down_input_bounds), down_weight, down_bias + ) + hidden_states = _clamp(hidden_states, down_output_bounds) + hidden_states = _gemma4_rms_norm( + hidden_states, post_ffn_norm_weight, spec.post_feedforward_norm_eps + ) + return residual + hidden_states + + return block_forward + + +def _prepare_vision_blocks(layers, num_heads, orig_head_dim, padded_head_dim): + compiled_by_spec = {} + state_signature_by_spec = {} + compiled_blocks = [] + for i, layer in enumerate(layers): + attn = layer.self_attn + if padded_head_dim != orig_head_dim: + for name, module in ( + ("q_proj", attn.q_proj), + ("k_proj", attn.k_proj), + ("v_proj", attn.v_proj), + ): + _assert_output_clamp_preserves_zero(module, f"layer {i} {name}") + + attn.q_proj.linear = _pad_qk_linear( + attn.q_proj, num_heads, orig_head_dim, padded_head_dim + ) + attn.k_proj.linear = _pad_qk_linear( + attn.k_proj, num_heads, orig_head_dim, padded_head_dim + ) + attn.v_proj.linear = _pad_proj_output_simple( + attn.v_proj.linear, num_heads, orig_head_dim, padded_head_dim + ) + attn.o_proj.linear = _pad_proj_input_simple( + attn.o_proj.linear, num_heads, orig_head_dim, padded_head_dim + ) + attn.q_norm.weight = _pad_norm_weight( + attn.q_norm, orig_head_dim, padded_head_dim + ) + attn.k_norm.weight = _pad_norm_weight( + attn.k_norm, orig_head_dim, padded_head_dim + ) + for module in ( + attn.q_proj, + attn.k_proj, + attn.v_proj, + attn.o_proj, + layer.mlp.gate_proj, + layer.mlp.up_proj, + layer.mlp.down_proj, + ): + _prepare_clip_bounds(module) + + spec = _vision_block_spec(layer, num_heads, orig_head_dim, padded_head_dim) + state = _vision_block_state(layer) + signature = tuple((tuple(tensor.shape), tensor.dtype) for tensor in state) + if ( + spec in state_signature_by_spec + and state_signature_by_spec[spec] != signature + ): + raise ValueError( + f"Gemma 4 vision layer {i} state does not match its compile group" + ) + state_signature_by_spec.setdefault(spec, signature) + if spec not in compiled_by_spec: + compiled_by_spec[spec] = torch.compile( + _make_vision_forward(spec), dynamic=False, fullgraph=True + ) + compiled_blocks.append(compiled_by_spec[spec]) + return compiled_blocks + + +def _build_rope_matrices(inv_freq, position_ids, padded_head_dim, dtype): + positions = position_ids.to("cpu").clamp(min=0).float() + angles = positions[..., None] * inv_freq.float() + cos = angles.cos() + sin = angles.sin() + bsz, seq_len, _, quarter = cos.shape + padded_half = padded_head_dim // 2 + matrices = torch.zeros(bsz, seq_len, 2, 2, padded_half) + matrices[:, :, 0, 0, :] = 1.0 + matrices[:, :, 1, 1, :] = 1.0 + for axis in range(2): + start = axis * quarter + end = start + quarter + matrices[:, :, 0, 0, start:end] = cos[:, :, axis] + matrices[:, :, 0, 1, start:end] = -sin[:, :, axis] + matrices[:, :, 1, 0, start:end] = sin[:, :, axis] + matrices[:, :, 1, 1, start:end] = cos[:, :, axis] + return matrices.to(dtype) + + +def _build_attention_mask(valid, padded_len, dtype): + bsz, seq_len = valid.shape + key_mask = F.pad(valid, (0, padded_len - seq_len), value=False) + mask = torch.zeros((bsz, 1, 1, padded_len), dtype=dtype) + return mask.masked_fill(~key_mask[:, None, None, :], -torch.inf) + + +def prepare_for_spyre(model): + tower = model.model.vision_tower + config = tower.config + layers = tower.encoder.layers + if config.num_key_value_heads != config.num_attention_heads: + raise SpyreUnsupportedFeatureError( + "Gemma 4 vision GQA is not supported on Spyre; num_key_value_heads " + "must equal num_attention_heads." + ) + if config.rope_parameters.get("rope_type", "default") != "default": + raise SpyreUnsupportedFeatureError( + "Gemma 4 vision supports only default, unscaled RoPE on Spyre." + ) + orig_head_dim = config.head_dim + padded_head_dim = math.ceil(orig_head_dim / (2 * BLOCK_SIZE)) * (2 * BLOCK_SIZE) + + orig_intermediate = config.intermediate_size + padded_intermediate = math.ceil(orig_intermediate / BLOCK_SIZE) * BLOCK_SIZE + if padded_intermediate > orig_intermediate: + _pad_mlp(layers, orig_intermediate, padded_intermediate) + + model._spyre_gemma4_vision_inv_freq = ( + tower.encoder.rotary_emb.inv_freq.detach().cpu() + ) + if tower.config.standardize: + model._spyre_gemma4_vision_std_bias = tower.std_bias.detach().cpu() + model._spyre_gemma4_vision_std_scale = tower.std_scale.detach().cpu() + model._spyre_gemma4_vision_head_dim = padded_head_dim + model._spyre_gemma4_vision_blocks = _prepare_vision_blocks( + layers, config.num_attention_heads, orig_head_dim, padded_head_dim + ) + + +def prefill_vision_tower(model, pixel_values, position_ids): + tower = model.model.vision_tower + pixel_values = pixel_values.to("cpu") + position_ids = position_ids.to("cpu") + padding_positions = (position_ids == -1).all(dim=-1) + hidden_states = tower.patch_embedder(pixel_values, position_ids, padding_positions) + output_dtype = hidden_states.dtype + + seq_len = hidden_states.shape[1] + padded_len = math.ceil(seq_len / BLOCK_SIZE) * BLOCK_SIZE + rope_matrices = _build_rope_matrices( + model._spyre_gemma4_vision_inv_freq, + position_ids, + model._spyre_gemma4_vision_head_dim, + hidden_states.dtype, + ) + if padded_len > seq_len: + hidden_states = F.pad(hidden_states, (0, 0, 0, padded_len - seq_len)) + identity = torch.zeros( + hidden_states.shape[0], + padded_len - seq_len, + 2, + 2, + model._spyre_gemma4_vision_head_dim // 2, + dtype=hidden_states.dtype, + ) + identity[:, :, 0, 0, :] = 1.0 + identity[:, :, 1, 1, :] = 1.0 + rope_matrices = torch.cat([rope_matrices, identity], dim=1) + + attn_mask = _build_attention_mask( + ~padding_positions, padded_len, hidden_states.dtype + ) + hidden_states = hidden_states.to(DEVICE) + rope_matrices = rope_matrices.to(DEVICE) + attn_mask = attn_mask.to(DEVICE) + for layer, block in zip(tower.encoder.layers, model._spyre_gemma4_vision_blocks): + hidden_states = block( + _vision_block_state(layer), hidden_states, rope_matrices, attn_mask + ).clone() + + hidden_states = hidden_states[:, :seq_len].to("cpu") + output_length = pixel_values.shape[-2] // (tower.config.pooling_kernel_size**2) + hidden_states, pooler_mask = tower.pooler( + hidden_states, + position_ids, + padding_positions, + output_length, + ) + hidden_states = hidden_states[pooler_mask] + if tower.config.standardize: + hidden_states = ( + hidden_states - model._spyre_gemma4_vision_std_bias.float() + ) * model._spyre_gemma4_vision_std_scale.float() + return hidden_states.to(output_dtype) diff --git a/hf_adapters/hf_granite_vision_mm.py b/hf_adapters/hf_granite_vision_mm.py index 82fcb946..0a39f28f 100644 --- a/hf_adapters/hf_granite_vision_mm.py +++ b/hf_adapters/hf_granite_vision_mm.py @@ -316,6 +316,7 @@ def _logits_from_embeds( cache_index, deepstack=None, vision_mask=None, + input_ids=None, ): """Run text backbone over embeds + LM head / logits scaling -> logits. diff --git a/hf_adapters/hf_mistral3_vision_mm.py b/hf_adapters/hf_mistral3_vision_mm.py index eef9e3ea..3ab9139c 100644 --- a/hf_adapters/hf_mistral3_vision_mm.py +++ b/hf_adapters/hf_mistral3_vision_mm.py @@ -256,6 +256,7 @@ def _logits_from_embeds( cache_index, image_features=None, vision_mask=None, + input_ids=None, ): """Run text backbone over embeds + LM head → logits.""" h = _run_text_backbone( diff --git a/tests/model_registry.py b/tests/model_registry.py index e7dd0ec7..e3290a47 100644 --- a/tests/model_registry.py +++ b/tests/model_registry.py @@ -605,17 +605,37 @@ def _include_gated() -> bool: "dtype": "bfloat16", # blocked-FP8 checkpoint, dequantized to bf16 "size": "3b", }, + # hf_gemma4_vision.py — Vision tower of encoder-based Gemma4 models + "gemma4_vision_tower": { + "name": "Gemma 4 26B-A4B (Vision tower)", + "path": "google/gemma-4-26B-A4B-it", + "adapter": "hf_gemma4_vision.py", + "kind": "tower", + }, # hf_gemma4_mm.py — unified encoder-free VLM (image + text -> text) - # Note: google/gemma-4-12b (base, no chat template) also resolves via this - # adapter (Gemma4UnifiedConfig -> hf_gemma4_mm) but is tested via the - # causal-LM path (gemma4_base in CAUSAL_LM_MODELS); the VLM harness requires - # apply_chat_template, which the base model does not provide. "gemma4_mm": { "name": "Gemma 4 12B IT (unified VLM)", "path": "google/gemma-4-12B-it", "adapter": "hf_gemma4_mm.py", "kind": "vlm", # multimodal: image + text -> generated text "size": "12b", + "always_test": True, + }, + "gemma4_e2b_mm": { + "name": "Gemma 4 E2B (VLM with PLE)", + "path": "google/gemma-4-E2B-it", + "adapter": "hf_gemma4_mm.py", + "kind": "vlm", + "size": "2b", + "always_test": True, + }, + "gemma4_moe_mm": { + "name": "Gemma 4 26B-A4B (MoE VLM)", + "path": "google/gemma-4-26B-A4B-it", + "adapter": "hf_gemma4_mm.py", + "kind": "vlm", + "size": "26b", + "always_test": True, }, # hf_clip.py — CLIP dual-encoder (image + text -> embeddings via ST backend) "clip_vit_b_32": { diff --git a/tests/spyre/test_vlm_e2e_spyre.py b/tests/spyre/test_vlm_e2e_spyre.py index 69c5f484..c903632c 100644 --- a/tests/spyre/test_vlm_e2e_spyre.py +++ b/tests/spyre/test_vlm_e2e_spyre.py @@ -204,6 +204,7 @@ def _write_token(tok_id): key_caches, value_caches, cache_index=make_cache_index(current_cache_len, 1, DEVICE), + input_ids=next_input, ) per_step_logits.append(logits.to("cpu")[0, -1, :].float()) current_cache_len += 1