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

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 7, 09:00 UTC

[BugFix] Add support for MTP num_speculative_tokens > 1 with sparse MLA (#34552)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

Lucas Wilkinson committed Mar 3, 2026 at 10:21 UTC 28ef9ba399340ea7013df8cd1c359b07acc0a302

7 files changed +258 -195

tests/v1/spec_decode/test_eagle.py

+24 -29

@@ -476,12 +476,12 @@ def test_set_inputs_first_pass_draft_model():

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

479 + # Mock draft_attn_groups 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

482 + mock_attn_group = mock.MagicMock()

483 + mock_attn_group.kv_cache_spec = mock_kv_cache_spec

484 + proposer.draft_attn_groups = [mock_attn_group]

485

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

487 batch_spec = BatchSpec(

@@ -616,12 +616,12 @@ def test_set_inputs_first_pass_parallel_drafting():

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

617 )

618

619 - # Mock the attn_metadata_builder

619 + # Mock draft_attn_groups

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

622 + mock_attn_group = mock.MagicMock()

623 + mock_attn_group.kv_cache_spec = mock_kv_cache_spec

624 + proposer.draft_attn_groups = [mock_attn_group]

625

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

627 batch_spec = BatchSpec(

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

916 proposer.model = model_mock

917

918 # Assign draft attn_layer_names since load_model is not invoked

919 - proposer.attn_layer_names = ["layer.0"]

919 + proposer._draft_attn_layer_names = {"layer.0"}

920

921 # Create input tensors

922 batch_spec = BatchSpec(

@@ -961,20 +961,18 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

961

962 attn_metadata_builder = attn_metadata_builder_cls(

963 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

964 - layer_names=proposer.attn_layer_names,

964 + layer_names=proposer._draft_attn_layer_names,

965 vllm_config=proposer.vllm_config,

966 device=device,

967 )

968

969 - # Mock runner for attention metadata building

969 + # Mock runner and draft_attn_groups for attention metadata building

970 proposer.runner = mock.MagicMock()

971 - proposer.runner.attn_groups.append([mock.MagicMock()])

972 - proposer.runner.attn_groups[0][

973 - 0

974 - ].get_metadata_builder.return_value = attn_metadata_builder

975 - proposer._get_attention_metadata_builder = mock.MagicMock(

976 - return_value=attn_metadata_builder

977 - )

971 + mock_attn_group = mock.MagicMock()

972 + mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder

973 + mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)

974 + mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec

975 + proposer.draft_attn_groups = [mock_attn_group]

976

977 result = proposer.propose(

978 target_token_ids=target_token_ids,

@@ -1089,7 +1087,7 @@ def test_propose_tree(spec_token_tree):

1087 proposer.model = model_mock

1088

1089 # Assign draft attn_layer_names since load_model is not invoked

1092 - proposer.attn_layer_names = ["layer.0"]

1090 + proposer._draft_attn_layer_names = {"layer.0"}

1091

1092 # Get the tree attention metadata builder.

1093 attn_metadata_builder_cls, _ = try_get_attention_backend(

@@ -1097,21 +1095,18 @@ def test_propose_tree(spec_token_tree):

1095 )

1096 attn_metadata_builder = attn_metadata_builder_cls(

1097 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

1100 - layer_names=proposer.attn_layer_names,

1098 + layer_names=proposer._draft_attn_layer_names,

1099 vllm_config=proposer.vllm_config,

1100 device=device,

1101 )

1102

1105 - # Mock runner for attention metadata building.

1103 + # Mock runner and draft_attn_groups for attention metadata building.

1104 proposer.runner = mock.MagicMock()

1107 - proposer.runner.attn_groups.append([mock.MagicMock()])

1108 - proposer.runner.attn_groups[0][0].metadata_builders = [attn_metadata_builder]

1109 - proposer.runner.attn_groups[0][

1110 - 0

1111 - ].get_metadata_builder.return_value = attn_metadata_builder

1112 - proposer._get_attention_metadata_builder = mock.MagicMock(

1113 - return_value=attn_metadata_builder

1114 - )

1105 + mock_attn_group = mock.MagicMock()

1106 + mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder

1107 + mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)

1108 + mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec

1109 + proposer.draft_attn_groups = [mock_attn_group]

1110

1111 # Setup inputs for the proposer.

1112 target_token_ids = torch.randint(0, vocab_size, (total_tokens,), device=device)

tests/v1/spec_decode/test_mtp.py

+7 -3

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

162 model_mock.compute_logits.side_effect = logits_returns

163

164 proposer.model = model_mock

165 - proposer.attn_layer_names = ["layer.0"]

165 + proposer._draft_attn_layer_names = {"layer.0"}

166

167 # Prepare inputs

168 batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=seq_lens)

@@ -190,13 +190,17 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):

