@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / af3162d3aa

Fork 0

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 8, 20:15 UTC

[Spec Decode] Unified Parallel Drafting (#32887)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>

Benjamin Chislett committed Feb 5, 2026 at 12:37 UTC af3162d3aaa559a738396baf5b5134c1ab0742f5

14 files changed +1085 -392

examples/offline_inference/spec_decode.py

@@ -75,6 +75,7 @@ def parse_args():

75 parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)

76 parser.add_argument("--disable-padded-drafter-batch", action="store_true")

77 parser.add_argument("--max-num-seqs", type=int, default=None)

78 + parser.add_argument("--parallel-drafting", action="store_true")

79 parser.add_argument("--allowed-local-media-path", type=str, default="")

80 return parser.parse_args()

@@ -121,6 +122,7 @@ def main(args):

122 "model": eagle_dir,

123 "num_speculative_tokens": args.num_spec_tokens,

124 "disable_padded_drafter_batch": args.disable_padded_drafter_batch,

125 + "parallel_drafting": args.parallel_drafting,

126 }

127 elif args.method == "ngram":

128 speculative_config = {

@@ -137,6 +139,7 @@ def main(args):

139 "num_speculative_tokens": args.num_spec_tokens,

140 "enforce_eager": args.enforce_eager,

141 "max_model_len": args.max_model_len,

142 + "parallel_drafting": args.parallel_drafting,

143 }

144 elif args.method == "mtp":

145 speculative_config = {

tests/v1/e2e/test_spec_decode.py

+32 -63

@@ -13,15 +13,12 @@ from vllm import LLM, SamplingParams

13 from vllm.assets.base import VLLM_S3_BUCKET_URL

14 from vllm.assets.image import VLM_IMAGES_DIR

15 from vllm.benchmarks.datasets import InstructCoderDataset

16 -from vllm.config.vllm import VllmConfig

16 +from vllm.config import VllmConfig

17 from vllm.distributed import cleanup_dist_env_and_memory

18 from vllm.engine.arg_utils import EngineArgs

19 from vllm.platforms import current_platform

20 from vllm.v1.metrics.reader import Metric

21 -from vllm.v1.spec_decode.draft_model import (

22 - create_vllm_config_for_draft_model,

23 - merge_toks_kernel,

24 -)

21 +from vllm.v1.spec_decode.utils import create_vllm_config_for_draft_model

23 MTP_SIMILARITY_RATE = 0.8

@@ -625,6 +622,8 @@ class ArgsTest:

622 expected_acceptance_rate: float

623 expected_acceptance_len: float

624 # Defaults

625 + enforce_eager: bool = True

626 + parallel_drafting: bool = False

627 target_tensor_parallel_size: int = 1

628 draft_tensor_parallel_size: int = 1

629 max_model_len: int = 1024

@@ -658,7 +657,8 @@ cases = [

657 @pytest.mark.parametrize("args", cases)

658 @pytest.mark.parametrize("enforce_eager", [True, False])

659 def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):

661 - assert_draft_model_correctness(args, enforce_eager)

660 + args.enforce_eager = enforce_eager

661 + assert_draft_model_correctness(args)

662

663

664 def test_draft_model_realistic_example():

@@ -668,11 +668,28 @@ def test_draft_model_realistic_example():

668 dataset="likaixin/InstructCoder",

669 num_speculative_tokens=3,

670 sampling_config=greedy_sampling(),

671 + enforce_eager=False,

672 # values below are not derived, but just prevent a regression

673 expected_acceptance_len=2.8,

674 expected_acceptance_rate=0.55,

675 )

675 - assert_draft_model_correctness(args, enforce_eager=False)

676 + assert_draft_model_correctness(args)

677 +

678 +

679 +def test_draft_model_parallel_drafting():

680 + args = ArgsTest(

681 + target_model="Qwen/Qwen3-1.7B",

682 + draft_model="amd/PARD-Qwen3-0.6B",

683 + dataset="likaixin/InstructCoder",

684 + num_speculative_tokens=3,

685 + sampling_config=greedy_sampling(),

686 + parallel_drafting=True,

687 + enforce_eager=False,

688 + # values below are collected from a stable run, with ~5% tolerance

689 + expected_acceptance_len=2.375,

690 + expected_acceptance_rate=0.45,

691 + )

692 + assert_draft_model_correctness(args)

693

694

695 @pytest.mark.parametrize(

@@ -691,8 +708,9 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):

708 target_model=tgt_model,

709 draft_model=draft_model,

710 **some_high_acceptance_metrics(),

711 + enforce_eager=enforce_eager,

712 )

695 - assert_draft_model_correctness(sd_case, enforce_eager)

713 + assert_draft_model_correctness(sd_case)

714

715

716 def test_draft_model_tensor_parallelism():

@@ -704,8 +722,9 @@ def test_draft_model_tensor_parallelism():

722 draft_model="Qwen/Qwen3-0.6B",

723 draft_tensor_parallel_size=2,

724 **some_high_acceptance_metrics(),

725 + enforce_eager=False,

726 )

708 - assert_draft_model_correctness(sd_case, enforce_eager=False)

727 + assert_draft_model_correctness(sd_case)

728

729

730 def test_draft_model_engine_args_tensor_parallelism():

@@ -750,7 +769,7 @@ def test_draft_model_engine_args_rejects_invalid_tp_argname():

769 engine_args.create_engine_config()

770

771

753 -def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool):

772 +def assert_draft_model_correctness(args: ArgsTest):

773 """Compare the outputs using and not using speculative decoding.

774 In the greedy decoding case, the outputs must match EXACTLY."""

775 test_prompts: list[Messages] = get_messages(

@@ -764,14 +783,15 @@ def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool):

783 "method": "draft_model",

784 "num_speculative_tokens": args.num_speculative_tokens,

785 "max_model_len": args.max_model_len,

767 - "enforce_eager": enforce_eager,

786 + "enforce_eager": args.enforce_eager,

787 "draft_tensor_parallel_size": args.draft_tensor_parallel_size,

788 + "parallel_drafting": args.parallel_drafting,

789 },

790 max_num_seqs=100, # limit cudagraph capture runtime

791 max_model_len=args.max_model_len,

792 gpu_memory_utilization=args.gpu_memory_utilization,

793 tensor_parallel_size=args.target_tensor_parallel_size,

774 - enforce_eager=enforce_eager,

794 + enforce_eager=args.enforce_eager,

795 disable_log_stats=False, # enables get_metrics()

796 )

797 # we don't check the outputs, only check the metrics

@@ -813,57 +833,6 @@ def some_high_acceptance_metrics() -> dict:

833 }

834

835

816 -def test_merge_toks_kernel():

817 - device = "cuda"

818 - merged_len = 5 + 2 # len(target_toks) = 5, batch_size = 2

819 - merged = torch.full((merged_len,), -100, device=device) # -100 is arbitrary

820 - is_rejected_tok = torch.full((merged_len,), True, device=device)

821 - grid = (2,)

822 - merge_toks_kernel[grid](

823 - target_toks_ptr=torch.tensor([0, 1, 2, 0, 1], device=device),

824 - next_toks_ptr=torch.tensor([3, 2], device=device),

825 - query_start_locs_ptr=torch.tensor([0, 3], device=device),

826 - query_end_locs_ptr=torch.tensor([2, 4], device=device),

827 - out_ptr_merged_toks=merged,

828 - out_ptr_is_rejected_tok=is_rejected_tok,

829 - target_toks_size=5,

830 - rejected_tok_fill=-1,

831 - )

832 - expected_merged = torch.tensor([0, 1, 2, 3, 0, 1, 2], device=device)

833 - assert torch.allclose(merged, expected_merged)

834 -

835 - expected_rejected_toks = torch.tensor([False] * merged_len, device=device)

836 - assert torch.allclose(is_rejected_tok, expected_rejected_toks)

837 -

838 -

839 -def test_merge_toks_kernel_with_rejected_tokens():

840 - device = "cuda"

841 - merged_size = 9 + 2 # len(target_toks) = 9, batch_size = 2

842 - merged = torch.full((merged_size,), -100, device=device)

843 - is_rejected_tok = torch.full((merged_size,), True, device=device)

844 - grid = (2,)

845 - merge_toks_kernel[grid](

846 - # rejected tokens

847 - # ↓ ↓ ↓ ↓

848 - target_toks_ptr=torch.tensor([0, 1, 2, 13, 14, 15, 0, 1, 22], device=device),

849 - next_toks_ptr=torch.tensor([3, 2], device=device),

850 - query_start_locs_ptr=torch.tensor([0, 6], device=device),

851 - query_end_locs_ptr=torch.tensor([2, 7], device=device),

852 - out_ptr_merged_toks=merged,

853 - out_ptr_is_rejected_tok=is_rejected_tok,

854 - target_toks_size=9,

855 - rejected_tok_fill=-1,

856 - )

857 - expected_merged = torch.tensor([0, 1, 2, 3, -1, -1, -1, 0, 1, 2, -1], device=device)

858 - assert torch.allclose(merged, expected_merged)

859 -

860 - expected_rejected_toks = torch.tensor(

861 - [False, False, False, False, True, True, True, False, False, False, True],

862 - device=device,

863 - )

864 - assert torch.allclose(is_rejected_tok, expected_rejected_toks)

865 -

866 -

836 def compute_acceptance_rate(metrics: list[Metric]) -> float:

837 name2metric = {metric.name: metric for metric in metrics}

838 n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value # type: ignore

tests/v1/spec_decode/test_eagle.py

+408 -6

@@ -27,6 +27,7 @@ from vllm.config.load import LoadConfig

27 from vllm.model_executor.models.llama import LlamaForCausalLM

28 from vllm.platforms import current_platform

29 from vllm.v1.attention.backends.registry import AttentionBackendEnum

30 +from vllm.v1.spec_decode.draft_model import DraftModelProposer

31 from vllm.v1.spec_decode.eagle import EagleProposer

32 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

33 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

@@ -34,6 +35,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

35 model_dir = "meta-llama/Llama-3.1-8B-Instruct"

36 eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"

37 eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"

38 +ar_draft_model_dir = "amd/PARD-Llama-3.2-1B" # Compatible with parallel and AR drafting

41 def _create_proposer(

@@ -41,11 +43,19 @@ def _create_proposer(

43 num_speculative_tokens: int,

44 attention_backend: str | None = None,

45 speculative_token_tree: list[tuple[int, ...]] | None = None,

46 + parallel_drafting: bool = False,

47 ) -> EagleProposer:

48 model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)

47 - # Choose model directory based on method

48 - draft_model_dir = eagle_dir if method == "eagle" else eagle3_dir

50 + # Method-dependent setup

51 + if method == "eagle":

52 + draft_model_dir = eagle_dir

53 + elif method == "eagle3":

