From a04afd76aa91106f3e59206f984fa68cfacd5d9f Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 28 May 2026 01:00:52 -0700 Subject: [PATCH] [DSV4] Remove AMD/XPU path in deepseek_v4/nvidia (#43829) Signed-off-by: Woosuk Kwon --- vllm/models/deepseek_v4/nvidia/model.py | 67 ++----------------------- vllm/models/deepseek_v4/nvidia/mtp.py | 19 ++----- 2 files changed, 8 insertions(+), 78 deletions(-) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 25f0a730fd..d71983884b 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -60,7 +60,6 @@ from vllm.models.deepseek_v4.attention import ( DeepseekV4MultiHeadLatentAttentionWrapper, ) from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs -from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import direct_register_custom_op @@ -262,13 +261,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): return (sf.to(torch.int32) << 23).view(torch.float32) def _check_runtime_supported(self) -> None: - if not torch.cuda.is_available(): - raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.") device = self.w13_weight.device - if device.type != "cuda": - raise NotImplementedError( - "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA." - ) if torch.cuda.get_device_capability(device)[0] != 10: raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: @@ -955,7 +948,7 @@ class DeepseekV4DecoderLayer(nn.Module): ): return self.mhc_post(x, residual, post, comb) - def _forward_cuda( + def forward( self, x: torch.Tensor, positions: torch.Tensor, @@ -1024,52 +1017,6 @@ class DeepseekV4DecoderLayer(nn.Module): x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix - def _forward_native( - self, - x: torch.Tensor, - positions: torch.Tensor, - input_ids: torch.Tensor | None, - post_mix: torch.Tensor | None = None, - res_mix: torch.Tensor | None = None, - residual: torch.Tensor | None = None, - ) -> tuple[ - torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None - ]: - residual = x - x, post, comb = self.hc_pre( - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base - ) - x = self.attn_norm(x) - x = self.attn(positions, x, None) - x = self.hc_post(x, residual, post, comb) - - residual = x - x, post, comb = self.hc_pre( - x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base - ) - x = self.ffn_norm(x) - x = self.ffn(x, input_ids) - x = self.hc_post(x, residual, post, comb) - return x, None, None, None - - def forward( - self, - x: torch.Tensor, - positions: torch.Tensor, - input_ids: torch.Tensor | None, - post_mix: torch.Tensor | None = None, - res_mix: torch.Tensor | None = None, - residual: torch.Tensor | None = None, - ) -> tuple[ - torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None - ]: - if current_platform.is_rocm() or current_platform.is_xpu(): - return self._forward_native( - x, positions, input_ids, post_mix, res_mix, residual - ) - - return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual) - @support_torch_compile class DeepseekV4Model(nn.Module): @@ -1098,20 +1045,13 @@ class DeepseekV4Model(nn.Module): # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute # (compressor kv_score, indexer.weights_proj, indexer.compressor # kv_score). fused_wqa_wkv stays on the default stream. - # Disable them on ROCm / XPU because of hang issues / no overlap. - aux_stream_list = ( - None - if current_platform.is_rocm() or current_platform.is_xpu() - else [torch.cuda.Stream() for _ in range(3)] - ) + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] - self.device = current_platform.device_type # Reserved topk indices buffer for all Indexer layers to reuse. self.topk_indices_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, config.index_topk, dtype=torch.int32, - device=self.device, ) if get_pp_group().is_first_rank: @@ -1171,7 +1111,6 @@ class DeepseekV4Model(nn.Module): vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, dtype=vllm_config.model_config.dtype, - device=self.device, ) else: self._mtp_hidden_buffer = None @@ -1229,7 +1168,7 @@ class DeepseekV4Model(nn.Module): res_mix, residual, ) - if layer is not None and current_platform.is_cuda(): + if layer is not None: hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) if not get_pp_group().is_last_rank: diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 9d0ac43530..d194f7a186 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -37,7 +37,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name from vllm.model_executor.models.utils import maybe_prefix -from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from .model import ( @@ -147,10 +146,9 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): hidden_states, residual, post_mix, res_mix = self.mtp_block( positions=positions, x=hidden_states, input_ids=None ) - if current_platform.is_cuda(): - hidden_states = self.mtp_block.hc_post( - hidden_states, residual, post_mix, res_mix - ) + hidden_states = self.mtp_block.hc_post( + hidden_states, residual, post_mix, res_mix + ) # Return the flat pre-hc_head residual so it can be re-fed as the # next spec step's `previous_hidden_states` when # num_speculative_tokens > 1. hc_head is deferred to compute_logits. @@ -163,23 +161,16 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): config = vllm_config.model_config.hf_config self.mtp_start_layer_idx = config.num_hidden_layers self.num_mtp_layers = config.num_nextn_predict_layers - self.device = current_platform.device_type topk_tokens = config.index_topk self.topk_indices_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, topk_tokens, dtype=torch.int32, - device=self.device, ) - # Three aux streams shared across all MTP layers, mirroring - # DeepseekV4Model. ROCm runs the same work serially for now. - aux_stream_list = ( - None - if current_platform.is_rocm() - else [torch.cuda.Stream() for _ in range(3)] - ) + # Three aux streams shared across all MTP layers, mirroring DeepseekV4Model. + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] # to map the exact layer index from weights self.layers = torch.nn.ModuleDict(