@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a311916a29

[MRV2][Spec] Fuse AR speculator multi-step decodes back into one CUDA graph (#46849)

Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>

Yizhou committed Aug 11, 2026 at 15:06 UTC a311916a291c1fed3dbfb72e60f74cd778c8419d

10 files changed +492 -49

docs/design/model_runner_v2.md

@@ -191,6 +191,14 @@ V1's CUDA graph handling is implicit and hard to reason about. MRV2 uses a `CUDA

191

192 This makes graph lifecycle and execution mode decisions more understandable and easier to extend. Example: MRV2 can capture multiple draft-model forward passes into one CUDA graph.

193

194 +### Fused Multi-Step Draft Decoding

195 +

196 +Autoregressive speculative decoding executes several dependent draft steps per scheduler step. In the fused path, MRV2 captures all post-prefill draft steps in one full CUDA graph instead of replaying a separate graph for each draft token. Attention metadata is built once before the loop, while common step-dependent tensors keep stable addresses and are updated in place between draft steps.

197 +

198 +Some attention backends also materialize derived state, such as scheduler metadata or sparse indices. Before opting into the fused path, these backends must implement `AttentionMetadataBuilder.update_draft_decode_metadata()` to update or invalidate that state after the draft inputs advance. The hook runs during CUDA graph capture, so only the GPU operations it issues are recorded and executed during replay; its Python body is not run again. Implementations must therefore use capture-safe operations and keep all replayed tensor state in persistent storage.

199 +

200 +For draft models that advance positions, the fused path is enabled only when every draft attention group declares `supports_draft_decode_metadata_update`. Otherwise, MRV2 falls back to rebuilding attention metadata between draft steps. Draft models that keep positions fixed do not require this update. Before enabling a backend, developers must audit all derived metadata, including state inherited from parent builders or owned by auxiliary attention backends.

201 +

202 ## Development Philosophy

203

204 MRV2 changes should meet a higher code quality bar. As feature gaps with V1 are filled, features should be reconsidered from first principles in the MRV2 design context instead of quickly porting V1 behavior.

tests/v1/worker/test_gpu_autoregressive_speculator.py

+166

@@ -3,6 +3,7 @@

4 from contextlib import nullcontext

5 from types import SimpleNamespace

6 +from unittest.mock import Mock

8 import pytest

9 import torch

@@ -16,6 +17,9 @@ from vllm.model_executor.models.mistral_eagle import EagleMistralForCausalLM

17 from vllm.model_executor.models.mistral_large_3_eagle import (

18 EagleMistralLarge3ForCausalLM,

19 )

20 +from vllm.v1.attention.backends import flash_attn as flash_attn_module

21 +from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata

22 +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor

23 from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module

24 from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module

25 from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (

@@ -278,3 +282,165 @@ def test_run_model_reuses_tensor_return_for_mtp(monkeypatch):

282

283 assert actual_logits_hidden is hidden

284 assert actual_feedback_hidden is hidden

285 +

286 +

287 +@pytest.mark.parametrize(

288 + (

289 + "method_name",

290 + "cg_mode",

291 + "expected_eager_calls",

292 + "expected_graph_replays",

293 + ),

294 + [

295 + ("_multi_step_decode", CUDAGraphMode.NONE, 3, 0),

296 + ("_multi_step_decode", CUDAGraphMode.FULL, 0, 3),

297 + ("_fused_multi_step_decode", CUDAGraphMode.NONE, 3, 0),

298 + ("_fused_multi_step_decode", CUDAGraphMode.FULL, 0, 1),

299 + ],

300 +)

301 +def test_multi_step_decode_replays_captured_graph_as_expected(

302 + method_name,

303 + cg_mode,

304 + expected_eager_calls,

305 + expected_graph_replays,

306 +):

307 + speculator = object.__new__(_TestSpeculator)

308 + speculator.num_speculative_steps = 4

309 + speculator.current_draft_step = torch.tensor(0)

310 + speculator.input_buffers = SimpleNamespace(

311 + positions=torch.arange(2),

312 + query_start_loc=torch.arange(3),

313 + )

314 + speculator.idx_mapping = torch.arange(2)

315 + generate_draft = Mock()

316 + speculator._generate_draft = generate_draft

317 + run_fullgraph = Mock()

318 + speculator.decode_cudagraph_manager = SimpleNamespace(run_fullgraph=run_fullgraph)

319 + batch_desc = BatchExecutionDescriptor(

320 + cg_mode=cg_mode,

321 + num_tokens=2,

322 + num_reqs=2,

323 + )

324 +

325 + getattr(speculator, method_name)(

326 + num_reqs=2,

327 + skip_attn=True,

328 + batch_desc=batch_desc,

329 + seq_lens_cpu_upper_bound=None,

330 + num_tokens_across_dp=None,

331 + )

332 +

333 + assert generate_draft.call_count == expected_eager_calls

334 + assert run_fullgraph.call_count == expected_graph_replays

335 +

336 +

337 +def test_update_draft_decode_metadata_updates_fa3_scheduler_metadata(

338 + monkeypatch,

339 +):

340 + builder = object.__new__(flash_attn_module.FlashAttentionMetadataBuilder)

341 + builder.aot_schedule = True

342 + builder.use_full_cuda_graph = True

343 + builder.scheduler_metadata = torch.zeros(8, dtype=torch.int32)

344 + builder.cache_config = SimpleNamespace(cache_dtype="bfloat16")

345 + builder.kv_cache_dtype = torch.bfloat16

346 + builder.num_heads_q = 2

347 + builder.num_heads_kv = 1

348 + builder.headdim = 128

349 + builder.block_size = 16

350 + builder.dcp_world_size = 1

351 + builder.dcp_rank = 0

352 + builder.cp_kv_cache_interleave_size = 1

353 + builder.aot_sliding_window = None

354 +

355 + expected = torch.tensor([7, 8, 9], dtype=torch.int32)

356 +

357 + def fake_get_scheduler_metadata(**kwargs):

358 + return expected

359 +

360 + monkeypatch.setattr(builder, "_get_scheduler_metadata", fake_get_scheduler_metadata)

361 +

362 + metadata = FlashAttentionMetadata(

363 + num_actual_tokens=3,

364 + max_query_len=2,

365 + query_start_loc=torch.tensor([0, 1, 3], dtype=torch.int32),

366 + max_seq_len=8,

367 + seq_lens=torch.tensor([5, 6], dtype=torch.int32),

368 + block_table=torch.zeros((2, 1), dtype=torch.int32),

369 + slot_mapping=torch.zeros(3, dtype=torch.int32),

370 + use_cascade=False,

371 + common_prefix_len=0,

372 + cu_prefix_query_lens=None,

373 + prefix_kv_lens=None,

374 + suffix_kv_lens=None,

375 + max_dcp_context_kv_len=None,

376 + dcp_context_kv_lens=None,

377 + num_decode_reqs=2,

378 + num_prefill_reqs=0,

379 + num_decode_tokens=3,

380 + num_prefill_tokens=0,

381 + scheduler_metadata=torch.tensor([-1, -1, -1], dtype=torch.int32),

382 + prefix_scheduler_metadata=None,

383 + max_num_splits=4,

384 + causal=True,

385 + sliding_window=None,

386 + mm_prefix_query_range_tensor=None,

387 + rswa_prefix_lens=None,

388 + rswa_window=None,

389 + rswa_window_tensor=None,

390 + )

391 +

392 + builder.update_draft_decode_metadata(metadata)

393 +

394 + assert torch.equal(metadata.scheduler_metadata, expected)

395 + assert torch.equal(builder.scheduler_metadata[:3], expected)

396 +

397 +

398 +def test_update_draft_decode_metadata_skips_without_scheduler_metadata(monkeypatch):

399 + builder = object.__new__(flash_attn_module.FlashAttentionMetadataBuilder)

400 + builder.aot_schedule = True

401 + builder.use_full_cuda_graph = True

402 + builder.scheduler_metadata = torch.zeros(4, dtype=torch.int32)

403 +

404 + called = False

405 +

406 + def fake_get_scheduler_metadata(**kwargs):

407 + nonlocal called

408 + called = True

409 + return torch.tensor([1], dtype=torch.int32)

410 +

411 + monkeypatch.setattr(builder, "_get_scheduler_metadata", fake_get_scheduler_metadata)

412 +

413 + metadata = FlashAttentionMetadata(

414 + num_actual_tokens=1,

415 + max_query_len=1,

416 + query_start_loc=torch.tensor([0, 1], dtype=torch.int32),

417 + max_seq_len=1,

418 + seq_lens=torch.tensor([1], dtype=torch.int32),

419 + block_table=torch.zeros((1, 1), dtype=torch.int32),

420 + slot_mapping=torch.zeros(1, dtype=torch.int32),

421 + use_cascade=False,

422 + common_prefix_len=0,

423 + cu_prefix_query_lens=None,

424 + prefix_kv_lens=None,

425 + suffix_kv_lens=None,

426 + max_dcp_context_kv_len=None,

427 + dcp_context_kv_lens=None,

428 + num_decode_reqs=1,

429 + num_prefill_reqs=0,

430 + num_decode_tokens=1,

431 + num_prefill_tokens=0,

432 + scheduler_metadata=None,

433 + prefix_scheduler_metadata=None,

434 + max_num_splits=1,

435 + causal=True,

436 + sliding_window=None,

437 + mm_prefix_query_range_tensor=None,

438 + rswa_prefix_lens=None,

439 + rswa_window=None,

440 + rswa_window_tensor=None,

441 + )

442 +

443 + builder.update_draft_decode_metadata(metadata)

444 +

445 + assert not called

446 + assert metadata.scheduler_metadata is None

vllm/models/deepseek_v4/amd/rocm.py

@@ -372,6 +372,10 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4FlashMLAMetadataBuil

372

373

374 class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):

375 + # Keep fused multi-step decode disabled until update_draft_decode_metadata()

376 + # also refreshes the ROCm-specific ragged SWA indices and indptrs.

377 + supports_draft_decode_metadata_update = False

378 +

379 def __init__(self, *args, **kwargs):

380 super().__init__(*args, **kwargs)

381 max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens

vllm/v1/attention/backend.py

+13

@@ -633,6 +633,9 @@ class AttentionMetadataBuilder(ABC, Generic[M]):

633 supports_update_block_table: bool = False

634 # Whether the builder constructor requires the block-table width.

635 requires_block_table_width: ClassVar[bool] = False

636 + # Whether all step-dependent draft decode metadata can be updated in place,

637 + # allowing one metadata build to be reused across autoregressive draft steps.

638 + supports_draft_decode_metadata_update: bool = False

639

640 @abstractmethod

641 def __init__(

@@ -757,6 +760,16 @@ class AttentionMetadataBuilder(ABC, Generic[M]):

760 fast_build=True,

761 )

762

763 + def update_draft_decode_metadata(self, metadata: M) -> None:

764 + """Update step-dependent draft decode metadata in place.

765 +

766 + The fused draft loop may call this method during full CUDA graph

767 + capture. CUDA graph replay does not run this Python method, so

768 + implementations must emit capture-safe operations and keep replayed

769 + tensor state in persistent storage.

770 + """

771 + raise NotImplementedError

772 +

773 def use_cascade_attention(

774 self,

775 common_prefix_len: int,

vllm/v1/attention/backends/flash_attn.py

+94 -42

@@ -370,6 +370,57 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

370 ) -> AttentionCGSupport:

371 return cls._cudagraph_support

372

373 + def _get_scheduler_metadata(

374 + self,

375 + *,

376 + aot_schedule: bool,

377 + batch_size: int,

378 + cu_query_lens: torch.Tensor,

379 + max_query_len: int,

380 + seqlens: torch.Tensor,

381 + max_seq_len: int,

382 + causal: bool | torch.Tensor,

383 + max_num_splits: int,

384 + ) -> torch.Tensor | None:

385 + if not aot_schedule:

386 + return None

387 +

388 + cache_dtype = self.cache_config.cache_dtype

389 + if is_quantized_kv_cache(cache_dtype):

390 + qkv_dtype = current_platform.fp8_dtype()

391 + else:

392 + qkv_dtype = self.kv_cache_dtype

393 + return get_scheduler_metadata(

394 + batch_size=batch_size,

395 + max_seqlen_q=max_query_len,

396 + max_seqlen_k=max_seq_len,

397 + num_heads_q=self.num_heads_q * self.dcp_world_size,

398 + num_heads_kv=self.num_heads_kv,

399 + headdim=self.headdim,

400 + cache_seqlens=seqlens,

401 + qkv_dtype=qkv_dtype,

402 + cu_seqlens_q=cu_query_lens,

403 + page_size=self.block_size,

404 + causal=causal,

405 + window_size=_maybe_symmetrize_window(self.aot_sliding_window, causal),

406 + num_splits=max_num_splits,

407 + )

408 +

409 + def _store_scheduler_metadata(

410 + self, scheduler_metadata: torch.Tensor | None

411 + ) -> torch.Tensor | None:

412 + if self.use_full_cuda_graph and scheduler_metadata is not None:

413 + n = scheduler_metadata.shape[0]

414 + assert self.scheduler_metadata is not None

415 + self.scheduler_metadata[:n] = scheduler_metadata

416 + # NOTE(woosuk): We should zero out the rest of the scheduler

417 + # metadata to guarantee the correctness. Otherwise, some thread

418 + # blocks may use the invalid scheduler metadata and overwrite the

419 + # output buffer.

420 + self.scheduler_metadata[n:] = 0

421 + return self.scheduler_metadata[:n]

422 + return scheduler_metadata

423 +

424 def __init__(

425 self,

426 kv_cache_spec: AttentionSpec,

@@ -405,6 +456,14 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

456 self.dcp_world_size = 1

457 self.dcp_rank = 0

458

459 + # Fused draft decode reuses the captured metadata object across draft

460 + # steps. For DCP, build-time host-side decisions such as

461 + # skip_dcp_context_attention() can change the metadata shape/control

462 + # path (for example max_dcp_context_kv_len), and those Python-side

463 + # fields are not refreshed in-place between graph replays. Keep the

464 + # fused path disabled until DCP gets a full replay-safe refresh model.

465 + self.supports_draft_decode_metadata_update = self.dcp_world_size == 1

466 +

467 self.cp_kv_cache_interleave_size = (

468 self.parallel_config.cp_kv_cache_interleave_size

469 )

@@ -533,34 +592,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

592 if envs.VLLM_BATCH_INVARIANT:

593 max_num_splits = 1

594

536 - def schedule(

537 - batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal

538 - ):

539 - cache_dtype = self.cache_config.cache_dtype

540 - if is_quantized_kv_cache(cache_dtype):

541 - qkv_dtype = current_platform.fp8_dtype()

542 - else:

543 - qkv_dtype = self.kv_cache_dtype

544 - if aot_schedule:

545 - return get_scheduler_metadata(

546 - batch_size=batch_size,

547 - max_seqlen_q=max_query_len,

548 - max_seqlen_k=max_seq_len,

549 - num_heads_q=self.num_heads_q * self.dcp_world_size,

550 - num_heads_kv=self.num_heads_kv,

551 - headdim=self.headdim,

552 - cache_seqlens=seqlens,

553 - qkv_dtype=qkv_dtype,

554 - cu_seqlens_q=cu_query_lens,

555 - page_size=self.block_size,

556 - causal=causal,

557 - window_size=_maybe_symmetrize_window(

558 - self.aot_sliding_window, causal

559 - ),

560 - num_splits=max_num_splits,

561 - )

562 - return None

563 -

595 use_cascade = common_prefix_len > 0

596 max_dcp_context_kv_len = 0

597 dcp_context_kv_lens = None

@@ -627,13 +658,15 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

658 (max_seq_len + num_partitions - 1) // num_partitions

659 ) * self.cp_kv_cache_interleave_size

660

630 - scheduler_metadata = schedule(

661 + scheduler_metadata = self._get_scheduler_metadata(

662 + aot_schedule=aot_schedule,

663 batch_size=num_reqs,

664 cu_query_lens=query_start_loc,

665 max_query_len=max_query_len,

666 seqlens=dcp_context_kv_lens,

667 max_seq_len=max_dcp_context_kv_len,

668 causal=False,

669 + max_num_splits=max_num_splits,

670 )

671 elif use_cascade:

672 cu_prefix_query_lens = torch.tensor(

@@ -644,41 +677,38 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

677 )

678 # Use GPU tensor directly - no CPU sync needed

679 suffix_kv_lens = seq_lens[:num_reqs] - common_prefix_len

647 - prefix_scheduler_metadata = schedule(

680 + prefix_scheduler_metadata = self._get_scheduler_metadata(

681 + aot_schedule=aot_schedule,

682 batch_size=1,

683 cu_query_lens=cu_prefix_query_lens,

684 max_query_len=num_actual_tokens,

685 seqlens=prefix_kv_lens,

686 max_seq_len=common_prefix_len,

687 causal=False,

688 + max_num_splits=max_num_splits,

689 )

655 - scheduler_metadata = schedule(

690 + scheduler_metadata = self._get_scheduler_metadata(

691 + aot_schedule=aot_schedule,

692 batch_size=num_reqs,

693 cu_query_lens=query_start_loc,

694 max_query_len=max_query_len,

695 seqlens=suffix_kv_lens,

696 max_seq_len=max_seq_len - common_prefix_len,

697 causal=True,

698 + max_num_splits=max_num_splits,

699 )

700 else:

664 - scheduler_metadata = schedule(

701 + scheduler_metadata = self._get_scheduler_metadata(

702 + aot_schedule=aot_schedule,

703 batch_size=num_reqs,

704 cu_query_lens=query_start_loc,

705 max_query_len=max_query_len,

706 seqlens=seq_lens,

707 max_seq_len=max_seq_len,

708 causal=causal,

709 + max_num_splits=max_num_splits,

710 )

672 - # For FA3 + full cudagraph

673 - if self.use_full_cuda_graph and scheduler_metadata is not None:

674 - n = scheduler_metadata.shape[0]

675 - self.scheduler_metadata[:n] = scheduler_metadata

676 - # NOTE(woosuk): We should zero out the rest of the scheduler

677 - # metadata to guarantee the correctness. Otherwise, some thread

678 - # blocks may use the invalid scheduler metadata and overwrite the

679 - # output buffer.

680 - self.scheduler_metadata[n:] = 0

681 - scheduler_metadata = self.scheduler_metadata[:n]

711 + scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)

712

713 if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32:

714 causal = causal.to(torch.int32)

@@ -772,6 +802,28 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

802 new_metadata.slot_mapping = slot_mapping

803 return new_metadata

804

805 + def update_draft_decode_metadata(self, metadata: FlashAttentionMetadata) -> None:

806 + if metadata.scheduler_metadata is None:

807 + return

808 +

809 + num_reqs = metadata.num_decode_reqs or metadata.seq_lens.shape[0]

810 +

811 + assert self.dcp_world_size == 1

812 + assert not metadata.use_cascade

813 +

814 + scheduler_metadata = self._get_scheduler_metadata(

815 + aot_schedule=True,

816 + batch_size=num_reqs,

817 + cu_query_lens=metadata.query_start_loc,

818 + max_query_len=metadata.max_query_len,

819 + seqlens=metadata.seq_lens,

820 + max_seq_len=metadata.max_seq_len,

821 + causal=metadata.causal,

822 + max_num_splits=metadata.max_num_splits,

823 + )

824 +

825 + metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)

826 +

827 def use_cascade_attention(self, *args, **kwargs) -> bool:

828 return use_cascade_attention(*args, **kwargs)

829

vllm/v1/attention/backends/mla/sparse_swa.py

+35

@@ -400,6 +400,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

400

401 reorder_batch_threshold: int | None = None

402 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

403 + supports_draft_decode_metadata_update = True

404

405 def __init__(self, *args, **kwargs):

406 super().__init__(*args, **kwargs)

@@ -649,6 +650,40 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

650 **deepseek_v4_fields, # type: ignore[arg-type]

651 )

652

653 + def update_draft_decode_metadata(

654 + self,

655 + metadata: DeepseekSparseSWAMetadata,

656 + ) -> None:

657 + if metadata.num_decode_tokens == 0:

658 + return

659 + assert metadata.query_start_loc is not None

660 + assert metadata.seq_lens is not None

661 + assert metadata.token_to_req_indices is not None

662 + assert metadata.is_valid_token is not None

663 + assert metadata.decode_swa_indices is not None

664 + assert metadata.decode_swa_lens is not None

665 +

666 + _compute_swa_indices_and_lens_kernel[(metadata.num_decode_tokens,)](

667 + metadata.decode_swa_indices,

668 + metadata.decode_swa_indices.stride(0),

669 + metadata.decode_swa_lens,

670 + metadata.decode_swa_indices.shape[-1],

671 + metadata.query_start_loc,

672 + metadata.seq_lens,

673 + metadata.token_to_req_indices,

674 + metadata.is_valid_token,

675 + metadata.block_table,

676 + metadata.block_table.stride(0),

677 + self.block_size,

678 + token_offset=0,

679 + TRITON_BLOCK_SIZE=1024,

680 + )

681 + tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens)

682 + metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]