54 + draft_model_dir = eagle3_dir

55 + elif method == "draft_model":

56 + draft_model_dir = ar_draft_model_dir

57 + else:

58 + raise ValueError(f"Unknown method: {method}")

60 spec_token_tree_str = None

61 if speculative_token_tree is not None:

@@ -59,13 +69,18 @@ def _create_proposer(

69 method=method,

70 num_speculative_tokens=num_speculative_tokens,

71 speculative_token_tree=spec_token_tree_str,

72 + parallel_drafting=parallel_drafting,

73 )

74 + if parallel_drafting:

75 + # Overwrite pard_token to avoid crash during init

76 + speculative_config.draft_model_config.hf_config.pard_token = 0

78 + device = current_platform.device_type

79 vllm_config = VllmConfig(

80 model_config=model_config,

81 cache_config=CacheConfig(),

82 speculative_config=speculative_config,

68 - device_config=DeviceConfig(device=current_platform.device_type),

83 + device_config=DeviceConfig(device=device),

84 parallel_config=ParallelConfig(),

85 load_config=LoadConfig(),

86 scheduler_config=SchedulerConfig(

@@ -75,7 +90,10 @@ def _create_proposer(

90 attention_config=AttentionConfig(backend=attention_backend),

91 )

78 - return EagleProposer(vllm_config=vllm_config, device=current_platform.device_type)

93 + if "eagle" in method:

94 + return EagleProposer(vllm_config=vllm_config, device=device)

95 + else:

96 + return DraftModelProposer(vllm_config=vllm_config, device=device)

99 def test_prepare_next_token_ids():

@@ -321,6 +339,390 @@ def test_prepare_inputs_padded():

339 assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)

340

341

342 +def test_set_inputs_first_pass_default_eagle():

343 + """

344 + Test for set_inputs_first_pass without extra input slots (default EAGLE).

345 +

346 + This tests the path where needs_extra_input_slots=False, which is the

347 + default EAGLE pathway. In this case:

348 + - Input IDs are rotated (shifted by one)

349 + - The next_token_ids are inserted at the last position of each request

350 + - Positions are copied as-is

351 + - Hidden states are copied as-is

352 + - The CommonAttentionMetadata is returned unchanged

353 +

354 + Setup:

355 + - 3 requests with query_lens [3, 2, 4]

356 + - Tokens: [a1, a2, a3, b1, b2, c1, c2, c3, c4]

357 + - After rotation: [a2, a3, -, b2, -, c2, c3, c4, -]

358 + - After inserting next_tokens [100, 200, 300]:

359 + [a2, a3, 100, b2, 200, c2, c3, c4, 300]

360 + """

361 + device = torch.device(current_platform.device_type)

362 +

363 + num_speculative_tokens = 3

364 + proposer = _create_proposer("eagle", num_speculative_tokens)

365 +

366 + # Setup batch with 3 requests

367 + batch_spec = BatchSpec(

368 + seq_lens=[10, 8, 12], # Arbitrary context lengths

369 + query_lens=[3, 2, 4],

370 + )

371 +

372 + common_attn_metadata = create_common_attn_metadata(

373 + batch_spec,

374 + block_size=16,

375 + device=device,

376 + )

377 +

378 + # Input tensors

379 + # Request 0: tokens [10, 11, 12] at positions [7, 8, 9]

380 + # Request 1: tokens [20, 21] at positions [6, 7]

381 + # Request 2: tokens [30, 31, 32, 33] at positions [8, 9, 10, 11]

382 + target_token_ids = torch.tensor(

383 + [10, 11, 12, 20, 21, 30, 31, 32, 33], dtype=torch.int32, device=device

384 + )

385 + target_positions = torch.tensor(

386 + [7, 8, 9, 6, 7, 8, 9, 10, 11], dtype=torch.int64, device=device

387 + )

388 + target_hidden_states = torch.randn(

389 + 9, proposer.hidden_size, dtype=proposer.dtype, device=device

390 + )

391 + next_token_ids = torch.tensor([100, 200, 300], dtype=torch.int32, device=device)

392 +

393 + num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass(

394 + target_token_ids=target_token_ids,

395 + next_token_ids=next_token_ids,

396 + target_positions=target_positions,

397 + target_hidden_states=target_hidden_states,

398 + token_indices_to_sample=None,

399 + cad=common_attn_metadata,

400 + num_rejected_tokens_gpu=None,

401 + )

402 +

403 + assert num_tokens == 9 # Total tokens unchanged

404 +

405 + expected_token_indices_to_sample = torch.tensor(

406 + [2, 4, 8], dtype=torch.int32, device=device

407 + )

408 + assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)

409 +

410 + assert output_cad is common_attn_metadata

411 +

412 + # Verify input_ids are rotated and next_tokens inserted

413 + # Original: [10, 11, 12, 20, 21, 30, 31, 32, 33]

414 + # After shift by 1: [11, 12, 12, 21, 21, 31, 32, 33, 33]

415 + # After inserting at last indices [2, 4, 8]: [11, 12, 100, 21, 200, 31, 32, 33, 300]

416 + expected_input_ids = torch.tensor(

417 + [11, 12, 100, 21, 200, 31, 32, 33, 300], dtype=torch.int32, device=device

418 + )

419 + assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids)

420 +

421 + # Verify positions are copied as-is

422 + assert torch.equal(proposer.positions[:num_tokens], target_positions)

423 +

424 + # Verify hidden states are copied as-is

425 + assert torch.equal(proposer.hidden_states[:num_tokens], target_hidden_states)

426 +

427 +

428 +def test_set_inputs_first_pass_draft_model():

429 + """

430 + Test for set_inputs_first_pass with a draft model (extra input slots,

431 + no shift).

432 +

433 + This tests the path where needs_extra_input_slots=True and

434 + shift_input_ids=False (draft model case). In this case:

435 + - Input IDs are NOT shifted

436 + - Each request gets extra_slots_per_request (1) new slots

437 + - The kernel handles copying tokens and inserting bonus/padding tokens

438 + - A new CommonAttentionMetadata is returned with updated query_start_loc

439 +

440 + Setup:

441 + - 2 requests

442 + - Request 0: tokens [10, 11, 12] at positions [0, 1, 2]

443 + - Only tokens [10, 11] are "valid" (query_end_loc=1),

444 + token 12 is a rejected token from previous speculation

445 + - Request 1: tokens [20, 21] at positions [0, 1], both valid.

446 + - Note: this is less than num_speculative_tokens (2) to ensure

447 + we handle variable lengths correctly.

448 + - next_token_ids: [100, 200] (bonus tokens)

449 +

450 + With extra_slots_per_request=1 and shift=False:

451 + Expected output layout:

452 + Request 0 (indices 0-3):

453 + - idx 0: token 10, pos 0

454 + - idx 1: token 11, pos 1

455 + - idx 2: token 100, pos 2 (bonus token)

456 + - idx 3: padding_token_id, is_rejected=True

457 + Request 1 (indices 4-6):

458 + - idx 4: token 20, pos 0

459 + - idx 5: token 21, pos 1

460 + - idx 6: token 200, pos 2 (bonus token)

461 + """

462 + device = torch.device(current_platform.device_type)

463 +

464 + num_speculative_tokens = 2

465 + block_size = 16

466 +

467 + # Create a proposer configured as a draft model (pass_hidden_states=False)

468 + # We need to mock this since _create_proposer defaults to EAGLE

469 + proposer = _create_proposer("draft_model", num_speculative_tokens)

470 +

471 + proposer.parallel_drafting_token_id = 0

472 + proposer.is_rejected_token_mask = torch.zeros(

473 + proposer.max_num_tokens, dtype=torch.bool, device=device

474 + )

475 + proposer.is_masked_token_mask = torch.zeros(

476 + proposer.max_num_tokens, dtype=torch.bool, device=device

477 + )

478 +

479 + # Mock the attn_metadata_builder to avoid needing the full model setup

480 + mock_kv_cache_spec = mock.MagicMock()

481 + mock_kv_cache_spec.block_size = block_size

482 + mock_builder = mock.MagicMock()

483 + mock_builder.kv_cache_spec = mock_kv_cache_spec

484 + proposer.attn_metadata_builder = mock_builder

485 +

486 + # Request 0: query_len=3 (but 1 rejected), Request 1: query_len=2

487 + batch_spec = BatchSpec(

488 + seq_lens=[3, 2],

489 + query_lens=[3, 2],

490 + )

491 +

492 + common_attn_metadata = create_common_attn_metadata(

493 + batch_spec,

494 + block_size=block_size,

495 + device=device,

496 + arange_block_indices=True, # Use predictable block indices

497 + )

498 +

499 + # Input tensors

500 + target_token_ids = torch.tensor(

501 + [10, 11, 12, 20, 21], dtype=torch.int32, device=device

502 + )

503 + target_positions = torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64, device=device)

504 + target_hidden_states = torch.randn(

505 + 5, proposer.hidden_size, dtype=proposer.dtype, device=device

506 + )

507 + next_token_ids = torch.tensor([100, 200], dtype=torch.int32, device=device)

508 +

509 + num_rejected_tokens_gpu = torch.tensor([1, 0], dtype=torch.int32, device=device)

510 +

511 + num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass(

512 + target_token_ids=target_token_ids,

513 + next_token_ids=next_token_ids,

514 + target_positions=target_positions,

515 + target_hidden_states=target_hidden_states,

516 + token_indices_to_sample=None,

517 + cad=common_attn_metadata,

518 + num_rejected_tokens_gpu=num_rejected_tokens_gpu,

519 + )

520 +

521 + assert proposer.net_num_new_slots_per_request == 1

522 + assert proposer.needs_extra_input_slots

523 +

524 + # total_output_tokens = total_input_tokens + net_num_new_slots * batch_size

525 + assert num_tokens == 7

526 +

527 + # Request 0: [10, 11, 100, padding_token (0)]

528 + # Request 1: [20, 21, 200]

529 + # Combined: [10, 11, 100, 0, 20, 21, 200]

530 + expected_input_ids = torch.tensor(

531 + [10, 11, 100, 0, 20, 21, 200], dtype=torch.int32, device=device

532 + )

533 + assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids)

534 +

535 + # Verify positions

536 + # Request 0: [0, 1, 2, 0 (don't care)]

537 + # Request 1: [0, 1, 2]

538 + # Combined: [0, 1, 2, 0, 0, 1, 2]

539 + expected_positions = torch.tensor(

540 + [0, 1, 2, 0, 0, 1, 2], dtype=torch.int64, device=device

541 + )

542 + assert torch.equal(

543 + proposer.positions[:num_tokens],

544 + expected_positions,

545 + )

546 +

547 + # Verify rejection mask

548 + expected_is_rejected = torch.zeros(7, dtype=torch.bool, device=device)

