Factor llama rotary_emb initialization out to protected method.

Override rotary_emb initialization for NemotronNAS Attention

Signed-off-by: Nave Assaf <[email protected]>

Author: Naveassaf

vllm/model_executor/models/llama.py

@@ -162,20 +162,9 @@ def __init__(
             prefix=f"{prefix}.o_proj",
         )
 
-        is_neox_style = True
-        is_gguf = quant_config and quant_config.get_name() == "gguf"
-        if is_gguf and config.model_type == "llama":
-            is_neox_style = False
-
-        self.rotary_emb = get_rope(
-            self.head_dim,
-            rotary_dim=self.head_dim,
-            max_position=max_position_embeddings,
-            base=rope_theta,
-            rope_scaling=rope_scaling,
-            is_neox_style=is_neox_style,
-            partial_rotary_factor=self.partial_rotary_factor,
-        )
+        self._init_rotary_emb(config,
+                              rope_scaling=rope_scaling,
+                              quant_config=quant_config)
 
         if hasattr(config, "interleaved_sliding_window"):
             interleaved_sliding_window = config.interleaved_sliding_window
@@ -214,6 +203,24 @@ def forward(
         output, _ = self.o_proj(attn_output)
         return output
 
+    def _init_rotary_emb(self, config: LlamaConfig,
+                         rope_scaling: Optional[dict[str, Any]],
+                         quant_config: Optional[QuantizationConfig]) -> None:
+        is_neox_style = True
+        is_gguf = quant_config and quant_config.get_name() == "gguf"
+        if is_gguf and self.config.model_type == "llama":
+            is_neox_style = False
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=self.max_position_embeddings,
+            base=self.rope_theta,
+            rope_scaling=rope_scaling,
+            is_neox_style=is_neox_style,
+            partial_rotary_factor=self.partial_rotary_factor,
+        )
+
 
 class LlamaDecoderLayer(nn.Module):
 

vllm/model_executor/models/nemotron_nas.py

@@ -23,7 +23,7 @@
 # limitations under the License.
 """Inference-only deci model compatible with HuggingFace weights."""
 from collections.abc import Iterable
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 from torch import nn
@@ -66,36 +66,43 @@ def _find_multiple(n: int, k: int) -> int:
 
 class DeciLMAttention(LlamaAttention):
 
-    def __init__(self,
-                 config,
-                 hidden_size,
-                 num_heads,
-                 num_kv_heads,
-                 rope_theta=10000,
-                 rope_scaling=None,
-                 max_position_embeddings=8192,
-                 quant_config=None,
-                 bias=False,
-                 bias_o_proj=False,
-                 cache_config=None,
-                 prefix="",
-                 attn_type=AttentionType.DECODER):
+    def __init__(
+        self,
+        config: LlamaConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_theta: float = 10000,
+        rope_scaling: Optional[dict[str, Any]] = None,
+        max_position_embeddings: int = 8192,
+        quant_config: Optional[QuantizationConfig] = None,
+        bias: bool = False,
+        bias_o_proj: bool = False,
+        cache_config: Optional[CacheConfig] = None,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ) -> None:
         super().__init__(config, hidden_size, num_heads, num_kv_heads,
                          rope_theta, rope_scaling, max_position_embeddings,
                          quant_config, bias, bias_o_proj, cache_config, prefix,
                          attn_type)
 
-        # Enable YARN by overriding rope
-        interleaved_rope = config.position_embedding_type in [
-            "mistral_yarn", "rope_llama4"
-        ]
+    def _init_rotary_emb(self, config, rope_scaling: Optional[dict[str, Any]],
+                         quant_config: Optional[QuantizationConfig]) -> None:
+        # Enables YARN for Mistral and LLaMA4 derivatives.
+        is_neox_style = True
+        if hasattr(config, "position_embedding_type"):
+            is_neox_style = config.position_embedding_type not in [
+                "mistral_yarn", "rope_llama4"
+            ]
+
         self.rotary_emb = get_rope(
             self.head_dim,
             rotary_dim=self.head_dim,
             max_position=self.max_position_embeddings,
             base=self.rope_theta,
             rope_scaling=rope_scaling,
-            is_neox_style=not interleaved_rope,
+            is_neox_style=is_neox_style,
             partial_rotary_factor=self.partial_rotary_factor)