e10c84e06a · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/core/test_mrope.py

+2 -6

@@ -113,12 +113,10 @@ def test_mrope(

113 is_neox_style = True

114

115 max_position = config.max_position_embeddings

116 - partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)

117 - rotary_dim = int(head_dim * partial_rotary_factor)

116

117 mrope_helper_class = get_rope(

118 head_size=head_dim,

121 - rotary_dim=rotary_dim,

119 + rotary_dim=head_dim,

120 max_position=max_position,

121 is_neox_style=is_neox_style,

122 rope_parameters=config.rope_parameters,

@@ -184,12 +182,10 @@ def test_mrope_torch_compile_tracing(

182 )

183 is_neox_style = True

184 max_position = config.max_position_embeddings

187 - partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)

188 - rotary_dim = int(head_dim * partial_rotary_factor)

185

186 mrope_helper_class = get_rope(

187 head_size=head_dim,

192 - rotary_dim=rotary_dim,

188 + rotary_dim=head_dim,

189 max_position=max_position,

190 is_neox_style=is_neox_style,

191 rope_parameters=config.rope_parameters,

vllm/model_executor/layers/rotary_embedding/__init__.py

+4 -1

@@ -30,7 +30,6 @@ def get_rope(

30 is_neox_style: bool = True,

31 rope_parameters: dict[str, Any] | None = None,

32 dtype: torch.dtype | None = None,

33 - partial_rotary_factor: float = 1.0,

33 dual_chunk_attention_config: dict[str, Any] | None = None,

34 ) -> RotaryEmbedding:

35 if dtype is None:

@@ -55,6 +54,10 @@ def get_rope(

54 else:

55 dual_chunk_attention_args = None

56

57 + partial_rotary_factor = 1.0

58 + if rope_parameters is not None:

59 + partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0)

60 +

61 if partial_rotary_factor < 1.0:

62 rotary_dim = int(rotary_dim * partial_rotary_factor)

63 key = (

vllm/model_executor/models/apertus.py

+1 -4

@@ -148,8 +148,6 @@ class ApertusAttention(nn.Module):

148 if head_dim is None:

149 head_dim = self.hidden_size // self.total_num_heads

150 self.head_dim = head_dim

151 - # Phi models introduced a partial_rotary_factor parameter in the config

152 - self.partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)

151 self.q_size = self.num_heads * self.head_dim

152 self.kv_size = self.num_kv_heads * self.head_dim

153 self.scaling = self.head_dim**-0.5

@@ -228,11 +226,10 @@ class ApertusAttention(nn.Module):

226

227 self.rotary_emb = get_rope(

228 self.head_dim,

231 - rotary_dim=int(self.partial_rotary_factor * self.head_dim),

229 + rotary_dim=self.head_dim,

230 max_position=self.max_position_embeddings,

231 rope_parameters=config.rope_parameters,

232 is_neox_style=is_neox_style,

235 - partial_rotary_factor=self.partial_rotary_factor,

233 )

234

235

vllm/model_executor/models/bailing_moe.py

-3

@@ -127,8 +127,6 @@ class BailingAttention(nn.Module):

127 prefix=f"{prefix}.dense",

128 )

129

130 - self.partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)

131 -

130 self.rotary_dim = getattr(config, "rotary_dim", self.head_dim)

131

132 self.rotary_emb = get_rope(

@@ -137,7 +135,6 @@ class BailingAttention(nn.Module):

135 max_position=config.max_position_embeddings,

136 rope_parameters=config.rope_parameters,

137 is_neox_style=True,

140 - partial_rotary_factor=self.partial_rotary_factor,

138 )

139

