[Perf] Remove redundant op for GLM 5.2 (#47198)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 4, 2026 at 13:25 UTC
2a9113f9985f10d2089a55e419c3b67b1b9731de
1 file changed
+16
-20
vllm/model_executor/models/deepseek_v2.py
+16
-20
@@ -713,6 +713,7 @@ class Indexer(nn.Module):
713
)
714
715
self.is_inplace_rope = is_inplace_rope
716
+ self.n_head_scale = self.n_head**-0.5
717
self.use_fused_indexer_q = (
718
current_platform.is_cuda()
719
and self.quant_block_size == self.head_dim
@@ -761,15 +762,16 @@ class Indexer(nn.Module):
762
rotary_emb.cos_sin_cache,
763
weights,
764
self.softmax_scale,
764
- self.n_head**-0.5,
765
+ self.n_head_scale,
766
rotary_emb.is_neox_style,
767
)
768
769
# rotate only the MQA K
769
- q_dummy = torch.empty_like(k_pe.unsqueeze(1))
770
- _, k_pe = rotary_emb(positions, q_dummy, k_pe.unsqueeze(1))
771
- k_pe = k_pe.reshape(-1, 1, self.rope_dim)
772
- k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
770
+ k_pe = k_pe.unsqueeze(1)
771
+ q_dummy = torch.empty_like(k_pe)
772
+ _, k_pe = rotary_emb(positions, q_dummy, k_pe)
773
+ k_pe = k_pe.reshape(-1, self.rope_dim)
774
+ k = torch.cat([k_pe, k_nope], dim=-1)
775
776
return self.indexer_op(hidden_states, q_fp8, k, weights)
777
else:
@@ -790,13 +792,13 @@ class Indexer(nn.Module):
792
# Note: RoPE (NeoX) can introduce extra leading dimensions during
793
# compilation so we need to reshape back to token-flattened shapes
794
q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
793
- k_pe = k_pe.reshape(-1, 1, self.rope_dim)
795
+ k_pe = k_pe.reshape(-1, self.rope_dim)
796
797
# `rotary_emb` is shape-preserving; `q_pe` is already
798
# [num_tokens, n_head, rope_dim].
799
q = torch.cat([q_pe, q_nope], dim=-1)
798
- # `k_pe` is [num_tokens, 1, rope_dim] (MQA).
799
- k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
800
+ # `k_pe` is [num_tokens, rope_dim] (MQA).
801
+ k = torch.cat([k_pe, k_nope], dim=-1)
802
803
# we only quant q here since k quant is fused with cache insertion
804
q = q.view(-1, self.head_dim)
@@ -807,12 +809,9 @@ class Indexer(nn.Module):
809
use_ue8m0=self.scale_fmt is not None,
810
)
811
q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
810
- q_scale = q_scale.view(-1, self.n_head, 1)
812
+ q_scale = q_scale.view(-1, self.n_head)
813
812
- weights = (
813
- weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
814
- )
815
- weights = weights.squeeze(-1)
814
+ weights = weights * q_scale * self.softmax_scale * self.n_head_scale
815
816
return self.indexer_op(hidden_states, q_fp8, k, weights)
817
@@ -1287,13 +1286,10 @@ class DeepseekV2DecoderLayer(nn.Module):
1286
hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
1287
hidden_states = hidden_states[:full_num_tokens]
1288
1290
- attn_kwargs = {
1291
- "positions": positions,
1292
- "hidden_states": hidden_states,
1293
- }
1294
- if not self.use_mha:
1295
- attn_kwargs["llama_4_scaling"] = llama_4_scaling
1296
- hidden_states = self.self_attn(**attn_kwargs)
1289
+ if self.use_mha:
1290
+ hidden_states = self.self_attn(positions, hidden_states)
1291
+ else:
1292
+ hidden_states = self.self_attn(positions, hidden_states, llama_4_scaling)
1293
1294
if (
1295
not isinstance(self.self_attn, DeepseekAttention)