549 + expected_is_rejected[3] = True # padding token at index 3

550 + assert torch.equal(

551 + proposer.is_rejected_token_mask[:num_tokens], expected_is_rejected

552 + )

553 +

554 + # Verify masked token mask (should all be False for non-parallel drafting)

555 + expected_is_masked = torch.zeros(7, dtype=torch.bool, device=device)

556 + assert torch.equal(proposer.is_masked_token_mask[:num_tokens], expected_is_masked)

557 +

558 + # Verify token_indices_to_sample (bonus tokens at indices 2 and 6)

559 + expected_token_indices_to_sample = torch.tensor(

560 + [2, 6], dtype=torch.int32, device=device

561 + )

562 + assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)

563 +

564 + # Verify the new CAD has updated query_start_loc

565 + # Original: [0, 3, 5] -> New: [0, 4, 7] (each request gains 1 slot)

566 + expected_query_start_loc = torch.tensor([0, 4, 7], dtype=torch.int32, device=device)

567 + assert torch.equal(output_cad.query_start_loc, expected_query_start_loc)

568 +

569 +

570 +def test_set_inputs_first_pass_parallel_drafting():

571 + """

572 + Test for set_inputs_first_pass with parallel drafting (extra input slots,

573 + with shift).

574 +

575 + This tests the path where needs_extra_input_slots=True and

576 + shift_input_ids=True (parallel drafting case). In this case:

577 + - Input IDs ARE shifted (like default EAGLE)

578 + - Each request gets extra_slots_per_request (3) new slots

579 + - Parallel drafting tokens are inserted and marked as masked

580 + - Hidden states are mapped correctly

581 +

582 + Setup:

583 + - 2 requests with query_lens [4, 4] (1 bonus + 3 spec tokens each)

584 + - Request 0: tokens [10, 11, 12, 13] at positions [5, 6, 7, 8]

585 + - Only tokens [10, 11, 12] are "valid", token 13 is rejected

586 + - Request 1: tokens [20, 21, 22, 23] at positions [10, 11, 12, 13], all valid.

587 + - next_token_ids: [100, 200] (bonus tokens)

588 +

589 + With shift_input_ids=True, extra_slots_per_request=3:

590 + Expected output layout:

591 + Request 0 (6 output slots = 4 - 1 + 3):

592 + - idx 0-2: shifted tokens [11, 12, 100]

593 + - idx 3-4: parallel_drafting_tokens, is_masked=True

594 + - idx 5: padding_token, is_rejected=True

595 + Request 1 (6 output slots = 4 - 1 + 3):

596 + - idx 6-8: shifted tokens [21, 22, 23]

597 + - idx 9: bonus token 200

598 + - idx 10-11: parallel_drafting_tokens, is_masked=True

599 + """

600 + device = torch.device(current_platform.device_type)

601 +

602 + num_speculative_tokens = 3

603 + block_size = 16

604 +

605 + proposer = _create_proposer("eagle", num_speculative_tokens, parallel_drafting=True)

606 +

607 + # Override to simulate parallel drafting behavior

608 + proposer.parallel_drafting_token_id = -2

609 + proposer.parallel_drafting_hidden_state_tensor = torch.zeros(

610 + proposer.hidden_size, dtype=proposer.dtype, device=device

611 + )

612 + proposer.is_rejected_token_mask = torch.zeros(

613 + proposer.max_num_tokens, dtype=torch.bool, device=device

614 + )

615 + proposer.is_masked_token_mask = torch.zeros(

616 + proposer.max_num_tokens, dtype=torch.bool, device=device

617 + )

618 +

619 + # Mock the attn_metadata_builder

620 + mock_kv_cache_spec = mock.MagicMock()

621 + mock_kv_cache_spec.block_size = block_size

622 + mock_builder = mock.MagicMock()

623 + mock_builder.kv_cache_spec = mock_kv_cache_spec

624 + proposer.attn_metadata_builder = mock_builder

625 +

626 + # Request 0: query_len=4 (1 rejected), Request 1: query_len=4 (all valid)

627 + batch_spec = BatchSpec(

628 + seq_lens=[9, 14],

629 + query_lens=[4, 4],

630 + )

631 +

632 + common_attn_metadata = create_common_attn_metadata(

633 + batch_spec,

634 + block_size=block_size,

635 + device=device,

636 + arange_block_indices=True,

637 + )

638 +

639 + # Input tensors

640 + target_token_ids = torch.tensor(

641 + [10, 11, 12, 13, 20, 21, 22, 23], dtype=torch.int32, device=device

642 + )

643 + target_positions = torch.tensor(

644 + [5, 6, 7, 8, 10, 11, 12, 13], dtype=torch.int64, device=device

645 + )

646 + target_hidden_states = torch.arange(

647 + 8 * proposer.hidden_size, dtype=proposer.dtype, device=device

648 + ).view(8, proposer.hidden_size)

649 + next_token_ids = torch.tensor([100, 200], dtype=torch.int32, device=device)

650 +

651 + num_rejected_tokens_gpu = torch.tensor([1, 0], dtype=torch.int32, device=device)

652 +

653 + num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass(

654 + target_token_ids=target_token_ids,

655 + next_token_ids=next_token_ids,

656 + target_positions=target_positions,

657 + target_hidden_states=target_hidden_states,

658 + token_indices_to_sample=None,

659 + cad=common_attn_metadata,

660 + num_rejected_tokens_gpu=num_rejected_tokens_gpu,

661 + )

662 +

663 + # total_output_tokens = total_input_tokens + net_num_new_slots * batch_size

664 + # = 8 + 2 * 2 = 12

665 + assert num_tokens == 12

666 +

667 + # Request 0: [11, 12, 100, -2, -2, 0(padding)]

668 + # Request 1: [21, 22, 23, 200, -2, -2]

669 + expected_input_ids = torch.tensor(

670 + [11, 12, 100, -2, -2, 0, 21, 22, 23, 200, -2, -2],

671 + dtype=torch.int32,

672 + device=device,

673 + )

674 + assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids)

675 +

676 + # Verify positions

677 + # Request 0: [5, 6, 7, 8, 9, 0 (don't care)]

678 + # Request 1: [10, 11, 12, 13, 14, 15]

679 + expected_positions = torch.tensor(

680 + [5, 6, 7, 8, 9, 0, 10, 11, 12, 13, 14, 15], dtype=torch.int64, device=device

681 + )

682 + assert torch.equal(

683 + proposer.positions[:num_tokens],

684 + expected_positions,

685 + )

686 +

687 + # Verify rejection mask

688 + expected_is_rejected = torch.zeros(12, dtype=torch.bool, device=device)

689 + expected_is_rejected[5] = True

690 + assert torch.equal(

691 + proposer.is_rejected_token_mask[:num_tokens], expected_is_rejected

692 + )

693 +

694 + # Verify masked token mask (parallel drafting slots should be masked)

695 + expected_is_masked = torch.zeros(12, dtype=torch.bool, device=device)

696 + expected_is_masked[3] = True

697 + expected_is_masked[4] = True

698 + expected_is_masked[10] = True

699 + expected_is_masked[11] = True

700 + assert torch.equal(proposer.is_masked_token_mask[:num_tokens], expected_is_masked)

701 +

702 + # Verify token_indices_to_sample (bonus + parallel drafting tokens)

703 + # Request 0: bonus at 2, parallel at 3, 4

704 + # Request 1: bonus at 9, parallel at 10, 11

705 + expected_token_indices_to_sample = torch.tensor(

706 + [2, 3, 4, 9, 10, 11], dtype=torch.int32, device=device

707 + )

708 + assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)

709 +

710 + # Verify the new CAD has updated query_start_loc

711 + # Original query_lens: [4, 4] -> Output: [6, 6]

712 + expected_query_start_loc = torch.tensor(

713 + [0, 6, 12], dtype=torch.int32, device=device

714 + )

715 + assert torch.equal(output_cad.query_start_loc, expected_query_start_loc)

716 +

717 + # Verify masked positions have the parallel drafting hidden state (zeros)

718 + parallel_drafting_hs = proposer.parallel_drafting_hidden_state_tensor

719 + for i in range(num_tokens):

720 + if expected_is_masked[i]:

721 + assert torch.equal(proposer.hidden_states[i], parallel_drafting_hs), (

722 + f"Masked position {i} should have parallel drafting hidden state"

723 + )

724 +

725 +

726 @pytest.mark.parametrize("method", ["eagle", "eagle3"])

727 @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())

728 @pytest.mark.parametrize("pp_size", [1, 2])

@@ -579,7 +981,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

981 target_positions=target_positions,

982 target_hidden_states=target_hidden_states,

983 next_token_ids=next_token_ids,

582 - last_token_indices=None,

984 + token_indices_to_sample=None,

985 common_attn_metadata=common_attn_metadata,

986 sampling_metadata=sampling_metadata,

987 )

@@ -737,7 +1139,7 @@ def test_propose_tree(spec_token_tree):

1139 target_positions=target_positions,

1140 target_hidden_states=target_hidden_states,

1141 next_token_ids=next_token_ids,

740 - last_token_indices=None,

1142 + token_indices_to_sample=None,

1143 common_attn_metadata=common_attn_metadata,

1144 sampling_metadata=sampling_metadata,

1145 )

tests/v1/spec_decode/test_mtp.py

+1 -1

@@ -204,7 +204,7 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):

204 target_positions=target_positions,

205 target_hidden_states=target_hidden_states,

206 next_token_ids=next_token_ids,

207 - last_token_indices=None,

207 + token_indices_to_sample=None,

208 common_attn_metadata=common_attn_metadata,

209 sampling_metadata=sampling_metadata,

210 )

vllm/config/speculative.py

@@ -116,9 +116,16 @@ class SpeculativeConfig:

116 """Minimum size of ngram token window when using Ngram proposer, if

117 provided. Defaults to 1."""

118

119 + # Alternative drafting strategies

120 speculative_token_tree: str | None = None

121 """Specifies the tree structure for speculative token generation.

122 """

123 + parallel_drafting: bool = False

124 + """Enable parallel drafting, where all speculative tokens are generated

125 + in parallel rather than sequentially. This can improve performance but

126 + requires the speculative model be trained to support parallel drafting.

127 + Only compatible with EAGLE and draft model methods."""

128 +

129 # required configuration params passed from engine

130 target_model_config: SkipValidation[ModelConfig] = None # type: ignore

131 """The configuration of the target model."""

vllm/config/vllm.py

+18 -10

@@ -604,10 +604,13 @@ class VllmConfig:

604 # Currently, async scheduling only support eagle speculative

605 # decoding.

606 if self.speculative_config is not None:

607 - if self.speculative_config.method not in get_args(EagleModelTypes):