140 self.attn = Attention(

vllm/model_executor/models/bamba.py

+1 -3

@@ -178,9 +178,7 @@ class BambaAttentionDecoderLayer(nn.Module):

178 self.scaling = self.head_dim**-0.5

179 self.max_position_embeddings = max_position_embeddings

180

181 - if hasattr(config, "partial_rotary_factor"):

182 - rotary_dim = int(self.head_dim * config.partial_rotary_factor)

183 - elif hasattr(config, "attn_rotary_emb"):

181 + if hasattr(config, "attn_rotary_emb"):

182 rotary_dim = config.attn_rotary_emb # for backward compatibility

183 else:

184 rotary_dim = self.head_dim # default

vllm/model_executor/models/config.py

-5

@@ -8,7 +8,6 @@ import vllm.envs as envs

8 from vllm.logger import init_logger

9 from vllm.model_executor.models import ModelRegistry

10 from vllm.platforms import current_platform

11 -from vllm.transformers_utils.config import set_default_rope_theta

11 from vllm.utils.math_utils import cdiv, round_up

12 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE

13 from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, MLAAttentionSpec

@@ -78,8 +77,6 @@ class JinaRobertaModelConfig(VerifyAndUpdateConfig):

77 if not model_config.enforce_eager:

78 max_position = round_up(max_position, 8)

79

81 - set_default_rope_theta(config, default_theta=config.rotary_emb_base)

82 -

80 config.rotary_kwargs = {

81 "head_size": head_dim,

82 "rotary_dim": getattr(config, "rotary_emb_dim", head_dim),

@@ -119,8 +116,6 @@ class NomicBertModelConfig(VerifyAndUpdateConfig):

116 rotary_emb_dim = int(head_dim * config.rotary_emb_fraction)

117 max_trained_positions = getattr(config, "max_trained_positions", 2048)

118

122 - set_default_rope_theta(config, default_theta=config.rotary_emb_base)

123 -

119 config.rotary_kwargs = {

120 "head_size": head_dim,

121 "rotary_dim": rotary_emb_dim,

vllm/model_executor/models/falcon_h1.py

+1 -3

@@ -242,9 +242,7 @@ class FalconH1AttentionDecoderLayer(nn.Module):

242 self.scaling = self.head_dim**-0.5

243 self.max_position_embeddings = max_position_embeddings

244

245 - if hasattr(config, "partial_rotary_factor"):

246 - rotary_dim = self.head_dim * config.partial_rotary_factor

247 - elif hasattr(config, "attn_rotary_emb"):

245 + if hasattr(config, "attn_rotary_emb"):

246 rotary_dim = config.attn_rotary_emb # for backward compatibility

247 else:

248 rotary_dim = self.head_dim # default

vllm/model_executor/models/glm.py

+2 -1

@@ -10,7 +10,8 @@ from .utils import PPMissingLayer

10

11 class GlmForCausalLM(LlamaForCausalLM):

12 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

13 - vllm_config.model_config.hf_config.partial_rotary_factor = 0.5

13 + hf_config = vllm_config.model_config.hf_config

14 + hf_config.rope_parameters["partial_rotary_factor"] = 0.5

15 super().__init__(vllm_config=vllm_config, prefix=prefix)

16 # Hack Llama model to fit HF format GLM implementation

17 # Attention difference between GLM and Llama:

vllm/model_executor/models/glm4.py

+1 -2

@@ -78,7 +78,7 @@ class Glm4Attention(nn.Module):

78 # Number of KV heads is less than TP size, so we replicate

79 # the KV heads across multiple tensor parallel GPUs.

80 assert tp_size % self.total_num_kv_heads == 0

81 - partial_rotary_factor = getattr(config, "partial_rotary_factor", 0.5)

81 + config.rope_parameters.setdefault("partial_rotary_factor", 0.5)

82 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

83 self.head_dim = head_dim or hidden_size // self.total_num_heads

84 self.rotary_dim = self.head_dim

@@ -106,7 +106,6 @@ class Glm4Attention(nn.Module):

106 rotary_dim=self.rotary_dim,

107 max_position=max_position,

108 rope_parameters=config.rope_parameters,

109 - partial_rotary_factor=partial_rotary_factor,

109 is_neox_style=False,

110 )

111 self.attn = Attention(

vllm/model_executor/models/glm4_moe.py

+1 -2

@@ -282,13 +282,12 @@ class Glm4MoeAttention(nn.Module):

282 prefix=f"{prefix}.o_proj",

283 )

284

285 - partial_rotary_factor = getattr(config, "partial_rotary_factor", 0.5)

285 + config.rope_parameters.setdefault("partial_rotary_factor", 0.5)

286 self.rotary_emb = get_rope(

287 self.head_dim,

288 rotary_dim=self.head_dim,

289 max_position=max_position_embeddings,

290 rope_parameters=config.rope_parameters,

291 - partial_rotary_factor=partial_rotary_factor,

291 )

292 self.attn = Attention(

293 self.num_heads,

vllm/model_executor/models/gpt_neox.py

+2 -4

@@ -89,16 +89,14 @@ class GPTNeoXAttention(nn.Module):

89 quant_config=quant_config,

90 prefix=f"{prefix}.dense",

91 )

92 - scaling = self.head_size**-0.5

93 - rotary_dim = int(self.head_size * config.rotary_pct)

94 - assert rotary_dim % 2 == 0

92 max_position_embeddings = getattr(config, "max_position_embeddings", 8192)

93 self.rotary_emb = get_rope(

94 self.head_size,

98 - rotary_dim=rotary_dim,

95 + rotary_dim=self.head_size,

96 max_position=max_position_embeddings,

97 rope_parameters=config.rope_parameters,

98 )

99 + scaling = self.head_size**-0.5

100 self.attn = Attention(

101 self.num_heads,

102 self.head_size,

vllm/model_executor/models/llama.py

-3

@@ -149,8 +149,6 @@ class LlamaAttention(nn.Module):

149 if head_dim is None:

150 head_dim = self.hidden_size // self.total_num_heads

151 self.head_dim = head_dim

152 - # Phi models introduced a partial_rotary_factor parameter in the config

153 - self.partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)

152 self.q_size = self.num_heads * self.head_dim

153 self.kv_size = self.num_kv_heads * self.head_dim

154 self.scaling = self.head_dim**-0.5

@@ -265,7 +263,6 @@ class LlamaAttention(nn.Module):

263 max_position=self.max_position_embeddings,

264 rope_parameters=getattr(config, "rope_parameters", None),

265 is_neox_style=is_neox_style,

268 - partial_rotary_factor=self.partial_rotary_factor,

266 )

267

268

vllm/model_executor/models/nemotron.py

-2

@@ -178,7 +178,6 @@ class NemotronAttention(nn.Module):

178 self.q_size = self.num_heads * self.head_dim

179 self.kv_size = self.num_kv_heads * self.head_dim

180 self.scaling = self.head_dim**-0.5

181 - self.partial_rotary_factor = config.partial_rotary_factor

181 self.max_position_embeddings = max_position_embeddings

182

183 self.qkv_proj = QKVParallelLinear(

@@ -203,7 +202,6 @@ class NemotronAttention(nn.Module):

202 rotary_dim=self.head_dim,

203 max_position=max_position_embeddings,

204 rope_parameters=config.rope_parameters,

206 - partial_rotary_factor=self.partial_rotary_factor,

205 )

206 self.attn = Attention(

207 self.num_heads,

vllm/model_executor/models/nemotron_nas.py

-1

@@ -122,7 +122,6 @@ class DeciLMAttention(LlamaAttention):

122 max_position=self.max_position_embeddings,

123 rope_parameters=config.rope_parameters,

124 is_neox_style=is_neox_style,

125 - partial_rotary_factor=self.partial_rotary_factor,

125 )

126

127

vllm/model_executor/models/persimmon.py

-2

@@ -106,7 +106,6 @@ class PersimmonAttention(nn.Module):

106 self.num_heads = self.total_num_heads // tensor_parallel_world_size

107 self.head_dim = self.hidden_size // self.total_num_heads

108 self.max_position_embeddings = config.max_position_embeddings

109 - self.partial_rotary_factor = config.partial_rotary_factor

109 self.is_causal = True

110

111 assert (self.head_dim * self.total_num_heads) == self.hidden_size

@@ -138,7 +137,6 @@ class PersimmonAttention(nn.Module):

137 rotary_dim=self.head_dim,

138 max_position=self.max_position_embeddings,

139 rope_parameters=config.rope_parameters,

141 - partial_rotary_factor=self.partial_rotary_factor,

140 )