190

191 attn_metadata_builder = attn_metadata_builder_cls(

192 kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),

193 - layer_names=proposer.attn_layer_names,

193 + layer_names=list(proposer._draft_attn_layer_names),

194 vllm_config=proposer.vllm_config,

195 device=device,

196 )

197

198 proposer.runner = mock.MagicMock()

199 - proposer.attn_metadata_builder = attn_metadata_builder

199 + mock_attn_group = mock.MagicMock()

200 + mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder

201 + mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)

202 + mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec

203 + proposer.draft_attn_groups = [mock_attn_group]

204

205 # Run propose

206 result = proposer.propose(

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -79,6 +79,12 @@ def sparse_attn_indexer(

79 has_prefill = attn_metadata.num_prefills > 0

80 num_decode_tokens = attn_metadata.num_decode_tokens

82 + # During speculative decoding, k may be padded to the CUDA graph batch

83 + # size while slot_mapping only covers actual tokens. Truncate k to avoid

84 + # out-of-bounds reads in the kernel.

85 + num_tokens = slot_mapping.shape[0]

86 + k = k[:num_tokens]

87 +

88 ops.indexer_k_quant_and_cache(

89 k,

90 kv_cache,

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

+108 -32

@@ -12,6 +12,7 @@ from vllm.utils.deep_gemm import (

12 get_paged_mqa_logits_metadata,

13 is_deep_gemm_supported,

14 )

15 +from vllm.utils.math_utils import cdiv

16 from vllm.utils.platform_utils import num_compute_units

17 from vllm.v1.attention.backend import (

18 AttentionBackend,

@@ -24,6 +25,7 @@ from vllm.v1.attention.backends.utils import (

25 split_decodes_and_prefills,

26 split_prefill_chunks,

27 )

28 +from vllm.v1.worker.cp_utils import get_total_cp_world_size

30 logger = init_logger(__name__)

@@ -214,20 +216,39 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

216 if self.vllm_config.speculative_config

217 else 0

218 )

217 - if self.num_speculative_tokens > 1:

218 - raise ValueError(

219 - "Sparse MLA only supports "

220 - "num_speculative_tokens <= 1 because the DeepGEMM "

221 - "fp8_paged_mqa_logits kernel does not support next_n > 2. "

222 - f"Got num_speculative_tokens={self.num_speculative_tokens}."

223 - )

219 self.reorder_batch_threshold += self.num_speculative_tokens

220

221 sm_count = num_compute_units(self.device.index)

222 self.num_sms = sm_count

223

224 self.decode_lens_buffer = torch.empty(

230 - (scheduler_config.max_num_seqs,), dtype=torch.int32, device=self.device

225 + (scheduler_config.max_num_batched_tokens,),

226 + dtype=torch.int32,

227 + device=self.device,

228 + )

229 +

230 + # Pre-allocated buffers for flattening (spec decode).

231 + self.arange_buffer = torch.arange(

232 + scheduler_config.max_num_seqs * (1 + self.num_speculative_tokens),

233 + dtype=torch.int32,

234 + device=self.device,

235 + )

236 + self.expanded_seq_lens_buffer = torch.zeros(

237 + (scheduler_config.max_num_batched_tokens,),

238 + dtype=torch.int32,

239 + device=self.device,

240 + )

241 + max_num_blocks_per_req = cdiv(

242 + self.vllm_config.model_config.max_model_len,

243 + self.kv_cache_spec.block_size * get_total_cp_world_size(),

244 + )

245 + self.expanded_block_table_buffer = torch.zeros(

246 + (

247 + scheduler_config.max_num_batched_tokens,

248 + max_num_blocks_per_req,

249 + ),

250 + dtype=torch.int32,

251 + device=self.device,

252 )

253

254 # See: DeepGMM/csrc/apis/attention.hpp

@@ -326,42 +347,97 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

347 common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]

348 )

349

329 - # Use CPU to avoid GPU sync; breaking async scheduling

330 - requires_padding = (decode_lens_cpu.max() > decode_lens_cpu.min()).item()

350 + seq_lens = common_attn_metadata.seq_lens[:num_decodes]

351 + block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]

