[Model Refactoring] Remove torch compile dependency in DSv4 (#43746)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

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 "

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__()

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__()

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",
 ]

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

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__()

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__()

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