683 + metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A]

684 + metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A]

685 + metadata.flashinfer_sparse_index_cache.clear()

686 +

687 def build_tile_scheduler(

688 self, num_decode_tokens: int

689 ) -> dict[str, FlashMLASchedMeta | None]:

vllm/v1/attention/backends/mla/triton_mla.py

@@ -57,6 +57,8 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

58 def __init__(self, kv_cache_spec, layer_names, vllm_config, device):

59 super().__init__(kv_cache_spec, layer_names, vllm_config, device)

60 + # DCP local sequence lengths are not advanced between draft steps.

61 + self.supports_draft_decode_metadata_update = self.dcp_world_size == 1

62 # Only the non-causal DSpark draft group serves multi-token blocks via

63 # the decode path; raise its reorder threshold to the spec block length

64 # so full-cudagraph capture admits it. Causal usage stays single-token.

@@ -64,6 +66,9 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

66 self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

67 self._reserve_attn_logits_workspace()

69 + def update_draft_decode_metadata(self, _metadata: MLACommonMetadata) -> None:

70 + pass

71 +

72 def _reserve_attn_logits_workspace(self) -> None:

73 """Pre-size the shared workspace for the decode split-KV attn logits.

vllm/v1/attention/backends/triton_attn.py

@@ -100,6 +100,8 @@ class TritonAttentionMetadata:

100

101 class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):

102 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS

103 + # Step-dependent fields reference persistent input buffers directly.

104 + supports_draft_decode_metadata_update = True

105

106 def __init__(

107 self,

@@ -267,6 +269,9 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet

269

270 return attn_metadata

271

272 + def update_draft_decode_metadata(self, _metadata: TritonAttentionMetadata) -> None:

273 + pass

274 +

275

276 class TritonAttentionBackend(AttentionBackend):

277 supported_dtypes: ClassVar[list[torch.dtype]] = [

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

+151 -7

@@ -10,17 +10,21 @@ from vllm.config.compilation import CUDAGraphMode

10 from vllm.forward_context import BatchDescriptor, set_forward_context

11 from vllm.logger import init_logger

12 from vllm.triton_utils import tl, triton

13 +from vllm.v1.kv_cache_interface import KVCacheConfig

14 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer

15 +from vllm.v1.worker.gpu.block_table import BlockTables

16 from vllm.v1.worker.gpu.cudagraph_utils import (

17 BatchExecutionDescriptor,

18 get_uniform_token_count,

19 )

20 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp

21 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

22 +from vllm.v1.worker.gpu.model_states.interface import ModelState

23 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (

24 SpeculatorCudaGraphManager,

25 )

26 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator

27 +from vllm.v1.worker.utils import AttentionGroup

29 logger = init_logger(__name__)

@@ -41,6 +45,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

46 self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None

47 self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None

48 + self.use_fused_multi_step_decode = False

50 def load_model(self, target_model: nn.Module) -> None:

51 super().load_model(target_model)

@@ -76,6 +81,49 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

81 """