352

332 - # Decide which top-k kernel to use based on batch size and sequence length

333 - batch_size = num_decodes

334 - _is_large_context = common_attn_metadata.max_seq_len > 8192

353 + # Padded CUDA graph requests have block_table entries of -1.

354 + # Clamp to 0 to prevent OOB access in the DeepGEMM kernel.

355 + # This is safe because padded requests have seq_lens=0, so the

356 + # kernel produces no meaningful output for those rows.

357 + block_table.clamp_(min=0)

358

336 - # Decision logic based on micro-benchmark results:

337 - # - large_context_topk wins for batch <= 128 and seq_len > 8K

338 - # - top_k_per_row_decode wins for batch > 128 or seq_len <= 8K

339 - use_large_context_topk = batch_size <= 128 and _is_large_context

359 + max_decode_len = int(decode_lens_cpu.max().item())

360 + if max_decode_len > 1:

361 + # Flatten multi-token decode requests into single-token

362 + # batch entries, expanding seq_lens and block tables so

363 + # the kernel always sees next_n=1.

364

341 - next_n = 1 + self.num_speculative_tokens

342 - if next_n > 1:

343 - offsets = torch.arange(next_n, device=self.device, dtype=torch.int32)

344 - else:

345 - offsets = None

365 + # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is

366 + # padding) and decode_lens [3, 1, 4, 0] in the below example comments.

367 + # The context lengths are therefore

368 + # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].

369

347 - seq_lens = common_attn_metadata.seq_lens[:num_decodes]

370 + # 3 + 1 + 4 + 0 = 8

371 + actual_expanded = int(decode_lens_cpu.sum().item())

372 +

373 + # [7, 6, 8, 0] -> [7, 7, 7, 6, 8, 8, 8, 8]

374 + expanded_base = torch.repeat_interleave(

375 + seq_lens - decode_lens, decode_lens

376 + )

377 +

378 + # [0, 3, 4, 8] -> [0, 0, 0, 3, 4, 4, 4, 4]

379 + expanded_starts = torch.repeat_interleave(

380 + common_attn_metadata.query_start_loc[:num_decodes], decode_lens

381 + )

382 +

383 + # [0, 1, 2, 0, 0, 1, 2, 3]

384 + positions_within = (

385 + self.arange_buffer[:actual_expanded] - expanded_starts

386 + )

387 +

388 + # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space

389 + self.expanded_seq_lens_buffer[:actual_expanded] = (

390 + expanded_base + positions_within + 1

391 + )

392 + self.expanded_seq_lens_buffer[actual_expanded:] = 0

393 + seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]

394 +

395 + # Give each of the flattened entries the same block table row as the

396 + # original request.

397 + self.expanded_block_table_buffer[:actual_expanded] = (

398 + torch.repeat_interleave(block_table, decode_lens, dim=0)

399 + )

400 + if actual_expanded < num_decode_tokens:

401 + self.expanded_block_table_buffer[

402 + actual_expanded:num_decode_tokens, 0

403 + ] = 0

404 + block_table = self.expanded_block_table_buffer[:num_decode_tokens]

405 +

406 + # All reqs now have decode_len=1

