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 5, 21:00 UTC
[Model Runner V2] Rebuild attention metadata before eagle decode full… (#38311)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Mar 27, 2026 at 13:46 UTC
384e4d5f48cea203b0fbcd41951d2d49775bf6bd
1 file changed
+97
-55
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+97
-55
@@ -19,6 +19,9 @@ from vllm.v1.worker.gpu.attn_utils import (
19
init_attn_backend,
20
)
21
from vllm.v1.worker.gpu.block_table import BlockTables
22
+from vllm.v1.worker.gpu.cudagraph_utils import (
23
+ BatchExecutionDescriptor,
24
+)
25
from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding
26
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
27
from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -239,6 +242,66 @@ class EagleSpeculator:
242
idx_mapping, query_start_loc, pos, num_tokens_padded
243
)
244
245
+ def _dispatch_and_sync_dp(
246
+ self,
247
+ cudagraph_manager: EagleCudaGraphManager,
248
+ num_reqs: int,
249
+ num_tokens: int,
250
+ uniform_token_count: int | None,
251
+ ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:
252
+ batch_desc = cudagraph_manager.dispatch(
253
+ num_reqs, num_tokens, uniform_token_count
254
+ )
255
+ num_tokens_across_dp = None
256
+ if self.dp_size > 1:
257
+ batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding(
258
+ cudagraph_manager,
259
+ batch_desc,
260
+ num_tokens,
261
+ num_reqs,
262
+ uniform_token_count,
263
+ self.dp_size,
264
+ self.dp_rank,
265
+ )
266
+ return batch_desc, num_tokens_across_dp
267
+
268
+ def _build_draft_attn_metadata(
269
+ self,
270
+ num_reqs: int,
271
+ num_reqs_padded: int,
272
+ num_tokens_padded: int,
273
+ max_query_len: int,
274
+ ) -> dict[str, Any] | None:
275
+ if not self.draft_attn_layer_names:
276
+ return None
277
+
278
+ query_start_loc_cpu = (
279
+ torch.arange(num_reqs_padded + 1, dtype=torch.int32, device="cpu").clamp_(
280
+ max=num_reqs
281
+ )
282
+ * max_query_len
283
+ )
284
+ block_tables = [
285
+ x[:num_reqs_padded] for x in self.block_tables.input_block_tables
286
+ ]
287
+ slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded]
288
+ attn_metadata = build_attn_metadata(
289
+ attn_groups=self.attn_groups,
290
+ num_reqs=num_reqs_padded,
291
+ num_tokens=num_tokens_padded,
292
+ query_start_loc_gpu=self.input_buffers.query_start_loc[
293
+ : num_reqs_padded + 1
294
+ ],
295
+ query_start_loc_cpu=query_start_loc_cpu,
296
+ max_query_len=max_query_len,
297
+ seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
298
+ max_seq_len=self.max_model_len,
299
+ block_tables=block_tables,
300
+ slot_mappings=slot_mappings,
301
+ kv_cache_config=self.kv_cache_config,
302
+ )
303
+ return attn_metadata
304
+
305
def capture_model(self) -> None:
306
if self.num_speculative_steps == 1:
307
return
@@ -319,7 +382,6 @@ class EagleSpeculator:
382
logits = self.model.compute_logits(sample_hidden_states)
383
384
num_reqs = input_batch.num_reqs
322
- num_reqs_padded = input_batch.num_reqs_after_padding
385
# NOTE(woosuk): For draft sampling, we only consider the temperature
386
# and ignore the other sampling parameters such as top_k and top_p,
387
# for simplicity and performance.
@@ -366,69 +428,49 @@ class EagleSpeculator:
428
self.max_num_reqs,
429
)
430
369
- # Get batch descriptor and sync across DP ranks.
370
- # Eagle uses FULL-only mode, dispatch with uniform_token_count=1 for decode
371
-
372
- batch_desc = self.cudagraph_manager.dispatch(num_reqs, num_reqs, 1)
373
- num_tokens_across_dp = None
374
-
375
- if self.dp_size > 1:
376
- batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding(
377
- self.cudagraph_manager,
378
- batch_desc,
379
- num_reqs,
380
- num_reqs,
381
- 1, # uniform_token_count
382
- self.dp_size,
383
- self.dp_rank,
384
- )
385
-
386
- if not (dummy_run and skip_attn_for_dummy_run):
387
- query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
388
- slot_mappings = self.block_tables.compute_slot_mappings(
389
- idx_mapping, query_start_loc, pos, batch_desc.num_tokens
390
- )
391
-
392
- if batch_desc.cg_mode == CUDAGraphMode.FULL:
393
- return self.cudagraph_manager.run_fullgraph(batch_desc)[:num_reqs]
431
+ # Each request produces exactly 1 token per draft decode step,
432
+ # enabling FULL cudagraph.
433
+ decode_batch_desc, num_tokens_across_dp = self._dispatch_and_sync_dp(
434
+ self.cudagraph_manager,
435
+ num_reqs,
436
+ num_reqs,
437
+ uniform_token_count=1,
438
+ )
439
395
- # Run eager or piecewise CUDA graph.
440
attn_metadata_updated = None
441
slot_mappings_updated = None
442
if not (dummy_run and skip_attn_for_dummy_run):
399
- query_start_loc_cpu = torch.arange(
400
- num_reqs_padded + 1, dtype=torch.int32, device="cpu"
401
- )
402
- block_tables = [
403
- x[:num_reqs_padded] for x in self.block_tables.input_block_tables
404
- ]
405
-
406
- # FIXME(woosuk): This is UNSAFE!!
407
- attn_metadata_updated = build_attn_metadata(
408
- attn_groups=self.attn_groups,
409
- num_reqs=num_reqs_padded,
410
- num_tokens=num_reqs_padded,
411
- query_start_loc_gpu=query_start_loc,
412
- query_start_loc_cpu=query_start_loc_cpu,
413
- max_query_len=1,
414
- seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
415
- max_seq_len=self.max_model_len,
416
- block_tables=block_tables,
417
- slot_mappings=slot_mappings,
418
- kv_cache_config=self.kv_cache_config,
443
+ # Build attention metadata and slot mappings for the draft
444
+ # decode steps. It is necessary to rebuild the attention
445
+ # metadata even when replaying the FULL cudagraph so that
446
+ # any attention metadata builder state is updated.
447
+ slot_mappings = self.block_tables.compute_slot_mappings(
448
+ idx_mapping,
449
+ self.input_buffers.query_start_loc[: num_reqs + 1],
450
+ pos,
451
+ decode_batch_desc.num_tokens,
452
)
453
slot_mappings_updated = build_slot_mappings_by_layer(
454
slot_mappings, self.kv_cache_config
455
)
456
+ attn_metadata_updated = self._build_draft_attn_metadata(
457
+ num_reqs=num_reqs,
458
+ num_reqs_padded=decode_batch_desc.num_reqs or num_reqs,
459
+ num_tokens_padded=decode_batch_desc.num_tokens,
460
+ max_query_len=1,
461
+ )
462
424
- self.generate_draft(
425
- num_reqs,
426
- batch_desc.num_tokens,
427
- attn_metadata_updated,
428
- slot_mappings_updated,
429
- num_tokens_across_dp=num_tokens_across_dp,
430
- cudagraph_runtime_mode=batch_desc.cg_mode,
431
- )
463
+ if decode_batch_desc.cg_mode == CUDAGraphMode.FULL:
464
+ self.cudagraph_manager.run_fullgraph(decode_batch_desc)
465
+ else:
466
+ self.generate_draft(
467
+ num_reqs,
468
+ decode_batch_desc.num_tokens,
469
+ attn_metadata_updated,
470
+ slot_mappings_updated,
471
+ num_tokens_across_dp=num_tokens_across_dp,
472
+ cudagraph_runtime_mode=decode_batch_desc.cg_mode,
473
+ )
474
return self.draft_tokens[:num_reqs]
475
476