fix(gemma4): cache rotary embeddings by layer type

QEfficient/transformers/models/gemma4/modeling_gemma4.py

@@ -23,11 +23,13 @@
     Gemma4TextAttention,
     Gemma4TextDecoderLayer,
     Gemma4TextModel,
+    Gemma4TextRotaryEmbedding,
     Gemma4TextRouter,
     Gemma4VisionAttention,
     apply_rotary_pos_emb,
     eager_attention_forward,
     repeat_kv,
+    rotate_half,
 )
 
 from QEfficient.customop.rms_norm import CustomRMSNormFunc
@@ -418,6 +420,36 @@ def forward(
         return attn_output, attn_weights
 
 
+class QEffGemma4TextRotaryEmbedding(Gemma4TextRotaryEmbedding):
+    """Gemma4 rotary embeddings with static caches for each text attention type."""
+
+    def __init__(self, config, device=None):
+        super().__init__(config=config, device=device)
+
+        for layer_type in sorted(self.layer_types):
+            self._set_cos_sin_cache(
+                layer_type=layer_type,
+                seq_len=self.original_max_seq_len,
+                device=getattr(self, f"{layer_type}_inv_freq").device,
+                dtype=config.dtype,
+            )
+
+    def _set_cos_sin_cache(self, layer_type, seq_len, device, dtype):
+        inv_freq = getattr(self, f"{layer_type}_inv_freq")
+        positions = torch.arange(seq_len, device=device, dtype=torch.int64).type_as(inv_freq)
+        freqs = torch.outer(positions, inv_freq)
+        embeddings = torch.cat((freqs, freqs), dim=-1)
+
+        self.register_buffer(f"{layer_type}_cos_cached", embeddings.cos().to(dtype), persistent=False)
+        self.register_buffer(f"{layer_type}_sin_cached", embeddings.sin().to(dtype), persistent=False)
+
+
+def qeff_apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Apply pre-shaped rotary embeddings and restore the input dtype."""
+    x_embed = (x * cos) + (rotate_half(x) * sin)
+    return x_embed.to(x.dtype)
+
+
 class QEffGemma4TextAttention(Gemma4TextAttention):
     def __qeff_init__(self):
         for norm_name in ("q_norm", "k_norm", "v_norm"):
@@ -447,7 +479,7 @@ def forward(
 
         query_states = self.q_proj(hidden_states).view(hidden_shape)
         query_states = self.q_norm(query_states)
-        query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2)
+        query_states = qeff_apply_rotary_pos_emb(query_states, cos, sin)
         query_states = query_states.transpose(1, 2)
 
         if self.is_kv_shared_layer and past_key_values is not None:
@@ -463,7 +495,7 @@ def forward(
             value_states = self.v_proj(hidden_states).view(hidden_shape) if self.v_proj is not None else key_states
 
             key_states = self.k_norm(key_states)
-            key_states = apply_rotary_pos_emb(key_states, cos, sin, unsqueeze_dim=2)
+            key_states = qeff_apply_rotary_pos_emb(key_states, cos, sin)
             key_states = key_states.transpose(1, 2)
 
             value_states = self.v_norm(value_states)
@@ -594,6 +626,15 @@ def forward(
 
 
 class QEffGemma4TextModel(Gemma4TextModel):
+    def __qeff_init__(self):
+        self.rotary_emb = QEffGemma4TextRotaryEmbedding(config=self.config)
+        for layer_type in sorted(self.rotary_emb.layer_types):
+            attention_scaling = getattr(self.rotary_emb, f"{layer_type}_attention_scaling")
+            sin_cached = getattr(self.rotary_emb, f"{layer_type}_sin_cached") * attention_scaling
+            cos_cached = getattr(self.rotary_emb, f"{layer_type}_cos_cached") * attention_scaling
+            setattr(self, f"{layer_type}_sin_cached", nn.Parameter(sin_cached))
+            setattr(self, f"{layer_type}_cos_cached", nn.Parameter(cos_cached))
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
@@ -636,8 +677,12 @@ def forward(
         hidden_states = inputs_embeds
 
         position_embeddings = {}
-        for layer_type in self.unique_layer_types:
-            position_embeddings[layer_type] = self.rotary_emb(hidden_states, position_ids, layer_type)
+        for layer_type in sorted(self.unique_layer_types):
+            sin_cached = getattr(self, f"{layer_type}_sin_cached")
+            cos_cached = getattr(self, f"{layer_type}_cos_cached")
+            sin = sin_cached[position_ids].unsqueeze(2)
+            cos = cos_cached[position_ids].unsqueeze(2)
+            position_embeddings[layer_type] = (cos, sin)
 
         for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
             per_layer_input = per_layer_inputs[:, :, i, :] if per_layer_inputs is not None else None

tests/unit_test/models/test_gemma4_accuracy.py

@@ -366,6 +366,42 @@ def test_qeff_model_is_eval_mode(self):
         qeff = _make_qeff_gemma4(model)
         assert not qeff.model.training
 
+    def test_text_rotary_caches_cover_supported_context_and_match_hf(self):
+        model, cfg = make_tiny_gemma4()
+        hf_text_model = model.model.language_model
+        position_ids = torch.tensor([[0, 1, 7, CTX_LEN - 1]], dtype=torch.long)
+        hidden_states = torch.zeros(1, position_ids.shape[1], cfg.hidden_size)
+        with torch.no_grad():
+            expected_embeddings = {
+                layer_type: hf_text_model.rotary_emb(hidden_states, position_ids, layer_type)
+                for layer_type in ("full_attention", "sliding_attention")
+            }
+
+        qeff = _make_qeff_gemma4(model)
+        text_model = qeff.lang_model.model.language_model
+
+        for layer_type in ("full_attention", "sliding_attention"):
+            expected_head_dim = cfg.global_head_dim if layer_type == "full_attention" else cfg.head_dim
+            for embedding_index, embedding_name in enumerate(("cos", "sin")):
+                cached_embedding = getattr(text_model, f"{layer_type}_{embedding_name}_cached")
+                assert cached_embedding.shape == (text_model.rotary_emb.original_max_seq_len, expected_head_dim)
+                torch.testing.assert_close(
+                    cached_embedding[position_ids],
+                    expected_embeddings[layer_type][embedding_index],
+                )
+
+    def test_text_forward_reuses_rotary_caches(self, mocker):
+        model, cfg = make_tiny_gemma4()
+        qeff = _make_qeff_gemma4(model)
+        text_model = qeff.lang_model.model.language_model
+        rotary_forward = mocker.spy(text_model.rotary_emb, "forward")
+        input_ids = torch.randint(0, VOCAB_SIZE, (1, PREFILL_LEN))
+
+        with torch.no_grad():
+            _qeff_forward(qeff, **_prefill_inputs(input_ids, cfg))
+
+        rotary_forward.assert_not_called()
+
 
 # ---------------------------------------------------------------------------
 # Tests: QEff Gemma4 logit shape (argmax-based extraction)