607 + if (

608 + self.speculative_config.method not in get_args(EagleModelTypes)

609 + and self.speculative_config.method != "draft_model"

610 + ):

611 raise ValueError(

612 "Currently, async scheduling is only supported "

610 - "with EAGLE/MTP kind of speculative decoding."

613 + "with EAGLE/MTP/Draft Model kind of speculative decoding."

614 )

615 if self.speculative_config.disable_padded_drafter_batch:

616 raise ValueError(

@@ -1298,16 +1301,21 @@ class VllmConfig:

1301 computed_compile_ranges_split_points = []

1302

1303 # The upper bound of the compile ranges is the max_num_batched_tokens.

1301 - # For speculative decoding with draft model, the compile range must be extended

1302 - # by 1 for each sequence.

1304 + # For speculative decoding, the compile range must be extended

1305 + # - Sequential: + 1 * max_num_seqs (one draft token per iteration)

1306 + # - Parallel draft: + num_speculative_tokens * max_num_seqs

1307 compile_range_end = self.scheduler_config.max_num_batched_tokens

1308 if compile_range_end is not None:

1305 - do_extend: bool = (

1306 - self.speculative_config is not None

1307 - and self.speculative_config.uses_draft_model()

1308 - )

1309 - if do_extend:

1310 - compile_range_end += self.scheduler_config.max_num_seqs

1309 + if self.speculative_config is not None and (

1310 + self.speculative_config.uses_draft_model()

1311 + or self.speculative_config.use_eagle()

1312 + ):

1313 + multiplier = (

1314 + self.speculative_config.num_speculative_tokens

1315 + if self.speculative_config.parallel_drafting

1316 + else 1

1317 + )

1318 + compile_range_end += multiplier * self.scheduler_config.max_num_seqs

1319

1320 computed_compile_ranges_split_points.append(compile_range_end)

1321

vllm/model_executor/models/llama_eagle3.py

+39 -3

@@ -52,13 +52,16 @@ class LlamaDecoderLayer(LlamaDecoderLayer):

52 # Subsequent layers use hidden_size (only hidden_states, no embeds)

53 qkv_input_size = 2 * self.hidden_size if layer_idx == 0 else self.hidden_size

55 - # override qkv

55 + # Parallel drafting checkpoints may have attention bias enabled

56 + qkv_bias = getattr(config, "attention_bias", False)

57 +

58 + # Override qkv_proj with correct input size and bias setting

59 self.self_attn.qkv_proj = QKVParallelLinear(

60 qkv_input_size,

61 self.self_attn.head_dim,

62 self.self_attn.total_num_heads,

63 self.self_attn.total_num_kv_heads,

61 - bias=False,

64 + bias=qkv_bias,

65 quant_config=quant_config,

66 prefix=maybe_prefix(prefix, "qkv_proj"),

67 )

@@ -293,6 +296,19 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

296 requires_grad=False,

297 )

298

299 + self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting

300 +

301 + if self.use_parallel_drafting:

302 + self.register_buffer(

303 + "mask_hidden",

304 + torch.zeros(

305 + 1,

306 + (3 if self.model.use_aux_hidden_state else 1)

307 + * self.config.hidden_size,

308 + ),

309 + persistent=False,

310 + )

311 +

312 def embed_input_ids(

313 self,

314 input_ids: torch.Tensor,

@@ -347,12 +363,25 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

363 model_weights = {}

364 includes_draft_id_mapping = False

365 includes_embed_tokens = False

366 + includes_mask_hidden = False

367 for name, loaded_weight in weights:

368 if "t2d" in name:

369 continue

370 if "d2t" in name:

371 name = name.replace("d2t", "draft_id_to_target_id")

372 includes_draft_id_mapping = True

373 + elif "mask_hidden" in name:

374 + # Load mask_hidden directly into buffer

375 + if not self.use_parallel_drafting:

376 + logger.warning(

377 + "mask_hidden found in weights but "

378 + "model is not configured for parallel drafting. "

379 + "Skipping loading mask_hidden."

380 + )

381 + continue

382 + self.mask_hidden.copy_(loaded_weight.view(1, -1))

383 + includes_mask_hidden = True

384 + continue

385 elif "lm_head" not in name:

386 name = "model." + name

387 if "embed_tokens" in name:

@@ -360,7 +389,14 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

389 model_weights[name] = loaded_weight

390 process_eagle_weight(self, name)

391

363 - skip_substrs = []

392 + if not includes_mask_hidden and self.use_parallel_drafting:

393 + raise ValueError(

394 + "mask_hidden not found in weights but "

395 + "model is configured for parallel drafting. "

396 + "Please provide mask_hidden in the weights."

397 + )

398 +

399 + skip_substrs = ["mask_hidden"]

400 if not includes_draft_id_mapping:

401 skip_substrs.append("draft_id_to_target_id")

402 if not includes_embed_tokens:

vllm/v1/attention/backend.py

+6 -1

@@ -480,9 +480,14 @@ class AttentionMetadataBuilder(ABC, Generic[M]):

480 speculative_config is not None

481 and speculative_config.num_speculative_tokens is not None

482 ):

483 + max_num_queries_for_spec = (

484 + 1

485 + + (2 if speculative_config.parallel_drafting else 1)

486 + * speculative_config.num_speculative_tokens

487 + )

488 self.reorder_batch_threshold = max(

489 self.reorder_batch_threshold,

485 - 1 + speculative_config.num_speculative_tokens,

490 + max_num_queries_for_spec,

491 )

492

493 if (

vllm/v1/attention/backends/flashinfer.py

+30 -6

@@ -60,7 +60,7 @@ from vllm.v1.attention.backends.utils import (

60 )

61 from vllm.v1.attention.ops.common import cp_lse_ag_out_rs

62 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states

63 -from vllm.v1.kv_cache_interface import AttentionSpec

63 +from vllm.v1.kv_cache_interface import AttentionSpec, UniformTypeKVCacheSpecs

64 from vllm.v1.utils import CpuGpuBuffer

66 FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024

@@ -658,12 +658,36 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

658 vllm_config: VllmConfig,

659 kv_cache_spec: AttentionSpec,

660 ) -> AttentionCGSupport:

661 - has_trtllm_support = can_use_trtllm_attention(

662 - num_qo_heads=vllm_config.model_config.get_num_attention_heads(

663 - vllm_config.parallel_config

664 - ),

665 - num_kv_heads=kv_cache_spec.num_kv_heads,

661 + """Get the cudagraph support level for FlashInfer attention.

662 +

663 + This depends on whether we can use TRTLLM attention for decodes, since we can

664 + only do UNIFORM_SINGLE_TOKEN_DECODE if it is unavailable.

665 + To check this, we must call can_use_trtllm_attention with the number of KV

666 + heads from the kv_cache_spec. We check all available KV cache specs and

667 + only return UNIFORM_BATCH if all of them support TRTLLM attention.

668 + """

669 + # For UniformTypeKVCacheSpecs, check all contained specs

670 + kv_specs = (

671 + kv_cache_spec.kv_cache_specs.values()

672 + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs)

673 + else [kv_cache_spec]

674 )

675 + num_qo_heads = vllm_config.model_config.get_num_attention_heads(

676 + vllm_config.parallel_config

677 + )

678 + has_trtllm_support: bool = len(kv_specs) > 0

679 + for spec in kv_specs:

680 + if not isinstance(spec, AttentionSpec):

681 + # FlashInfer only applies to attention, so we don't consider other types

682 + # of KV spec (e.g. Mamba) here. This is mostly for type checking.

683 + continue

684 + if not can_use_trtllm_attention(

685 + num_qo_heads=num_qo_heads,

686 + num_kv_heads=spec.num_kv_heads,

687 + ):

688 + has_trtllm_support = False

689 + break

690 +

691 if has_trtllm_support:

692 return AttentionCGSupport.UNIFORM_BATCH

693 else:

vllm/v1/attention/backends/utils.py

-32

@@ -825,38 +825,6 @@ def get_dcp_local_seq_lens(

825 return dcp_local_seq_lens.squeeze(1)

826

827

828 -def extend_all_queries_by_1(

829 - common_attn_metadata: CommonAttentionMetadata,

830 - arange: torch.Tensor,

831 - new_slot_mapping: torch.Tensor,

832 -) -> CommonAttentionMetadata:

833 - """

834 - Creates a new CommonAttentionMetadata with all query lengths increased by 1.

835 - Also all seq lens are increased by 1.

836 - This is useful e.g. in speculative decoding with draft models, where we

837 - extend each sequence by 1 token.

838 - The slot mapping is computed externally, as it requires more information.

839 - """

840 - cad = common_attn_metadata

841 - # query start loc must be increased by [+0, +1, +2, ..., +batch_size]

842 - new_query_start_loc = cad.query_start_loc + arange[: len(cad.query_start_loc)]

843 - new_query_start_loc_cpu = cad.query_start_loc_cpu + torch.arange(

844 - len(cad.query_start_loc_cpu), dtype=torch.int32

845 - )

846 - new_cad = cad.replace(

847 - query_start_loc=new_query_start_loc,

848 - query_start_loc_cpu=new_query_start_loc_cpu,

849 - seq_lens=cad.seq_lens + 1,

850 - # each request is extended by 1 token -> batch_size tokens are added

851 - num_actual_tokens=cad.num_actual_tokens + cad.batch_size(),

852 - # All query lens increase by 1, so max query len increases by 1

853 - max_query_len=cad.max_query_len + 1,

854 - max_seq_len=cad.max_seq_len + 1,

855 - slot_mapping=new_slot_mapping,

856 - )

857 - return new_cad

858 -

859 -

828 def mamba_get_block_table_tensor(

829 block_table: torch.Tensor,

830 seq_lens: torch.Tensor,

vllm/v1/spec_decode/draft_model.py

+23 -222

@@ -1,19 +1,15 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Any

4 import torch

5 +import torch.nn as nn

6 +from typing_extensions import override

7 -from vllm.config import VllmConfig, get_layers_from_vllm_config, replace

8 +from vllm.config import VllmConfig

9 from vllm.logger import init_logger

9 -from vllm.model_executor.layers.attention import Attention

10 from vllm.model_executor.model_loader import get_model

11 -from vllm.triton_utils import tl, triton

12 -from vllm.v1.attention.backends.utils import (

13 - CommonAttentionMetadata,

14 - extend_all_queries_by_1,

15 -)

16 -from vllm.v1.spec_decode.eagle import PADDING_SLOT_ID, SpecDecodeBaseProposer

11 +from vllm.v1.spec_decode.eagle import SpecDecodeBaseProposer

12 +from vllm.v1.spec_decode.utils import create_vllm_config_for_draft_model

14 logger = init_logger(__name__)

@@ -31,37 +27,9 @@ class DraftModelProposer(SpecDecodeBaseProposer):

27 pass_hidden_states_to_model=False,

28 runner=runner,

29 )

34 - self._raise_if_multimodal()

35 - self._raise_if_mrope()

36 - self._raise_if_padded_drafter_batch_disabled()

30 self._raise_if_vocab_size_mismatch()

31 self._raise_if_draft_tp_mismatch()

40 - def _block_size(self) -> int:

41 - builder = self._get_attention_metadata_builder()

42 - return builder.kv_cache_spec.block_size

43 -

44 - def _raise_if_multimodal(self):

45 - if self.supports_mm_inputs:

46 - raise NotImplementedError(

47 - "Speculative Decoding with draft models "

48 - "does not support multimodal models yet"

49 - )

50 -

51 - def _raise_if_mrope(self):

52 - if self.draft_model_config.uses_mrope:

53 - raise NotImplementedError(

54 - "Speculative Decoding with draft models does not support M-RoPE yet"

55 - )

56 -

57 - def _raise_if_padded_drafter_batch_disabled(self):

58 - if self.speculative_config.disable_padded_drafter_batch:

59 - raise NotImplementedError(

60 - "Speculative Decoding with draft models only supports "

61 - "padded drafter batch. Please don't pass --disable-padded-drafter-batch"

62 - " in the speculative_config."

63 - )

64 -

33 def _raise_if_vocab_size_mismatch(self):

34 self.speculative_config.verify_equal_vocab_size_if_draft_model()

@@ -82,193 +50,26 @@ class DraftModelProposer(SpecDecodeBaseProposer):

50 "Please pass 'draft_tensor_parallel_size' in the speculative_config."

51 )

