address comments

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

Author: shivghai

tensorrt_llm/layers/attention.py

@@ -702,7 +702,9 @@ def create_attention_const_params(model_cls, config):
                           is_buffer=True))
         else:
 
-            def register_rope_params(rotary_base, scale, scale_type, scaling,
+            def register_rope_params(rotary_base, rotary_embedding_scale,
+                                     rotary_embedding_scale_type,
+                                     rotary_embedding_scaling,
                                      names_to_register):
                 # Rotary const weights.
                 embed_positions = RopeEmbeddingUtils.create_sinusoidal_positions(
@@ -712,7 +714,8 @@ def register_rope_params(rotary_base, scale, scale_type, scaling,
 
                 rotary_inv_freq, embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin(
                     max_position_embeddings, rotary_embedding_dim, rotary_base,
-                    scale, scale_type, scaling)
+                    rotary_embedding_scale, rotary_embedding_scale_type,
+                    rotary_embedding_scaling)
                 model_cls.register_parameter(
                     names_to_register[0],
                     Parameter(embed_positions, dtype='float32', is_buffer=True))
@@ -725,24 +728,25 @@ def register_rope_params(rotary_base, scale, scale_type, scaling,
                               dtype='float32',
                               is_buffer=True))
 
-            register_rope_params(rotary_base=rotary_embedding_base,
-                                 scale=rotary_embedding_scale,
-                                 scale_type=rotary_embedding_scale_type,
-                                 scaling=rotary_embedding_scaling,
-                                 names_to_register=[
-                                     'embed_positions', 'rotary_inv_freq',
-                                     'embed_positions_for_gpt_attention'
-                                 ])
+            register_rope_params(
+                rotary_base=rotary_embedding_base,
+                rotary_embedding_scale=rotary_embedding_scale,
+                rotary_embedding_scale_type=rotary_embedding_scale_type,
+                rotary_embedding_scaling=rotary_embedding_scaling,
+                names_to_register=[
+                    'embed_positions', 'rotary_inv_freq',
+                    'embed_positions_for_gpt_attention'
+                ])
 
             # For models with non-homegeneous attention layers requiring a second set of rope params. e.g. Gemma3.
             rotary_embedding_base_local = getattr(config,
                                                   'rope_local_base_freq', None)
             if rotary_embedding_base_local is not None:
                 register_rope_params(
                     rotary_base=rotary_embedding_base_local,
-                    scale=1.0,
-                    scale_type=RotaryScalingType.none,
-                    scaling=None,
+                    rotary_embedding_scale=1.0,
+                    rotary_embedding_scale_type=RotaryScalingType.none,
+                    rotary_embedding_scaling=None,
                     names_to_register=[
                         'embed_positions_local', 'rotary_inv_freq_local',
                         'embed_positions_for_gpt_attention_local'

tests/unittest/others/test_layer.py

@@ -2184,7 +2184,6 @@ def register_parameter(cls, name, param):
         # The global and local inv_freq should be different because:
         # 1. Global uses rope_scaling with factor=8.0 (linear scaling applies 1/8 to inv_freq)
         # 2. Local uses scale=1.0 (no scaling)
-        # Also they use different base frequencies (1000000 vs 10000)
         self.assertFalse(
             np.allclose(global_inv_freq, local_inv_freq),
             "Global and local rotary_inv_freq should be different "
@@ -2197,8 +2196,8 @@ def register_parameter(cls, name, param):
             "(global has scaling, local does not)")
 
         # Additional verification: Check that local inv_freq matches unscaled calculation
-        # For local attention with scale=1.0 and base=10000:
-        # inv_freq = 1.0 / (10000 ** (arange(0, dim, 2) / dim))
+        # For local attention with scale=1.0 and base=10:
+        # inv_freq = 1.0 / (10 ** (arange(0, dim, 2) / dim))
         dim = config.head_size  # rotary_embedding_dim = head_size * rotary_pct = 128
         expected_local_inv_freq = 1.0 / (config.rope_local_base_freq
                                          **(np.arange(0, dim, 2) / dim))
@@ -2211,7 +2210,7 @@ def register_parameter(cls, name, param):
 
         # For global attention with linear scaling (factor=8.0):
         # scale = 1.0 / 8.0 = 0.125
-        # inv_freq = 0.125 / (1000000 ** (arange(0, dim, 2) / dim))
+        # inv_freq = 0.125 / (100 ** (arange(0, dim, 2) / dim))
         expected_global_inv_freq = (1.0 / 8.0) / (config.rotary_base**
                                                   (np.arange(0, dim, 2) / dim))