82 return True

84 + def set_attn(

85 + self,

86 + model_state: ModelState,

87 + kv_cache_config: KVCacheConfig,

88 + block_tables: BlockTables,

89 + target_input_buffers: InputBuffers,

90 + target_attn_groups: list[list[AttentionGroup]],

91 + ) -> None:

92 + super().set_attn(

93 + model_state,

94 + kv_cache_config,

95 + block_tables,

96 + target_input_buffers,

97 + target_attn_groups,

98 + )

99 + self._configure_fused_multi_step_decode()

100 +

101 + def _configure_fused_multi_step_decode(self) -> None:

102 + if self.num_speculative_steps == 1:

103 + self.use_fused_multi_step_decode = False

104 + return

105 +

106 + if not self.advance_draft_positions:

107 + self.use_fused_multi_step_decode = True

108 + return

109 +

110 + unsupported_backends = sorted(

111 + {

112 + attn_group.backend.get_name()

113 + for attn_groups in self.attn_groups

114 + for attn_group in attn_groups

115 + if not attn_group.supports_draft_decode_metadata_update

116 + }

117 + )

118 + self.use_fused_multi_step_decode = not unsupported_backends

119 + if unsupported_backends:

120 + logger.info_once(

121 + "Fused multi-step draft decode is not supported by attention "

122 + "backend(s) %s; falling back to rebuilding attention metadata "

123 + "between draft steps.",

124 + ", ".join(unsupported_backends),

125 + )