141 self.scaling = self.head_dim**-0.5

142 self.attn = Attention(

vllm/model_executor/models/phi.py

+1 -4

@@ -109,10 +109,7 @@ class PhiAttention(nn.Module):

109 )

110

111 scaling = self.head_size**-0.5

112 - rotary_dim = int(

113 - config.partial_rotary_factor

114 - * (config.hidden_size // config.num_attention_heads)

115 - )

112 + rotary_dim = config.hidden_size // config.num_attention_heads

113 assert rotary_dim % 2 == 0

114

115 max_position_embeddings = getattr(config, "max_position_embeddings", 2048)

vllm/model_executor/models/qwen3_next.py

-1

@@ -750,7 +750,6 @@ class Qwen3NextAttention(nn.Module):

750 rotary_dim=self.head_dim,

751 max_position=config.max_position_embeddings,

752 rope_parameters=config.rope_parameters,

753 - partial_rotary_factor=config.partial_rotary_factor,

753 dual_chunk_attention_config=self.dual_chunk_attention_config,

754 )

755

vllm/model_executor/models/stablelm.py

-4

@@ -119,9 +119,6 @@ class StablelmAttention(nn.Module):

119 self.num_key_value_heads = max(1, self.total_num_key_value_heads // tp_size)

120 self.head_dim = self.hidden_size // self.total_num_heads

121 self.max_position_embeddings = config.max_position_embeddings

122 - self.partial_rotary_factor = getattr(

123 - config, "rope_pct", getattr(config, "partial_rotary_factor", 1)

124 - )

122 self.scaling = self.head_dim**-0.5

123 self.q_size = self.num_heads * self.head_dim

124 self.kv_size = self.num_key_value_heads * self.head_dim

@@ -154,7 +151,6 @@ class StablelmAttention(nn.Module):

151 rotary_dim=self.head_dim,

152 max_position=self.config.max_position_embeddings,

153 rope_parameters=self.config.rope_parameters,

157 - partial_rotary_factor=self.partial_rotary_factor,

154 )

