iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[GLM Perf] DSv32/glm use skip topk for MTP case, 2.0x kernel performance improvement (#50904)

Browse Source
This commit is contained in:
Wentao Ye
2026-08-06 01:48:36 -07:00
committed by GitHub
parent 2dfb8ba590
commit ad5280255b
2 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/deepseek_v32/amd/rocm.py
+4 -3
View File
@@ -167,7 +167,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
assert isinstance(slot_mapping, dict)
mla_slot = slot_mapping.get(self.layer_name)
if self.indexer is not None:
if self.indexer is not None and not self.skip_topk:
has_indexer = True
indexer_k_norm_w = self.indexer.k_norm.weight
indexer_k_norm_bias = self.indexer.k_norm.bias
@@ -222,7 +222,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
ql_nope, q_pe = self._compute_ql_nope(q_c)
if self.indexer is not None:
if self.indexer is not None and not self.skip_topk:
index_q = self.indexer.wq_b(q_c)[0]
index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim)
else:
@@ -244,7 +244,8 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
quantize_mqa=self._fp8_kv,
)
self._run_indexer(q_c, index_q_fp8, index_weights_out)
if self.indexer is not None and not self.skip_topk:
self._run_indexer(q_c, index_q_fp8, index_weights_out)
if attn_metadata is None:
output.zero_()
vllm/models/deepseek_v32/attention.py
+3 -3
View File
@@ -372,7 +372,7 @@ class DeepseekV32Attention(MLAAttention):
assert isinstance(slot_mapping, dict)
mla_slot = slot_mapping.get(self.layer_name)
if self.indexer is not None:
if self.indexer is not None and not self.skip_topk:
has_indexer = True
indexer_k_norm_w = self.indexer.k_norm.weight
indexer_k_norm_bias = self.indexer.k_norm.bias
@@ -430,7 +430,7 @@ class DeepseekV32Attention(MLAAttention):
q_nope = q_nope.transpose(0, 1)
ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1)
if self.indexer is not None:
if self.indexer is not None and not self.skip_topk:
index_q = self.indexer.wq_b(q_c)[0]
index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim)
else:
@@ -452,7 +452,7 @@ class DeepseekV32Attention(MLAAttention):
quantize_mqa=self._fp8_query,
)
if self.indexer is not None:
if self.indexer is not None and not self.skip_topk:
sparse_attn_indexer(
q_c,
self.indexer.k_cache.prefix,