126 +

127 def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:

128 # Initialize cudagraph manager for draft prefill (draft position 0).

129 self.prefill_cudagraph_manager = SpeculatorCudaGraphManager(

@@ -104,6 +152,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

152 # Reset indices to zeros to prevent stale values from prior

153 # dummy runs to cause out-of-bounds indexing during capture.

154 self.last_token_indices.zero_()

155 + self.idx_mapping.zero_()

156

157 # Capture the prefill routine (model forward + compute_logits +

158 # sample).

@@ -133,12 +182,15 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

182 return

183

184 self.on_multi_step_decode_begin(self.max_num_reqs)

136 - # Capture the decode draft generation routine (model forward +

137 - # sample + update_draft_inputs) for a single

138 - # step.

185 + # Capture either the fused decode loop or one decode step per graph.

186 assert self.decode_cudagraph_manager is not None

187 + decode_fn = (

188 + self._generate_fused_drafts

189 + if self.use_fused_multi_step_decode

190 + else self._generate_draft

191 + )

192 self.decode_cudagraph_manager.capture(

141 - self._generate_draft,

193 + decode_fn,

194 self.model_state,

195 self.input_buffers,

196 self.block_tables,

@@ -290,7 +342,12 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

342

343 self.on_multi_step_decode_begin(num_reqs)

344 # Generate the remaining num_speculative_steps - 1 draft tokens.

293 - self._multi_step_decode(

345 + decode_fn = (

346 + self._fused_multi_step_decode

347 + if self.use_fused_multi_step_decode

348 + else self._multi_step_decode

349 + )

350 + decode_fn(

351 num_reqs,

352 dummy_run and skip_attn_for_dummy_run,

353 decode_batch_desc,

@@ -435,10 +492,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

492 step=step,

493 )

494

438 - # Update the current draft step.

495 self.current_draft_step.fill_(step)

496

441 - # Generate draft tokens for the current step.

497 if batch_desc.cg_mode == CUDAGraphMode.FULL:

498 assert self.decode_cudagraph_manager is not None

499 self.decode_cudagraph_manager.run_fullgraph(batch_desc)

@@ -452,6 +507,95 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

507 cudagraph_runtime_mode=batch_desc.cg_mode,

508 )

509

510 + def _fused_multi_step_decode(

511 + self,

512 + num_reqs: int,

513 + skip_attn: bool,

514 + batch_desc: BatchExecutionDescriptor,

515 + num_tokens_across_dp: torch.Tensor | None,

516 + seq_lens_cpu_upper_bound: torch.Tensor,

517 + ) -> None:

518 + positions = self.input_buffers.positions[:num_reqs]

519 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]

520 + idx_mapping = self.idx_mapping[:num_reqs]

521 +

522 + attn_metadata = None

523 + slot_mappings_by_layer = None

524 + if not skip_attn:

525 + slot_mappings = self.block_tables.compute_slot_mappings(

526 + idx_mapping,

527 + query_start_loc,

528 + positions,

529 + batch_desc.num_tokens,

530 + )

531 + if batch_desc.cg_mode != CUDAGraphMode.FULL:

532 + slot_mappings_by_layer = build_slot_mappings_by_layer(

533 + slot_mappings, self.kv_cache_config

534 + )

535 + attn_metadata = self._build_draft_attn_metadata(

536 + num_reqs=num_reqs,

537 + num_reqs_padded=batch_desc.num_reqs or num_reqs,

538 + num_tokens_padded=batch_desc.num_tokens,

539 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

540 + step=1,

541 + )

542 +

543 + if batch_desc.cg_mode == CUDAGraphMode.FULL:

544 + assert self.decode_cudagraph_manager is not None

545 + self.decode_cudagraph_manager.run_fullgraph(batch_desc)

546 + return

547 +

548 + self._generate_fused_drafts(

549 + num_reqs,

550 + batch_desc.num_tokens,

551 + attn_metadata,

552 + slot_mappings_by_layer,

553 + num_tokens_across_dp,

554 + batch_desc.cg_mode,

555 + )

556 +

557 + def _generate_fused_drafts(

558 + self,

559 + num_reqs: int,

560 + num_tokens_padded: int,

561 + attn_metadata: dict[str, Any] | None,

562 + slot_mappings: dict[str, torch.Tensor] | None,

563 + num_tokens_across_dp: torch.Tensor | None,

564 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

565 + ) -> None:

566 + idx_mapping = self.idx_mapping[:num_reqs]

567 + positions = self.input_buffers.positions[:num_reqs]

568 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]