85 - def set_inputs_first_pass(

86 - self,

87 - target_token_ids: torch.Tensor,

88 - next_token_ids: torch.Tensor,

89 - target_positions: torch.Tensor,

90 - last_token_indices: torch.Tensor | None,

91 - cad: CommonAttentionMetadata,

92 - num_rejected_tokens_gpu: torch.Tensor | None,

93 - ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:

94 - batch_size = cad.batch_size()

95 - grid = (batch_size,)

96 - start_locs = cad.query_start_loc[:-1]

97 - end_locs = cad.query_start_loc[1:] - 1

98 - if num_rejected_tokens_gpu is not None:

99 - end_locs -= num_rejected_tokens_gpu

100 -

101 - num_tokens = target_token_ids.shape[0] + batch_size

102 - is_rejected_tok = torch.empty(

103 - (num_tokens,), device=self.input_ids.device, dtype=torch.bool

104 - )

105 - merge_toks_kernel[grid](

106 - target_toks_ptr=target_token_ids,

107 - next_toks_ptr=next_token_ids,

108 - query_start_locs_ptr=start_locs,

109 - query_end_locs_ptr=end_locs,

110 - out_ptr_merged_toks=self.input_ids,

111 - out_ptr_is_rejected_tok=is_rejected_tok,

112 - target_toks_size=target_token_ids.shape[0],

113 - # passing a negative rejected_tok_fill value will raise an error

114 - # when the value is used to index into embeddings.

115 - # Therefore, we pass a valid integer, e.g. 0.

116 - rejected_tok_fill=0,

117 - )

118 - merge_toks_kernel[grid](

119 - target_toks_ptr=target_positions,

120 - next_toks_ptr=target_positions[end_locs] + 1,

121 - query_start_locs_ptr=start_locs,

122 - query_end_locs_ptr=end_locs,

123 - out_ptr_merged_toks=self.positions,

124 - out_ptr_is_rejected_tok=is_rejected_tok,

125 - target_toks_size=target_positions.shape[0],

126 - rejected_tok_fill=0,

127 - )

128 -

129 - # recompute slot mapping

130 - new_slot_mapping = compute_new_slot_mapping(

131 - cad=cad,

132 - new_positions=self.positions[:num_tokens],

133 - is_rejected_token_mask=is_rejected_tok,

134 - block_size=self._block_size(),

135 - max_model_len=self.max_model_len,

136 - )

137 - # update common_attn_metadata

138 - new_cad: CommonAttentionMetadata = extend_all_queries_by_1(

139 - cad,

140 - arange=self.arange,

141 - new_slot_mapping=new_slot_mapping,

142 - )

143 -

144 - new_last_token_indices = new_cad.query_start_loc[1:] - 1

145 - if num_rejected_tokens_gpu is not None:

146 - new_last_token_indices -= num_rejected_tokens_gpu

147 -

148 - return num_tokens, new_last_token_indices, new_cad

149 -

150 - def load_model(self, target_model: Any) -> None:

151 - """Takes target_model to satisfy the type checker."""

152 -

153 - # This must be computed before loading the draft model

154 - # because that mutates the forward_context of the vllm_config

155 - target_attn_layer_names = set(

156 - get_layers_from_vllm_config(self.vllm_config, Attention).keys()

157 - )

158 -

53 + @override

54 + def _get_model(self) -> nn.Module:

55 + # Draft models may be quantized or on different parallelism,

56 + # so we load them with a modified vllm config

57 from vllm.compilation.backends import set_model_tag

161 - draft_vllm_config: VllmConfig = create_vllm_config_for_draft_model(

162 - target_model_vllm_config=self.vllm_config

163 - )

164 - logger.info(

165 - "Starting to load draft model %s. TP=%d, rank=%d",

166 - draft_vllm_config.model_config.model,

167 - draft_vllm_config.parallel_config.tensor_parallel_size,

168 - draft_vllm_config.parallel_config.rank,

169 - )

59 + temp_vllm_config = create_vllm_config_for_draft_model(self.vllm_config)

60 with set_model_tag("draft_model"):

171 - self.model = get_model(vllm_config=draft_vllm_config, prefix="draft_model")

172 -

173 - # This must be computed after loading the draft model

174 - # because that mutates the forward_context of the vllm_config

175 - draft_attn_layer_names = (

176 - get_layers_from_vllm_config(self.vllm_config, Attention).keys()

177 - - target_attn_layer_names

178 - )

179 - self.attn_layer_names = list(draft_attn_layer_names)

180 -

181 -

182 -def create_vllm_config_for_draft_model(

183 - target_model_vllm_config: VllmConfig,

184 -) -> VllmConfig:

185 - """The vllm_config is configured for the target model, e.g.

186 - its quant_config and parallel_config. But the draft model is potentially

187 - quantized differently, and has potentially different tensor_parallel_size.

188 - This function creates a new vllm_config configured for the draft model.

189 - The vllm_config is useful when loading the draft model with get_model().

190 - """

191 - old = target_model_vllm_config

192 - assert old.speculative_config is not None, "speculative_config is not set"

193 - old_spec_config = old.speculative_config

194 - new_parallel_config = replace(

195 - old_spec_config.draft_parallel_config,

196 - rank=old.parallel_config.rank,

197 - )

198 - new: VllmConfig = replace(

199 - old,

200 - quant_config=None, # quant_config is recomputed in __init__()

201 - model_config=old_spec_config.draft_model_config,

202 - parallel_config=new_parallel_config,

203 - )

204 - return new

205 -

206 -

207 -def compute_new_slot_mapping(

208 - cad: CommonAttentionMetadata,

209 - new_positions: torch.Tensor,

210 - is_rejected_token_mask: torch.Tensor,

211 - block_size: int,

212 - max_model_len: int,

213 -):

214 - batch_size, n_blocks_per_req = cad.block_table_tensor.shape

215 - req_indices = torch.arange(batch_size, device=cad.query_start_loc.device)

216 - req_indices = torch.repeat_interleave(

217 - req_indices, cad.naive_query_lens() + 1, output_size=len(new_positions)

218 - )

219 - # Clamp the positions to prevent an out-of-bounds error when indexing

220 - # into block_table_tensor.

221 - clamped_positions = torch.clamp(new_positions, max=max_model_len - 1)

222 - block_table_indices = (

223 - req_indices * n_blocks_per_req + clamped_positions // block_size

224 - )

225 - block_nums = cad.block_table_tensor.view(-1)[block_table_indices]

226 - block_offsets = clamped_positions % block_size

227 - new_slot_mapping = block_nums * block_size + block_offsets

228 - # Mask out the position ids that exceed the max model length.

229 - exceeds_max_model_len = new_positions >= max_model_len

230 - new_slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID)

231 - # Mask out rejected tokens to prevent saves to the KV cache.

232 - new_slot_mapping.masked_fill_(is_rejected_token_mask, PADDING_SLOT_ID)

233 - return new_slot_mapping

234 -

61 + model = get_model(

62 + vllm_config=temp_vllm_config,

63 + prefix="draft_model",

64 + )

65 + return model

236 -@triton.jit

237 -def merge_toks_kernel(

238 - target_toks_ptr,

239 - next_toks_ptr,

240 - query_start_locs_ptr,

241 - query_end_locs_ptr,

242 - out_ptr_merged_toks,

243 - out_ptr_is_rejected_tok,

244 - target_toks_size,

245 - rejected_tok_fill,

246 -):

247 - """

248 - Merges the `target_toks_ptr` and the `next_toks_ptr` into a new tensor

249 - called `out_ptr_merged_toks`. Rejected tokens are those after the

250 - `query_end_locs_ptr` and before the next `query_start_locs_ptr`. Fills the

251 - rejected tokens positions with the value `rejected_tok_fill`. Also fills a mask

252 - of the rejected tokens in `out_ptr_is_rejected_tok`.

253 - """

254 - pid = tl.program_id(0)

255 - start_loc = tl.load(query_start_locs_ptr + pid)

256 - is_last_program = pid == tl.num_programs(0) - 1

257 - if is_last_program:

258 - next_start_loc = target_toks_size.to(tl.int32)

259 - else:

260 - next_start_loc = tl.load(query_start_locs_ptr + pid + 1).to(tl.int32)

67 + @override

68 + def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None:

69 + # Draft models don't share embeddings with the target model

70 + pass

262 - end_loc = tl.load(query_end_locs_ptr + pid)

263 - new_val = tl.load(next_toks_ptr + pid)

264 - for i in range(start_loc, next_start_loc + 1):

265 - if i <= end_loc: # copy existing tokens

266 - old_val = tl.load(target_toks_ptr + i)

267 - tl.store(out_ptr_merged_toks + pid + i, old_val)

268 - tl.store(out_ptr_is_rejected_tok + pid + i, False)

269 - elif i == end_loc + 1: # copy bonus token

270 - tl.store(out_ptr_merged_toks + pid + i, new_val)

271 - tl.store(out_ptr_is_rejected_tok + pid + i, False)

272 - else: # fill rejected tokens

273 - tl.store(out_ptr_merged_toks + pid + i, rejected_tok_fill)

274 - tl.store(out_ptr_is_rejected_tok + pid + i, True)

72 + @override

73 + def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:

