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

[DSV4] Remove AMD/XPU path in deepseek_v4/nvidia (#43829)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
This commit is contained in:
Woosuk Kwon
2026-05-28 08:00:52 +00:00
committed by GitHub
parent 6cc8577421
commit a04afd76aa
2 changed files with 8 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/deepseek_v4/nvidia/model.py
+3 -64
View File
@@ -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:
vllm/models/deepseek_v4/nvidia/mtp.py
+5 -14
View File
@@ -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(