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, 13:30 UTC
[Minor][Models] Pass partial_rotary_factor parameter to rope (#17266)
Signed-off-by: evian <eviantai@u.nus.edu> Co-authored-by: evian <eviantai@u.nus.edu>
Wanrui Dai committed
Apr 28, 2025 at 12:28 UTC
7fcc4223dca53ea1531093cbbfcf59ddfea1800d
3 files changed
+10
-8
vllm/model_executor/models/llama.py
+4
-3
@@ -130,8 +130,8 @@ class LlamaAttention(nn.Module):
130
self.head_dim = getattr(config, "head_dim",
131
self.hidden_size // self.total_num_heads)
132
# Phi models introduced a partial_rotary_factor parameter in the config
133
- partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)
134
- self.rotary_dim = int(partial_rotary_factor * self.head_dim)
133
+ self.partial_rotary_factor = getattr(config, "partial_rotary_factor",
134
+ 1)
135
self.q_size = self.num_heads * self.head_dim
136
self.kv_size = self.num_kv_heads * self.head_dim
137
self.scaling = self.head_dim**-0.5
@@ -163,11 +163,12 @@ class LlamaAttention(nn.Module):
163
164
self.rotary_emb = get_rope(
165
self.head_dim,
166
- rotary_dim=self.rotary_dim,
166
+ rotary_dim=self.head_dim,
167
max_position=max_position_embeddings,
168
base=rope_theta,
169
rope_scaling=rope_scaling,
170
is_neox_style=is_neox_style,
171
+ partial_rotary_factor=self.partial_rotary_factor,
172
)
173
174
if hasattr(config, "interleaved_sliding_window"):
vllm/model_executor/models/persimmon.py
+2
-1
@@ -115,9 +115,10 @@ class PersimmonAttention(nn.Module):
115
116
self.rotary_emb = get_rope(
117
self.head_dim,
118
- rotary_dim=int(self.partial_rotary_factor * self.head_dim),
118
+ rotary_dim=self.head_dim,
119
max_position=self.max_position_embeddings,
120
base=self.rope_theta,
121
+ partial_rotary_factor=self.partial_rotary_factor,
122
)
123
self.scaling = self.head_dim**-0.5
124
self.attn = Attention(self.num_heads,
vllm/model_executor/models/stablelm.py
+4
-4
@@ -104,9 +104,8 @@ class StablelmAttention(nn.Module):
104
1, self.total_num_key_value_heads // tp_size)
105
self.head_dim = self.hidden_size // self.total_num_heads
106
self.max_position_embeddings = config.max_position_embeddings
107
- rope_pct = getattr(config, "rope_pct",
108
- getattr(config, "partial_rotary_factor", 1))
109
- self.rotary_ndims = int(self.head_dim * rope_pct)
107
+ self.partial_rotary_factor = getattr(
108
+ config, "rope_pct", getattr(config, "partial_rotary_factor", 1))
109
self.scaling = self.head_dim**-0.5
110
self.q_size = self.num_heads * self.head_dim
111
self.kv_size = self.num_key_value_heads * self.head_dim
@@ -130,9 +129,10 @@ class StablelmAttention(nn.Module):
129
prefix=f"{prefix}.o_proj")
130
self.rotary_emb = get_rope(
131
self.head_dim,
133
- rotary_dim=self.rotary_ndims,
132
+ rotary_dim=self.head_dim,
133
max_position=self.config.max_position_embeddings,
134
base=self.config.rope_theta,
135
+ partial_rotary_factor=self.partial_rotary_factor,
136
)
137
self.attn = Attention(self.num_heads,
138
self.head_dim,