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 14, 05:30 UTC
[Spec-Decode] Support piecewise cudagraphs for Eagle head (#25109)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>
Lucas Wilkinson committed
Oct 10, 2025 at 01:20 UTC
29255cfc3b2d094336b5deb65ea555cbe97fcfb5
6 files changed
+84
-16
vllm/config/compilation.py
+7
-4
@@ -50,11 +50,14 @@ class CUDAGraphMode(enum.Enum):
50
def mixed_mode(self) -> "CUDAGraphMode":
51
return CUDAGraphMode(self.value[1]) if self.separate_routine() else self
52
53
+ def has_mode(self, mode: "CUDAGraphMode") -> bool:
54
+ assert not mode.separate_routine()
55
+ if self.separate_routine():
56
+ return mode.value in self.value
57
+ return self == mode
58
+
59
def requires_piecewise_compilation(self) -> bool:
54
- return (
55
- self.decode_mode() == CUDAGraphMode.PIECEWISE
56
- or self.mixed_mode() == CUDAGraphMode.PIECEWISE
57
- )
60
+ return self.has_mode(CUDAGraphMode.PIECEWISE)
61
62
def max_cudagraph_mode(self) -> "CUDAGraphMode":
63
return CUDAGraphMode(max(self.value)) if self.separate_routine() else self
vllm/forward_context.py
+6
@@ -283,6 +283,12 @@ def set_forward_context(
283
vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp
284
)
285
286
+ # Convenience: if cudagraph is used and num_tokens is given, we can just
287
+ # create a batch descriptor here if not given (there's no harm since if it
288
+ # doesn't match in the wrapper it'll fall through).
289
+ if cudagraph_runtime_mode != CUDAGraphMode.NONE and num_tokens is not None:
290
+ batch_descriptor = batch_descriptor or BatchDescriptor(num_tokens=num_tokens)
291
+
292
forward_context = create_forward_context(
293
attn_metadata,
294
vllm_config,
vllm/model_executor/models/deepseek_mtp.py
+2
@@ -7,6 +7,7 @@ import torch
7
import torch.nn as nn
8
from transformers import PretrainedConfig
9
10
+from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
from vllm.model_executor.layers.fused_moe import FusedMoE
13
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -162,6 +163,7 @@ class DeepSeekMultiTokenPredictor(nn.Module):
163
return logits
164
165
166
+@support_torch_compile
167
class DeepSeekMTP(nn.Module, SupportsPP):
168
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
169
super().__init__()
vllm/model_executor/models/llama_eagle3.py
+11
@@ -8,6 +8,7 @@ import torch
8
import torch.nn as nn
9
from transformers import LlamaConfig
10
11
+from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig, get_current_vllm_config
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -21,6 +22,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
22
)
23
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
24
from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM
25
+from vllm.multimodal import MULTIMODAL_REGISTRY
26
from vllm.multimodal.inputs import NestedTensors
27
28
from .utils import AutoWeightsLoader, maybe_prefix
@@ -119,6 +121,15 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
121
return hidden_states, residual
122
123
124
+@support_torch_compile(
125
+ # torch.compile is disabled for multimodal EAGLE3 models due to constraint
126
+ # violations with dynamic shapes during tensor concatenation operations.
127
+ # See: https://github.com/vllm-project/vllm/pull/22872/files#r2362028132
128
+ # Non-multimodal EAGLE3 models can still use torch.compile safely.
129
+ enable_if=lambda vllm_config: not MULTIMODAL_REGISTRY.supports_multimodal_inputs(
130
+ vllm_config.model_config
131
+ ),
132
+)
133
class LlamaModel(nn.Module):
134
def __init__(
135
self,
vllm/v1/spec_decode/eagle.py
+56
-11
@@ -9,7 +9,12 @@ import numpy as np
9
import torch
10
import torch.nn as nn
11
12
-from vllm.config import CompilationLevel, VllmConfig, get_layers_from_vllm_config
12
+from vllm.config import (
13
+ CompilationLevel,
14
+ CUDAGraphMode,
15
+ VllmConfig,
16
+ get_layers_from_vllm_config,
17
+)
18
from vllm.distributed.parallel_state import get_pp_group
19
from vllm.forward_context import set_forward_context
20
from vllm.logger import init_logger
@@ -80,12 +85,25 @@ class EagleProposer:
85
self.attn_layer_names: list[str] = []
86
self.indexer_layer_names: list[str] = []
87
83
- self.use_cuda_graph = (
84
- not current_platform.is_xpu()
85
- and self.vllm_config.compilation_config.level == CompilationLevel.PIECEWISE
86
- and not self.vllm_config.model_config.enforce_eager
87
- and not self.speculative_config.enforce_eager
88
- )
88
+ self.use_cuda_graph = False
89
+
90
+ compilation_config = self.vllm_config.compilation_config
91
+ if compilation_config.level == CompilationLevel.PIECEWISE:
92
+ cudagraph_mode = compilation_config.cudagraph_mode
93
+ if cudagraph_mode != CUDAGraphMode.NONE and not cudagraph_mode.has_mode(
94
+ CUDAGraphMode.PIECEWISE
95
+ ):
96
+ logger.warning(
97
+ "Currently the eagle proposer only supports cudagraph_mode "
98
+ "PIECEWISE, if you want the drafter to use cuda graphs, "
99
+ "please set compilation_config.cudagraph_mode to PIECEWISE "
100
+ "or FULL_AND_PIECEWISE"
101
+ )
102
+ self.use_cuda_graph = (
103
+ cudagraph_mode.has_mode(CUDAGraphMode.PIECEWISE)
104
+ and not self.speculative_config.enforce_eager
105
+ )
106
+
107
self.cudagraph_batch_sizes = (
108
list(reversed(self.vllm_config.compilation_config.cudagraph_capture_sizes))
109
if self.use_cuda_graph
@@ -239,12 +257,15 @@ class EagleProposer:
257
per_layer_attn_metadata = {}
258
for layer_name in self.attn_layer_names:
259
per_layer_attn_metadata[layer_name] = attn_metadata
260
+
261
for layer_name in self.indexer_layer_names:
262
assert draft_indexer_metadata is not None
263
per_layer_attn_metadata[layer_name] = draft_indexer_metadata
264
265
+ cudagraph_runtime_mode = CUDAGraphMode.NONE
266
if self.use_cuda_graph and num_tokens <= self.cudagraph_batch_sizes[-1]:
267
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
268
+ cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
269
else:
270
num_input_tokens = num_tokens
271
# copy inputs to buffer for cudagraph
@@ -267,7 +288,10 @@ class EagleProposer:
288
inputs_embeds = None
289
290
with set_forward_context(
270
- per_layer_attn_metadata, self.vllm_config, num_tokens=num_input_tokens
291
+ per_layer_attn_metadata,
292
+ self.vllm_config,
293
+ num_tokens=num_input_tokens,
294
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
295
):
296
ret_hidden_states = self.model(
297
input_ids=input_ids,
@@ -326,8 +350,10 @@ class EagleProposer:
350
351
if self.use_cuda_graph and batch_size <= self.cudagraph_batch_sizes[-1]:
352
input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
353
+ cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
354
else:
355
input_batch_size = batch_size
356
+ cudagraph_runtime_mode = CUDAGraphMode.NONE
357
358
common_attn_metadata.num_actual_tokens = batch_size
359
common_attn_metadata.max_query_len = 1
@@ -424,7 +450,10 @@ class EagleProposer:
450
451
# Run the model.
452
with set_forward_context(
427
- per_layer_attn_metadata, self.vllm_config, num_tokens=input_batch_size
453
+ per_layer_attn_metadata,
454
+ self.vllm_config,
455
+ num_tokens=input_batch_size,
456
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
457
):
458
ret_hidden_states = self.model(
459
input_ids=input_ids,
@@ -731,11 +760,16 @@ class EagleProposer:
760
761
if self.use_cuda_graph and num_tokens <= self.cudagraph_batch_sizes[-1]:
762
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
763
+ cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
764
else:
765
num_input_tokens = num_tokens
766
+ cudagraph_runtime_mode = CUDAGraphMode.NONE
767
# Run the model.
768
with set_forward_context(
738
- per_layer_attn_metadata, self.vllm_config, num_tokens=num_input_tokens
769
+ per_layer_attn_metadata,
770
+ self.vllm_config,
771
+ num_tokens=num_input_tokens,
772
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
773
):
774
last_hidden_states, hidden_states = self.model(
775
input_ids=self.input_ids[:num_input_tokens],
@@ -1015,8 +1049,19 @@ class EagleProposer:
1049
def dummy_run(
1050
self,
1051
num_tokens: int,
1052
+ use_cudagraphs=True,
1053
) -> None:
1019
- with set_forward_context(None, self.vllm_config, num_tokens=num_tokens):
1054
+ if use_cudagraphs and num_tokens <= self.cudagraph_batch_sizes[-1]:
1055
+ num_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
1056
+
1057
+ with set_forward_context(
1058
+ None,
1059
+ self.vllm_config,
1060
+ num_tokens=num_tokens,
1061
+ cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE
1062
+ if use_cudagraphs
1063
+ else CUDAGraphMode.NONE,
1064
+ ):
1065
if self.supports_mm_inputs:
1066
input_ids = None
1067
inputs_embeds = self.inputs_embeds[:num_tokens]
vllm/v1/worker/gpu_model_runner.py
+2
-1
@@ -3441,7 +3441,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3441
3442
if self.speculative_config and self.speculative_config.use_eagle():
3443
assert isinstance(self.drafter, EagleProposer)
3444
- self.drafter.dummy_run(num_tokens)
3444
+ use_cudagraphs = cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE
3445
+ self.drafter.dummy_run(num_tokens, use_cudagraphs=use_cudagraphs)
3446
3447
# This is necessary to avoid blocking DP.
3448
# For dummy runs, we typically skip EPLB since we don't have any real