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 2, 14:45 UTC
[DSV4] Remove AMD/XPU path in deepseek_v4/nvidia (#43829)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 28, 2026 at 01:00 UTC
a04afd76aa91106f3e59206f984fa68cfacd5d9f
2 files changed
+8
-78
vllm/models/deepseek_v4/nvidia/model.py
+3
-64
@@ -60,7 +60,6 @@ from vllm.models.deepseek_v4.attention import (
60
DeepseekV4MultiHeadLatentAttentionWrapper,
61
)
62
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
63
-from vllm.platforms import current_platform
63
from vllm.sequence import IntermediateTensors
64
from vllm.utils.torch_utils import direct_register_custom_op
65
@@ -262,13 +261,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
261
return (sf.to(torch.int32) << 23).view(torch.float32)
262
263
def _check_runtime_supported(self) -> None:
265
- if not torch.cuda.is_available():
266
- raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.")
264
device = self.w13_weight.device
268
- if device.type != "cuda":
269
- raise NotImplementedError(
270
- "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA."
271
- )
265
if torch.cuda.get_device_capability(device)[0] != 10:
266
raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.")
267
if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0:
@@ -955,7 +948,7 @@ class DeepseekV4DecoderLayer(nn.Module):
948
):
949
return self.mhc_post(x, residual, post, comb)
950
958
- def _forward_cuda(
951
+ def forward(
952
self,
953
x: torch.Tensor,
954
positions: torch.Tensor,
@@ -1024,52 +1017,6 @@ class DeepseekV4DecoderLayer(nn.Module):
1017
x = self.ffn(x, input_ids)
1018
return x, residual, post_mix, res_mix
1019
1027
- def _forward_native(
1028
- self,
1029
- x: torch.Tensor,
1030
- positions: torch.Tensor,
1031
- input_ids: torch.Tensor | None,
1032
- post_mix: torch.Tensor | None = None,
1033
- res_mix: torch.Tensor | None = None,
1034
- residual: torch.Tensor | None = None,
1035
- ) -> tuple[
1036
- torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
1037
- ]:
1038
- residual = x
1039
- x, post, comb = self.hc_pre(
1040
- x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1041
- )
1042
- x = self.attn_norm(x)
1043
- x = self.attn(positions, x, None)
1044
- x = self.hc_post(x, residual, post, comb)
1045
-
1046
- residual = x
1047
- x, post, comb = self.hc_pre(
1048
- x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
1049
- )
1050
- x = self.ffn_norm(x)
1051
- x = self.ffn(x, input_ids)
1052
- x = self.hc_post(x, residual, post, comb)
1053
- return x, None, None, None
1054
-
1055
- def forward(
1056
- self,
1057
- x: torch.Tensor,
1058
- positions: torch.Tensor,
1059
- input_ids: torch.Tensor | None,
1060
- post_mix: torch.Tensor | None = None,
1061
- res_mix: torch.Tensor | None = None,
1062
- residual: torch.Tensor | None = None,
1063
- ) -> tuple[
1064
- torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
1065
- ]:
1066
- if current_platform.is_rocm() or current_platform.is_xpu():
1067
- return self._forward_native(
1068
- x, positions, input_ids, post_mix, res_mix, residual
1069
- )
1070
-
1071
- return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual)
1072
-
1020
1021
@support_torch_compile
1022
class DeepseekV4Model(nn.Module):
@@ -1098,20 +1045,13 @@ class DeepseekV4Model(nn.Module):
1045
# DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute
1046
# (compressor kv_score, indexer.weights_proj, indexer.compressor
1047
# kv_score). fused_wqa_wkv stays on the default stream.
1101
- # Disable them on ROCm / XPU because of hang issues / no overlap.
1102
- aux_stream_list = (
1103
- None
1104
- if current_platform.is_rocm() or current_platform.is_xpu()
1105
- else [torch.cuda.Stream() for _ in range(3)]
1106
- )
1048
+ aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
1049
1108
- self.device = current_platform.device_type
1050
# Reserved topk indices buffer for all Indexer layers to reuse.
1051
self.topk_indices_buffer = torch.empty(
1052
vllm_config.scheduler_config.max_num_batched_tokens,
1053
config.index_topk,
1054
dtype=torch.int32,
1114
- device=self.device,
1055
)
1056
1057
if get_pp_group().is_first_rank:
@@ -1171,7 +1111,6 @@ class DeepseekV4Model(nn.Module):
1111
vllm_config.scheduler_config.max_num_batched_tokens,
1112
self.hc_dim,
1113
dtype=vllm_config.model_config.dtype,
1174
- device=self.device,
1114
)
1115
else:
1116
self._mtp_hidden_buffer = None
@@ -1229,7 +1168,7 @@ class DeepseekV4Model(nn.Module):
1168
res_mix,
1169
residual,
1170
)
1232
- if layer is not None and current_platform.is_cuda():
1171
+ if layer is not None:
1172
hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
1173
1174
if not get_pp_group().is_last_rank:
vllm/models/deepseek_v4/nvidia/mtp.py
+5
-14
@@ -37,7 +37,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37
from vllm.model_executor.models.deepseek_mtp import SharedHead
38
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
39
from vllm.model_executor.models.utils import maybe_prefix
40
-from vllm.platforms import current_platform
40
from vllm.sequence import IntermediateTensors
41
42
from .model import (
@@ -147,10 +146,9 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
146
hidden_states, residual, post_mix, res_mix = self.mtp_block(
147
positions=positions, x=hidden_states, input_ids=None
148
)
150
- if current_platform.is_cuda():
151
- hidden_states = self.mtp_block.hc_post(
152
- hidden_states, residual, post_mix, res_mix
153
- )
149
+ hidden_states = self.mtp_block.hc_post(
150
+ hidden_states, residual, post_mix, res_mix
151
+ )
152
# Return the flat pre-hc_head residual so it can be re-fed as the
153
# next spec step's `previous_hidden_states` when
154
# num_speculative_tokens > 1. hc_head is deferred to compute_logits.
@@ -163,23 +161,16 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
161
config = vllm_config.model_config.hf_config
162
self.mtp_start_layer_idx = config.num_hidden_layers
163
self.num_mtp_layers = config.num_nextn_predict_layers
166
- self.device = current_platform.device_type
164
165
topk_tokens = config.index_topk
166
self.topk_indices_buffer = torch.empty(
167
vllm_config.scheduler_config.max_num_batched_tokens,
168
topk_tokens,
169
dtype=torch.int32,
173
- device=self.device,
170
)
171
176
- # Three aux streams shared across all MTP layers, mirroring
177
- # DeepseekV4Model. ROCm runs the same work serially for now.
178
- aux_stream_list = (
179
- None
180
- if current_platform.is_rocm()
181
- else [torch.cuda.Stream() for _ in range(3)]
182
- )
172
+ # Three aux streams shared across all MTP layers, mirroring DeepseekV4Model.
173
+ aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
174
175
# to map the exact layer index from weights
176
self.layers = torch.nn.ModuleDict(