569 + attn_groups = (

570 + [group for groups in self.attn_groups for group in groups]

571 + if attn_metadata is not None

572 + else []

573 + )

574 +

575 + for step in range(1, self.num_speculative_steps):

576 + self.current_draft_step.fill_(step)

577 + self._generate_draft(

578 + num_reqs,

579 + num_tokens_padded,

580 + attn_metadata,

581 + slot_mappings,

582 + num_tokens_across_dp,

583 + cudagraph_runtime_mode,

584 + )

585 + if (

586 + step < self.num_speculative_steps - 1

587 + and attn_metadata is not None

588 + and self.advance_draft_positions

589 + ):

590 + self.block_tables.compute_slot_mappings(

591 + idx_mapping,

592 + query_start_loc,

593 + positions,

594 + num_tokens_padded,

595 + )

596 + for attn_group in attn_groups:

597 + attn_group.update_draft_decode_metadata(attn_metadata)

598 +

599 def _generate_draft(

600 self,

601 num_reqs: int,

vllm/v1/worker/utils.py

+11

@@ -299,6 +299,17 @@ class AttentionGroup:

299 assert len(self.metadata_builders) > ubatch_id

300 return self.metadata_builders[ubatch_id]

301

302 + @property

303 + def supports_draft_decode_metadata_update(self) -> bool:

304 + return self.get_metadata_builder().supports_draft_decode_metadata_update

305 +

306 + def update_draft_decode_metadata(

307 + self,

308 + attn_metadata: Mapping[str, Any],

309 + ) -> None:

310 + metadata = attn_metadata[self.layer_names[0]]

311 + self.get_metadata_builder().update_draft_decode_metadata(metadata)

312 +

313

314 def select_common_block_size(

315 kv_manager_block_size: int,