407 + self.decode_lens_buffer[:num_decode_tokens] = 1

408 + decode_lens = self.decode_lens_buffer[:num_decode_tokens]

409 + offsets = None

410 + batch_size = num_decode_tokens

411 + else:

412 + next_n = 1 + self.num_speculative_tokens

413 + if next_n > 1:

414 + offsets = torch.arange(

415 + next_n, device=self.device, dtype=torch.int32

416 + )

417 + else:

418 + offsets = None

419 + batch_size = num_decodes

420

421 # DeepGEMM is required for the paged MQA logits on CUDA devices

422 if current_platform.is_cuda() and is_deep_gemm_supported():

423 self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(

352 - seq_lens, self.kv_cache_spec.block_size, self.num_sms

424 + seq_lens,

425 + self.kv_cache_spec.block_size,

426 + self.num_sms,

427 )

354 - block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]

355 - # Padded CUDA graph requests have block_table entries of -1.

356 - # Clamp to 0 to prevent OOB access in the DeepGEMM kernel.

357 - # This is safe because padded requests have seq_lens=0, so the

358 - # kernel produces no meaningful output for those rows.

359 - block_table.clamp_(min=0)

428 +

429 + # Decide which top-k kernel to use based on batch size and sequence length

430 + # Decision logic based on micro-benchmark results:

431 + # - large_context_topk wins for batch <= 128 and seq_len > 8K

432 + # - top_k_per_row_decode wins for batch > 128 or seq_len <= 8K

433 + _is_large_context = common_attn_metadata.max_seq_len > 8192

434 + use_large_context_topk = batch_size <= 128 and _is_large_context

435 +

