[Gemma3] Fix RoPE for local attention for Gemma3

Signed-off-by: Shiv Ghai <[email protected]>

Author: shivghai

tensorrt_llm/layers/attention.py

@@ -702,16 +702,20 @@ def create_attention_const_params(model_cls, config):
                           is_buffer=True))
         else:
 
-            def register_rope_params(rotary_base, names_to_register):
+            def register_rope_params(rotary_base, names_to_register, is_local=False):
                 # Rotary const weights.
                 embed_positions = RopeEmbeddingUtils.create_sinusoidal_positions(
                     max_position_embeddings,
                     rotary_embedding_dim,
                 )
+                # For local attention, use no scaling (consistent with forward pass)
+                local_scale = 1.0 if is_local else rotary_embedding_scale
+                local_scale_type = RotaryScalingType.none if is_local else rotary_embedding_scale_type
+                local_scaling = None if is_local else rotary_embedding_scaling
+
                 rotary_inv_freq, embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin(
                     max_position_embeddings, rotary_embedding_dim, rotary_base,
-                    rotary_embedding_scale, rotary_embedding_scale_type,
-                    rotary_embedding_scaling)
+                    local_scale, local_scale_type, local_scaling)
                 model_cls.register_parameter(
                     names_to_register[0],
                     Parameter(embed_positions, dtype='float32', is_buffer=True))
@@ -739,7 +743,8 @@ def register_rope_params(rotary_base, names_to_register):
                     names_to_register=[
                         'embed_positions_local', 'rotary_inv_freq_local',
                         'embed_positions_for_gpt_attention_local'
-                    ])
+                    ],
+                    is_local=True)
 
     @staticmethod
     def fill_attention_params(model_cls, attention_params):
@@ -1141,10 +1146,10 @@ def compute_cross_kv(encoder_output):
                 rotary_embedding_dim=self.rotary_embedding_dim,
                 rotary_embedding_base=self.rotary_embedding_base
                 if not self.is_local else self.rotary_embedding_base_local,
-                rotary_embedding_scale_type=self.rotary_embedding_scale_type,
+                rotary_embedding_scale_type=self.rotary_embedding_scale_type if not self.is_local else RotaryScalingType.none,
                 rotary_embedding_short_m_scale=attention_params.short_mscale,
                 rotary_embedding_long_m_scale=attention_params.long_mscale,
-                rotary_embedding_scale=self.rotary_embedding_scale,
+                rotary_embedding_scale=self.rotary_embedding_scale if not self.is_local else 1.0,
                 rotary_embedding_max_positions=self.max_position_embeddings,
                 rotary_embedding_original_max_positions=self.
                 original_max_position_embeddings,
@@ -2792,4 +2797,4 @@ def forward(self,
             attention_mask=attention_mask,
             max_input_length=max_input_length,
             *args,
-            **kwargs)
+            **kwargs)