[GLM Perf] DSv32/glm use skip topk for MTP case, 2.0x kernel performance improvement (#50904)
Wentao Ye committed
Aug 6, 2026 at 04:48 UTC
ad5280255ba6e37dd11cf697bd7145d2a90213ff
2 files changed
+7
-6
vllm/models/deepseek_v32/amd/rocm.py
+4
-3
@@ -167,7 +167,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
167
assert isinstance(slot_mapping, dict)
168
mla_slot = slot_mapping.get(self.layer_name)
169
170
- if self.indexer is not None:
170
+ if self.indexer is not None and not self.skip_topk:
171
has_indexer = True
172
indexer_k_norm_w = self.indexer.k_norm.weight
173
indexer_k_norm_bias = self.indexer.k_norm.bias
@@ -222,7 +222,7 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
222
223
ql_nope, q_pe = self._compute_ql_nope(q_c)
224
225
- if self.indexer is not None:
225
+ if self.indexer is not None and not self.skip_topk:
226
index_q = self.indexer.wq_b(q_c)[0]
227
index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim)
228
else:
@@ -244,7 +244,8 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
244
quantize_mqa=self._fp8_kv,
245
)
246
247
- self._run_indexer(q_c, index_q_fp8, index_weights_out)
247
+ if self.indexer is not None and not self.skip_topk:
248
+ self._run_indexer(q_c, index_q_fp8, index_weights_out)
249
250
if attn_metadata is None:
251
output.zero_()
vllm/models/deepseek_v32/attention.py
+3
-3
@@ -372,7 +372,7 @@ class DeepseekV32Attention(MLAAttention):
372
assert isinstance(slot_mapping, dict)
373
mla_slot = slot_mapping.get(self.layer_name)
374
375
- if self.indexer is not None:
375
+ if self.indexer is not None and not self.skip_topk:
376
has_indexer = True
377
indexer_k_norm_w = self.indexer.k_norm.weight
378
indexer_k_norm_bias = self.indexer.k_norm.bias
@@ -430,7 +430,7 @@ class DeepseekV32Attention(MLAAttention):
430
q_nope = q_nope.transpose(0, 1)
431
ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1)
432
433
- if self.indexer is not None:
433
+ if self.indexer is not None and not self.skip_topk:
434
index_q = self.indexer.wq_b(q_c)[0]
435
index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim)
436
else:
@@ -452,7 +452,7 @@ class DeepseekV32Attention(MLAAttention):
452
quantize_mqa=self._fp8_query,
453
)
454
455
- if self.indexer is not None:
455
+ if self.indexer is not None and not self.skip_topk:
456
sparse_attn_indexer(
457
q_c,
458
self.indexer.k_cache.prefix,