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 12, 14:45 UTC
EAGLE Support DP>1 (#26086)
Signed-off-by: Rémi Delacourt <remi@mistral.ai> Signed-off-by: Rémi Delacourt <54138269+Flechman@users.noreply.github.com> Signed-off-by: remi <remi@mistral.ai>
Rémi Delacourt committed
Nov 25, 2025 at 08:32 UTC
12c007e288bf5c0ae3bd438036fbafbad88e706b
4 files changed
+176
-31
.buildkite/test-pipeline.yaml
+2
@@ -192,6 +192,7 @@ steps:
192
# test with internal dp
193
- python3 ../examples/offline_inference/data_parallel.py --enforce-eager
194
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
195
+ - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
196
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
197
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
198
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
@@ -1116,6 +1117,7 @@ steps:
1117
# https://github.com/NVIDIA/nccl/issues/1838
1118
- export NCCL_CUMEM_HOST_ENABLE=0
1119
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
1120
+ - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
1121
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
1122
- DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
1123
- pytest -v -s entrypoints/llm/test_collective_rpc.py
tests/v1/distributed/test_eagle_dp.py
new
+77
@@ -0,0 +1,77 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import asyncio
4
+import os
5
+from contextlib import AsyncExitStack
6
+from dataclasses import replace
7
+
8
+import pytest
9
+
10
+from vllm import SamplingParams
11
+from vllm.engine.arg_utils import AsyncEngineArgs
12
+from vllm.sampling_params import RequestOutputKind
13
+from vllm.v1.engine.async_llm import AsyncLLM
14
+
15
+DP_SIZE = int(os.getenv("DP_SIZE", 2))
16
+
17
+
18
+@pytest.mark.asyncio
19
+async def test_run_eagle_dp():
20
+ target_model = "meta-llama/Llama-3.1-8B-Instruct"
21
+ draft_model = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
22
+
23
+ engine_args = AsyncEngineArgs(
24
+ model=target_model,
25
+ tokenizer_mode="auto",
26
+ enforce_eager=False,
27
+ tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
28
+ data_parallel_size=DP_SIZE,
29
+ data_parallel_backend="mp", # ray takes more time
30
+ trust_remote_code=True,
31
+ max_model_len=16384,
32
+ )
33
+
34
+ eagle_engine_args = replace(
35
+ engine_args,
36
+ speculative_config={
37
+ "model": draft_model,
38
+ "method": "eagle",
39
+ "num_speculative_tokens": 3,
40
+ },
41
+ )
42
+
43
+ prompt = "This is a test of data parallel with eagle"
44
+ num_expected_tokens = 100
45
+ sampling_params = SamplingParams(
46
+ min_tokens=num_expected_tokens,
47
+ max_tokens=num_expected_tokens,
48
+ ignore_eos=True,
49
+ output_kind=RequestOutputKind.FINAL_ONLY,
50
+ temperature=0,
51
+ )
52
+
53
+ async def generate_with_timeout(given_engine: AsyncLLM):
54
+ async for out in given_engine.generate(
55
+ request_id="test-eagle-dp", prompt=prompt, sampling_params=sampling_params
56
+ ):
57
+ token_ids = out.outputs[0].token_ids
58
+ assert len(token_ids) == num_expected_tokens
59
+ return token_ids
60
+
61
+ async def engine_create_and_generate(engine_args: AsyncEngineArgs):
62
+ async with AsyncExitStack() as after:
63
+ engine = AsyncLLM.from_engine_args(engine_args)
64
+ after.callback(engine.shutdown)
65
+
66
+ token_ids = await asyncio.wait_for(
67
+ generate_with_timeout(engine), timeout=30
68
+ )
69
+
70
+ assert not engine.output_processor.has_unfinished_requests()
71
+ return token_ids
72
+
73
+ token_ids_with_eagle = await engine_create_and_generate(eagle_engine_args)
74
+ token_ids_no_eagle = await engine_create_and_generate(engine_args)
75
+
76
+ # Test for correctness
77
+ assert token_ids_with_eagle == token_ids_no_eagle
vllm/v1/spec_decode/eagle.py
+93
-30
@@ -40,6 +40,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
40
from vllm.v1.sample.sampler import _SAMPLING_EPS
41
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
42
from vllm.v1.utils import CpuGpuBuffer
43
+from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
44
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
45
46
logger = init_logger(__name__)
@@ -65,6 +66,7 @@ class EagleProposer:
66
self.dtype = vllm_config.model_config.dtype
67
self.max_model_len = vllm_config.model_config.max_model_len
68
self.block_size = vllm_config.cache_config.block_size
69
+ self.dp_rank = vllm_config.parallel_config.data_parallel_rank
70
self.num_speculative_tokens = self.speculative_config.num_speculative_tokens
71
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
72
self.token_arange_np = np.arange(self.max_num_tokens)
@@ -271,15 +273,24 @@ class EagleProposer:
273
assert draft_indexer_metadata is not None
274
per_layer_attn_metadata[layer_name] = draft_indexer_metadata
275
276
+ num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(
277
+ num_tokens_unpadded=num_tokens,
278
+ num_tokens_padded=num_tokens,
279
+ )
280
+
281
cudagraph_runtime_mode = CUDAGraphMode.NONE
282
if (
283
self.use_cuda_graph
277
- and num_tokens <= self.compilation_config.max_cudagraph_capture_size
284
+ and num_tokens_dp_padded
285
+ <= self.compilation_config.max_cudagraph_capture_size
286
):
279
- num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
287
+ num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens_dp_padded)
288
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
289
else:
282
- num_input_tokens = num_tokens
290
+ num_input_tokens = num_tokens_dp_padded
291
+ if num_tokens_across_dp is not None:
292
+ num_tokens_across_dp[self.dp_rank] = num_input_tokens
293
+
294
# copy inputs to buffer for cudagraph
295
self._set_positions(num_tokens, target_positions)
296
self.hidden_states[:num_tokens] = target_hidden_states
@@ -303,6 +314,7 @@ class EagleProposer:
314
per_layer_attn_metadata,
315
self.vllm_config,
316
num_tokens=num_input_tokens,
317
+ num_tokens_across_dp=num_tokens_across_dp,
318
cudagraph_runtime_mode=cudagraph_runtime_mode,
319
):
320
ret_hidden_states = self.model(
@@ -365,15 +377,23 @@ class EagleProposer:
377
# Generate the remaining draft tokens.
378
draft_token_ids_list = [draft_token_ids]
379
380
+ batch_size_dp_padded, batch_size_across_dp = self._pad_batch_across_dp(
381
+ num_tokens_unpadded=batch_size,
382
+ num_tokens_padded=batch_size,
383
+ )
384
+
385
if (
386
self.use_cuda_graph
370
- and batch_size <= self.compilation_config.max_cudagraph_capture_size
387
+ and batch_size_dp_padded
388
+ <= self.compilation_config.max_cudagraph_capture_size
389
):
372
- input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
390
+ input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size_dp_padded)
391
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
392
else:
375
- input_batch_size = batch_size
393
+ input_batch_size = batch_size_dp_padded
394
cudagraph_runtime_mode = CUDAGraphMode.NONE
395
+ if batch_size_across_dp is not None:
396
+ batch_size_across_dp[self.dp_rank] = input_batch_size
397
398
common_attn_metadata.num_actual_tokens = batch_size
399
common_attn_metadata.max_query_len = 1
@@ -474,6 +494,7 @@ class EagleProposer:
494
per_layer_attn_metadata,
495
self.vllm_config,
496
num_tokens=input_batch_size,
497
+ num_tokens_across_dp=batch_size_across_dp,
498
cudagraph_runtime_mode=cudagraph_runtime_mode,
499
):
500
ret_hidden_states = self.model(
@@ -1116,36 +1137,56 @@ class EagleProposer:
1137
self,
1138
num_tokens: int,
1139
use_cudagraphs=True,
1140
+ is_graph_capturing=False,
1141
) -> None:
1142
# Determine if CUDA graphs should be used for this run.
1143
cudagraphs_enabled = use_cudagraphs and self.use_cuda_graph
1122
- if (
1123
- cudagraphs_enabled
1124
- and num_tokens <= self.compilation_config.max_cudagraph_capture_size
1125
- ):
1126
- num_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
1144
1128
- with set_forward_context(
1129
- None,
1130
- self.vllm_config,
1131
- num_tokens=num_tokens,
1132
- cudagraph_runtime_mode=(
1133
- CUDAGraphMode.PIECEWISE if cudagraphs_enabled else CUDAGraphMode.NONE
1134
- ),
1145
+ # FIXME: when using tree-based specdec, adjust number of forward-passes
1146
+ # according to the depth of the tree.
1147
+ for fwd_idx in range(
1148
+ self.num_speculative_tokens if not is_graph_capturing else 1
1149
):
1136
- if self.supports_mm_inputs:
1137
- input_ids = None
1138
- inputs_embeds = self.inputs_embeds[:num_tokens]
1139
- else:
1140
- input_ids = self.input_ids[:num_tokens]
1141
- inputs_embeds = None
1150
+ if fwd_idx <= 1:
1151
+ num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(
1152
+ num_tokens_unpadded=num_tokens,
1153
+ num_tokens_padded=num_tokens,
1154
+ )
1155
+ if (
1156
+ cudagraphs_enabled
1157
+ and num_tokens_dp_padded
1158
+ <= self.compilation_config.max_cudagraph_capture_size
1159
+ ):
1160
+ num_input_tokens = self.vllm_config.pad_for_cudagraph(
1161
+ num_tokens_dp_padded
1162
+ )
1163
+ else:
1164
+ num_input_tokens = num_tokens_dp_padded
1165
+ if num_tokens_across_dp is not None:
1166
+ num_tokens_across_dp[self.dp_rank] = num_input_tokens
1167
1143
- self.model(
1144
- input_ids=input_ids,
1145
- positions=self._get_positions(num_tokens),
1146
- hidden_states=self.hidden_states[:num_tokens],
1147
- inputs_embeds=inputs_embeds,
1148
- )
1168
+ with set_forward_context(
1169
+ None,
1170
+ self.vllm_config,
1171
+ num_tokens=num_input_tokens,
1172
+ num_tokens_across_dp=num_tokens_across_dp,
1173
+ cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE
1174
+ if cudagraphs_enabled
1175
+ else CUDAGraphMode.NONE,
1176
+ ):
1177
+ if self.supports_mm_inputs:
1178
+ input_ids = None
1179
+ inputs_embeds = self.inputs_embeds[:num_input_tokens]
1180
+ else:
1181
+ input_ids = self.input_ids[:num_input_tokens]
1182
+ inputs_embeds = None
1183
+
1184
+ self.model(
1185
+ input_ids=input_ids,
1186
+ positions=self._get_positions(num_input_tokens),
1187
+ hidden_states=self.hidden_states[:num_input_tokens],
1188
+ inputs_embeds=inputs_embeds,
1189
+ )
1190
1191
def _get_attention_metadata_builder(self) -> AttentionMetadataBuilder:
1192
"""Find and return the attention metadata builders for EAGLE layers.
@@ -1211,6 +1252,28 @@ class EagleProposer:
1252
== 1
1253
), "All eagle layers should belong to the same kv cache group"
1254
1255
+ def _pad_batch_across_dp(
1256
+ self,
1257
+ num_tokens_unpadded: int,
1258
+ num_tokens_padded: int,
1259
+ ) -> tuple[int, torch.Tensor]:
1260
+ # TODO(Flechman): support DBO ubatching
1261
+ ubatch_slices, num_toks_across_dp = coordinate_batch_across_dp(
1262
+ num_tokens_unpadded=num_tokens_unpadded,
1263
+ parallel_config=self.vllm_config.parallel_config,
1264
+ allow_microbatching=False,
1265
+ allow_dp_padding=self.use_cuda_graph,
1266
+ num_tokens_padded=num_tokens_padded,
1267
+ uniform_decode=None,
1268
+ num_scheduled_tokens_per_request=None,
1269
+ )
1270
+ assert ubatch_slices is None, "DBO ubatching not implemented for EAGLE"
1271
+
1272
+ num_tokens_dp_padded = num_tokens_padded
1273
+ if num_toks_across_dp is not None:
1274
+ num_tokens_dp_padded = int(num_toks_across_dp[self.dp_rank].item())
1275
+ return num_tokens_dp_padded, num_toks_across_dp
1276
+
1277
1278
# NOTE(woosuk): Currently, the below code is not used and we always use argmax
1279
# to sample the draft tokens. We will use this after we find a way to manage
vllm/v1/worker/gpu_model_runner.py
+4
-1
@@ -3746,6 +3746,7 @@ class GPUModelRunner(
3746
create_mixed_batch: bool = False,
3747
remove_lora: bool = True,
3748
activate_lora: bool = False,
3749
+ is_graph_capturing: bool = False,
3750
) -> tuple[torch.Tensor, torch.Tensor]:
3751
"""
3752
Run a dummy forward pass to warm up/profile run or capture the
@@ -3981,7 +3982,7 @@ class GPUModelRunner(
3982
if self.speculative_config and self.speculative_config.use_eagle():
3983
assert isinstance(self.drafter, EagleProposer)
3984
use_cudagraphs = (
3984
- cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE
3985
+ cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE)
3986
and not self.speculative_config.enforce_eager
3987
)
3988
@@ -3995,6 +3996,7 @@ class GPUModelRunner(
3996
self.drafter.dummy_run(
3997
num_tokens,
3998
use_cudagraphs=use_cudagraphs,
3999
+ is_graph_capturing=is_graph_capturing,
4000
)
4001
4002
# This is necessary to avoid blocking DP.
@@ -4427,6 +4429,7 @@ class GPUModelRunner(
4429
skip_eplb=True,
4430
remove_lora=False,
4431
activate_lora=activate_lora,
4432
+ is_graph_capturing=True,
4433
)
4434
self.maybe_remove_all_loras(self.lora_config)
4435