155 self.attn = Attention(

156 self.num_heads,

vllm/transformers_utils/config.py

+9 -1

@@ -25,6 +25,7 @@ from transformers.models.auto.tokenization_auto import get_tokenizer_config

25 from transformers.utils import CONFIG_NAME as HF_CONFIG_NAME

26

27 from vllm import envs

28 +from vllm.config.utils import getattr_iter

29 from vllm.logger import init_logger

30 from vllm.transformers_utils.utils import parse_safetensors_file_metadata

31

@@ -304,7 +305,8 @@ def set_default_rope_theta(config: PretrainedConfig, default_theta: float) -> No

305

306 def patch_rope_parameters(config: PretrainedConfig) -> None:

307 """Provide backwards compatibility for RoPE."""

307 - rope_theta = getattr(config, "rope_theta", None)

308 + rope_theta_names = ("rope_theta", "rotary_emb_base")

309 + rope_theta = getattr_iter(config, rope_theta_names, None)

310 if Version(version("transformers")) < Version("5.0.0.dev0"):

311 # Transformers v4 installed, legacy config fields may be present

312 if (rope_scaling := getattr(config, "rope_scaling", None)) is not None:

@@ -313,6 +315,12 @@ def patch_rope_parameters(config: PretrainedConfig) -> None:

315 if not hasattr(config, "rope_parameters"):

316 config.rope_parameters = {"rope_type": "default"}

317 config.rope_parameters["rope_theta"] = rope_theta

318 + partial_rotary_factor_names = ("partial_rotary_factor", "rotary_pct")

319 + partial_rotary_factor = getattr_iter(config, partial_rotary_factor_names, None)

320 + if partial_rotary_factor is not None:

321 + if not hasattr(config, "rope_parameters"):

322 + config.rope_parameters = {"rope_type": "default"}

323 + config.rope_parameters["partial_rotary_factor"] = partial_rotary_factor

324 elif rope_theta is not None or hasattr(config, "rope_parameters"):

325 # Transformers v5 installed

326 config.standardize_rope_params()

vllm/transformers_utils/configs/nemotron.py

+13 -7

@@ -89,9 +89,14 @@ class NemotronConfig(PretrainedConfig):

89 tie_word_embeddings (`bool`, *optional*, defaults to `False`):

90 Whether to tie weight embeddings

91 rope_parameters (`dict`, *optional*):

92 - The parameters of the RoPE embeddings.

93 - partial_rotary_factor (`float`, *optional*, defaults to 0.5):

94 - Percentage of the query and keys which will have rotary embedding.

92 + The parameters of the RoPE embeddings. Expected contents:

93 + `rope_theta` (`float`): The base period of the RoPE embeddings.

94 + `rope_type` (`str`):

95 + The sub-variant of RoPE to use. Can be one of ['default', 'linear',

96 + 'dynamic', 'yarn', 'longrope', 'llama3'], with 'default' being the

97 + original RoPE implementation.

98 + `partial_rotary_factor` (`float`, *optional*, defaults to 0.5):

99 + Percentage of the query and keys which will have rotary embedding.

100 attention_bias (`bool`, *optional*, defaults to `False`):

101 Whether to use a bias in the query, key, value and output

102 projection layers during self-attention.

@@ -133,7 +138,6 @@ class NemotronConfig(PretrainedConfig):

138 eos_token_id=3,

139 tie_word_embeddings=False,

140 rope_parameters=None,

136 - partial_rotary_factor=0.5,

141 attention_bias=False,

142 attention_dropout=0.0,

143 mlp_bias=False,

@@ -165,14 +169,16 @@ class NemotronConfig(PretrainedConfig):

169 rope_theta = kwargs.pop("rope_theta", 10000.0)

170 if "rope_theta" not in rope_parameters:

171 rope_parameters["rope_theta"] = rope_theta

168 - self.rope_parameters = rope_parameters

172 # for backward compatibility

173 partial_rotary_factor = (

174 kwargs.get("rope_percent")

175 or kwargs.get("rope_percentage")

173 - or partial_rotary_factor

176 + or kwargs.get("partial_rotary_factor")

177 + or 0.5

178 )

175 - self.partial_rotary_factor = partial_rotary_factor

179 + if "partial_rotary_factor" not in rope_parameters:

180 + rope_parameters["partial_rotary_factor"] = partial_rotary_factor

181 + self.rope_parameters = rope_parameters

182 self._rope_parameters_validation()

183 self.attention_bias = attention_bias

184 self.attention_dropout = attention_dropout

vllm/transformers_utils/configs/qwen3_next.py

+5 -3

@@ -103,8 +103,8 @@ class Qwen3NextConfig(PretrainedConfig):

103 Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE

104 `high_freq_factor` (`float`, *optional*):

105 Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE

106 - partial_rotary_factor (`float`, *optional*, defaults to 0.25):

107 - Percentage of the query and keys which will have rotary embedding.

106 + `partial_rotary_factor` (`float`, *optional*, defaults to 0.25):

107 + Percentage of the query and keys which will have rotary embedding.

108 attention_bias (`bool`, *optional*, defaults to `False`):

109 Whether to use a bias in the query, key, value and output projection layers during self-attention.

110 attention_dropout (`float`, *optional*, defaults to 0.0):

@@ -198,7 +198,6 @@ class Qwen3NextConfig(PretrainedConfig):

198 use_cache=True,

199 tie_word_embeddings=False,

200 rope_parameters=None,

201 - partial_rotary_factor=0.25,

201 attention_bias=False,

202 attention_dropout=0.0,

203 head_dim=256,

@@ -239,6 +238,9 @@ class Qwen3NextConfig(PretrainedConfig):

238 rope_theta = kwargs.pop("rope_theta", 10000.0)

239 if "rope_theta" not in rope_parameters:

240 rope_parameters["rope_theta"] = rope_theta

241 + partial_rotary_factor = kwargs.pop("partial_rotary_factor", 0.25)

242 + if "partial_rotary_factor" not in rope_parameters:

243 + rope_parameters["partial_rotary_factor"] = partial_rotary_factor

244 self.rope_parameters = rope_parameters

245 self.partial_rotary_factor = partial_rotary_factor

246 self.attention_bias = attention_bias

Access `partial_rotary_factor` from `rope_parameters` (#29966)