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 7, 09:00 UTC
[Model Runner V2] support dp & ep for spec decoding (#35294)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai>
zhrrr committed
Mar 4, 2026 at 07:19 UTC
97286a20ed5803583c50af3dd1f45268346be0e8
3 files changed
+124
-57
vllm/v1/worker/gpu/model_runner.py
+39
-17
@@ -57,10 +57,7 @@ from vllm.v1.worker.gpu.block_table import BlockTables
57
from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
58
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
59
from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
60
-from vllm.v1.worker.gpu.dp_utils import (
61
- get_cudagraph_and_dp_padding,
62
- make_num_tokens_across_dp,
63
-)
60
+from vllm.v1.worker.gpu.dp_utils import get_cudagraph_and_dp_padding
61
from vllm.v1.worker.gpu.input_batch import (
62
InputBatch,
63
InputBuffers,
@@ -265,7 +262,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
262
263
prepare_communication_buffer_for_model(self.model)
264
if self.speculator is not None:
268
- prepare_communication_buffer_for_model(self.speculator)
265
+ prepare_communication_buffer_for_model(self.speculator.model)
266
267
# Initialize the components that require the model.
268
self.model_state = init_model_state(
@@ -382,8 +379,41 @@ class GPUModelRunner(LoRAModelRunnerMixin):
379
return None, None
380
381
assert self.execute_model_state is not None
385
- input_batch, _, _, _, hidden_states, _, _ = self.execute_model_state
382
+ (
383
+ input_batch,
384
+ model_inputs,
385
+ attn_metadata,
386
+ slot_mappings_by_layer,
387
+ hidden_states,
388
+ aux_hidden_states,
389
+ kv_connector_output,
390
+ num_tokens_across_dp,
391
+ ) = self.execute_model_state
392
self.execute_model_state = None
393
+
394
+ # dummy run the eagle speculator's propose to ensure DP/EP sync.
395
+ if self.speculator is not None:
396
+ self.speculator.propose(
397
+ input_batch=input_batch,
398
+ attn_metadata=attn_metadata,
399
+ slot_mappings=slot_mappings_by_layer,
400
+ last_hidden_states=hidden_states,
401
+ aux_hidden_states=aux_hidden_states,
402
+ num_sampled=torch.ones(
403
+ input_batch.num_reqs, dtype=torch.int32, device=self.device
404
+ ),
405
+ num_rejected=torch.zeros(
406
+ input_batch.num_reqs, dtype=torch.int32, device=self.device
407
+ ),
408
+ last_sampled=self.req_states.last_sampled_tokens,
409
+ next_prefill_tokens=self.req_states.next_prefill_tokens,
410
+ temperature=self.sampler.sampling_states.temperature.gpu,
411
+ seeds=self.sampler.sampling_states.seeds.gpu,
412
+ num_tokens_across_dp=num_tokens_across_dp,
413
+ dummy_run=True,
414
+ skip_attn_for_dummy_run=skip_attn,
415
+ )
416
+
417
assert hidden_states is not None # Last PP rank always has hidden_states
418
sample_hidden_states = hidden_states[input_batch.logits_indices]
419
return hidden_states, sample_hidden_states
@@ -431,17 +461,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
461
else:
462
self._dummy_pooler_run(hidden_states)
463
434
- if self.speculator is not None:
435
- num_tokens_across_dp = make_num_tokens_across_dp(
436
- self.parallel_config.data_parallel_size, self.max_num_tokens
437
- )
438
- self.speculator.run_model(
439
- self.max_num_tokens,
440
- attn_metadata=None,
441
- slot_mappings=None,
442
- num_tokens_across_dp=num_tokens_across_dp,
443
- )
444
-
464
torch.cuda.synchronize()
465
del hidden_states, sample_hidden_states
466
gc.collect()
@@ -979,6 +998,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
998
hidden_states,
999
aux_hidden_states,
1000
kv_connector_output,
1001
+ num_tokens_across_dp,
1002
)
1003
1004
if not self.is_last_pp_rank:
@@ -1005,6 +1025,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1025
hidden_states,
1026
aux_hidden_states,
1027
kv_connector_output,
1028
+ num_tokens_across_dp,
1029
) = self.execute_model_state
1030
self.execute_model_state = None
1031
@@ -1078,6 +1099,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1099
self.req_states.next_prefill_tokens,
1100
self.sampler.sampling_states.temperature.gpu,
1101
self.sampler.sampling_states.seeds.gpu,
1102
+ num_tokens_across_dp=num_tokens_across_dp,
1103
)
1104
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1105
self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+20
@@ -55,6 +55,26 @@ class EagleCudaGraphManager:
55
def get_cudagraph_size(self, num_tokens: int) -> int | None:
56
return self.cudagraph_sizes.get(num_tokens)
57
58
+ def get_cudagraph_runtime_mode(
59
+ self, num_tokens: int
60
+ ) -> tuple[CUDAGraphMode, int | None]:
61
+ cudagraph_size = self.get_cudagraph_size(num_tokens)
62
+ if cudagraph_size is None:
63
+ cudagraph_mode = CUDAGraphMode.NONE
64
+ else:
65
+ cudagraph_mode = self.cudagraph_mode
66
+
67
+ if (
68
+ cudagraph_mode == CUDAGraphMode.FULL
69
+ and cudagraph_size is not None
70
+ and cudagraph_size not in self.graphs
71
+ ):
72
+ # If graph wasn't captured yet, fall back to eager.
73
+ # This might happen when the dummy run is called before capture.
74
+ cudagraph_mode = CUDAGraphMode.NONE
75
+ cudagraph_size = None
76
+ return cudagraph_mode, cudagraph_size
77
+
78
def capture_graph(
79
self,
80
num_tokens: int,
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+65
-40
@@ -16,6 +16,7 @@ from vllm.v1.worker.gpu.attn_utils import (
16
build_slot_mappings_by_layer,
17
)
18
from vllm.v1.worker.gpu.block_table import BlockTables
19
+from vllm.v1.worker.gpu.dp_utils import get_cudagraph_and_dp_padding
20
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
21
from vllm.v1.worker.gpu.model_states.interface import ModelState
22
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
@@ -48,6 +49,10 @@ class EagleSpeculator:
49
self.vocab_size = self.draft_model_config.get_vocab_size()
50
self.dtype = vllm_config.model_config.dtype
51
52
+ # DP configuration
53
+ self.dp_size = vllm_config.parallel_config.data_parallel_size
54
+ self.dp_rank = vllm_config.parallel_config.data_parallel_rank
55
+
56
self.input_buffers = InputBuffers(
57
max_num_reqs=self.max_num_reqs,
58
max_num_tokens=self.max_num_tokens,
@@ -122,8 +127,8 @@ class EagleSpeculator:
127
self,
128
num_reqs: int,
129
num_tokens_padded: int,
125
- attn_metadata: dict[str, Any],
126
- slot_mappings: dict[str, torch.Tensor],
130
+ attn_metadata: dict[str, Any] | None,
131
+ slot_mappings: dict[str, torch.Tensor] | None,
132
num_tokens_across_dp: torch.Tensor | None,
133
cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
134
) -> None:
@@ -164,9 +169,10 @@ class EagleSpeculator:
169
self.hidden_states,
170
self.max_model_len,
171
)
167
- self.block_tables.compute_slot_mappings(
168
- idx_mapping, query_start_loc, pos
169
- )
172
+ if attn_metadata is not None:
173
+ self.block_tables.compute_slot_mappings(
174
+ idx_mapping, query_start_loc, pos
175
+ )
176
177
def capture_model(self) -> None:
178
if self.num_speculative_steps == 1:
@@ -203,6 +209,9 @@ class EagleSpeculator:
209
temperature: torch.Tensor,
210
# [max_num_reqs]
211
seeds: torch.Tensor,
212
+ num_tokens_across_dp: torch.Tensor | None = None,
213
+ dummy_run: bool = False,
214
+ skip_attn_for_dummy_run: bool = False,
215
) -> torch.Tensor:
216
# NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
217
# number of rejected tokens, we maintain the size of eagle's input_ids and
@@ -236,7 +245,7 @@ class EagleSpeculator:
245
num_tokens,
246
attn_metadata,
247
slot_mappings,
239
- num_tokens_across_dp=None, # FIXME
248
+ num_tokens_across_dp=num_tokens_across_dp,
249
)
250
sample_hidden_states = last_hidden_states[last_token_indices]
251
logits = self.model.compute_logits(sample_hidden_states)
@@ -282,48 +291,64 @@ class EagleSpeculator:
291
self.max_model_len,
292
self.max_num_reqs,
293
)
285
- query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
286
- slot_mappings = self.block_tables.compute_slot_mappings(
287
- idx_mapping, query_start_loc, pos
288
- )
294
290
- cudagraph_size = self.cudagraph_manager.get_cudagraph_size(num_reqs)
291
- cudagraph_mode = self.cudagraph_manager.cudagraph_mode
292
- if cudagraph_size is not None and cudagraph_mode == CUDAGraphMode.FULL:
295
+ if not (dummy_run and skip_attn_for_dummy_run):
296
+ query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
297
+ slot_mappings = self.block_tables.compute_slot_mappings(
298
+ idx_mapping, query_start_loc, pos
299
+ )
300
+
301
+ cudagraph_mode, cudagraph_size = (
302
+ self.cudagraph_manager.get_cudagraph_runtime_mode(num_reqs)
303
+ )
304
+ num_tokens_padded, num_tokens_across_dp, synced_cudagraph_mode = (
305
+ get_cudagraph_and_dp_padding(
306
+ num_reqs,
307
+ cudagraph_size,
308
+ cudagraph_mode.value,
309
+ self.dp_size,
310
+ self.dp_rank,
311
+ )
312
+ )
313
+ cudagraph_mode = CUDAGraphMode(synced_cudagraph_mode)
314
+ if cudagraph_mode == CUDAGraphMode.FULL:
315
# Run full CUDA graph.
294
- self.cudagraph_manager.run_fullgraph(cudagraph_size)
316
+ self.cudagraph_manager.run_fullgraph(num_tokens_padded)
317
return self.draft_tokens[:num_reqs]
318
319
# Run eager or piecewise CUDA graph.
298
- num_tokens_padded = cudagraph_size if cudagraph_size is not None else num_reqs
299
- query_start_loc_cpu = torch.arange(
300
- num_reqs + 1, dtype=torch.int32, device="cpu"
301
- )
302
- block_tables = [x[:num_reqs] for x in self.block_tables.input_block_tables]
303
-
304
- # FIXME(woosuk): This is UNSAFE!!
305
- attn_metadata = build_attn_metadata(
306
- attn_groups=self.attn_groups,
307
- num_reqs=num_reqs,
308
- num_tokens=num_reqs,
309
- query_start_loc_gpu=query_start_loc,
310
- query_start_loc_cpu=query_start_loc_cpu,
311
- max_query_len=1,
312
- seq_lens=self.input_buffers.seq_lens[:num_reqs],
313
- max_seq_len=self.max_model_len,
314
- block_tables=block_tables,
315
- slot_mappings=slot_mappings,
316
- kv_cache_config=self.kv_cache_config,
317
- )
318
- slot_mappings_by_layer = build_slot_mappings_by_layer(
319
- slot_mappings, self.kv_cache_config
320
- )
320
+ attn_metadata_updated = None
321
+ slot_mappings_updated = None
322
+ if not (dummy_run and skip_attn_for_dummy_run):
323
+ query_start_loc_cpu = torch.arange(
324
+ num_reqs + 1, dtype=torch.int32, device="cpu"
325
+ )
326
+ block_tables = [x[:num_reqs] for x in self.block_tables.input_block_tables]
327
+
328
+ # FIXME(woosuk): This is UNSAFE!!
329
+ attn_metadata_updated = build_attn_metadata(
330
+ attn_groups=self.attn_groups,
331
+ num_reqs=num_reqs,
332
+ num_tokens=num_reqs,
333
+ query_start_loc_gpu=query_start_loc,
334
+ query_start_loc_cpu=query_start_loc_cpu,
335
+ max_query_len=1,
336
+ seq_lens=self.input_buffers.seq_lens[:num_reqs],
337
+ max_seq_len=self.max_model_len,
338
+ block_tables=block_tables,
339
+ slot_mappings=slot_mappings,
340
+ kv_cache_config=self.kv_cache_config,
341
+ )
342
+ slot_mappings_updated = build_slot_mappings_by_layer(
343
+ slot_mappings, self.kv_cache_config
344
+ )
345
+
346
self.generate_draft(
347
num_reqs,
348
num_tokens_padded,
324
- attn_metadata,
325
- slot_mappings_by_layer,
326
- num_tokens_across_dp=None, # FIXME
349
+ attn_metadata_updated,
350
+ slot_mappings_updated,
351
+ num_tokens_across_dp=num_tokens_across_dp,
352
cudagraph_runtime_mode=cudagraph_mode,
353
)
354
return self.draft_tokens[:num_reqs]