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 10, 17:45 UTC
[Bugfix] Properly apply v_scale for mimo_v2_flash (#31175)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jan 5, 2026 at 18:20 UTC
95130298981428053a23376e89a771afc7fc37af
1 file changed
+10
-13
vllm/model_executor/models/mimo_v2_flash.py
+10
-13
@@ -211,6 +211,7 @@ class MiMoV2Attention(nn.Module):
211
num_kv_heads: int,
212
head_dim: int,
213
v_head_dim: int | None = None,
214
+ v_scale: float | None = None,
215
sliding_window_size: int = -1,
216
attention_bias: bool = False,
217
add_swa_attention_sink_bias: bool = False,
@@ -241,6 +242,7 @@ class MiMoV2Attention(nn.Module):
242
self.k_size = self.num_kv_heads * self.head_dim
243
self.v_size = self.num_kv_heads * self.v_head_dim
244
245
+ self.v_scale = v_scale
246
self.scaling = self.head_dim**-0.5
247
self.rope_theta = rope_theta
248
self.max_position_embeddings = max_position_embeddings
@@ -304,6 +306,10 @@ class MiMoV2Attention(nn.Module):
306
q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
307
q, k = self.rotary_emb(positions, q, k)
308
309
+ # Apply v_scale before attention
310
+ if self.v_scale is not None:
311
+ v = v * self.v_scale
312
+
313
v = v.view(-1, self.num_kv_heads, self.v_head_dim)
314
v = torch.nn.functional.pad(v, [0, self.head_dim - self.v_head_dim], value=0)
315
v = v.view(-1, self.num_kv_heads * self.head_dim)
@@ -332,6 +338,8 @@ class MiMoV2FlashDecoderLayer(nn.Module):
338
rope_theta = getattr(config, "rope_theta", 1000000)
339
max_position_embeddings = getattr(config, "max_position_embeddings", 32768)
340
341
+ v_scale = getattr(config, "attention_value_scale", None)
342
+
343
if self.is_compressed_softmax_layer():
344
self.self_attn = MiMoV2Attention(
345
hidden_size=self.hidden_size,
@@ -339,6 +347,7 @@ class MiMoV2FlashDecoderLayer(nn.Module):
347
num_kv_heads=config.swa_num_key_value_heads,
348
head_dim=config.swa_head_dim,
349
v_head_dim=getattr(config, "swa_v_head_dim", None),
350
+ v_scale=v_scale,
351
sliding_window_size=config.sliding_window_size,
352
attention_bias=config.attention_bias,
353
add_swa_attention_sink_bias=getattr(
@@ -358,6 +367,7 @@ class MiMoV2FlashDecoderLayer(nn.Module):
367
num_kv_heads=config.num_key_value_heads,
368
head_dim=config.head_dim,
369
v_head_dim=getattr(config, "v_head_dim", None),
370
+ v_scale=v_scale,
371
sliding_window_size=-1, # normal attention
372
attention_bias=config.attention_bias,
373
layer_id=layer_id,
@@ -433,7 +443,6 @@ class MiMoV2Model(nn.Module):
443
self.quant_config = quant_config
444
self.vocab_size = config.vocab_size
445
self.num_redundant_experts = eplb_config.num_redundant_experts
436
- self.v_scale = getattr(config, "attention_value_scale", None)
446
447
if get_pp_group().is_first_rank or (
448
config.tie_word_embeddings and get_pp_group().is_last_rank
@@ -605,18 +614,6 @@ class MiMoV2Model(nn.Module):
614
615
param = params_dict[name_rewritten]
616
weight_loader = getattr(param, "weight_loader", default_weight_loader)
608
-
609
- if param_name == "qkv_proj" and shard_id == "v":
610
- v_scale = (
611
- self.v_scale
612
- if self.v_scale is not None
613
- else getattr(self.config, "attention_value_scale", None)
614
- )
615
- if v_scale is not None and (
616
- name.endswith("weight_scale_inv") or name.endswith(".bias")
617
- ):
618
- loaded_weight *= float(v_scale)
619
-
617
weight_loader(param, loaded_weight, shard_id)
618
loaded_params.add(name_rewritten)
619