74 + # Draft models don't share lm_head with the target model

75 + pass

vllm/v1/spec_decode/eagle.py

+269 -47

@@ -43,8 +43,12 @@ from vllm.v1.sample.metadata import SamplingMetadata

43 from vllm.v1.sample.sampler import _SAMPLING_EPS

44 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

45 from vllm.v1.spec_decode.utils import (

46 + PADDING_SLOT_ID,

47 + compute_new_slot_mapping,

48 + copy_and_expand_eagle_inputs_kernel,

49 eagle_prepare_inputs_padded_kernel,

50 eagle_prepare_next_token_padded_kernel,

51 + extend_all_queries_by_N,

52 )

53 from vllm.v1.utils import CpuGpuBuffer

54 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp

@@ -52,8 +56,6 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

57 logger = init_logger(__name__)

55 -PADDING_SLOT_ID = -1

56 -

60 class SpecDecodeBaseProposer:

61 def __init__(

@@ -76,18 +78,35 @@ class SpecDecodeBaseProposer:

78 self.max_model_len = vllm_config.model_config.max_model_len

79 self.dp_rank = vllm_config.parallel_config.data_parallel_rank

80 self.num_speculative_tokens = self.speculative_config.num_speculative_tokens

79 - # The drafter can get longer sequences than the target model.

80 - max_batch_size = vllm_config.scheduler_config.max_num_seqs

81 - self.max_num_tokens = (

82 - vllm_config.scheduler_config.max_num_batched_tokens + max_batch_size

83 - )

84 - self.token_arange_np = np.arange(self.max_num_tokens)

81 +

82 # We need to get the hidden size from the draft model config because

83 # the draft model's hidden size can be different from the target model's

84 # hidden size (e.g., Llama 3.3 70B).

85 self.hidden_size = self.draft_model_config.get_hidden_size()

86 self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()

88 + # Unifying eagle, draft model, and parallel drafting support

89 + self.parallel_drafting: bool = self.speculative_config.parallel_drafting

90 + self.extra_slots_per_request = (

91 + 1 if not self.parallel_drafting else self.num_speculative_tokens

92 + )

93 + self.net_num_new_slots_per_request = self.extra_slots_per_request - (

94 + 1 if self.pass_hidden_states_to_model else 0

95 + )

96 + self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0

97 +

98 + self.parallel_drafting_token_id: int = 0

99 + self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None

100 + if self.parallel_drafting:

101 + self._init_parallel_drafting_params()

102 +

103 + # The drafter can get longer sequences than the target model.

104 + max_batch_size = vllm_config.scheduler_config.max_num_seqs

105 + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + (

106 + self.net_num_new_slots_per_request * max_batch_size

107 + )

108 + self.token_arange_np = np.arange(self.max_num_tokens)

109 +

110 # Multi-modal data support

111 self.mm_registry = MULTIMODAL_REGISTRY

112 self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(

@@ -155,6 +174,26 @@ class SpecDecodeBaseProposer:

174 max_num_slots_for_arange, device=device, dtype=torch.int32

175 )

176

177 + if self.needs_extra_input_slots:

178 + self._raise_if_padded_drafter_batch_disabled()

179 + self._raise_if_multimodal()

180 + self._raise_if_mrope()

181 +

182 + self.is_rejected_token_mask: torch.Tensor | None = None

183 + self.is_masked_token_mask: torch.Tensor | None = None

184 + if self.needs_extra_input_slots:

185 + # For draft models and parallel drafting, we need to keep track of

186 + # which tokens are rejected to update the slot mapping with padding slots.

187 + self.is_rejected_token_mask = torch.zeros(

188 + (self.max_num_tokens,), dtype=torch.bool, device=device

189 + )

190 + # For parallel drafting, we also need to keep track of which tokens

191 + # are parallel-padding tokens used to sample at later positions.

192 + # We populate this tensor even when using draft models for simplicity.

193 + self.is_masked_token_mask = torch.zeros(

194 + (self.max_num_tokens,), dtype=torch.bool, device=device

195 + )

196 +

197 self.inputs_embeds = torch.zeros(

198 (self.max_num_tokens, self.inputs_embeds_size),

199 dtype=self.dtype,

@@ -231,6 +270,49 @@ class SpecDecodeBaseProposer:

270 1, len(self.tree_choices) + 1, device=device, dtype=torch.int32

271 ).repeat(max_batch_size, 1)

272

273 + def _raise_if_padded_drafter_batch_disabled(self):

274 + if self.speculative_config.disable_padded_drafter_batch:

275 + raise NotImplementedError(

276 + "Speculative Decoding with draft models or parallel drafting only "

277 + "supports padded drafter batch. Please unset "

278 + "disable_padded_drafter_batch in the speculative_config."

279 + )

280 +

281 + def _raise_if_multimodal(self):

282 + if self.supports_mm_inputs:

283 + raise NotImplementedError(

284 + "Speculative Decoding with draft models or parallel drafting "

285 + "does not support multimodal models yet"

286 + )

287 +

288 + def _raise_if_mrope(self):

289 + if self.draft_model_config.uses_mrope:

290 + raise NotImplementedError(

291 + "Speculative Decoding with draft models or parallel drafting "

292 + "does not support M-RoPE yet"

293 + )

294 +

295 + def _init_parallel_drafting_params(self):

296 + # For parallel drafting, we need the token ID to use for masked slots

297 + # And for EAGLE + parallel drafting, we need the hidden state tensor to use

298 + # for those masked slots.

299 +

300 + model_hf_config = self.draft_model_config.hf_config

301 + if hasattr(model_hf_config, "pard_token"):

302 + self.parallel_drafting_token_id = model_hf_config.pard_token

303 + elif hasattr(model_hf_config, "ptd_token_id"):

304 + self.parallel_drafting_token_id = model_hf_config.ptd_token_id

305 + else:

306 + raise ValueError(

307 + "For parallel drafting, the draft model config must have "

308 + "`pard_token` or `ptd_token_id` specified in its config.json."

309 + )

310 +

311 + if self.pass_hidden_states_to_model:

312 + self.parallel_drafting_hidden_state_tensor = torch.empty(

313 + self.hidden_size, dtype=self.dtype, device=self.device

314 + )

315 +

316 def _get_positions(self, num_tokens: int):

317 if self.uses_mrope:

318 return self.mrope_positions[:, :num_tokens]

@@ -296,7 +378,7 @@ class SpecDecodeBaseProposer:

378 target_hidden_states: torch.Tensor,

379 # [batch_size]

380 next_token_ids: torch.Tensor,

299 - last_token_indices: torch.Tensor | None,

381 + token_indices_to_sample: torch.Tensor | None,

382 common_attn_metadata: CommonAttentionMetadata,

383 sampling_metadata: SamplingMetadata,

384 mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,

@@ -314,12 +396,13 @@ class SpecDecodeBaseProposer:

396 )

397 assert target_hidden_states.shape[-1] == self.hidden_size

398