436 decode_metadata = DeepSeekV32IndexerDecodeMetadata(

437 block_table=block_table,

362 - seq_lens=common_attn_metadata.seq_lens[:num_decodes],

438 + seq_lens=seq_lens,

439 decode_lens=decode_lens,

364 - requires_padding=requires_padding,

440 + requires_padding=False,

441 schedule_metadata=self.scheduler_metadata_buffer,

442 use_large_context_topk=use_large_context_topk,

443 offsets=offsets,

vllm/v1/spec_decode/eagle.py

+99 -121

@@ -20,17 +20,13 @@ from vllm.logger import init_logger

20 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

21 from vllm.model_executor.model_loader import get_model

22 from vllm.model_executor.models import supports_multimodal

23 -from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache

23 from vllm.model_executor.models.interfaces import SupportsMultiModal

24 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

25 from vllm.multimodal import MULTIMODAL_REGISTRY

26 from vllm.platforms import current_platform

27 from vllm.triton_utils import triton

28 from vllm.utils.platform_utils import is_pin_memory_available

30 -from vllm.v1.attention.backend import (

31 - AttentionMetadataBuilder,

32 - CommonAttentionMetadata,

33 -)

29 +from vllm.v1.attention.backend import CommonAttentionMetadata

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

31 from vllm.v1.attention.backends.tree_attn import (

32 TreeAttentionMetadata,

@@ -38,7 +34,7 @@ from vllm.v1.attention.backends.tree_attn import (

34 )

35 from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata

36 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

41 -from vllm.v1.kv_cache_interface import KVCacheConfig

37 +from vllm.v1.kv_cache_interface import KVCacheConfig, UniformTypeKVCacheSpecs

38 from vllm.v1.sample.metadata import SamplingMetadata

39 from vllm.v1.sample.sampler import _SAMPLING_EPS

40 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

@@ -53,6 +49,7 @@ from vllm.v1.spec_decode.utils import (

49 from vllm.v1.utils import CpuGpuBuffer

50 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp

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

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

54 logger = init_logger(__name__)

@@ -113,10 +110,8 @@ class SpecDecodeBaseProposer:

110 vllm_config.model_config

111 )

112

116 - self.attn_metadata_builder: AttentionMetadataBuilder | None = None

117 - self.draft_indexer_metadata_builder: AttentionMetadataBuilder | None = None

118 - self.attn_layer_names: list[str] = []

119 - self.indexer_layer_names: list[str] = []

113 + self.draft_attn_groups: list[AttentionGroup] = []

114 + self.kv_cache_gid: int = -1

115 self.eagle3_use_aux_hidden_state: bool = (

116 self._get_eagle3_use_aux_hidden_state_from_config()

117 )

@@ -353,7 +348,7 @@ class SpecDecodeBaseProposer:

348 self._slot_mapping_buffer[num_actual:num_tokens].fill_(PADDING_SLOT_ID)

349

350 view = self._slot_mapping_buffer[:num_tokens]

356 - return {name: view for name in self.attn_layer_names + self.indexer_layer_names}

351 + return {name: view for name in self._draft_attn_layer_names}

352

353 def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:

354 """Initialize cudagraph dispatcher keys for eagle.

@@ -420,33 +415,13 @@ class SpecDecodeBaseProposer:

415

416 assert self.runner is not None

417

423 - if self.attn_metadata_builder is None:

424 - attn_metadata_builder = self._get_attention_metadata_builder()

425 - else:

426 - attn_metadata_builder = self.attn_metadata_builder

427 -

428 - attn_metadata = attn_metadata_builder.build_for_drafting(

429 - common_attn_metadata=common_attn_metadata, draft_index=0

430 - )

431 - # FIXME: support hybrid kv for draft model (remove separate indexer)

432 - if self.draft_indexer_metadata_builder:

433 - draft_indexer_metadata = (

434 - self.draft_indexer_metadata_builder.build_for_drafting(

435 - common_attn_metadata=common_attn_metadata,

436 - draft_index=0,

437 - )

418 + per_layer_attn_metadata: dict[str, object] = {}

419 + for attn_group in self.draft_attn_groups:

420 + attn_metadata = attn_group.get_metadata_builder().build_for_drafting(

421 + common_attn_metadata=common_attn_metadata, draft_index=0

422 )

439 - else:

440 - draft_indexer_metadata = None

441 - # At this moment, we assume all eagle layers belong to the same KV

442 - # cache group, thus using the same attention metadata.

443 - per_layer_attn_metadata = {}

444 - for layer_name in self.attn_layer_names:

445 - per_layer_attn_metadata[layer_name] = attn_metadata

446 -

447 - for layer_name in self.indexer_layer_names:

448 - assert draft_indexer_metadata is not None

449 - per_layer_attn_metadata[layer_name] = draft_indexer_metadata

423 + for layer_name in attn_group.layer_names:

424 + per_layer_attn_metadata[layer_name] = attn_metadata

425

426 cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

427 self._determine_batch_execution_and_padding(num_tokens)

@@ -503,12 +478,7 @@ class SpecDecodeBaseProposer:

478 positions = self.mrope_positions[:, token_indices_to_sample]

479 else:

480 positions = self.positions[token_indices_to_sample]

506 - if self.method in (

507 - "deepseek_mtp",

508 - "ernie_mtp",

509 - "longcat_flash_mtp",

510 - "pangu_ultra_moe_mtp",

511 - ):

481 + if self.method == "mtp":

482 hidden_states = self.hidden_states[token_indices_to_sample]

483 else:

484 hidden_states = hidden_states[token_indices_to_sample]

@@ -613,7 +583,8 @@ class SpecDecodeBaseProposer:

583 common_attn_metadata._num_computed_tokens_cpu += 1

584

585 # Compute the slot mapping.

616 - block_size = attn_metadata_builder.kv_cache_spec.block_size

586 + # Use the first draft attention group's kv_cache_spec for block_size

587 + block_size = self.draft_attn_groups[0].kv_cache_spec.block_size

588 if self.uses_mrope:

589 # all dimensions of positions are the same

590 block_numbers = clamped_positions[0] // block_size

@@ -639,11 +610,13 @@ class SpecDecodeBaseProposer:

610 )

611

612 # Rebuild attention metadata

642 - attn_metadata = attn_metadata_builder.build_for_drafting( # type: ignore

643 - common_attn_metadata=common_attn_metadata, draft_index=token_index + 1

644 - )

645 - for layer_name in self.attn_layer_names:

646 - per_layer_attn_metadata[layer_name] = attn_metadata

613 + for attn_group in self.draft_attn_groups:

614 + attn_metadata = attn_group.get_metadata_builder().build_for_drafting(

615 + common_attn_metadata=common_attn_metadata,

616 + draft_index=token_index + 1,

617 + )

618 + for layer_name in attn_group.layer_names:

619 + per_layer_attn_metadata[layer_name] = attn_metadata

620

621 # copy inputs to buffer for cudagraph

622 self.input_ids[:batch_size] = input_ids

@@ -805,18 +778,17 @@ class SpecDecodeBaseProposer:

778 # 2.

779 # Recompute the slot mapping based on the new positions and

780 # rejection mask.

808 - builder = (

809 - self._get_attention_metadata_builder()

810 - if self.attn_metadata_builder is None

811 - else self.attn_metadata_builder

812 - )

781 + # Use the first draft attention group's kv_cache_spec for block_size

782 + # (all draft layers share the same kv-cache group)

783 + assert len(self.draft_attn_groups) > 0

784 + block_size = self.draft_attn_groups[0].kv_cache_spec.block_size

785 new_slot_mapping = compute_new_slot_mapping(

786 cad=cad,

787 new_positions=self.positions[:total_num_output_tokens],

788 is_rejected_token_mask=self.is_rejected_token_mask[

789 :total_num_output_tokens

790 ],

819 - block_size=builder.kv_cache_spec.block_size,

791 + block_size=block_size,

792 num_new_tokens=self.net_num_new_slots_per_request,

793 max_model_len=self.max_model_len,

794 )

@@ -1000,9 +972,7 @@ class SpecDecodeBaseProposer:

972 | list[dict[str, torch.Tensor]]

973 | None = None,

974 ) -> list[torch.Tensor]:

1003 - tree_attn_metadata_builder = self.runner.attn_groups[0][

1004 - 0

1005 - ].get_metadata_builder()

975 + tree_attn_metadata_builder = self.draft_attn_groups[0].get_metadata_builder()

976 assert isinstance(tree_attn_metadata_builder, TreeAttentionMetadataBuilder)

977

978 total_num_drafts = self.cu_drafts_per_level[0]

@@ -1078,10 +1048,11 @@ class SpecDecodeBaseProposer:

1048 common_attn_metadata=common_attn_metadata, draft_index=level + 1

1049 )

1050

1081 - # Apply new attention metadata to all layers.

1051 + # Apply new attention metadata to all draft layers.

1052 per_layer_attn_metadata = {}

1083 - for layer_name in self.attn_layer_names:

1084 - per_layer_attn_metadata[layer_name] = attn_metadata

1053 + for attn_group in self.draft_attn_groups:

1054 + for layer_name in attn_group.layer_names:

1055 + per_layer_attn_metadata[layer_name] = attn_metadata

1056

1057 # Consider max model length.

1058 attn_metadata.max_seq_len = min(

@@ -1288,43 +1259,17 @@ class SpecDecodeBaseProposer:

1259 AttentionLayerBase, # type: ignore[type-abstract]

1260 ).keys()

1261 )

1291 - # FIXME: support hybrid kv for draft model

1292 - target_indexer_layer_names = set(

1293 - get_layers_from_vllm_config(

1294 - self.vllm_config, DeepseekV32IndexerCache

1295 - ).keys()

1296 - )

1262

1263 self.model = self._get_model()

1264

1300 - draft_attn_layer_names = (

1301 - get_layers_from_vllm_config(

1302 - self.vllm_config,

1303 - AttentionLayerBase, # type: ignore[type-abstract]

1304 - ).keys()

1305 - - target_attn_layer_names

1265 + # Find draft layers (attention layers added by draft model)

1266 + all_attn_layers = get_layers_from_vllm_config(

1267 + self.vllm_config,

1268 + AttentionLayerBase, # type: ignore[type-abstract]

1269 )

1307 - indexer_layers = get_layers_from_vllm_config(

1308 - self.vllm_config, DeepseekV32IndexerCache

1270 + self._draft_attn_layer_names = (

1271 + set(all_attn_layers.keys()) - target_attn_layer_names

1272 )

1310 - draft_indexer_layer_names = indexer_layers.keys() - target_indexer_layer_names

1311 - self.attn_layer_names = list(draft_attn_layer_names - draft_indexer_layer_names)

1312 - self.indexer_layer_names = list(draft_indexer_layer_names)

1313 -

1314 - if self.indexer_layer_names:

1315 - first_layer = self.indexer_layer_names[0]

1316 - self.draft_indexer_metadata_builder = (

1317 - indexer_layers[first_layer]

1318 - .get_attn_backend()

1319 - .get_builder_cls()(

1320 - indexer_layers[first_layer].get_kv_cache_spec(self.vllm_config),

1321 - self.indexer_layer_names,

1322 - self.vllm_config,

1323 - self.device,

1324 - )

1325 - )

1326 - else:

1327 - self.draft_indexer_metadata_builder = None

1273

1274 if self.supports_mm_inputs:

1275 # Even if the target model is multimodal, we can also use

@@ -1562,9 +1507,9 @@ class SpecDecodeBaseProposer:

1507

1508 # Make sure to use EAGLE's own buffer during cudagraph capture.

1509 if (

1565 - self.attn_layer_names

1510 + self._draft_attn_layer_names

1511 and slot_mappings is not None

1567 - and self.attn_layer_names[0] in slot_mappings

1512 + and next(iter(self._draft_attn_layer_names)) in slot_mappings

1513 ):

1514 slot_mapping_dict = self._get_slot_mapping(num_input_tokens)

1515 else:

@@ -1594,31 +1539,6 @@ class SpecDecodeBaseProposer:

1539 kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]

1540 self.model(**kwargs)

1541

1597 - def _get_attention_metadata_builder(self) -> AttentionMetadataBuilder:

1598 - """Find and return the attention metadata builders for EAGLE layers.

1599 -

1600 - Returns:

1601 - The metadata builders for EAGLE layers.

1602 -

1603 - Raises:

1604 - AssertionError: If no metadata builders are found for EAGLE layers.

1605 - """

1606 - builder = None

1607 - chosen_layer = self.attn_layer_names[0]

1608 -

1609 - for kv_cache_group in self.runner.attn_groups:

1610 - for attn_group in kv_cache_group:

1611 - if chosen_layer in attn_group.layer_names:

1612 - builder = attn_group.get_metadata_builder()

1613 - break

1614 - if builder is not None:

1615 - break

1616 -

1617 - assert builder is not None, (

1618 - "Failed to find attention metadata builder for EAGLE layers."

1619 - )

1620 - return builder

1621 -

1542 def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:

1543 """

1544 Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary

@@ -1651,13 +1571,71 @@ class SpecDecodeBaseProposer:

1571 set(

1572 [

1573 kv_cache_groups[layer_name]

1654 - for layer_name in self.attn_layer_names

1574 + for layer_name in self._draft_attn_layer_names

1575 ]

1576 )

1577 )

1578 == 1

1579 ), "All drafting layers should belong to the same kv cache group"

