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 14, 05:30 UTC
[DSA][MLA] Tiny refactor on DeepSeek to make it reusable for different backends (#26656)
Signed-off-by: MengqingCao <cmq0113@163.com>
Mengqing Cao committed
Oct 15, 2025 at 15:16 UTC
302ef403a2305e9158064f8e386d1b5284d12cb2
3 files changed
+12
-3
vllm/attention/layer.py
+2
@@ -587,6 +587,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
587
prefix: str = "",
588
use_sparse: bool = False,
589
indexer: object | None = None,
590
+ **extra_impl_args,
591
):
592
super().__init__()
593
self.num_heads = num_heads
@@ -639,6 +640,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
640
v_head_dim=self.v_head_dim,
641
kv_b_proj=kv_b_proj,
642
indexer=indexer,
643
+ **extra_impl_args,
644
)
645
646
self.use_direct_call = not current_platform.opaque_attention_op()
vllm/model_executor/models/deepseek_mtp.py
+8
-2
@@ -17,9 +17,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
17
VocabParallelEmbedding,
18
)
19
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
20
+from vllm.platforms import current_platform
21
from vllm.sequence import IntermediateTensors
22
22
-from .deepseek_v2 import DeepseekV2DecoderLayer, get_spec_layer_idx_from_weight_name
23
+from .deepseek_v2 import (
24
+ DeepseekV2DecoderLayer,
25
+ get_spec_layer_idx_from_weight_name,
26
+)
27
from .interfaces import SupportsPP
28
from .utils import maybe_prefix
29
@@ -56,6 +60,8 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
60
self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
61
self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
62
63
+ self.device = current_platform.device_type
64
+
65
self.is_v32 = hasattr(config, "index_topk")
66
if self.is_v32:
67
topk_tokens = config.index_topk
@@ -63,7 +69,7 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
69
vllm_config.scheduler_config.max_num_batched_tokens,
70
topk_tokens,
71
dtype=torch.int32,
66
- device="cuda",
72
+ device=self.device,
73
)
74
else:
75
topk_indices_buffer = None
vllm/model_executor/models/deepseek_v2.py
+2
-1
@@ -1165,6 +1165,7 @@ class DeepseekV2Model(nn.Module):
1165
config = vllm_config.model_config.hf_config
1166
quant_config = vllm_config.quant_config
1167
self.config = config
1168
+ self.device = current_platform.device_type
1169
1170
self.vocab_size = config.vocab_size
1171
self.is_v32 = hasattr(config, "index_topk")
@@ -1174,7 +1175,7 @@ class DeepseekV2Model(nn.Module):
1175
vllm_config.scheduler_config.max_num_batched_tokens,
1176
topk_tokens,
1177
dtype=torch.int32,
1177
- device="cuda",
1178
+ device=self.device,
1179
)
1180
else:
1181
topk_indices_buffer = None