317 - num_tokens, last_token_indices, common_attn_metadata = (

399 + num_tokens, token_indices_to_sample, common_attn_metadata = (

400 self.set_inputs_first_pass(

401 target_token_ids=target_token_ids,

402 next_token_ids=next_token_ids,

403 target_positions=target_positions,

322 - last_token_indices=last_token_indices,

404 + target_hidden_states=target_hidden_states,

405 + token_indices_to_sample=token_indices_to_sample,

406 cad=common_attn_metadata,

407 num_rejected_tokens_gpu=num_rejected_tokens_gpu,

408 )

@@ -366,11 +449,6 @@ class SpecDecodeBaseProposer:

449 if num_tokens_across_dp is not None:

450 num_tokens_across_dp[self.dp_rank] = num_input_tokens

451

369 - if self.pass_hidden_states_to_model:

370 - # target_hidden_states and self.hidden_states can have different

371 - # hidden dims. E.g. large target model and small draft model.

372 - self.hidden_states[:num_tokens] = target_hidden_states

373 -

452 if self.supports_mm_inputs:

453 mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

454

@@ -411,27 +489,27 @@ class SpecDecodeBaseProposer:

489 else:

490 last_hidden_states, hidden_states = ret_hidden_states

491

414 - sample_hidden_states = last_hidden_states[last_token_indices]

492 + sample_hidden_states = last_hidden_states[token_indices_to_sample]

493 logits = self.model.compute_logits(sample_hidden_states)

494

495 # Early exit if there is only one draft token to be generated.

418 - if self.num_speculative_tokens == 1:

496 + if self.num_speculative_tokens == 1 or self.parallel_drafting:

497 draft_token_ids = logits.argmax(dim=-1)

420 - return draft_token_ids.view(-1, 1)

498 + return draft_token_ids.view(-1, self.num_speculative_tokens)

499

500 if self.uses_mrope:

423 - positions = self.mrope_positions[:, last_token_indices]

501 + positions = self.mrope_positions[:, token_indices_to_sample]

502 else:

425 - positions = self.positions[last_token_indices]

503 + positions = self.positions[token_indices_to_sample]

504 if self.method in (

505 "deepseek_mtp",

506 "ernie_mtp",

507 "longcat_flash_mtp",

508 "pangu_ultra_moe_mtp",

509 ):

432 - hidden_states = self.hidden_states[last_token_indices]

510 + hidden_states = self.hidden_states[token_indices_to_sample]

511 else:

434 - hidden_states = hidden_states[last_token_indices]

512 + hidden_states = hidden_states[token_indices_to_sample]

513

514 if isinstance(attn_metadata, TreeAttentionMetadata):

515 # Draft using tree attention.

@@ -624,27 +702,139 @@ class SpecDecodeBaseProposer:

702 target_token_ids: torch.Tensor,

703 next_token_ids: torch.Tensor,

704 target_positions: torch.Tensor,

627 - last_token_indices: torch.Tensor | None,

705 + target_hidden_states: torch.Tensor,

706 + token_indices_to_sample: torch.Tensor | None,

707 cad: CommonAttentionMetadata,

708 num_rejected_tokens_gpu: torch.Tensor | None,

709 ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:

631 - if last_token_indices is None:

632 - last_token_indices = cad.query_start_loc[1:] - 1

710 + if not self.needs_extra_input_slots:

711 + # Default EAGLE pathway: no reshaping of input tensors needed.

712 + # Simply rotate the input ids and leave the positions unchanged,

713 + # Inserting the next token ids at the last slot in each request.

714 + if token_indices_to_sample is None:

715 + token_indices_to_sample = cad.query_start_loc[1:] - 1

716 +

717 + num_tokens = target_token_ids.shape[0]

718 + # Shift the input ids by one token.

719 + # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]

720 + self.input_ids[: num_tokens - 1] = target_token_ids[1:]

721 + # Replace the last token with the next token.

722 + # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]

723 + self.input_ids[token_indices_to_sample] = next_token_ids

724 +

725 + # copy inputs to buffer for cudagraph

726 + if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim == 0:

727 + target_positions = target_positions[0]

728 + self._set_positions(num_tokens, target_positions)

729 +

730 + self.hidden_states[:num_tokens] = target_hidden_states

731 +

732 + return num_tokens, token_indices_to_sample, cad

733 + else:

734 + assert self.is_rejected_token_mask is not None

735 + assert self.is_masked_token_mask is not None

736 + # 1.

737 + # Call a custom triton kernel to copy input_ids and positions

738 + # into the correct slots in the preallocated buffers self.input_ids,

739 + # self.positions.

740 + batch_size = cad.batch_size()

741 + # Since we might have to copy a lot of data for prefills, we select the

742 + # block size based on the max query length and limit to max 256 slots/block.

743 + max_num_tokens_per_request = (

744 + cad.max_query_len + self.net_num_new_slots_per_request

745 + )

746 + BLOCK_SIZE_TOKENS = min(

747 + 256, triton.next_power_of_2(max_num_tokens_per_request)

748 + )

749 + num_blocks = (

750 + max_num_tokens_per_request + BLOCK_SIZE_TOKENS - 1

751 + ) // BLOCK_SIZE_TOKENS

752 + total_num_input_tokens = target_token_ids.shape[0]

753 + total_num_output_tokens = total_num_input_tokens + (

754 + self.net_num_new_slots_per_request * batch_size

755 + )

756

634 - num_tokens = target_token_ids.shape[0]

635 - # Shift the input ids by one token.

636 - # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]

637 - self.input_ids[: num_tokens - 1] = target_token_ids[1:]

638 - # Replace the last token with the next token.

639 - # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]

640 - self.input_ids[last_token_indices] = next_token_ids

757 + token_indices_to_sample = torch.empty(

758 + batch_size * self.extra_slots_per_request,

759 + dtype=torch.int32,

760 + device=self.device,

761 + )

762

642 - # copy inputs to buffer for cudagraph

643 - if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim == 0:

644 - target_positions = target_positions[0]

645 - self._set_positions(num_tokens, target_positions)

763 + # Destination indices to write target_hidden_states into drafting buffer.

764 + out_hidden_state_mapping = torch.empty(

765 + total_num_input_tokens, dtype=torch.int32, device=self.device

766 + )

767 +

768 + # Kernel grid: one program per request (row)

769 + grid = (batch_size, num_blocks)

770 + query_start_loc = cad.query_start_loc

771 + query_end_loc = cad.query_start_loc[1:] - 1

772 + if num_rejected_tokens_gpu is not None:

773 + query_end_loc = query_end_loc - num_rejected_tokens_gpu

774 + copy_and_expand_eagle_inputs_kernel[grid](

775 + # (Padded) Inputs from the target model

776 + target_token_ids_ptr=target_token_ids,

777 + target_positions_ptr=target_positions,

778 + next_token_ids_ptr=next_token_ids, # sampled tokens, one per request

779 + # Outputs to the drafting buffers

780 + out_input_ids_ptr=self.input_ids,

781 + out_positions_ptr=self.positions, # Doesn't support mrope for now

782 + out_is_rejected_token_mask_ptr=self.is_rejected_token_mask,

783 + out_is_masked_token_mask_ptr=self.is_masked_token_mask,

784 + out_new_token_indices_ptr=token_indices_to_sample,

785 + out_hidden_state_mapping_ptr=out_hidden_state_mapping,

786 + # Input metadata

787 + query_start_loc_ptr=query_start_loc,

788 + query_end_loc_ptr=query_end_loc,

789 + padding_token_id=0,

790 + parallel_drafting_token_id=self.parallel_drafting_token_id,

791 + # Sizing info

792 + # Note that we can deduce batch_size for free from the grid size

793 + total_input_tokens=total_num_input_tokens,

794 + num_padding_slots_per_request=self.extra_slots_per_request,

795 + shift_input_ids=self.pass_hidden_states_to_model,

796 + BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,

797 + )

798 + if self.pass_hidden_states_to_model:

799 + assert self.parallel_drafting_hidden_state_tensor is not None

800 + self.hidden_states[out_hidden_state_mapping] = target_hidden_states

801 + # Use torch.where to avoid DtoH sync from boolean indexing

802 + mask = self.is_masked_token_mask[:total_num_output_tokens]

803 + torch.where(

804 + mask.unsqueeze(1),

805 + self.parallel_drafting_hidden_state_tensor,

806 + self.hidden_states[:total_num_output_tokens],

807 + out=self.hidden_states[:total_num_output_tokens],

808 + )

809 +

810 + # 2.

811 + # Recompute the slot mapping based on the new positions and

812 + # rejection mask.

813 + builder = (

814 + self._get_attention_metadata_builder()

815 + if self.attn_metadata_builder is None

816 + else self.attn_metadata_builder

817 + )

818 + new_slot_mapping = compute_new_slot_mapping(

819 + cad=cad,

820 + new_positions=self.positions[:total_num_output_tokens],

821 + is_rejected_token_mask=self.is_rejected_token_mask[

822 + :total_num_output_tokens

823 + ],

824 + block_size=builder.kv_cache_spec.block_size,

825 + num_new_tokens=self.net_num_new_slots_per_request,

826 + max_model_len=self.max_model_len,

827 + )

828 +

829 + # 3. Update the common attention metadata with the new (meta)data

830 + new_cad = extend_all_queries_by_N(

831 + cad,

832 + N=self.net_num_new_slots_per_request,

833 + arange=self.arange,

834 + new_slot_mapping=new_slot_mapping,

835 + )

836

647 - return num_tokens, last_token_indices, cad

837 + return total_num_output_tokens, token_indices_to_sample, new_cad

838

839 def model_returns_tuple(self) -> bool:

840 return self.method not in ("mtp", "draft_model")

@@ -1081,8 +1271,21 @@ class SpecDecodeBaseProposer:

1271 model = model.module

1272 return model.__class__.__name__

1273

1274 + def _get_model(self) -> nn.Module:

1275 + """

1276 + Default method to call get_model(). Can be overridden by subclasses which

1277 + need to customize model loading.

1278 + """

1279 + from vllm.compilation.backends import set_model_tag

1280 +

1281 + with set_model_tag("eagle_head"):

1282 + model = get_model(

1283 + vllm_config=self.vllm_config,

1284 + model_config=self.speculative_config.draft_model_config,

1285 + )

1286 + return model

1287 +

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

1085 - draft_model_config = self.speculative_config.draft_model_config

1289 target_attn_layer_names = set(

1290 get_layers_from_vllm_config(

1291 self.vllm_config,

@@ -1096,12 +1299,7 @@ class SpecDecodeBaseProposer:

1299 ).keys()

1300 )

1301

1099 - from vllm.compilation.backends import set_model_tag

1100 -

1101 - with set_model_tag("eagle_head"):

1102 - self.model = get_model(

1103 - vllm_config=self.vllm_config, model_config=draft_model_config

1104 - )

1302 + self.model = self._get_model()

1303

1304 draft_attn_layer_names = (

1305 get_layers_from_vllm_config(

@@ -1170,7 +1368,26 @@ class SpecDecodeBaseProposer:

1368 else:

1369 target_language_model = target_model

1370

1173 - # share embed_tokens with the target model if needed

1371 + self._maybe_share_embeddings(target_language_model)

1372 + self._maybe_share_lm_head(target_language_model)

1373 +

1374 + if self.parallel_drafting and self.pass_hidden_states_to_model:

1375 + assert self.parallel_drafting_hidden_state_tensor is not None

1376 + self.parallel_drafting_hidden_state_tensor.copy_(

1377 + self.model.combine_hidden_states(

1378 + self.model.mask_hidden.view(3 * self.hidden_size)

1379 + )

1380 + if self.eagle3_use_aux_hidden_state

1381 + else self.model.mask_hidden.view(self.hidden_size)

1382 + )

1383 +

1384 + def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None:

1385 + """

1386 + Some draft models may not have their own embedding layers, and some may

1387 + have a duplicate copy of the target model's embedding layers. In these cases,

1388 + we share the target model's embedding layers with the draft model to save

1389 + memory.

1390 + """

1391 if get_pp_group().world_size == 1:

1392 inner_model = getattr(target_language_model, "model", None)

1393 if inner_model is None:

@@ -1233,7 +1450,12 @@ class SpecDecodeBaseProposer:

1450 " from the target model."

1451 )

1452

1236 - # share lm_head with the target model if needed

1453 + def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:

1454 + """

1455 + Some draft models may not have their own LM head, and some may have a

1456 + duplicate copy of the target model's LM head. In these cases, we share

1457 + the target model's LM head with the draft model to save memory.

1458 + """

1459 share_lm_head = False

1460 if hasattr(self.model, "has_own_lm_head"):

1461 # EAGLE model

vllm/v1/spec_decode/utils.py

+248

@@ -1,6 +1,14 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import torch

4 +

5 +from vllm.config import VllmConfig, replace

6 from vllm.triton_utils import tl, triton

7 +from vllm.v1.attention.backends.utils import (

8 + CommonAttentionMetadata,

9 +)

10 +

11 +PADDING_SLOT_ID = -1

14 @triton.jit

@@ -107,3 +115,243 @@ def eagle_prepare_next_token_padded_kernel(

115 tl.store(next_token_ids_ptr + req_idx, backup_token)

116

117 tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)

118 +

119 +

120 +def compute_new_slot_mapping(

121 + cad: CommonAttentionMetadata,

122 + new_positions: torch.Tensor,

123 + is_rejected_token_mask: torch.Tensor,

124 + block_size: int,

125 + num_new_tokens: int,

126 + max_model_len: int,

127 +):

128 + batch_size, n_blocks_per_req = cad.block_table_tensor.shape

129 + req_indices = torch.arange(batch_size, device=cad.query_start_loc.device)

130 + req_indices = torch.repeat_interleave(

131 + req_indices,

132 + cad.naive_query_lens() + num_new_tokens,

133 + output_size=len(new_positions),

134 + )

135 + # Clamp the positions to prevent an out-of-bounds error when indexing

136 + # into block_table_tensor.

137 + clamped_positions = torch.clamp(new_positions, max=max_model_len - 1)

138 + block_table_indices = (

139 + req_indices * n_blocks_per_req + clamped_positions // block_size

140 + )

141 + block_nums = cad.block_table_tensor.view(-1)[block_table_indices]

142 + block_offsets = clamped_positions % block_size

143 + new_slot_mapping = block_nums * block_size + block_offsets

144 + # Mask out the position ids that exceed the max model length.

145 + exceeds_max_model_len = new_positions >= max_model_len

146 + new_slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID)

147 + # Mask out rejected tokens to prevent saves to the KV cache.

148 + new_slot_mapping.masked_fill_(is_rejected_token_mask, PADDING_SLOT_ID)

149 + return new_slot_mapping

150 +

151 +

152 +def create_vllm_config_for_draft_model(

153 + target_model_vllm_config: VllmConfig,

154 +) -> VllmConfig:

155 + """The vllm_config is configured for the target model, e.g.

156 + its quant_config and parallel_config. But the draft model is potentially

157 + quantized differently, and has potentially different tensor_parallel_size.

158 + This function creates a new vllm_config configured for the drafter.

159 + The vllm_config is useful when loading the draft model with get_model().

160 + """

161 + old = target_model_vllm_config

162 + assert old.speculative_config is not None, "speculative_config is not set"

163 + old_spec_config = old.speculative_config

164 + new_parallel_config = replace(

165 + old_spec_config.draft_parallel_config, rank=old.parallel_config.rank

166 + )

167 + new: VllmConfig = replace(

168 + old,

169 + quant_config=None,

170 + parallel_config=new_parallel_config,

171 + model_config=old_spec_config.draft_model_config,

172 + )

173 + return new

174 +

175 +

176 +def extend_all_queries_by_N(

177 + common_attn_metadata: CommonAttentionMetadata,

178 + N: int,

179 + arange: torch.Tensor,

180 + new_slot_mapping: torch.Tensor,

181 +) -> CommonAttentionMetadata:

182 + """

183 + Creates a new CommonAttentionMetadata with all query lengths increased by N.

184 + Also all seq lens are increased by N.

185 + This is useful e.g. in speculative decoding with parallel drafting, where we

186 + extend each sequence by N tokens and predict all tokens in one pass.

187 + The slot mapping is computed externally, as it requires more information.

188 + """

189 + cad = common_attn_metadata

190 + # query start loc must be increased by [+0, +N, +2N, ..., +batch_size * N]

191 + new_query_start_loc = cad.query_start_loc + N * arange[: len(cad.query_start_loc)]

192 + new_query_start_loc_cpu = cad.query_start_loc_cpu + N * torch.arange(

193 + len(cad.query_start_loc_cpu), dtype=torch.int32

194 + )

195 + new_cad = cad.replace(

196 + query_start_loc=new_query_start_loc,

197 + query_start_loc_cpu=new_query_start_loc_cpu,

198 + seq_lens=cad.seq_lens + N,

199 + # each request is extended by N tokens -> batch_size * N tokens are added

200 + num_actual_tokens=cad.num_actual_tokens + cad.batch_size() * N,

201 + # All query lens increase by N, so max query len increases by N

202 + max_query_len=cad.max_query_len + N,

203 + max_seq_len=cad.max_seq_len + N,

204 + slot_mapping=new_slot_mapping,

205 + )

206 + return new_cad

207 +

208 +

209 +# Unified copy/expand kernel

210 +@triton.jit

211 +def copy_and_expand_eagle_inputs_kernel(

212 + # (Padded) Inputs from the target model

213 + target_token_ids_ptr, # [total_tokens_in_batch]

214 + target_positions_ptr, # [total_tokens_in_batch]

215 + next_token_ids_ptr, # [num_reqs]

216 + # Outputs to the drafting buffers

217 + out_input_ids_ptr, # [total_draft_tokens_in_batch] (output)

218 + out_positions_ptr, # [total_draft_tokens_in_batch] (output)

219 + out_is_rejected_token_mask_ptr, # [total_draft_tokens_in_batch] (output)

220 + out_is_masked_token_mask_ptr, # [total_draft_tokens_in_batch] (output)

221 + out_new_token_indices_ptr, # [num_padding_slots_per_request * num_reqs] (output)

222 + out_hidden_state_mapping_ptr, # [total_tokens_in_batch]

223 + # Input metadata

224 + query_start_loc_ptr, # [num_reqs + 1], last value is the total num input tokens

225 + query_end_loc_ptr, # [num_reqs]

226 + padding_token_id, # tl.int32

227 + parallel_drafting_token_id, # tl.int32

228 + # Sizing info

229 + total_input_tokens, # tl.int32

230 + num_padding_slots_per_request, # tl.int32

231 + shift_input_ids, # tl.bool

232 + BLOCK_SIZE_TOKENS: tl.constexpr, # Blocks along token dim to handle prefills

233 +):

234 + """

235 + Copy and expand inputs from the target model to the drafting buffers for Eagle

236 + speculative decoding. This kernel handles padding slots and parallel drafting

237 + tokens, if enabled.

238 + """

239 + request_idx = tl.program_id(axis=0)

240 + token_batch_idx = tl.program_id(axis=1)

241 +

242 + # Load query locations

243 + query_start_loc = tl.load(query_start_loc_ptr + request_idx)

244 + next_query_start_loc = tl.load(query_start_loc_ptr + request_idx + 1)

245 + query_end_loc = tl.load(query_end_loc_ptr + request_idx)

246 +

247 + # Calculate number of valid tokens to copy and input offset

248 + # With shift_input_ids=True, we skip the first token

249 + # Output layout: each request gets (input_len + num_padding_slots_per_request) slots

250 + # But with shift, we lose one token per request

251 + if shift_input_ids:

252 + num_valid_tokens = query_end_loc - query_start_loc

253 + input_offset = 1

254 + output_start = query_start_loc + request_idx * (

255 + num_padding_slots_per_request - 1

256 + )

257 + else:

258 + num_valid_tokens = query_end_loc - query_start_loc + 1

259 + input_offset = 0

260 + output_start = query_start_loc + request_idx * num_padding_slots_per_request

261 +

262 + # Number of rejected tokens from previous speculation

263 + num_rejected = next_query_start_loc - query_end_loc - 1

264 +

265 + # Total output tokens for this request

266 + total_output_tokens = (

267 + num_valid_tokens + num_padding_slots_per_request + num_rejected

268 + )

269 +

270 + # Process tokens in this block

271 + j = token_batch_idx * BLOCK_SIZE_TOKENS + tl.arange(0, BLOCK_SIZE_TOKENS)

272 +

273 + # Compute masks for different output regions:

274 + # [0, num_valid_tokens): valid tokens copied from input

275 + # [num_valid_tokens]: bonus token from next_token_ids

276 + # (num_valid_tokens, num_valid_tokens + num_padding_slots_per_request):

277 + # parallel drafting slots

278 + # [num_valid_tokens + num_padding_slots_per_request, total_output_tokens):

279 + # rejected slots

280 + in_bounds = j < total_output_tokens

281 + is_valid_region = j < num_valid_tokens

282 + is_bonus_region = j == num_valid_tokens

283 + is_parallel_draft_region = (j > num_valid_tokens) & (

284 + j < num_valid_tokens + num_padding_slots_per_request

285 + )

286 + is_rejected_region = j >= num_valid_tokens + num_padding_slots_per_request

287 +

288 + # Compute output indices

289 + out_idx = output_start + j

290 +

291 + # For valid tokens, compute input index

292 + in_idx = query_start_loc + input_offset + j

293 + # Clamp to avoid out-of-bounds access (masked loads still need valid addresses)

294 + in_idx_clamped = tl.minimum(in_idx, total_input_tokens - 1)

295 +

296 + # Load input tokens (masked to valid region)

297 + token_ids = tl.load(

298 + target_token_ids_ptr + in_idx_clamped, mask=is_valid_region & in_bounds, other=0

299 + )

300 +

301 + # Load the starting position for this request (first position in the sequence)

302 + start_pos = tl.load(target_positions_ptr + query_start_loc)

303 +

304 + # Load bonus token for this request

305 + bonus_token = tl.load(next_token_ids_ptr + request_idx)

306 +

307 + # Build final token_ids based on region

308 + token_ids = tl.where(is_bonus_region, bonus_token, token_ids)

309 + token_ids = tl.where(

310 + is_parallel_draft_region, parallel_drafting_token_id, token_ids

311 + )

312 + token_ids = tl.where(is_rejected_region, padding_token_id, token_ids)

313 +

314 + # Build final positions:

315 + # Positions are NOT shifted - they start from the first input position and increment

316 + # Output position j gets start_pos + j

317 + # (e.g., input positions [5,6,7] -> output [5,6,7,8,9,...])

318 + positions = start_pos + j

319 + # Rejected positions are don't-care, set to 0

320 + positions = tl.where(is_rejected_region, 0, positions)

321 +

322 + # Compute output masks

323 + is_rejected_out = is_rejected_region & in_bounds

324 + is_masked_out = is_parallel_draft_region & in_bounds

325 +

326 + # Compute indices of new tokens (bonus + parallel drafting) for sampling

327 + # New tokens are at positions

328 + # [num_valid_tokens, num_valid_tokens + num_padding_slots_per_request)

329 + is_new_token_region = (j >= num_valid_tokens) & (

330 + j < num_valid_tokens + num_padding_slots_per_request

331 + )

332 + new_token_local_idx = (

333 + j - num_valid_tokens

334 + ) # 0 for bonus, 1, 2, ... for parallel drafting

335 + new_token_out_idx = (

336 + request_idx * num_padding_slots_per_request + new_token_local_idx

337 + )

338 +

339 + # Compute hidden state mapping (source index -> destination index)

340 + # This maps each input position to its corresponding output position

341 + # Hidden states don't get shifted, so we map all input tokens (including rejected)

342 + if shift_input_ids:

343 + num_input_tokens_this_request = next_query_start_loc - query_start_loc

344 + is_input_region = j < num_input_tokens_this_request

345 + src_idx = query_start_loc + j

346 + tl.store(out_hidden_state_mapping_ptr + src_idx, out_idx, mask=is_input_region)

347 +

348 + # Store outputs

349 + tl.store(out_input_ids_ptr + out_idx, token_ids, mask=in_bounds)

350 + tl.store(out_positions_ptr + out_idx, positions, mask=in_bounds)

351 + tl.store(out_is_rejected_token_mask_ptr + out_idx, is_rejected_out, mask=in_bounds)

352 + tl.store(out_is_masked_token_mask_ptr + out_idx, is_masked_out, mask=in_bounds)

353 + tl.store(

354 + out_new_token_indices_ptr + new_token_out_idx,

355 + out_idx,

356 + mask=is_new_token_region & in_bounds,

357 + )

vllm/v1/worker/gpu_model_runner.py

+1 -1

@@ -4090,7 +4090,7 @@ class GPUModelRunner(

4090 target_positions=target_positions,

4091 target_hidden_states=target_hidden_states,

4092 next_token_ids=next_token_ids,

4093 - last_token_indices=token_indices_to_sample,

4093 + token_indices_to_sample=token_indices_to_sample,

4094 sampling_metadata=sampling_metadata,

4095 common_attn_metadata=common_attn_metadata,

4096 mm_embed_inputs=mm_embed_inputs,