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, 13:30 UTC
[Model Refactoring] Remove torch compile dependency in DSv4 (#43746)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 28, 2026 at 07:26 UTC
9957e4d240aa5fc48014d4ed7d7ada614ba53302
8 files changed
+270
-24
vllm/config/vllm.py
+17
@@ -1044,6 +1044,23 @@ class VllmConfig:
1044
)
1045
self.compilation_config.mode = CompilationMode.NONE
1046
1047
+ # DeepSeek V4's model classes don't carry @support_torch_compile —
1048
+ # the breakable cudagraph is the supported PIECEWISE path. Auto-enable
1049
+ # it unless the user has explicitly opted out via the env var.
1050
+ if (
1051
+ self.model_config is not None
1052
+ and "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
1053
+ and any(
1054
+ a in ("DeepseekV4ForCausalLM", "DeepSeekV4MTPModel")
1055
+ for a in self.model_config.architectures
1056
+ )
1057
+ ):
1058
+ os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
1059
+ logger.info_once(
1060
+ "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1 for DeepSeek V4. "
1061
+ "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
1062
+ )
1063
+
1064
if envs.VLLM_USE_BREAKABLE_CUDAGRAPH:
1065
logger.warning_once(
1066
"VLLM_USE_BREAKABLE_CUDAGRAPH is set, disabling vLLM's "
vllm/models/deepseek_v4/amd/model.py
-2
@@ -8,7 +8,6 @@ import regex as re
8
import torch
9
import torch.nn as nn
10
11
-from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
from vllm.distributed import (
13
get_pp_group,
@@ -605,7 +604,6 @@ class DeepseekV4DecoderLayer(nn.Module):
604
)
605
606
608
-@support_torch_compile
607
class DeepseekV4Model(nn.Module):
608
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
609
super().__init__()
vllm/models/deepseek_v4/amd/mtp.py
+21
-10
@@ -18,7 +18,6 @@ import regex as re
18
import torch
19
import torch.nn as nn
20
21
-from vllm.compilation.decorators import support_torch_compile
21
from vllm.config import VllmConfig
22
from vllm.distributed import (
23
get_tensor_model_parallel_rank,
@@ -37,6 +36,10 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36
from vllm.model_executor.models.deepseek_mtp import SharedHead
37
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
38
from vllm.model_executor.models.utils import maybe_prefix
39
+from vllm.models.deepseek_v4.common.ops import (
40
+ fused_mtp_input_rmsnorm,
41
+ mtp_shared_head_rmsnorm,
42
+)
43
from vllm.platforms import current_platform
44
from vllm.sequence import IntermediateTensors
45
from vllm.utils.import_utils import has_tilelang
@@ -130,16 +133,22 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
133
spec_step_index: int = 0,
134
) -> torch.Tensor:
135
assert inputs_embeds is not None
133
- # masking inputs at position 0, as not needed by MTP
134
- inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
135
- inputs_embeds = self.enorm(inputs_embeds)
136
-
136
# Target stashes pre-hc_head residual as flat (T, hc_mult * D);
138
- # reshape to (T, hc_mult, D) — the training-time layout.
137
+ # reshape to (T, hc_mult, D) — the training-time layout — before
138
+ # the fused norm pass so both inputs are 3D-friendly.
139
previous_hidden_states = previous_hidden_states.view(
140
-1, self.hc_mult, self.config.hidden_size
141
)
142
- previous_hidden_states = self.hnorm(previous_hidden_states)
142
+ # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm.
143
+ inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm(
144
+ inputs_embeds,
145
+ positions,
146
+ previous_hidden_states,
147
+ self.enorm.weight.data,
148
+ self.hnorm.weight.data,
149
+ self.enorm.variance_epsilon,
150
+ self.hc_mult,
151
+ )
152
hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
153
inputs_embeds
154
).unsqueeze(-2)
@@ -244,13 +253,15 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
253
mtp_layer.rms_norm_eps,
254
mtp_layer.hc_eps,
255
)
247
- logits = self.logits_processor(
248
- mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
256
+ hidden_states = mtp_shared_head_rmsnorm(
257
+ hidden_states,
258
+ mtp_layer.shared_head.norm.weight.data,
259
+ mtp_layer.shared_head.norm.variance_epsilon,
260
)
261
+ logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
262
return logits
263
264
253
-@support_torch_compile
265
class DeepSeekV4MTP(nn.Module):
266
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
267
super().__init__()
vllm/models/deepseek_v4/common/ops/__init__.py
+3
@@ -9,6 +9,7 @@ from .cache_utils import (
9
)
10
from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
11
from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
12
+from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm
13
from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
14
from .save_partial_states import save_partial_states
15
@@ -19,7 +20,9 @@ __all__ = [
20
"dequantize_and_gather_k_cache",
21
"fused_indexer_q_rope_quant",
22
"fused_inv_rope_fp8_quant",
23
+ "fused_mtp_input_rmsnorm",
24
"fused_q_kv_rmsnorm",
25
+ "mtp_shared_head_rmsnorm",
26
"quantize_and_insert_k_cache",
27
"save_partial_states",
28
]
vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py
new
+203
@@ -0,0 +1,203 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Fused MTP-input RMSNorm: enorm (with mask-zero at position 0) + hnorm.
4
+
5
+Replaces the eager sequence at the top of the MTP draft forward:
6
+ inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
7
+ inputs_embeds = self.enorm(inputs_embeds)
8
+ previous_hidden_states = previous_hidden_states.view(-1, hc_mult, H)
9
+ previous_hidden_states = self.hnorm(previous_hidden_states)
10
+
11
+which lowers to ~6 small kernels (CompareEq, where, Fill, enorm rms_norm,
12
+hnorm rms_norm, plus aten elementwise helpers) on the breakable-cudagraph
13
+path. Math is preserved: positions==0 → masked row → zero RMS output
14
+regardless of weight.
15
+
16
+A single grid (T, hc_mult+1) drives both norms: task 0 is enorm on
17
+inputs_embeds[token, :], task k+1 is hnorm on previous_hidden_states[token, k, :].
18
+"""
19
+
20
+import torch
21
+
22
+from vllm.triton_utils import tl, triton
23
+
24
+
25
+@triton.jit
26
+def _rmsnorm_row(
27
+ x,
28
+ w_ptr,
29
+ out_row_ptr,
30
+ block,
31
+ mask,
32
+ eps,
33
+ HIDDEN: tl.constexpr,
34
+):
35
+ x = x.to(tl.float32)
36
+ variance = tl.sum(x * x, axis=0) / HIDDEN
37
+ rrms = tl.rsqrt(variance + eps)
38
+ w = tl.load(w_ptr + block, mask=mask, other=0.0).to(tl.float32)
39
+ y = x * rrms * w
40
+ tl.store(out_row_ptr + block, y.to(out_row_ptr.dtype.element_ty), mask=mask)
41
+
42
+
43
+@triton.jit
44
+def _fused_mtp_input_rmsnorm_kernel(
45
+ inputs_embeds_ptr,
46
+ positions_ptr,
47
+ prev_hidden_ptr,
48
+ enorm_weight_ptr,
49
+ hnorm_weight_ptr,
50
+ enorm_out_ptr,
51
+ hnorm_out_ptr,
52
+ eps,
53
+ HIDDEN: tl.constexpr,
54
+ HC_MULT: tl.constexpr,
55
+ BLOCK_SIZE: tl.constexpr,
56
+):
57
+ # int64 token index so per-token offsets don't overflow int32 at
58
+ # large num_tokens (matches the convention in fused_q_kv_rmsnorm).
59
+ token_idx = tl.program_id(0).to(tl.int64)
60
+ pid_task = tl.program_id(1)
61
+
62
+ block = tl.arange(0, BLOCK_SIZE)
63
+ mask = block < HIDDEN
64
+
65
+ if pid_task == 0:
66
+ # enorm path: load inputs_embeds[token, :] then zero-mask at pos==0.
67
+ # Math is preserved: pos==0 → x=0 → variance=0 → RMSNorm output is 0
68
+ # regardless of weight, matching torch.where(pos==0, 0, x) + RMSNorm.
69
+ pos = tl.load(positions_ptr + token_idx)
70
+ keep = pos != 0
71
+ x = tl.load(
72
+ inputs_embeds_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0
73
+ )
74
+ x = tl.where(keep, x, 0.0)
75
+ _rmsnorm_row(
76
+ x,
77
+ enorm_weight_ptr,
78
+ enorm_out_ptr + token_idx * HIDDEN,
79
+ block,
80
+ mask,
81
+ eps,
82
+ HIDDEN,
83
+ )
84
+ else:
85
+ # hnorm path: load prev_hidden[token, slot, :].
86
+ slot = pid_task - 1
87
+ row_offset = (token_idx * HC_MULT + slot) * HIDDEN
88
+ x = tl.load(prev_hidden_ptr + row_offset + block, mask=mask, other=0.0)
89
+ _rmsnorm_row(
90
+ x,
91
+ hnorm_weight_ptr,
92
+ hnorm_out_ptr + row_offset,
93
+ block,
94
+ mask,
95
+ eps,
96
+ HIDDEN,
97
+ )
98
+
99
+
100
+@triton.jit
101
+def _mtp_shared_head_rmsnorm_kernel(
102
+ x_ptr,
103
+ weight_ptr,
104
+ out_ptr,
105
+ eps,
106
+ HIDDEN: tl.constexpr,
107
+ BLOCK_SIZE: tl.constexpr,
108
+):
109
+ token_idx = tl.program_id(0).to(tl.int64)
110
+ block = tl.arange(0, BLOCK_SIZE)
111
+ mask = block < HIDDEN
112
+ x = tl.load(x_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0)
113
+ _rmsnorm_row(
114
+ x,
115
+ weight_ptr,
116
+ out_ptr + token_idx * HIDDEN,
117
+ block,
118
+ mask,
119
+ eps,
120
+ HIDDEN,
121
+ )
122
+
123
+
124
+def mtp_shared_head_rmsnorm(
125
+ hidden_states: torch.Tensor,
126
+ weight: torch.Tensor,
127
+ eps: float,
128
+) -> torch.Tensor:
129
+ """RMSNorm for MTP's SharedHead.norm, on (T, H) bf16 input.
130
+
131
+ Uses the same ``_rmsnorm_row`` body as ``fused_mtp_input_rmsnorm`` so the
132
+ MTP draft path runs one consistent RMSNorm implementation end to end.
133
+ """
134
+ assert hidden_states.ndim == 2
135
+ assert hidden_states.is_contiguous()
136
+ assert weight.is_contiguous()
137
+ num_tokens, hidden = hidden_states.shape
138
+ out = torch.empty_like(hidden_states)
139
+ if num_tokens == 0:
140
+ return out
141
+ block_size = triton.next_power_of_2(hidden)
142
+ _mtp_shared_head_rmsnorm_kernel[(num_tokens,)](
143
+ hidden_states,
144
+ weight,
145
+ out,
146
+ eps,
147
+ HIDDEN=hidden,
148
+ BLOCK_SIZE=block_size,
149
+ )
150
+ return out
151
+
152
+
153
+def fused_mtp_input_rmsnorm(
154
+ inputs_embeds: torch.Tensor,
155
+ positions: torch.Tensor,
156
+ previous_hidden_states: torch.Tensor,
157
+ enorm_weight: torch.Tensor,
158
+ hnorm_weight: torch.Tensor,
159
+ eps: float,
160
+ hc_mult: int,
161
+) -> tuple[torch.Tensor, torch.Tensor]:
162
+ """Returns (enorm_out, hnorm_out).
163
+
164
+ enorm_out has the same shape as inputs_embeds (2D, [T, H]).
165
+ hnorm_out has the same shape as previous_hidden_states (3D, [T, hc_mult, H]).
166
+ previous_hidden_states must already be reshaped to 3D.
167
+ """
168
+ assert inputs_embeds.ndim == 2
169
+ assert previous_hidden_states.ndim == 3
170
+ assert previous_hidden_states.shape[1] == hc_mult
171
+ assert inputs_embeds.shape[0] == previous_hidden_states.shape[0], (
172
+ "token dim mismatch"
173
+ )
174
+ assert (
175
+ inputs_embeds.shape[1]
176
+ == previous_hidden_states.shape[2]
177
+ == enorm_weight.shape[0]
178
+ == hnorm_weight.shape[0]
179
+ )
180
+ assert inputs_embeds.is_contiguous() and previous_hidden_states.is_contiguous()
181
+ assert enorm_weight.is_contiguous() and hnorm_weight.is_contiguous()
182
+
183
+ num_tokens, hidden = inputs_embeds.shape
184
+ enorm_out = torch.empty_like(inputs_embeds)
185
+ hnorm_out = torch.empty_like(previous_hidden_states)
186
+ if num_tokens == 0:
187
+ return enorm_out, hnorm_out
188
+
189
+ block_size = triton.next_power_of_2(hidden)
190
+ _fused_mtp_input_rmsnorm_kernel[(num_tokens, hc_mult + 1)](
191
+ inputs_embeds,
192
+ positions,
193
+ previous_hidden_states,
194
+ enorm_weight,
195
+ hnorm_weight,
196
+ enorm_out,
197
+ hnorm_out,
198
+ eps,
199
+ HIDDEN=hidden,
200
+ HC_MULT=hc_mult,
201
+ BLOCK_SIZE=block_size,
202
+ )
203
+ return enorm_out, hnorm_out
vllm/models/deepseek_v4/nvidia/model.py
-2
@@ -8,7 +8,6 @@ import regex as re
8
import torch
9
import torch.nn as nn
10
11
-from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
from vllm.distributed import (
13
get_ep_group,
@@ -1018,7 +1017,6 @@ class DeepseekV4DecoderLayer(nn.Module):
1017
return x, residual, post_mix, res_mix
1018
1019
1021
-@support_torch_compile
1020
class DeepseekV4Model(nn.Module):
1021
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1022
super().__init__()
vllm/models/deepseek_v4/nvidia/mtp.py
+21
-10
@@ -18,7 +18,6 @@ import regex as re
18
import torch
19
import torch.nn as nn
20
21
-from vllm.compilation.decorators import support_torch_compile
21
from vllm.config import VllmConfig
22
from vllm.distributed import (
23
get_tensor_model_parallel_rank,
@@ -37,6 +36,10 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36
from vllm.model_executor.models.deepseek_mtp import SharedHead
37
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
38
from vllm.model_executor.models.utils import maybe_prefix
39
+from vllm.models.deepseek_v4.common.ops import (
40
+ fused_mtp_input_rmsnorm,
41
+ mtp_shared_head_rmsnorm,
42
+)
43
from vllm.sequence import IntermediateTensors
44
45
from .model import (
@@ -130,16 +133,22 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
133
spec_step_index: int = 0,
134
) -> torch.Tensor:
135
assert inputs_embeds is not None
133
- # masking inputs at position 0, as not needed by MTP
134
- inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
135
- inputs_embeds = self.enorm(inputs_embeds)
136
-
136
# Target stashes pre-hc_head residual as flat (T, hc_mult * D);
138
- # reshape to (T, hc_mult, D) — the training-time layout.
137
+ # reshape to (T, hc_mult, D) — the training-time layout — before
138
+ # the fused norm pass so both inputs are 3D-friendly.
139
previous_hidden_states = previous_hidden_states.view(
140
-1, self.hc_mult, self.config.hidden_size
141
)
142
- previous_hidden_states = self.hnorm(previous_hidden_states)
142
+ # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm.
143
+ inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm(
144
+ inputs_embeds,
145
+ positions,
146
+ previous_hidden_states,
147
+ self.enorm.weight.data,
148
+ self.hnorm.weight.data,
149
+ self.enorm.variance_epsilon,
150
+ self.hc_mult,
151
+ )
152
hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
153
inputs_embeds
154
).unsqueeze(-2)
@@ -236,13 +245,15 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
245
mtp_layer.rms_norm_eps,
246
mtp_layer.hc_eps,
247
)
239
- logits = self.logits_processor(
240
- mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
248
+ hidden_states = mtp_shared_head_rmsnorm(
249
+ hidden_states,
250
+ mtp_layer.shared_head.norm.weight.data,
251
+ mtp_layer.shared_head.norm.variance_epsilon,
252
)
253
+ logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states)
254
return logits
255
256
245
-@support_torch_compile
257
class DeepSeekV4MTP(nn.Module):
258
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
259
super().__init__()
vllm/v1/worker/gpu_model_runner.py
+5
@@ -5182,6 +5182,11 @@ class GPUModelRunner(
5182
and not self.parallel_config.use_ubatching
5183
):
5184
self.model = BreakableCUDAGraphWrapper(self.model, self.vllm_config)
5185
+ drafter = getattr(self, "drafter", None)
5186
+ if drafter is not None and hasattr(drafter, "model"):
5187
+ drafter.model = BreakableCUDAGraphWrapper(
5188
+ drafter.model, self.vllm_config
5189
+ )
5190
elif (
5191
cudagraph_mode.has_full_cudagraphs()
5192
and not self.parallel_config.use_ubatching