lint/fmt using pre-commit

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, names_to_register, is_local=False):
+            def register_rope_params(rotary_base,
+                                     names_to_register,
+                                     is_local=False):
                 # Rotary const weights.
                 embed_positions = RopeEmbeddingUtils.create_sinusoidal_positions(
                     max_position_embeddings,
@@ -1146,10 +1148,12 @@ 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 if not self.is_local else RotaryScalingType.none,
+                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 if not self.is_local else 1.0,
+                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,
@@ -2797,4 +2801,4 @@ def forward(self,
             attention_mask=attention_mask,
             max_input_length=max_input_length,
             *args,
-            **kwargs)
+            **kwargs)

tests/unittest/others/test_layer.py

@@ -2115,7 +2115,6 @@ def fuse_rg_lru(recurrent_layer):
                                    atol=atol,
                                    rtol=rtol)
 
-
     def test_gemma3_local_attention_rope_scaling(self):
         """
         Test that local attention layers in Gemma3 do NOT apply rope scaling,
@@ -2126,8 +2125,7 @@ def test_gemma3_local_attention_rope_scaling(self):
         ensures that local attention layers get scale=1.0 and scale_type=none,
         while global layers get the configured scaling.
         """
-        from tensorrt_llm.functional import (PositionEmbeddingType,
-                                             RotaryScalingType)
+        from tensorrt_llm.functional import PositionEmbeddingType
         from tensorrt_llm.layers.attention import Attention
 
         # Create a mock config similar to Gemma3 27B with rope_scaling
@@ -2138,10 +2136,7 @@ class MockGemma3Config:
             max_position_embeddings = 32768
             position_embedding_type = PositionEmbeddingType.rope_gpt_neox
             rotary_base = 1000000.0
-            rotary_scaling = {
-                "factor": 8.0,
-                "rope_type": "linear"
-            }
+            rotary_scaling = {"factor": 8.0, "rope_type": "linear"}
             rotary_pct = 1.0
             # Local attention uses a different base frequency
             rope_local_base_freq = 10000.0
@@ -2202,8 +2197,8 @@ def register_parameter(cls, name, param):
         # For local attention with scale=1.0 and base=10000:
         # inv_freq = 1.0 / (10000 ** (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))
+        expected_local_inv_freq = 1.0 / (config.rope_local_base_freq
+                                         **(np.arange(0, dim, 2) / dim))
 
         np.testing.assert_allclose(
             local_inv_freq,
@@ -2214,14 +2209,15 @@ 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))
-        expected_global_inv_freq = (1.0 / 8.0) / (config.rotary_base**(
-            np.arange(0, dim, 2) / dim))
+        expected_global_inv_freq = (1.0 / 8.0) / (config.rotary_base**
+                                                  (np.arange(0, dim, 2) / dim))
 
         np.testing.assert_allclose(
             global_inv_freq,
             expected_global_inv_freq,
             rtol=1e-5,
-            err_msg="Global rotary_inv_freq should be computed WITH linear scaling")
+            err_msg=
+            "Global rotary_inv_freq should be computed WITH linear scaling")
 
 
 if __name__ == '__main__':