1580

1581 + def initialize_attn_backend(

1582 + self,

1583 + kv_cache_config: KVCacheConfig,

1584 + kernel_block_sizes: list[int] | None = None,

1585 + ) -> None:

1586 + """

1587 + Initialize AttentionGroups for draft layers using kv_cache_config.

1588 + Called from the model runner's initialize_metadata_builders.

1589 + """

1590 + all_attn_layers = get_layers_from_vllm_config(

1591 + self.vllm_config,

1592 + AttentionLayerBase, # type: ignore[type-abstract]

1593 + )

1594 +

1595 + # Find which kv_cache_group the draft layers belong to

1596 + self.validate_same_kv_cache_group(kv_cache_config)

1597 + kv_cache_spec = None

1598 + for gid, group in enumerate(kv_cache_config.kv_cache_groups):

1599 + if self._draft_attn_layer_names & set(group.layer_names):

1600 + self.kv_cache_gid = gid

1601 + kv_cache_spec = group.kv_cache_spec

1602 + break

1603 +

1604 + attention_groups: dict[tuple[str, str], AttentionGroup] = {}

1605 + if kv_cache_spec is not None:

1606 + for layer_name in self._draft_attn_layer_names:

1607 + attn_backend = all_attn_layers[layer_name].get_attn_backend()

