From 9957e4d240aa5fc48014d4ed7d7ada614ba53302 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 28 May 2026 07:26:25 -0700 Subject: [PATCH] [Model Refactoring] Remove torch compile dependency in DSv4 (#43746) Signed-off-by: Woosuk Kwon --- vllm/config/vllm.py | 17 ++ vllm/models/deepseek_v4/amd/model.py | 2 - vllm/models/deepseek_v4/amd/mtp.py | 31 ++- .../models/deepseek_v4/common/ops/__init__.py | 3 + .../common/ops/fused_mtp_input_rmsnorm.py | 203 ++++++++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 2 - vllm/models/deepseek_v4/nvidia/mtp.py | 31 ++- vllm/v1/worker/gpu_model_runner.py | 5 + 8 files changed, 270 insertions(+), 24 deletions(-) create mode 100644 vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 9e71d7b47e..4dacf4980e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1044,6 +1044,23 @@ class VllmConfig: ) self.compilation_config.mode = CompilationMode.NONE + # DeepSeek V4's model classes don't carry @support_torch_compile — + # the breakable cudagraph is the supported PIECEWISE path. Auto-enable + # it unless the user has explicitly opted out via the env var. + if ( + self.model_config is not None + and "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ + and any( + a in ("DeepseekV4ForCausalLM", "DeepSeekV4MTPModel") + for a in self.model_config.architectures + ) + ): + os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1" + logger.info_once( + "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1 for DeepSeek V4. " + "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out." + ) + if envs.VLLM_USE_BREAKABLE_CUDAGRAPH: logger.warning_once( "VLLM_USE_BREAKABLE_CUDAGRAPH is set, disabling vLLM's " diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 0a829c117e..28836a2b14 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -8,7 +8,6 @@ import regex as re import torch import torch.nn as nn -from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( get_pp_group, @@ -605,7 +604,6 @@ class DeepseekV4DecoderLayer(nn.Module): ) -@support_torch_compile class DeepseekV4Model(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index 2f83af5ad9..5938cde695 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -18,7 +18,6 @@ import regex as re import torch import torch.nn as nn -from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -37,6 +36,10 @@ 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.models.deepseek_v4.common.ops import ( + fused_mtp_input_rmsnorm, + mtp_shared_head_rmsnorm, +) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.import_utils import has_tilelang @@ -130,16 +133,22 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - # masking inputs at position 0, as not needed by MTP - inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) - inputs_embeds = self.enorm(inputs_embeds) - # Target stashes pre-hc_head residual as flat (T, hc_mult * D); - # reshape to (T, hc_mult, D) — the training-time layout. + # reshape to (T, hc_mult, D) — the training-time layout — before + # the fused norm pass so both inputs are 3D-friendly. previous_hidden_states = previous_hidden_states.view( -1, self.hc_mult, self.config.hidden_size ) - previous_hidden_states = self.hnorm(previous_hidden_states) + # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm. + inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm( + inputs_embeds, + positions, + previous_hidden_states, + self.enorm.weight.data, + self.hnorm.weight.data, + self.enorm.variance_epsilon, + self.hc_mult, + ) hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( inputs_embeds ).unsqueeze(-2) @@ -244,13 +253,15 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): mtp_layer.rms_norm_eps, mtp_layer.hc_eps, ) - logits = self.logits_processor( - mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states) + hidden_states = mtp_shared_head_rmsnorm( + hidden_states, + mtp_layer.shared_head.norm.weight.data, + mtp_layer.shared_head.norm.variance_epsilon, ) + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) return logits -@support_torch_compile class DeepSeekV4MTP(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() diff --git a/vllm/models/deepseek_v4/common/ops/__init__.py b/vllm/models/deepseek_v4/common/ops/__init__.py index 0f80c32974..dc6f3c608d 100644 --- a/vllm/models/deepseek_v4/common/ops/__init__.py +++ b/vllm/models/deepseek_v4/common/ops/__init__.py @@ -9,6 +9,7 @@ from .cache_utils import ( ) from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant +from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm from .fused_qk_rmsnorm import fused_q_kv_rmsnorm from .save_partial_states import save_partial_states @@ -19,7 +20,9 @@ __all__ = [ "dequantize_and_gather_k_cache", "fused_indexer_q_rope_quant", "fused_inv_rope_fp8_quant", + "fused_mtp_input_rmsnorm", "fused_q_kv_rmsnorm", + "mtp_shared_head_rmsnorm", "quantize_and_insert_k_cache", "save_partial_states", ] diff --git a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py new file mode 100644 index 0000000000..874b4f6079 --- /dev/null +++ b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused MTP-input RMSNorm: enorm (with mask-zero at position 0) + hnorm. + +Replaces the eager sequence at the top of the MTP draft forward: + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + inputs_embeds = self.enorm(inputs_embeds) + previous_hidden_states = previous_hidden_states.view(-1, hc_mult, H) + previous_hidden_states = self.hnorm(previous_hidden_states) + +which lowers to ~6 small kernels (CompareEq, where, Fill, enorm rms_norm, +hnorm rms_norm, plus aten elementwise helpers) on the breakable-cudagraph +path. Math is preserved: positions==0 → masked row → zero RMS output +regardless of weight. + +A single grid (T, hc_mult+1) drives both norms: task 0 is enorm on +inputs_embeds[token, :], task k+1 is hnorm on previous_hidden_states[token, k, :]. +""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _rmsnorm_row( + x, + w_ptr, + out_row_ptr, + block, + mask, + eps, + HIDDEN: tl.constexpr, +): + x = x.to(tl.float32) + variance = tl.sum(x * x, axis=0) / HIDDEN + rrms = tl.rsqrt(variance + eps) + w = tl.load(w_ptr + block, mask=mask, other=0.0).to(tl.float32) + y = x * rrms * w + tl.store(out_row_ptr + block, y.to(out_row_ptr.dtype.element_ty), mask=mask) + + +@triton.jit +def _fused_mtp_input_rmsnorm_kernel( + inputs_embeds_ptr, + positions_ptr, + prev_hidden_ptr, + enorm_weight_ptr, + hnorm_weight_ptr, + enorm_out_ptr, + hnorm_out_ptr, + eps, + HIDDEN: tl.constexpr, + HC_MULT: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # int64 token index so per-token offsets don't overflow int32 at + # large num_tokens (matches the convention in fused_q_kv_rmsnorm). + token_idx = tl.program_id(0).to(tl.int64) + pid_task = tl.program_id(1) + + block = tl.arange(0, BLOCK_SIZE) + mask = block < HIDDEN + + if pid_task == 0: + # enorm path: load inputs_embeds[token, :] then zero-mask at pos==0. + # Math is preserved: pos==0 → x=0 → variance=0 → RMSNorm output is 0 + # regardless of weight, matching torch.where(pos==0, 0, x) + RMSNorm. + pos = tl.load(positions_ptr + token_idx) + keep = pos != 0 + x = tl.load( + inputs_embeds_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0 + ) + x = tl.where(keep, x, 0.0) + _rmsnorm_row( + x, + enorm_weight_ptr, + enorm_out_ptr + token_idx * HIDDEN, + block, + mask, + eps, + HIDDEN, + ) + else: + # hnorm path: load prev_hidden[token, slot, :]. + slot = pid_task - 1 + row_offset = (token_idx * HC_MULT + slot) * HIDDEN + x = tl.load(prev_hidden_ptr + row_offset + block, mask=mask, other=0.0) + _rmsnorm_row( + x, + hnorm_weight_ptr, + hnorm_out_ptr + row_offset, + block, + mask, + eps, + HIDDEN, + ) + + +@triton.jit +def _mtp_shared_head_rmsnorm_kernel( + x_ptr, + weight_ptr, + out_ptr, + eps, + HIDDEN: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + block = tl.arange(0, BLOCK_SIZE) + mask = block < HIDDEN + x = tl.load(x_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0) + _rmsnorm_row( + x, + weight_ptr, + out_ptr + token_idx * HIDDEN, + block, + mask, + eps, + HIDDEN, + ) + + +def mtp_shared_head_rmsnorm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """RMSNorm for MTP's SharedHead.norm, on (T, H) bf16 input. + + Uses the same ``_rmsnorm_row`` body as ``fused_mtp_input_rmsnorm`` so the + MTP draft path runs one consistent RMSNorm implementation end to end. + """ + assert hidden_states.ndim == 2 + assert hidden_states.is_contiguous() + assert weight.is_contiguous() + num_tokens, hidden = hidden_states.shape + out = torch.empty_like(hidden_states) + if num_tokens == 0: + return out + block_size = triton.next_power_of_2(hidden) + _mtp_shared_head_rmsnorm_kernel[(num_tokens,)]( + hidden_states, + weight, + out, + eps, + HIDDEN=hidden, + BLOCK_SIZE=block_size, + ) + return out + + +def fused_mtp_input_rmsnorm( + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, + hc_mult: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Returns (enorm_out, hnorm_out). + + enorm_out has the same shape as inputs_embeds (2D, [T, H]). + hnorm_out has the same shape as previous_hidden_states (3D, [T, hc_mult, H]). + previous_hidden_states must already be reshaped to 3D. + """ + assert inputs_embeds.ndim == 2 + assert previous_hidden_states.ndim == 3 + assert previous_hidden_states.shape[1] == hc_mult + assert inputs_embeds.shape[0] == previous_hidden_states.shape[0], ( + "token dim mismatch" + ) + assert ( + inputs_embeds.shape[1] + == previous_hidden_states.shape[2] + == enorm_weight.shape[0] + == hnorm_weight.shape[0] + ) + assert inputs_embeds.is_contiguous() and previous_hidden_states.is_contiguous() + assert enorm_weight.is_contiguous() and hnorm_weight.is_contiguous() + + num_tokens, hidden = inputs_embeds.shape + enorm_out = torch.empty_like(inputs_embeds) + hnorm_out = torch.empty_like(previous_hidden_states) + if num_tokens == 0: + return enorm_out, hnorm_out + + block_size = triton.next_power_of_2(hidden) + _fused_mtp_input_rmsnorm_kernel[(num_tokens, hc_mult + 1)]( + inputs_embeds, + positions, + previous_hidden_states, + enorm_weight, + hnorm_weight, + enorm_out, + hnorm_out, + eps, + HIDDEN=hidden, + HC_MULT=hc_mult, + BLOCK_SIZE=block_size, + ) + return enorm_out, hnorm_out diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index d71983884b..a5f4bae842 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -8,7 +8,6 @@ import regex as re import torch import torch.nn as nn -from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( get_ep_group, @@ -1018,7 +1017,6 @@ class DeepseekV4DecoderLayer(nn.Module): return x, residual, post_mix, res_mix -@support_torch_compile class DeepseekV4Model(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index d194f7a186..3db831f026 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -18,7 +18,6 @@ import regex as re import torch import torch.nn as nn -from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -37,6 +36,10 @@ 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.models.deepseek_v4.common.ops import ( + fused_mtp_input_rmsnorm, + mtp_shared_head_rmsnorm, +) from vllm.sequence import IntermediateTensors from .model import ( @@ -130,16 +133,22 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - # masking inputs at position 0, as not needed by MTP - inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) - inputs_embeds = self.enorm(inputs_embeds) - # Target stashes pre-hc_head residual as flat (T, hc_mult * D); - # reshape to (T, hc_mult, D) — the training-time layout. + # reshape to (T, hc_mult, D) — the training-time layout — before + # the fused norm pass so both inputs are 3D-friendly. previous_hidden_states = previous_hidden_states.view( -1, self.hc_mult, self.config.hidden_size ) - previous_hidden_states = self.hnorm(previous_hidden_states) + # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm. + inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm( + inputs_embeds, + positions, + previous_hidden_states, + self.enorm.weight.data, + self.hnorm.weight.data, + self.enorm.variance_epsilon, + self.hc_mult, + ) hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( inputs_embeds ).unsqueeze(-2) @@ -236,13 +245,15 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): mtp_layer.rms_norm_eps, mtp_layer.hc_eps, ) - logits = self.logits_processor( - mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states) + hidden_states = mtp_shared_head_rmsnorm( + hidden_states, + mtp_layer.shared_head.norm.weight.data, + mtp_layer.shared_head.norm.variance_epsilon, ) + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) return logits -@support_torch_compile class DeepSeekV4MTP(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index f8f256944a..890abd53fa 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5182,6 +5182,11 @@ class GPUModelRunner( and not self.parallel_config.use_ubatching ): self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config) + drafter = getattr(self, "drafter", None) + if drafter is not None and hasattr(drafter, "model"): + drafter.model = BreakableCUDAGraphWrapper( + drafter.model, self.vllm_config + ) elif ( cudagraph_mode.has_full_cudagraphs() and not self.parallel_config.use_ubatching