A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 01:00 UTC
Let GraniteMoeAttention use YaRN (#21174)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Jul 18, 2025 at 14:52 UTC
ed8cbfedf84f1b1fc1d3eadf3622d9903e906cb0
2 files changed
+7
-1
vllm/model_executor/models/granitemoe.py
+5
-1
@@ -24,7 +24,7 @@
24
# limitations under the License.
25
"""Inference-only GraniteMoe model."""
26
from collections.abc import Iterable
27
-from typing import Optional
27
+from typing import Any, Optional
28
29
import torch
30
from torch import nn
@@ -113,6 +113,7 @@ class GraniteMoeAttention(nn.Module):
113
num_kv_heads: int,
114
max_position: int = 4096 * 32,
115
rope_theta: float = 10000,
116
+ rope_scaling: Optional[dict[str, Any]] = None,
117
cache_config: Optional[CacheConfig] = None,
118
quant_config: Optional[QuantizationConfig] = None,
119
attention_multiplier: Optional[float] = None,
@@ -163,6 +164,7 @@ class GraniteMoeAttention(nn.Module):
164
max_position=max_position,
165
base=int(self.rope_theta),
166
is_neox_style=True,
167
+ rope_scaling=rope_scaling,
168
)
169
self.attn = Attention(self.num_heads,
170
self.head_dim,
@@ -198,12 +200,14 @@ class GraniteMoeDecoderLayer(nn.Module):
200
self.hidden_size = config.hidden_size
201
# Requires transformers > 4.32.0
202
rope_theta = getattr(config, "rope_theta", 10000)
203
+ rope_scaling = getattr(config, "rope_scaling", None)
204
self.self_attn = GraniteMoeAttention(
205
hidden_size=self.hidden_size,
206
num_heads=config.num_attention_heads,
207
max_position=config.max_position_embeddings,
208
num_kv_heads=config.num_key_value_heads,
209
rope_theta=rope_theta,
210
+ rope_scaling=rope_scaling,
211
cache_config=cache_config,
212
quant_config=quant_config,
213
prefix=f"{prefix}.self_attn",
vllm/model_executor/models/granitemoeshared.py
+2
@@ -81,12 +81,14 @@ class GraniteMoeSharedDecoderLayer(nn.Module):
81
self.hidden_size = config.hidden_size
82
# Requires transformers > 4.32.0
83
rope_theta = getattr(config, "rope_theta", 10000)
84
+ rope_scaling = getattr(config, "rope_scaling", None)
85
self.self_attn = GraniteMoeAttention(
86
hidden_size=self.hidden_size,
87
num_heads=config.num_attention_heads,
88
max_position=config.max_position_embeddings,
89
num_kv_heads=config.num_key_value_heads,
90
rope_theta=rope_theta,
91
+ rope_scaling=rope_scaling,
92
cache_config=cache_config,
93
quant_config=quant_config,
94
prefix=f"{prefix}.self_attn",