1608 + backend_key = attn_backend.full_cls_name()

1609 + if backend_key not in attention_groups:

1610 + layer_kv_cache_spec = kv_cache_spec

1611 + if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):

1612 + layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[

1613 + layer_name

1614 + ]

1615 +

1616 + kernel_block_size = (

1617 + kernel_block_sizes[self.kv_cache_gid]

1618 + if kernel_block_sizes is not None

1619 + and self.kv_cache_gid < len(kernel_block_sizes)

1620 + else None

1621 + )

1622 + attn_group = AttentionGroup(

1623 + backend=attn_backend,

1624 + layer_names=[layer_name],

1625 + kv_cache_spec=layer_kv_cache_spec,

1626 + kv_cache_group_id=self.kv_cache_gid,

1627 + )

1628 + attn_group.create_metadata_builders(

1629 + self.vllm_config,

1630 + self.device,

1631 + kernel_block_size=kernel_block_size,

1632 + )

1633 + attention_groups[backend_key] = attn_group

1634 + else:

1635 + attention_groups[backend_key].layer_names.append(layer_name)

1636 +

1637 + self.draft_attn_groups = list(attention_groups.values())

1638 +

1639 def _determine_batch_execution_and_padding(

1640 self,

1641 num_tokens: int,

vllm/v1/worker/gpu_model_runner.py

+13 -9

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

1936

1937 if self.speculative_config and spec_decode_common_attn_metadata is None:

1938 if isinstance(self.drafter, EagleProposer):

1939 - if self.drafter.attn_layer_names[0] in kv_cache_group.layer_names:

1939 + if self.drafter.kv_cache_gid == kv_cache_gid:

1940 spec_decode_common_attn_metadata = cm

1941 else:

1942 spec_decode_common_attn_metadata = cm

@@ -5559,6 +5559,14 @@ class GPUModelRunner(

5559 # because some of them change the threshold at init time.

5560 self.calculate_reorder_batch_threshold()

5561

5562 + # Initialize drafter attention backend

5563 + if self.speculative_config and (

5564 + self.speculative_config.use_eagle()

5565 + or self.speculative_config.uses_draft_model()

5566 + ):

5567 + assert isinstance(self.drafter, EagleProposer | DraftModelProposer)

5568 + self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes)

5569 +

5570 def _check_and_update_cudagraph_mode(

5571 self,

5572 attention_backends: list[set[type[AttentionBackend]]],

@@ -6079,15 +6087,11 @@ class GPUModelRunner(

6087 kv_cache_config, kernel_block_sizes

6088 )

6089

6082 - if self.speculative_config and (

6083 - self.speculative_config.use_eagle()

6084 - or self.speculative_config.uses_draft_model()

6085 - or self.speculative_config.uses_extract_hidden_states()

6090 + if (

6091 + self.speculative_config

6092 + and self.speculative_config.uses_extract_hidden_states()

6093 ):

6087 - assert isinstance(

6088 - self.drafter,

6089 - EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,

6090 - )

6094 + assert isinstance(self.drafter, ExtractHiddenStatesProposer)

6095 # validate all draft model layers belong to the same kv cache

6096 # group

6097 self.drafter.validate_same_kv_cache_group(kv_cache_config)

vllm/v1/worker/utils.py

+1 -1

@@ -48,7 +48,7 @@ class AttentionGroup:

48 self,

49 vllm_config,

50 device,

51 - kernel_block_size: int | None,

51 + kernel_block_size: int | None = None,

52 num_metadata_builders: int = 1,

53 ):

54 kv_cache_spec_builder = (