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

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 4, 17:30 UTC

[Spec Decode] Move `SpecDecodeBaseProposer` out of `eagle.py` (#40732)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Apr 23, 2026 at 18:28 UTC cde8d247102652713eca27101aae731d8f041cbd

7 files changed +1792 -1785

tests/v1/spec_decode/test_eagle.py

+3 -3

@@ -741,9 +741,9 @@ def test_set_inputs_first_pass_parallel_drafting():

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

742 @pytest.mark.parametrize("use_distinct_embed_tokens", [True, False])

743 @pytest.mark.parametrize("use_distinct_lm_head", [True, False])

744 -@mock.patch("vllm.v1.spec_decode.eagle.get_pp_group")

745 -@mock.patch("vllm.v1.spec_decode.eagle.get_layers_from_vllm_config")

746 -@mock.patch("vllm.v1.spec_decode.eagle.get_model")

744 +@mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_pp_group")

745 +@mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_layers_from_vllm_config")

746 +@mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_model")

747 def test_load_model(

748 mock_get_model,

749 mock_get_layers,

tests/v1/spec_decode/test_mtp.py

+3 -3

@@ -61,9 +61,9 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer:

61 return EagleProposer(vllm_config=vllm_config, device=DEVICE_TYPE)

64 -@mock.patch("vllm.v1.spec_decode.eagle.get_pp_group")

65 -@mock.patch("vllm.v1.spec_decode.eagle.get_layers_from_vllm_config")

66 -@mock.patch("vllm.v1.spec_decode.eagle.get_model")

64 +@mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_pp_group")

65 +@mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_layers_from_vllm_config")

66 +@mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_model")

67 def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_group):

68 """Test MTP-specific model loading with unified model approach."""

vllm/v1/spec_decode/dflash.py

+1 -1

@@ -11,7 +11,7 @@ from vllm.forward_context import set_forward_context

11 from vllm.logger import init_logger

12 from vllm.triton_utils import triton

13 from vllm.v1.attention.backend import CommonAttentionMetadata

14 -from vllm.v1.spec_decode.eagle import SpecDecodeBaseProposer

14 +from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer

15 from vllm.v1.spec_decode.utils import copy_and_expand_dflash_inputs_kernel

17 logger = init_logger(__name__)

vllm/v1/spec_decode/draft_model.py

+1 -1

@@ -9,7 +9,7 @@ from vllm.config import VllmConfig

9 from vllm.config.utils import replace

10 from vllm.logger import init_logger

11 from vllm.model_executor.model_loader import get_model

12 -from vllm.v1.spec_decode.eagle import SpecDecodeBaseProposer

12 +from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer

14 logger = init_logger(__name__)

vllm/v1/spec_decode/eagle.py

+2 -1773

@@ -1,1735 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -import ast

4 -from importlib.util import find_spec

5 -from typing import Any, cast

7 -import numpy as np

4 import torch

9 -import torch.nn as nn

11 -from vllm.config import (

12 - CUDAGraphMode,

13 - VllmConfig,

14 - get_layers_from_vllm_config,

15 - replace,

16 -)

17 -from vllm.distributed.parallel_state import get_pp_group

18 -from vllm.forward_context import set_forward_context

19 -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_eagle3 import Eagle3DeepseekV2ForCausalLM

24 -from vllm.model_executor.models.interfaces import SupportsMultiModal

25 -from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

26 -from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM

27 -from vllm.multimodal import MULTIMODAL_REGISTRY

28 -from vllm.platforms import current_platform

29 -from vllm.utils.platform_utils import is_pin_memory_available

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

31 -from vllm.v1.attention.backends.registry import AttentionBackendEnum

32 -from vllm.v1.attention.backends.tree_attn import (

33 - TreeAttentionMetadata,

34 - TreeAttentionMetadataBuilder,

35 -)

36 -from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata

37 -from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

38 -from vllm.v1.kv_cache_interface import KVCacheConfig, UniformTypeKVCacheSpecs

39 -from vllm.v1.sample.metadata import SamplingMetadata

40 -from vllm.v1.sample.sampler import _SAMPLING_EPS

41 -from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

42 -from vllm.v1.spec_decode.utils import (

43 - PADDING_SLOT_ID,

44 - compute_new_slot_mapping,

45 - copy_and_expand_eagle_inputs_kernel,

46 - eagle_prepare_inputs_padded_kernel,

47 - eagle_prepare_next_token_padded_kernel,

48 - eagle_step_update_slot_mapping_and_metadata,

49 - extend_all_queries_by_N,

50 - next_power_of_2,

51 -)

52 -from vllm.v1.utils import CpuGpuBuffer

53 -from vllm.v1.worker.dp_utils import coordinate_batch_across_dp

54 -from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

55 -from vllm.v1.worker.utils import AttentionGroup

56 -

57 -logger = init_logger(__name__)

58 -

59 -

60 -class SpecDecodeBaseProposer:

61 - def __init__(

62 - self,

63 - vllm_config: VllmConfig,

64 - device: torch.device,

65 - pass_hidden_states_to_model: bool,

66 - runner=None,

67 - ):

68 - self.vllm_config = vllm_config

69 - assert vllm_config.speculative_config is not None

70 - self.speculative_config = vllm_config.speculative_config

71 - self.draft_model_config = self.speculative_config.draft_model_config

72 - self.method = self.speculative_config.method

73 - self.pass_hidden_states_to_model = pass_hidden_states_to_model

74 -

75 - self.device = device

76 - self.dtype = vllm_config.model_config.dtype

77 - self.max_model_len = vllm_config.model_config.max_model_len

78 - self.dp_rank = vllm_config.parallel_config.data_parallel_rank

79 - self.num_speculative_tokens = self.speculative_config.num_speculative_tokens

80 -

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

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

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

84 - self.hidden_size = self.draft_model_config.get_hidden_size()

85 - self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()

86 -

87 - # Unifying eagle, draft model, and parallel drafting support.

88 - # DFlash always uses parallel drafting (all tokens in one pass),

89 - # but has an additional slot for the next_token_id (does not shift like EAGLE)

90 - self.parallel_drafting: bool = self.speculative_config.parallel_drafting

91 - self.extra_slots_per_request = (

92 - 1 if not self.parallel_drafting else self.num_speculative_tokens

93 - )

94 - self.net_num_new_slots_per_request = self.extra_slots_per_request - (

95 - 1 if (self.pass_hidden_states_to_model and self.method != "dflash") else 0

96 - )

97 - self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0

98 -

99 - self.parallel_drafting_token_id: int = 0

100 - self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None

101 - if self.parallel_drafting:

102 - self._init_parallel_drafting_params()

103 - self.use_local_argmax_reduction: bool = (

104 - self.speculative_config.use_local_argmax_reduction

105 - )

106 -

107 - self.max_batch_size = vllm_config.scheduler_config.max_num_seqs

108 - self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

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

110 -

111 - # Can be specialized by methods like DFlash to reduce the limit

112 - self.max_query_tokens = self.max_num_tokens

113 - self.max_positions = self.max_num_tokens

114 -

115 - # Multi-modal data support

116 - self.mm_registry = MULTIMODAL_REGISTRY

117 - self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(

118 - vllm_config.model_config

119 - )

120 -

121 - self.draft_attn_groups: list[AttentionGroup] = []

122 - self.kv_cache_gid: int = -1

123 - self.eagle3_use_aux_hidden_state: bool = (

124 - self._get_eagle3_use_aux_hidden_state_from_config()

125 - )

126 -

127 - self.compilation_config = self.vllm_config.compilation_config

128 -

129 - # Cudagraph dispatcher for PIECEWISE-only dispatching in eagle.

130 - # Keys are initialized later via initialize_cudagraph_keys() called from

131 - # gpu_model_runner._check_and_update_cudagraph_mode after

132 - # adjust_cudagraph_sizes_for_spec_decode is called.

133 - self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)

134 -

135 - # persistent buffers for cuda graph

136 - self.input_ids = torch.zeros(

137 - self.max_num_tokens, dtype=torch.int32, device=device

138 - )

139 - # Use draft model's M-RoPE setting, not target model's

140 - # Draft models may be text-only even if target is multimodal

141 - self.uses_mrope = self.draft_model_config.uses_mrope

142 - self.uses_xdrope_dim = self.vllm_config.model_config.uses_xdrope_dim

143 - self.draft_uses_xdrope_dim = self.draft_model_config.uses_xdrope_dim

144 - if self.uses_mrope:

145 - # NOTE: `mrope_positions` is implemented with one additional dummy

146 - # position on purpose to make it non-contiguous so that it can work

147 - # with torch compile.

148 - # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923

149 -

150 - # NOTE: When M-RoPE is enabled, position ids are 3D regardless of

151 - # the modality of inputs. For text-only inputs, each dimension has

152 - # identical position IDs, making M-RoPE functionally equivalent to

153 - # 1D-RoPE.

154 - # See page 5 of https://arxiv.org/abs/2409.12191

155 - self.mrope_positions = torch.zeros(

156 - (3, self.max_positions + 1), dtype=torch.int64, device=device

157 - )

158 - elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

159 - self.xdrope_positions = torch.zeros(

160 - (self.uses_xdrope_dim, self.max_positions + 1),

161 - dtype=torch.int64,

162 - device=device,

163 - )

164 - else:

165 - # RoPE need (max_num_tokens,)

166 - self.positions = torch.zeros(

167 - self.max_positions,

168 - dtype=torch.int64,

169 - device=device,

170 - )

171 - self.hidden_states = torch.zeros(

172 - (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device

173 - )

174 -

175 - # Will be set when we initialize the attention backend

176 - self.block_size: int = -1

177 -

178 - # We need +1 here because the arange is used to set query_start_loc,

179 - # which has one more element than batch_size.

180 - max_num_slots_for_arange = max(self.max_batch_size + 1, self.max_num_tokens)

181 - self.arange = torch.arange(

182 - max_num_slots_for_arange, device=device, dtype=torch.int32

183 - )

184 -

185 - if self.needs_extra_input_slots:

186 - self._raise_if_padded_drafter_batch_disabled()

187 - self._raise_if_multimodal()

188 - self._raise_if_mrope()

189 -

190 - self.is_rejected_token_mask: torch.Tensor | None = None

191 - self.is_masked_token_mask: torch.Tensor | None = None

192 - if self.needs_extra_input_slots:

193 - # For draft models and parallel drafting, we need to keep track of

194 - # which tokens are rejected to update the slot mapping with padding slots.

195 - self.is_rejected_token_mask = torch.zeros(

196 - (self.max_num_tokens,), dtype=torch.bool, device=device

197 - )

198 - # For parallel drafting, we also need to keep track of which tokens

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

200 - # We populate this tensor even when using draft models for simplicity.

201 - self.is_masked_token_mask = torch.zeros(

202 - (self.max_num_tokens,), dtype=torch.bool, device=device

203 - )

204 -

205 - self.inputs_embeds = torch.zeros(

206 - (self.max_num_tokens, self.inputs_embeds_size),

207 - dtype=self.dtype,

208 - device=device,

209 - )

210 -

211 - self.backup_next_token_ids = CpuGpuBuffer(

212 - self.max_batch_size,

213 - dtype=torch.int32,

214 - pin_memory=is_pin_memory_available(),

215 - device=device,

216 - with_numpy=True,

217 - )

218 -

219 - self._slot_mapping_buffer = torch.zeros(

220 - self.max_positions,

221 - dtype=torch.int64,

222 - device=device,

223 - )

224 -

225 - # Determine allowed attention backends once during initialization.

226 - self.allowed_attn_types: tuple | None = None

227 - if current_platform.is_rocm():

228 - from vllm.v1.attention.backends.mla.indexer import (

229 - DeepseekV32IndexerMetadata,

230 - )

231 - from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (

232 - ROCMAiterMLASparseMetadata,

233 - )

234 - from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata

235 -

236 - rocm_types = [

237 - TritonAttentionMetadata,

238 - RocmAttentionMetadata,

239 - ROCMAiterMLASparseMetadata,

240 - DeepseekV32IndexerMetadata,

241 - ]

242 - # ROCM_AITER_FA is an optional backend

243 - # We check is_enabled() here to avoid importing the backend module during

244 - # auto-discovery when VLLM_ROCM_USE_AITER=0, which would trigger aiter

245 - # import and JIT compilation warnings. Explicit backend selection via

246 - # attention_config still works because the backend module is loaded

247 - # directly when selected, not through this auto-discovery path.

248 - # Check if backend module exists to allow explicit selection

249 - if find_spec(

250 - AttentionBackendEnum.ROCM_AITER_FA.get_path(include_classname=False)

251 - ):

252 - from vllm.v1.attention.backends.rocm_aiter_fa import (

253 - AiterFlashAttentionMetadata,

254 - )

255 -

256 - rocm_types.append(AiterFlashAttentionMetadata)

257 -

258 - # TRITON_MLA backend support for MLA models (e.g., DeepSeek)

259 - from vllm.model_executor.layers.attention.mla_attention import (

260 - MLACommonMetadata,

261 - )

262 -

263 - rocm_types.append(MLACommonMetadata)

264 -

265 - # FlexAttention backend support

266 - from vllm.v1.attention.backends.flex_attention import FlexAttentionMetadata

267 -

268 - rocm_types.append(FlexAttentionMetadata)

269 -

270 - self.allowed_attn_types = tuple(rocm_types)

271 -

272 - # Parse the speculative token tree.

273 - spec_token_tree = self.speculative_config.speculative_token_tree

274 - assert spec_token_tree is not None

275 - self.tree_choices: list[tuple[int, ...]] = ast.literal_eval(spec_token_tree)

276 - tree_depth = len(self.tree_choices[-1])

277 - # Precompute per-level properties of the tree.

278 - num_drafts_per_level = [0] * tree_depth

279 - for node in self.tree_choices:

280 - num_drafts_per_level[len(node) - 1] += 1

281 - self.cu_drafts_per_level = [num_drafts_per_level[0]]

282 - self.child_drafts_per_level = [num_drafts_per_level[0]]

283 - for level in range(1, tree_depth):

284 - self.cu_drafts_per_level.append(

285 - self.cu_drafts_per_level[-1] + num_drafts_per_level[level]

286 - )

287 - self.child_drafts_per_level.append(

288 - num_drafts_per_level[level] // num_drafts_per_level[level - 1]

289 - )

290 - # Precompute draft position offsets in flattened tree.

291 - self.tree_draft_pos_offsets = torch.arange(

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

293 - ).repeat(self.max_batch_size, 1)

294 -

295 - def _raise_if_padded_drafter_batch_disabled(self):

296 - if self.speculative_config.disable_padded_drafter_batch:

297 - raise NotImplementedError(

298 - "Speculative Decoding with draft models or parallel drafting only "

299 - "supports padded drafter batch. Please unset "

300 - "disable_padded_drafter_batch in the speculative_config."

301 - )

302 -

303 - def _raise_if_multimodal(self):

304 - if self.supports_mm_inputs:

305 - raise NotImplementedError(

306 - "Speculative Decoding with draft models or parallel drafting "

307 - "does not support multimodal models yet"

308 - )

309 -

310 - def _raise_if_mrope(self):

311 - if self.draft_model_config.uses_mrope:

312 - raise NotImplementedError(

313 - "Speculative Decoding with draft models or parallel drafting "

314 - "does not support M-RoPE yet"

315 - )

316 -

317 - def _init_parallel_drafting_params(self):

318 - # For parallel drafting, we need the token ID to use for masked slots

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

320 - # for those masked slots.

321 -

322 - model_hf_config = self.draft_model_config.hf_config

323 - # DFlash stores mask_token_id in dflash_config

324 - dflash_config = getattr(model_hf_config, "dflash_config", None)

325 - if dflash_config and "mask_token_id" in dflash_config:

326 - self.parallel_drafting_token_id = dflash_config["mask_token_id"]

327 - elif hasattr(model_hf_config, "pard_token"):

328 - self.parallel_drafting_token_id = model_hf_config.pard_token

329 - elif hasattr(model_hf_config, "ptd_token_id"):

330 - self.parallel_drafting_token_id = model_hf_config.ptd_token_id

331 - else:

332 - raise ValueError(

333 - "For parallel drafting, the draft model config must have "

334 - "`pard_token`, `ptd_token_id`, or "

335 - "`dflash_config.mask_token_id` specified in its config.json."

336 - )

337 -

338 - if self.pass_hidden_states_to_model:

339 - self.parallel_drafting_hidden_state_tensor = torch.empty(

340 - self.hidden_size, dtype=self.dtype, device=self.device

341 - )

342 -

343 - def _get_positions(self, num_tokens: int):

344 - if self.uses_mrope:

345 - return self.mrope_positions[:, :num_tokens]

346 - if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

347 - return self.xdrope_positions[:, :num_tokens]

348 - return self.positions[:num_tokens]

349 -

350 - def _set_positions(self, num_tokens: int, positions: torch.Tensor):

351 - if self.uses_mrope:

352 - self.mrope_positions[:, :num_tokens] = positions

353 - elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

354 - self.xdrope_positions[:, :num_tokens] = positions

355 - else:

356 - # Convert M-RoPE positions if target model uses M-RoPE

357 - # but draft doesn't, For text inputs, all M-RoPE

358 - # dimensions are identical

359 - if self.vllm_config.model_config.uses_mrope:

360 - positions = positions[0]

361 - self.positions[:num_tokens] = positions

362 -

363 - def _get_slot_mapping(

364 - self,

365 - num_tokens: int,

366 - slot_mapping: torch.Tensor | None = None,

367 - ) -> dict[str, torch.Tensor]:

368 - """Return slot_mapping dict for EAGLE layers.

369 -

370 - If slot_mapping is provided, copies it into the buffer first.

371 - """

372 - if slot_mapping is not None:

373 - num_actual = slot_mapping.shape[0]

374 - self._slot_mapping_buffer[:num_actual].copy_(slot_mapping)

375 - if num_tokens > num_actual:

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

377 -

378 - view = self._slot_mapping_buffer[:num_tokens]

379 - return {name: view for name in self._draft_attn_layer_names}

380 -

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

382 - """Initialize cudagraph dispatcher keys for eagle.

383 -

384 - Eagle only supports PIECEWISE cudagraphs (via mixed_mode).

385 - This should be called after adjust_cudagraph_sizes_for_spec_decode.

386 - """

387 - if (

388 - not self.speculative_config.enforce_eager

389 - and cudagraph_mode.mixed_mode()

390 - in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]

391 - ):

392 - eagle_cudagraph_mode = CUDAGraphMode.PIECEWISE

393 - else:

394 - eagle_cudagraph_mode = CUDAGraphMode.NONE

395 -

396 - self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode)

397 -

398 - def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor:

399 - """Greedy-sample draft tokens from hidden states."""

400 - if self.use_local_argmax_reduction:

401 - return self.model.get_top_tokens(hidden_states)

402 - return self.model.compute_logits(hidden_states).argmax(dim=-1)

403 -

404 - def propose(

405 - self,

406 - # [num_tokens]

407 - target_token_ids: torch.Tensor,

408 - # [num_tokens] or [3, num_tokens] when M-RoPE is enabled

409 - target_positions: torch.Tensor,

410 - # [num_tokens, hidden_size]

411 - target_hidden_states: torch.Tensor,

412 - # [batch_size]

413 - next_token_ids: torch.Tensor,

414 - token_indices_to_sample: torch.Tensor | None,

415 - common_attn_metadata: CommonAttentionMetadata,

416 - sampling_metadata: SamplingMetadata,

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

418 - num_rejected_tokens_gpu: torch.Tensor | None = None,

419 - slot_mappings: dict[str, torch.Tensor]

420 - | list[dict[str, torch.Tensor]]

421 - | None = None,

422 - ) -> torch.Tensor:

423 - batch_size = common_attn_metadata.batch_size()

424 -

425 - if self.method in ("eagle3", "dflash"):

426 - assert isinstance(

427 - self.model,

428 - (

429 - Eagle3LlamaForCausalLM,

430 - Eagle3DeepseekV2ForCausalLM,

431 - DFlashQwen3ForCausalLM,

432 - ),

433 - )

434 - target_hidden_states = self.model.combine_hidden_states(

435 - target_hidden_states

436 - )

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

438 -

439 - num_tokens, token_indices_to_sample, common_attn_metadata = (

440 - self.set_inputs_first_pass(

441 - target_token_ids=target_token_ids,

442 - next_token_ids=next_token_ids,

443 - target_positions=target_positions,

444 - target_hidden_states=target_hidden_states,

445 - token_indices_to_sample=token_indices_to_sample,

446 - cad=common_attn_metadata,

447 - num_rejected_tokens_gpu=num_rejected_tokens_gpu,

448 - )

449 - )

450 -

451 - per_group_attn_metadata, per_layer_attn_metadata = (

452 - self.build_per_group_and_layer_attn_metadata(common_attn_metadata)

453 - )

454 -

455 - cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

456 - self._determine_batch_execution_and_padding(num_tokens)

457 - )

458 -

459 - model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass(

460 - num_tokens, num_input_tokens, mm_embed_inputs

461 - )

462 -

463 - with set_forward_context(

464 - per_layer_attn_metadata,

465 - self.vllm_config,

466 - num_tokens=num_input_tokens,

467 - num_tokens_across_dp=num_tokens_across_dp,

468 - cudagraph_runtime_mode=cudagraph_runtime_mode,

469 - slot_mapping=self._get_slot_mapping(

470 - slot_mapping_size, common_attn_metadata.slot_mapping

471 - ),

472 - ):

473 - ret_hidden_states = self.model(**model_kwargs)

474 - if not self.model_returns_tuple():

475 - last_hidden_states = ret_hidden_states

476 - hidden_states = last_hidden_states

477 - else:

478 - last_hidden_states, hidden_states = ret_hidden_states

479 -

480 - sample_hidden_states = last_hidden_states[token_indices_to_sample]

481 -

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

483 - if self.num_speculative_tokens == 1 or self.parallel_drafting:

484 - draft_token_ids = self._greedy_sample(sample_hidden_states)

485 - return draft_token_ids.view(-1, self.num_speculative_tokens)

486 -

487 - if self.uses_mrope:

488 - positions = self.mrope_positions[:, token_indices_to_sample]

489 - else:

490 - positions = self.positions[token_indices_to_sample]

491 - hidden_states = hidden_states[token_indices_to_sample]

492 -

493 - if any(isinstance(md, TreeAttentionMetadata) for md in per_group_attn_metadata):

494 - # Draft using tree attention - requires full logits for top-k

495 - logits = self.model.compute_logits(sample_hidden_states)

496 - draft_token_ids_list = self.propose_tree(

497 - batch_size=batch_size,

498 - logits=logits,

499 - positions=positions,

500 - hidden_states=hidden_states,

501 - common_attn_metadata=common_attn_metadata,

502 - slot_mappings=slot_mappings,

503 - )

504 - # [batch_size, num_tree_tokens]

505 - return torch.cat(draft_token_ids_list, dim=1)

506 -

507 - draft_token_ids = self._greedy_sample(sample_hidden_states)

508 -

509 - if self.allowed_attn_types is not None:

510 - for group_md in per_group_attn_metadata:

511 - if not isinstance(group_md, self.allowed_attn_types):

512 - raise ValueError(

513 - f"Unsupported attention metadata type for speculative "

514 - "decoding with num_speculative_tokens > 1: "

515 - f"{type(group_md)}. Supported types are: "

516 - f"{self.allowed_attn_types}"

517 - )

518 -

519 - # Generate the remaining draft tokens.

520 - draft_token_ids_list = [draft_token_ids]

521 -

522 - cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = (

523 - self._determine_batch_execution_and_padding(batch_size)

524 - )

525 -

526 - common_attn_metadata.num_actual_tokens = batch_size

527 - common_attn_metadata.max_query_len = 1

528 - common_attn_metadata.query_start_loc = self.arange[: batch_size + 1]

529 - common_attn_metadata.query_start_loc_cpu = torch.from_numpy(

530 - self.token_arange_np[: batch_size + 1]

531 - ).clone()

532 -

533 - # In padded drafter batch, we need to adjust the sequence lengths

534 - # to remove the "padding" (i.e. rejected tokens).

535 - # Only apply this adjustment when we have rejected tokens

536 - # (i.e., not the first proposal).

537 - if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None:

538 - common_attn_metadata.seq_lens -= num_rejected_tokens_gpu

539 - # Invalidate the CPU-side shadows to avoid H<>D sync.

540 - common_attn_metadata._seq_lens_cpu = None

541 - common_attn_metadata._num_computed_tokens_cpu = None

542 -

543 - block_size = self.block_size

544 - assert block_size > 0, "block_size has not been initialized."

545 - for token_index in range(self.num_speculative_tokens - 1):

546 - # Update the inputs.

547 - # cast to int32 is crucial when eagle model is compiled.

548 - # tensor.argmax() returns int64 by default.

549 - input_ids = draft_token_ids_list[-1].int()

550 - # Use fused kernel for slot mapping and metadata updates.

551 - # Write clamped positions directly into the positions buffer to

552 - # avoid an extra D2D copy for the common (non-mrope) case.

553 - positions_1d = positions[0] if self.uses_mrope else positions

554 - if self.uses_mrope:

555 - out_pos = self.mrope_positions[0, :batch_size]

556 - elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

557 - out_pos = self.xdrope_positions[0, :batch_size]

558 - else:

559 - out_pos = self.positions[:batch_size]

560 - eagle_step_update_slot_mapping_and_metadata(

561 - positions_1d=positions_1d,

562 - block_table_tensor=common_attn_metadata.block_table_tensor,

563 - seq_lens=common_attn_metadata.seq_lens,

564 - block_size=block_size,

565 - max_model_len=self.max_model_len,

566 - out_clamped_positions=out_pos,

567 - out_slot_mapping=self._slot_mapping_buffer[:input_batch_size],

568 - input_batch_size=input_batch_size,

569 - )

570 - common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]

571 - if self.uses_mrope:

572 - self.mrope_positions[1:, :batch_size] = self.mrope_positions[

573 - 0, :batch_size

574 - ]

575 - positions = self.mrope_positions[:, :batch_size]

576 - elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

577 - self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[

578 - 0, :batch_size

579 - ]

580 - positions = self.xdrope_positions[0, :batch_size]

581 - else:

582 - positions = self.positions[:batch_size]

583 - # Increment the maximum sequence length. We increment max_seq_len

584 - # unconditionally even though some seq_lens may have been capped above,

585 - # as max_seq_len serves as an upper bound for sequence lengths.

586 - common_attn_metadata.max_seq_len = min(

587 - common_attn_metadata.max_seq_len + 1, self.max_model_len

588 - )

589 -

590 - # Also update the CPU-side shadow; NOTE: this is hacky and should be

591 - # removed in when common_attn_metadata.seq_lens_cpu is deprecated.

592 - if common_attn_metadata._seq_lens_cpu is not None:

593 - common_attn_metadata._seq_lens_cpu += 1

594 - if common_attn_metadata._num_computed_tokens_cpu is not None:

595 - common_attn_metadata._num_computed_tokens_cpu += 1

596 -

597 - # Rebuild attention metadata

598 - _, per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata(

599 - common_attn_metadata, draft_index=token_index + 1

600 - )

601 -

602 - # copy inputs to buffer for cudagraph

603 - self.input_ids[:batch_size] = input_ids

604 - self.hidden_states[:batch_size] = hidden_states

605 - if self.supports_mm_inputs:

606 - self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids)

607 -

608 - input_ids = None

609 - inputs_embeds = self.inputs_embeds[:input_batch_size]

610 - else:

611 - input_ids = self.input_ids[:input_batch_size]

612 - inputs_embeds = None

613 -

614 - # Run the model.

615 - model_kwargs = {

616 - "input_ids": input_ids,

617 - "positions": self._get_positions(input_batch_size),

618 - "inputs_embeds": inputs_embeds,

619 - }

620 - if self.pass_hidden_states_to_model:

621 - model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size]

622 -

623 - with set_forward_context(

624 - per_layer_attn_metadata,

625 - self.vllm_config,

626 - num_tokens=input_batch_size,

627 - num_tokens_across_dp=batch_size_across_dp,

628 - cudagraph_runtime_mode=cudagraph_runtime_mode,

629 - slot_mapping=self._get_slot_mapping(input_batch_size),

630 - ):

631 - ret_hidden_states = self.model(**model_kwargs)

632 - if not self.model_returns_tuple():

633 - last_hidden_states = ret_hidden_states

634 - hidden_states = ret_hidden_states

635 - else:

636 - last_hidden_states, hidden_states = ret_hidden_states

637 -

638 - hidden_states = hidden_states[:batch_size]

639 - draft_token_ids = self._greedy_sample(last_hidden_states[:batch_size])

640 - draft_token_ids_list.append(draft_token_ids)

641 -

642 - # [batch_size, num_speculative_tokens]

643 - draft_token_ids = torch.stack(draft_token_ids_list, dim=1)

644 - return draft_token_ids

645 -

646 - def set_inputs_first_pass(

647 - self,

648 - target_token_ids: torch.Tensor,

649 - next_token_ids: torch.Tensor,

650 - target_positions: torch.Tensor,

651 - target_hidden_states: torch.Tensor,

652 - token_indices_to_sample: torch.Tensor | None,

653 - cad: CommonAttentionMetadata,

654 - num_rejected_tokens_gpu: torch.Tensor | None,

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

656 - if not self.needs_extra_input_slots:

657 - # Default EAGLE pathway: no reshaping of input tensors needed.

658 - # Simply rotate the input ids and leave the positions unchanged,

659 - # Inserting the next token ids at the last slot in each request.

660 - if token_indices_to_sample is None:

661 - token_indices_to_sample = cad.query_start_loc[1:] - 1

662 -

663 - num_tokens = target_token_ids.shape[0]

664 - # Shift the input ids by one token.

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

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

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

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

669 - self.input_ids[token_indices_to_sample] = next_token_ids

670 -

671 - # copy inputs to buffer for cudagraph

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

673 - target_positions = target_positions[0]

674 - self._set_positions(num_tokens, target_positions)

675 -

676 - self.hidden_states[:num_tokens] = target_hidden_states

677 -

678 - return num_tokens, token_indices_to_sample, cad

679 - else:

680 - assert self.is_rejected_token_mask is not None

681 - assert self.is_masked_token_mask is not None

682 - # 1.

683 - # Call a custom triton kernel to copy input_ids and positions

684 - # into the correct slots in the preallocated buffers self.input_ids,

685 - # self.positions.

686 - batch_size = cad.batch_size()

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

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

689 - max_num_tokens_per_request = (

690 - cad.max_query_len + self.net_num_new_slots_per_request

691 - )

692 - BLOCK_SIZE_TOKENS = min(256, next_power_of_2(max_num_tokens_per_request))

693 - num_blocks = (

694 - max_num_tokens_per_request + BLOCK_SIZE_TOKENS - 1

695 - ) // BLOCK_SIZE_TOKENS

696 - total_num_input_tokens = target_token_ids.shape[0]

697 - total_num_output_tokens = total_num_input_tokens + (

698 - self.net_num_new_slots_per_request * batch_size

699 - )

700 -

701 - token_indices_to_sample = torch.empty(

702 - batch_size * self.extra_slots_per_request,

703 - dtype=torch.int32,

704 - device=self.device,

705 - )

706 -

707 - # Destination indices to write target_hidden_states into drafting buffer.

708 - out_hidden_state_mapping = torch.empty(

709 - total_num_input_tokens, dtype=torch.int32, device=self.device

710 - )

711 -

712 - # Kernel grid: one program per request (row)

713 - grid = (batch_size, num_blocks)

714 - query_start_loc = cad.query_start_loc

715 - query_end_loc = cad.query_start_loc[1:] - 1

716 - if num_rejected_tokens_gpu is not None:

717 - query_end_loc = query_end_loc - num_rejected_tokens_gpu

718 -

719 - copy_and_expand_eagle_inputs_kernel[grid](

720 - # (Padded) Inputs from the target model

721 - target_token_ids_ptr=target_token_ids,

722 - target_positions_ptr=target_positions,

723 - next_token_ids_ptr=next_token_ids, # sampled tokens, one per request

724 - # Outputs to the drafting buffers

725 - out_input_ids_ptr=self.input_ids,

726 - out_positions_ptr=self.positions, # Doesn't support mrope for now

727 - out_is_rejected_token_mask_ptr=self.is_rejected_token_mask,

728 - out_is_masked_token_mask_ptr=self.is_masked_token_mask,

729 - out_new_token_indices_ptr=token_indices_to_sample,

730 - out_hidden_state_mapping_ptr=out_hidden_state_mapping,

731 - # Input metadata

732 - query_start_loc_ptr=query_start_loc,

733 - query_end_loc_ptr=query_end_loc,

734 - padding_token_id=0,

735 - parallel_drafting_token_id=self.parallel_drafting_token_id,

736 - # Sizing info

737 - # Note that we can deduce batch_size for free from the grid size

738 - total_input_tokens=total_num_input_tokens,

739 - num_padding_slots_per_request=self.extra_slots_per_request,

740 - shift_input_ids=self.pass_hidden_states_to_model,

741 - BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,

742 - )

743 - if self.pass_hidden_states_to_model:

744 - assert self.parallel_drafting_hidden_state_tensor is not None

745 - self.hidden_states[out_hidden_state_mapping] = target_hidden_states

746 - # Use torch.where to avoid DtoH sync from boolean indexing

747 - mask = self.is_masked_token_mask[:total_num_output_tokens]

748 - torch.where(

749 - mask.unsqueeze(1),

750 - self.parallel_drafting_hidden_state_tensor,

751 - self.hidden_states[:total_num_output_tokens],

752 - out=self.hidden_states[:total_num_output_tokens],

753 - )

754 -

755 - # 2.

756 - # Recompute the slot mapping based on the new positions and

757 - # rejection mask.

758 - assert self.block_size > 0, "block_size has not been initialized."

759 - new_slot_mapping = compute_new_slot_mapping(

760 - cad=cad,

761 - new_positions=self.positions[:total_num_output_tokens],

762 - is_rejected_token_mask=self.is_rejected_token_mask[

763 - :total_num_output_tokens

764 - ],

765 - block_size=self.block_size,

766 - num_new_tokens=self.net_num_new_slots_per_request,

767 - max_model_len=self.max_model_len,

768 - )

769 -

770 - # 3. Update the common attention metadata with the new (meta)data

771 - new_cad = extend_all_queries_by_N(

772 - cad,

773 - N=self.net_num_new_slots_per_request,

774 - arange=self.arange,

775 - new_slot_mapping=new_slot_mapping,

776 - )

777 -

778 - return total_num_output_tokens, token_indices_to_sample, new_cad

779 -

780 - def build_model_inputs_first_pass(

781 - self,

782 - num_tokens: int,

783 - num_input_tokens: int,

784 - mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,

785 - ) -> tuple[dict[str, Any], int]:

786 - if self.supports_mm_inputs:

787 - mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

788 -

789 - self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(

790 - self.input_ids[:num_tokens],

791 - multimodal_embeddings=mm_embeds,

792 - is_multimodal=is_mm_embed,

793 - )

794 -

795 - input_ids = None

796 - inputs_embeds = self.inputs_embeds[:num_input_tokens]

797 - else:

798 - input_ids = self.input_ids[:num_input_tokens]

799 - inputs_embeds = None

800 -

801 - model_kwargs = {

802 - "input_ids": input_ids,

803 - "positions": self._get_positions(num_input_tokens),

804 - "inputs_embeds": inputs_embeds,

805 - }

806 - if self.pass_hidden_states_to_model:

807 - model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]

808 -

809 - return model_kwargs, num_input_tokens

810 -

811 - def build_per_group_and_layer_attn_metadata(

812 - self, common_attn_metadata: CommonAttentionMetadata, draft_index: int = 0

813 - ) -> tuple[list[object], dict[str, object]]:

814 - per_group_attn_metadata: list[object] = []

815 - per_layer_attn_metadata: dict[str, object] = {}

816 - for attn_group in self.draft_attn_groups:

817 - attn_metadata = attn_group.get_metadata_builder().build_for_drafting(

818 - common_attn_metadata=common_attn_metadata, draft_index=draft_index

819 - )

820 - per_group_attn_metadata.append(attn_metadata)

821 - for layer_name in attn_group.layer_names:

822 - per_layer_attn_metadata[layer_name] = attn_metadata

823 - return per_group_attn_metadata, per_layer_attn_metadata

824 -

825 - def model_returns_tuple(self) -> bool:

826 - return self.method not in ("mtp", "draft_model", "dflash")

827 -

828 - def prepare_next_token_ids_cpu(

829 - self,

830 - sampled_token_ids: list[list[int]],

831 - requests: dict[str, CachedRequestState],

832 - gpu_input_batch: InputBatch,

833 - num_scheduled_tokens: dict[str, int],

834 - ) -> torch.Tensor:

835 - """

836 - This function is used to prepare the inputs for speculative decoding.

837 - It calculates the next token ids for each request based on the sampled

838 - token ids from the CPU. If a request has no sampled token ids (e.g.,

839 - during the initial decoding steps), it falls back to using the request

840 - state to get the next token id.

841 - """

842 - req_ids = gpu_input_batch.req_ids

843 - next_token_ids: list[int] = []

844 - for i, token_ids in enumerate(sampled_token_ids):

845 - if token_ids:

846 - # Common case.

847 - next_token_id = token_ids[-1]

848 - else:

849 - # Partial prefill (rare case).

850 - # Get the next token id from the request state.

851 - req_id = req_ids[i]

852 - req_state = requests[req_id]

853 - seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]

854 - next_token_id = req_state.get_token_id(seq_len)

855 - next_token_ids.append(next_token_id)

856 - next_token_ids = torch.tensor(

857 - next_token_ids, dtype=torch.int32, device=self.input_ids.device

858 - )

859 - return next_token_ids

860 -

861 - def prepare_next_token_ids_padded(

862 - self,

863 - sampled_token_ids: torch.Tensor,

864 - requests: dict[str, CachedRequestState],

865 - gpu_input_batch: InputBatch,

866 - discard_request_mask: torch.Tensor,

867 - ) -> tuple[torch.Tensor, torch.Tensor]:

868 - """

869 - This function is used to prepare the inputs for speculative decoding.

870 - It calculates the next token ids and the number of valid sampled tokens

871 - for each request, considering the "discarded" requests whose next token

872 - is not sampled and comes from `request.get_token_id()` instead. This is denoted

873 - the "backup" token id. It also counts rejected tokens via `sampled_token_ids`.

874 - """

875 - # Precompute get_token_id for when there is no valid next token

876 - num_reqs = gpu_input_batch.num_reqs

877 - seq_lens_list = (gpu_input_batch.num_tokens_no_spec[:num_reqs] - 1).tolist()

878 - self.backup_next_token_ids.np[:num_reqs] = np.array(

879 - [

880 - requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])

881 - for i in range(num_reqs)

882 - ],

883 - dtype=np.int32,

884 - )

885 - self.backup_next_token_ids.copy_to_gpu(num_reqs)

886 - backup_tokens_gpu = self.backup_next_token_ids.gpu

887 -

888 - batch_size, num_tokens = sampled_token_ids.shape

889 - device = sampled_token_ids.device

890 -

891 - assert discard_request_mask.dtype == torch.bool

892 - assert backup_tokens_gpu.dtype == torch.int32

893 -

894 - next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=device)

895 - valid_sampled_tokens_count = next_token_ids.new_empty(batch_size)

896 -

897 - # Kernel grid: one program per request (row)

898 - grid = (batch_size,)

899 -

900 - # Find the next power of 2 for block sizes

901 - BLOCK_SIZE_TOKENS = next_power_of_2(num_tokens)

902 - eagle_prepare_next_token_padded_kernel[grid](

903 - sampled_token_ids,

904 - discard_request_mask,

905 - backup_tokens_gpu,

906 - next_token_ids,

907 - valid_sampled_tokens_count,

908 - gpu_input_batch.vocab_size,

909 - num_tokens,

910 - batch_size,

911 - sampled_token_ids.stride(0),

912 - BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,

913 - )

914 -

915 - return next_token_ids, valid_sampled_tokens_count

916 -

917 - def prepare_inputs_padded(

918 - self,

919 - common_attn_metadata: CommonAttentionMetadata,

920 - spec_decode_metadata: SpecDecodeMetadata,

921 - valid_sampled_tokens_count: torch.Tensor,

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

923 - """

924 - This function is used to prepare the inputs for speculative decoding

925 - It updates the common_attn_metadata for speculative decoding,

926 - but does not consider the rejected tokens. Instead, all tokens

927 - are included as inputs to the speculator, with the rejected tokens

928 - used as padding and filtered out later by `token_indices_to_sample`.

929 - No blocking CPU operations should be introduced in this function.

930 - """

931 - num_reqs = common_attn_metadata.num_reqs

932 - device = valid_sampled_tokens_count.device

933 -

934 - token_indices_to_sample = torch.empty(

935 - (num_reqs,), dtype=torch.int32, device=device

936 - )

937 - num_rejected_tokens_gpu = torch.empty(

938 - (num_reqs,), dtype=torch.int32, device=device

939 - )

940 -

941 - grid = (num_reqs,)

942 - eagle_prepare_inputs_padded_kernel[grid](

943 - spec_decode_metadata.cu_num_draft_tokens,

944 - valid_sampled_tokens_count,

945 - common_attn_metadata.query_start_loc,

946 - token_indices_to_sample,

947 - num_rejected_tokens_gpu,

948 - num_reqs,

949 - )

950 -

951 - query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

952 - new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

953 -

954 - total_num_tokens = query_start_loc_cpu[-1].item()

955 -

956 - spec_common_attn_metadata = CommonAttentionMetadata(

957 - query_start_loc=common_attn_metadata.query_start_loc,

958 - seq_lens=common_attn_metadata.seq_lens,

959 - query_start_loc_cpu=query_start_loc_cpu,

960 - _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,

961 - _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

962 - num_reqs=common_attn_metadata.num_reqs,

963 - num_actual_tokens=total_num_tokens,

964 - max_query_len=new_query_len_per_req.max().item(),

965 - max_seq_len=common_attn_metadata.max_seq_len,

966 - block_table_tensor=common_attn_metadata.block_table_tensor,

967 - slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens],

968 - causal=True,

969 - dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,

970 - )

971 -

972 - return (

973 - spec_common_attn_metadata,

974 - token_indices_to_sample,

975 - num_rejected_tokens_gpu,

976 - )

977 -

978 - def propose_tree(

979 - self,

980 - batch_size: int,

981 - # [num_tokens, vocab_size]

982 - logits: torch.Tensor,

983 - # [num_tokens]

984 - positions: torch.Tensor,

985 - # [num_tokens, hidden_size]

986 - hidden_states: torch.Tensor,

987 - common_attn_metadata: CommonAttentionMetadata,

988 - slot_mappings: dict[str, torch.Tensor]

989 - | list[dict[str, torch.Tensor]]

990 - | None = None,

991 - ) -> list[torch.Tensor]:

992 - tree_attn_metadata_builder = self.draft_attn_groups[0].get_metadata_builder()

993 - assert isinstance(tree_attn_metadata_builder, TreeAttentionMetadataBuilder)

994 -

995 - total_num_drafts = self.cu_drafts_per_level[0]

996 - level_num_drafts = total_num_drafts

997 - # Sample a draft token for each child at the tree root level.

998 - num_children = self.child_drafts_per_level[0]

999 - if num_children == 1:

1000 - draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)

1001 - else:

1002 - draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(

1003 - batch_size, -1

1004 - )

1005 - draft_token_ids_list = [draft_token_ids]

1006 - draft_hidden_states = hidden_states.view(batch_size, 1, -1)

1007 -

1008 - # Initialize empty tensors for concatenation with the level outputs.

1009 - tree_input_ids = torch.empty(

1010 - 0, device=self.input_ids.device, dtype=self.input_ids.dtype

1011 - )

1012 - tree_positions = torch.empty(

1013 - 0, device=self.positions.device, dtype=self.positions.dtype

1014 - )

1015 - tree_hidden_states = torch.empty(

1016 - 0, device=self.hidden_states.device, dtype=self.hidden_states.dtype

1017 - )

1018 - # Precompute the draft token positions.

1019 - flattened_draft_positions = (

1020 - positions.view(batch_size, -1) + self.tree_draft_pos_offsets[:batch_size, :]

1021 - )

1022 - tree_depth = len(self.cu_drafts_per_level)

1023 - for level in range(tree_depth - 1):

1024 - # Get draft positions for RoPE.

1025 - draft_positions = positions + (level + 1)

1026 - exceeds_max_model_len = (positions + total_num_drafts) >= self.max_model_len

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

1028 - # Otherwise, we may get out-of-range error in RoPE.

1029 - draft_positions = torch.where(

1030 - exceeds_max_model_len,

1031 - 0,

1032 - draft_positions,

1033 - ).view(batch_size, -1)

1034 -

1035 - if level_num_drafts > 1:

1036 - # Repeat the positions for each draft at this level.

1037 - draft_positions = draft_positions.repeat_interleave(

1038 - level_num_drafts, dim=1

1039 - )

1040 -

1041 - if num_children > 1:

1042 - # Repeat draft hidden states for each child.

1043 - draft_hidden_states = draft_hidden_states.repeat_interleave(

1044 - num_children, dim=1

1045 - )

1046 -

1047 - # Concatenate the draft tokens, positions, and hidden states.

1048 - tree_input_ids = torch.cat([tree_input_ids, draft_token_ids], dim=1)

1049 - tree_positions = torch.cat([tree_positions, draft_positions], dim=1)

1050 - tree_hidden_states = torch.cat(

1051 - [tree_hidden_states, draft_hidden_states], dim=1

1052 - )

1053 -

1054 - # Build new attention metadata for the next level of drafts.

1055 - # This is necessary to support tree attention.

1056 - query_len = total_num_drafts

1057 - common_attn_metadata = replace(

1058 - common_attn_metadata,

1059 - query_start_loc=query_len * self.arange[: batch_size + 1],

1060 - seq_lens=common_attn_metadata.seq_lens + level_num_drafts,

1061 - num_actual_tokens=batch_size * query_len,

1062 - max_query_len=query_len,

1063 - )

1064 - attn_metadata = tree_attn_metadata_builder.build_for_drafting(

1065 - common_attn_metadata=common_attn_metadata, draft_index=level + 1

1066 - )

1067 -

1068 - # Apply new attention metadata to all draft layers.

1069 - per_layer_attn_metadata = {}

1070 - for attn_group in self.draft_attn_groups:

1071 - for layer_name in attn_group.layer_names:

1072 - per_layer_attn_metadata[layer_name] = attn_metadata

1073 -

1074 - # Consider max model length.

1075 - attn_metadata.max_seq_len = min(

1076 - attn_metadata.max_seq_len, self.max_model_len

1077 - )

1078 - # For the requests that exceed the max model length, we set the

1079 - # sequence length to 1 to minimize their overheads in attention.

1080 - attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)

1081 -

1082 - # Compute the slot mapping.

1083 - block_size = tree_attn_metadata_builder.kv_cache_spec.block_size

1084 - query_positions = flattened_draft_positions[:, level : level + query_len]

1085 - block_numbers = query_positions // block_size

1086 - block_ids = attn_metadata.block_table.gather(dim=1, index=block_numbers)

1087 - slot_mapping = block_ids * block_size + query_positions % block_size

1088 - # Mask out the slot mappings that exceed the max model length.

1089 - # Otherwise, the KV cache will be inadvertently updated with the

1090 - # padding tokens.

1091 - slot_mapping[exceeds_max_model_len] = PADDING_SLOT_ID

1092 - attn_metadata.slot_mapping = slot_mapping.view(-1)

1093 -

1094 - # Copy inputs to buffer for cudagraph.

1095 - num_tokens = attn_metadata.num_actual_tokens

1096 - input_ids = tree_input_ids.view(-1)

1097 - self.input_ids[:num_tokens] = input_ids

1098 - self.positions[:num_tokens] = tree_positions.view(-1)

1099 - self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1)

1100 -

1101 - cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

1102 - num_tokens

1103 - )

1104 - num_input_tokens = batch_desc.num_tokens

1105 - # Run the model.

1106 - with set_forward_context(

1107 - per_layer_attn_metadata,

1108 - self.vllm_config,

1109 - num_tokens=num_input_tokens,

1110 - cudagraph_runtime_mode=cudagraph_runtime_mode,

1111 - slot_mapping=self._get_slot_mapping(

1112 - num_input_tokens, attn_metadata.slot_mapping

1113 - ),

1114 - ):

1115 - last_hidden_states, hidden_states = self.model(

1116 - input_ids=self.input_ids[:num_input_tokens],

1117 - positions=self.positions[:num_input_tokens],

1118 - hidden_states=self.hidden_states[:num_input_tokens],

1119 - inputs_embeds=None,

1120 - )

1121 -

1122 - # Get the output hidden states for the draft tokens.

1123 - draft_hidden_states = hidden_states[:num_tokens].view(

1124 - batch_size, query_len, -1

1125 - )[:, -level_num_drafts:]

1126 - draft_last_hidden_states = last_hidden_states[:num_tokens].view(

1127 - batch_size, query_len, -1

1128 - )[:, -level_num_drafts:]

1129 -

1130 - # Get the output logits for the draft tokens.

1131 - logits = self.model.compute_logits(

1132 - draft_last_hidden_states.reshape(batch_size * level_num_drafts, -1)

1133 - )

1134 -

1135 - # Sample a draft token for each child at the next tree level.

1136 - num_children = self.child_drafts_per_level[level + 1]

1137 - if num_children == 1:

1138 - draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)

1139 - else:

1140 - draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(

1141 - batch_size, -1

1142 - )

1143 - draft_token_ids_list.append(draft_token_ids)

1144 -

1145 - # Update the # drafts counters for the next tree level.

1146 - level_num_drafts = self.cu_drafts_per_level[level + 1] - total_num_drafts

1147 - total_num_drafts = self.cu_drafts_per_level[level + 1]

1148 - return draft_token_ids_list

1149 -

1150 - def prepare_inputs(

1151 - self,

1152 - common_attn_metadata: CommonAttentionMetadata,

1153 - sampled_token_ids: list[list[int]],

1154 - num_draft_tokens: list[int],

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

1156 - """

1157 - This function is used to prepare the inputs for speculative decoding.

1158 - It updates to the common_attn_metadata to account for the rejected

1159 - tokens (and newly sampled tokens). It also returns the token indices

1160 - of the tokens that should be fed to the speculator.

1161 - """

1162 - # E.g.

1163 - # common_attn_metadata.query_start_loc{_cpu}:

1164 - # [0, q1, q1 + q2, q1 + q2 + q3]

1165 - # common_attn_metadata.seq_lens{_cpu}: [s1, s2, s3]

1166 - # num_rejected_tokens: [n1, n2, n3]

1167 - # This function computes the intermediate values:

1168 - # num_tokens_per_req: [q1 - n1, q2 - n2, q3 - n3]

1169 - # And returns:

1170 - # common_attn_metadata.query_start_loc{_cpu}:

1171 - # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]

1172 - # common_attn_metadata.seq_lens{_cpu}:

1173 - # [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1]

1174 - # token_indices: [0, 1, ..., q1 - n1 - 1,

1175 - # q1, q1 + 1, ..., q1 + q2 - n2 - 1,

1176 - # q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]

1177 -

1178 - num_rejected_tokens = [

1179 - n + 1 - len(sampled_token_ids[i]) if n > 0 else 0

1180 - for i, n in enumerate(num_draft_tokens)

1181 - ]

1182 - num_rejected_tokens = torch.tensor(num_rejected_tokens, dtype=torch.int32)

1183 -

1184 - device = common_attn_metadata.query_start_loc.device

1185 - query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

1186 - new_seq_lens_cpu = common_attn_metadata.seq_lens_cpu - num_rejected_tokens

1187 -

1188 - # [0, q1, q1 + q2, q1 + q2 + q3] -> [q1, q2, q3]

1189 - new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

1190 - # [q1, q2, q3] -> [q1 - n1, q2 - n2, q3 - n3]

1191 - new_num_tokens_per_req = new_query_len_per_req - num_rejected_tokens

1192 - new_num_tokens_per_req_np = new_num_tokens_per_req.numpy()

1193 -

1194 - # [q1 - n1, q2 - n2, q3 - n3] ->

1195 - # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]

1196 - new_query_start_loc_cpu = torch.zeros(

1197 - query_start_loc_cpu.shape,

1198 - dtype=torch.int32,

1199 - pin_memory=is_pin_memory_available(),

1200 - )

1201 - new_query_start_loc_np = new_query_start_loc_cpu.numpy()

1202 - np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:])

1203 -

1204 - total_num_tokens = new_query_start_loc_np[-1]

1205 - # Example assuming num_tokens_per_req_np = [2, 4, 3]

1206 - # this implies that `new_query_start_locs` is:

1207 - # [0, 2, 6, 9] ->

1208 - # [0, 0, 2, 2, 2, 2, 6, 6, 6]

1209 - # _r1_ ____r2____ ___r3__

1210 - new_query_start_locs_expanded = np.repeat(

1211 - new_query_start_loc_np[:-1], new_num_tokens_per_req_np

1212 - )

1213 - # [0, 1, 2, 3, 4, 5, 6, 7, 8] ->

1214 - # [0, 1, 0, 1, 2, 3, 0, 1, 2]

1215 - # _r1_ ____r2____ ___r3__

1216 - token_offsets = (

1217 - self.token_arange_np[:total_num_tokens] - new_query_start_locs_expanded

1218 - )

1219 -

1220 - # Expand starting positions to match token pattern

1221 - # [0, q1, q1 + q2] ->

1222 - # [0, 0, q1, q1, q1, q1, q1 + q2, q1 + q2, q1 + q2]

1223 - # _r1_ _____r2_______ ___________r3____________

1224 - old_query_start_locs_expanded = np.repeat(

1225 - query_start_loc_cpu[:-1].numpy(), new_num_tokens_per_req_np

1226 - )

1227 - # Final token indices are:

1228 - # [0, 1, // req 1

1229 - # q1 + 0, q1 + 1, q1 + 2, q1 + 3, // req 2

1230 - # q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3

1231 - token_indices_np = token_offsets + old_query_start_locs_expanded

1232 - token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True)

1233 -

1234 - spec_common_attn_metadata = CommonAttentionMetadata(

1235 - query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True),

1236 - seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),

1237 - query_start_loc_cpu=new_query_start_loc_cpu,

1238 - _seq_lens_cpu=new_seq_lens_cpu,

1239 - _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

1240 - num_reqs=common_attn_metadata.num_reqs,

1241 - num_actual_tokens=total_num_tokens,

1242 - max_query_len=new_query_len_per_req.max().item(),

1243 - max_seq_len=new_seq_lens_cpu.max().item(),

1244 - block_table_tensor=common_attn_metadata.block_table_tensor,

1245 - slot_mapping=common_attn_metadata.slot_mapping[token_indices],

1246 - causal=True,

1247 - dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,

1248 - )

1249 -

1250 - return spec_common_attn_metadata, token_indices

1251 -

1252 - def get_model_name(self, model: nn.Module) -> str:

1253 - if hasattr(model, "module"): # multi-GPU

1254 - model = model.module

1255 - return model.__class__.__name__

1256 -

1257 - def _create_draft_vllm_config(self) -> VllmConfig:

1258 - """Return a VllmConfig with kernel-level overrides for the proposer.

1259 - Subclasses may override to apply additional config changes.

1260 - """

1261 - spec_cfg = self.speculative_config

1262 - if spec_cfg.moe_backend is not None:

1263 - return replace(

1264 - self.vllm_config,

1265 - kernel_config=replace(

1266 - self.vllm_config.kernel_config,

1267 - moe_backend=spec_cfg.moe_backend,

1268 - ),

1269 - )

1270 - return self.vllm_config

1271 -

1272 - def _get_model(self) -> nn.Module:

1273 - """

1274 - Default method to call get_model(). Can be overridden by subclasses which

1275 - need to customize model loading.

1276 - """

1277 - from vllm.compilation.backends import set_model_tag

1278 -

1279 - draft_vllm_config = self._create_draft_vllm_config()

1280 - with set_model_tag("eagle_head"):

1281 - model = get_model(

1282 - vllm_config=draft_vllm_config,

1283 - model_config=self.speculative_config.draft_model_config,

1284 - load_config=self.speculative_config.draft_load_config,

1285 - )

1286 - return model

1287 -

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

1289 - target_attn_layer_names = set(

1290 - get_layers_from_vllm_config(

1291 - self.vllm_config,

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

1293 - ).keys()

1294 - )

1295 -

1296 - self.model = self._get_model()

1297 -

1298 - # Find draft layers (attention layers added by draft model)

1299 - all_attn_layers = get_layers_from_vllm_config(

1300 - self.vllm_config,

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

1302 - )

1303 - self._draft_attn_layer_names = (

1304 - set(all_attn_layers.keys()) - target_attn_layer_names

1305 - )

1306 -

1307 - if self.supports_mm_inputs:

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

1309 - # text-only draft models

1310 - try:

1311 - dummy_input_ids = torch.tensor([[1]], device=self.input_ids.device)

1312 - self.model.embed_input_ids(dummy_input_ids, multimodal_embeddings=None)

1313 - except (NotImplementedError, AttributeError, TypeError):

1314 - logger.warning(

1315 - "Draft model does not support multimodal inputs, "

1316 - "falling back to text-only mode"

1317 - )

1318 - self.supports_mm_inputs = False

1319 -

1320 - if supports_multimodal(target_model):

1321 - # handle multimodality

1322 - assert hasattr(target_model, "config")

1323 - if self.get_model_name(target_model) in [

1324 - "Exaone4_5_ForConditionalGeneration",

1325 - "GlmOcrForConditionalGeneration",

1326 - "HunYuanVLForConditionalGeneration",

1327 - "Qwen2_5_VLForConditionalGeneration",

1328 - "Qwen3_5ForConditionalGeneration",

1329 - "Qwen3_5MoeForConditionalGeneration",

1330 - "Qwen3VLForConditionalGeneration",

1331 - "Qwen3VLMoeForConditionalGeneration",

1332 - "Gemma4ForConditionalGeneration",

1333 - ]:

1334 - self.model.config.image_token_index = target_model.config.image_token_id

1335 - elif self.get_model_name(target_model) == "PixtralForConditionalGeneration":

1336 - self.model.config.image_token_index = (

1337 - target_model.config.vision_config.image_token_id

1338 - )

1339 - elif self.get_model_name(target_model) == "KimiK25ForConditionalGeneration":

1340 - self.model.config.image_token_index = (

1341 - target_model.config.media_placeholder_token_id

1342 - )

1343 - else:

1344 - self.model.config.image_token_index = (

1345 - target_model.config.image_token_index

1346 - )

1347 - target_language_model = cast(

1348 - SupportsMultiModal, target_model

1349 - ).get_language_model()

1350 - else:

1351 - target_language_model = target_model

1352 -

1353 - self._maybe_share_embeddings(target_language_model)

1354 - self._maybe_share_lm_head(target_language_model)

1355 -

1356 - if (

1357 - self.parallel_drafting

1358 - and self.pass_hidden_states_to_model

1359 - and self.parallel_drafting_hidden_state_tensor is not None

1360 - ):

1361 - flat_mask = self.model.mask_hidden.view(-1)

1362 - if self.eagle3_use_aux_hidden_state:

1363 - # EAGLE3: mask_hidden stores all aux hidden states,

1364 - # project through combine_hidden_states

1365 - self.parallel_drafting_hidden_state_tensor.copy_(

1366 - self.model.combine_hidden_states(flat_mask)

1367 - )

1368 - else:

1369 - self.parallel_drafting_hidden_state_tensor.copy_(flat_mask)

1370 -

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

1372 - """

1373 - Some draft models may not have their own embedding layers, and some may

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

1375 - we share the target model's embedding layers with the draft model to save

1376 - memory.

1377 - """

1378 - if get_pp_group().world_size == 1:

1379 - inner_model = getattr(target_language_model, "model", None)

1380 - if inner_model is None:

1381 - raise AttributeError("Target model does not have 'model' attribute")

1382 - if hasattr(inner_model, "embed_tokens"):

1383 - target_embed_tokens = inner_model.embed_tokens

1384 - elif hasattr(inner_model, "embedding"):

1385 - target_embed_tokens = inner_model.embedding

1386 - else:

1387 - raise AttributeError(

1388 - "Target model does not have 'embed_tokens' or 'embedding' attribute"

1389 - )

1390 -

1391 - share_embeddings = False

1392 - if hasattr(self.model, "has_own_embed_tokens"):

1393 - # EAGLE model

1394 - if not self.model.has_own_embed_tokens:

1395 - share_embeddings = True

1396 - logger.info(

1397 - "Detected EAGLE model without its own embed_tokens in the"

1398 - " checkpoint. Sharing target model embedding weights with the"

1399 - " draft model."

1400 - )

1401 - elif (

1402 - isinstance(target_embed_tokens.weight, torch.Tensor)

1403 - and isinstance(self.model.model.embed_tokens.weight, torch.Tensor)

1404 - # TODO: Offload to CPU for comparison to avoid extra GPU memory

1405 - # usage in CI testing environments with limited GPU memory

1406 - and torch.equal(

1407 - target_embed_tokens.weight.cpu(),

1408 - self.model.model.embed_tokens.weight.cpu(),

1409 - )

1410 - ):

1411 - share_embeddings = True

1412 - logger.info(

1413 - "Detected EAGLE model with embed_tokens identical to the target"

1414 - " model. Sharing target model embedding weights with the draft"

1415 - " model."

1416 - )

1417 - else:

1418 - logger.info(

1419 - "Detected EAGLE model with distinct embed_tokens weights. "

1420 - "Keeping separate embedding weights from the target model."

1421 - )

1422 - else:

1423 - # MTP model

1424 - share_embeddings = True

1425 - logger.info(

1426 - "Detected MTP model. "

1427 - "Sharing target model embedding weights with the draft model."

1428 - )

1429 -

1430 - if share_embeddings:

1431 - if hasattr(self.model.model, "embed_tokens"):

1432 - del self.model.model.embed_tokens

1433 - self.model.model.embed_tokens = target_embed_tokens

1434 - else:

1435 - logger.info(

1436 - "The draft model's vocab embedding will be loaded separately"

1437 - " from the target model."

1438 - )

1439 -

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

1441 - """

1442 - Some draft models may not have their own LM head, and some may have a

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

1444 - the target model's LM head with the draft model to save memory.

1445 - """

1446 - share_lm_head = False

1447 - if hasattr(self.model, "has_own_lm_head"):

1448 - # EAGLE model

1449 - if not self.model.has_own_lm_head:

1450 - share_lm_head = True

1451 - logger.info(

1452 - "Detected EAGLE model without its own lm_head in the checkpoint. "

1453 - "Sharing target model lm_head weights with the draft model."

1454 - )

1455 - elif (

1456 - hasattr(target_language_model, "lm_head")

1457 - and hasattr(target_language_model.lm_head, "weight")

1458 - and hasattr(self.model.lm_head, "weight")

1459 - and isinstance(target_language_model.lm_head.weight, torch.Tensor)

1460 - and isinstance(self.model.lm_head.weight, torch.Tensor)

1461 - # TODO: Offload to CPU for comparison to avoid extra GPU memory

1462 - # usage in CI testing environments with limited GPU memory

1463 - and torch.equal(

1464 - target_language_model.lm_head.weight.cpu(),

1465 - self.model.lm_head.weight.cpu(),

1466 - )

1467 - ):

1468 - share_lm_head = True

1469 - logger.info(

1470 - "Detected EAGLE model with lm_head identical to the target model. "

1471 - "Sharing target model lm_head weights with the draft model."

1472 - )

1473 - else:

1474 - logger.info(

1475 - "Detected EAGLE model with distinct lm_head weights. "

1476 - "Keeping separate lm_head weights from the target model."

1477 - )

1478 - else:

1479 - # MTP model

1480 - share_lm_head = True

1481 - logger.info(

1482 - "Detected MTP model. "

1483 - "Sharing target model lm_head weights with the draft model."

1484 - )

1485 -

1486 - if share_lm_head and hasattr(target_language_model, "lm_head"):

1487 - if hasattr(self.model, "lm_head"):

1488 - del self.model.lm_head

1489 - self.model.lm_head = target_language_model.lm_head

1490 -

1491 - # MTP models call compute_logits via shared_head.head (a

1492 - # ParallelLMHead inside each MTP layer), not self.model.lm_head.

1493 - # If the checkpoint omits a copy of the lm_head weights at the

1494 - # MTP layer path, shared_head.head stays uninitialised and

1495 - # produces NaN logits. Always share it explicitly.

1496 - inner = getattr(self.model, "model", None)

1497 - layers = getattr(inner, "layers", None) if inner else None

1498 - if layers is not None:

1499 - items = layers.values() if isinstance(layers, nn.ModuleDict) else layers

1500 - for layer in items:

1501 - sh = getattr(layer, "shared_head", None)

1502 - if sh is not None and hasattr(sh, "head"):

1503 - del sh.head

1504 - sh.head = target_language_model.lm_head

1505 - logger.info(

1506 - "Shared target model lm_head with MTP shared_head.head."

1507 - )

1508 -

1509 - if self.use_local_argmax_reduction:

1510 - if not hasattr(self.model, "get_top_tokens"):

1511 - raise ValueError(

1512 - "use_local_argmax_reduction is enabled but draft model "

1513 - f"{self.model.__class__.__name__} does not implement "

1514 - "get_top_tokens()."

1515 - )

1516 - # Warn if draft model has vocab remapping, which forces fallback

1517 - # to the full-logits path (negating the optimization).

1518 - if (

1519 - hasattr(self.model, "draft_id_to_target_id")

1520 - and self.model.draft_id_to_target_id is not None

1521 - ):

1522 - logger.warning(

1523 - "use_local_argmax_reduction is enabled but draft model "

1524 - "uses draft_id_to_target_id vocab remapping. The "

1525 - "optimization will be bypassed (falling back to full "

1526 - "logits gather + argmax)."

1527 - )

1528 - else:

1529 - logger.info(

1530 - "Using local argmax reduction for draft token generation "

1531 - "(communication: O(2*tp_size) vs O(vocab_size))."

1532 - )

1533 -

1534 - @torch.inference_mode()

1535 - def dummy_run(

1536 - self,

1537 - num_tokens: int,

1538 - use_cudagraphs: bool = True,

1539 - is_graph_capturing: bool = False,

1540 - slot_mappings: dict[str, torch.Tensor] | None = None,

1541 - ) -> None:

1542 - # FIXME: when using tree-based specdec, adjust number of forward-passes

1543 - # according to the depth of the tree.

1544 - only_one_forward_pass = is_graph_capturing or self.parallel_drafting

1545 - for fwd_idx in range(

1546 - 1 if only_one_forward_pass else self.num_speculative_tokens

1547 - ):

1548 - if fwd_idx <= 1:

1549 - cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

1550 - self._determine_batch_execution_and_padding(

1551 - num_tokens, use_cudagraphs=use_cudagraphs

1552 - )

1553 - )

1554 -

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

1556 - if (

1557 - self._draft_attn_layer_names

1558 - and slot_mappings is not None

1559 - and next(iter(self._draft_attn_layer_names)) in slot_mappings

1560 - ):

1561 - slot_mapping_dict = self._get_slot_mapping(num_input_tokens)

1562 - else:

1563 - slot_mapping_dict = slot_mappings or {}

1564 -

1565 - with set_forward_context(

1566 - None,

1567 - self.vllm_config,

1568 - num_tokens=num_input_tokens,

1569 - num_tokens_across_dp=num_tokens_across_dp,

1570 - cudagraph_runtime_mode=cudagraph_runtime_mode,

1571 - slot_mapping=slot_mapping_dict,

1572 - ):

1573 - if self.supports_mm_inputs:

1574 - input_ids = None

1575 - inputs_embeds = self.inputs_embeds[:num_input_tokens]

1576 - else:

1577 - input_ids = self.input_ids[:num_input_tokens]

1578 - inputs_embeds = None

1579 -

1580 - kwargs = dict(

1581 - input_ids=input_ids,

1582 - positions=self._get_positions(num_input_tokens),

1583 - inputs_embeds=inputs_embeds,

1584 - )

1585 - if self.pass_hidden_states_to_model:

1586 - kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]

1587 - self.model(**kwargs)

1588 -

1589 - def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:

1590 - """

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

1592 - hidden states and directly uses the last layer output just like eagle1.

1593 - They might indicate this by setting "use_aux_hidden_state" to False

1594 - inside the "eagle_config" dict of their hf_config.

1595 - """

1596 - if self.method != "eagle3":

1597 - return False

1598 - # Assume that eagle3 heads use aux hidden states by default

1599 - use_aux_hidden_state = True

1600 - eagle_config = getattr(self.draft_model_config.hf_config, "eagle_config", None)

1601 - if eagle_config is not None:

1602 - use_aux_hidden_state = eagle_config.get("use_aux_hidden_state", True)

1603 - return use_aux_hidden_state

1604 -

1605 - def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:

1606 - """

1607 - Validate that all drafting layers belong to the same KVCacheGroup.

1608 - Need this assumption to ensure all drafting layers can use the

1609 - same AttentionMetadata.

1610 - May extend to multiple AttentionMetadata in the future.

1611 - """

1612 - kv_cache_groups: dict[str, int] = {}

1613 - for id, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):

1614 - for layer_name in kv_cache_group.layer_names:

1615 - kv_cache_groups[layer_name] = id

1616 - assert (

1617 - len(

1618 - set(

1619 - [

1620 - kv_cache_groups[layer_name]

1621 - for layer_name in self._draft_attn_layer_names

1622 - ]

1623 - )

1624 - )

1625 - == 1

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

1627 -

1628 - def initialize_attn_backend(

1629 - self,

1630 - kv_cache_config: KVCacheConfig,

1631 - kernel_block_sizes: list[int] | None = None,

1632 - ) -> None:

1633 - """

1634 - Initialize AttentionGroups for draft layers using kv_cache_config.

1635 - Called from the model runner's initialize_metadata_builders.

1636 - """

1637 - all_attn_layers = get_layers_from_vllm_config(

1638 - self.vllm_config,

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

1640 - )

1641 -

1642 - # Find which kv_cache_group the draft layers belong to

1643 - self.validate_same_kv_cache_group(kv_cache_config)

1644 - kv_cache_spec = None

1645 - for gid, group in enumerate(kv_cache_config.kv_cache_groups):

1646 - if self._draft_attn_layer_names & set(group.layer_names):

1647 - self.kv_cache_gid = gid

1648 - kv_cache_spec = group.kv_cache_spec

1649 - break

1650 -

1651 - attention_groups: dict[tuple[str, str], AttentionGroup] = {}

1652 - if kv_cache_spec is not None:

1653 - for layer_name in self._draft_attn_layer_names:

1654 - attn_backend = all_attn_layers[layer_name].get_attn_backend()

1655 - backend_key = attn_backend.full_cls_name()

1656 - if backend_key not in attention_groups:

1657 - layer_kv_cache_spec = kv_cache_spec

1658 - if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):

1659 - layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[

1660 - layer_name

1661 - ]

1662 -

1663 - kernel_block_size = (

1664 - kernel_block_sizes[self.kv_cache_gid]

1665 - if kernel_block_sizes is not None

1666 - and self.kv_cache_gid < len(kernel_block_sizes)

1667 - else None

1668 - )

1669 - attn_group = AttentionGroup(

1670 - backend=attn_backend,

1671 - layer_names=[layer_name],

1672 - kv_cache_spec=layer_kv_cache_spec,

1673 - kv_cache_group_id=self.kv_cache_gid,

1674 - )

1675 - attn_group.create_metadata_builders(

1676 - self.vllm_config,

1677 - self.device,

1678 - kernel_block_size=kernel_block_size,

1679 - )

1680 - attention_groups[backend_key] = attn_group

1681 - else:

1682 - attention_groups[backend_key].layer_names.append(layer_name)

1683 -

1684 - self.draft_attn_groups = list(attention_groups.values())

1685 - self.block_size = (

1686 - self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size

1687 - )

1688 - logger.debug("Using block size %d for drafting layers", self.block_size)

1689 -

1690 - def _determine_batch_execution_and_padding(

1691 - self,

1692 - num_tokens: int,

1693 - use_cudagraphs: bool = True,

1694 - ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:

1695 - cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

1696 - num_tokens,

1697 - valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None),

1698 - )

1699 - num_tokens_padded = batch_desc.num_tokens

1700 -

1701 - # Extra coordination when running data-parallel since we need to

1702 - # coordinate across ranks

1703 - # TODO(Flechman): support DBO ubatching

1704 - should_ubatch, num_tokens_across_dp = False, None

1705 - if self.vllm_config.parallel_config.data_parallel_size > 1:

1706 - should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = (

1707 - coordinate_batch_across_dp(

1708 - num_tokens_unpadded=num_tokens,

1709 - parallel_config=self.vllm_config.parallel_config,

1710 - allow_microbatching=False,

1711 - num_tokens_padded=num_tokens_padded,

1712 - cudagraph_mode=cudagraph_mode.value,

1713 - )

1714 - )

1715 - assert not should_ubatch, "DBO ubatching not implemented for EAGLE"

1716 -

1717 - # Extract DP-synced values

1718 - if num_tokens_across_dp is not None:

1719 - dp_rank = self.dp_rank

1720 - num_tokens_padded = int(num_tokens_across_dp[dp_rank].item())

1721 - # Re-dispatch with DP padding so we have the correct

1722 - # batch_descriptor

1723 - cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

1724 - num_tokens_padded,

1725 - valid_modes={CUDAGraphMode(synced_cudagraph_mode)},

1726 - )

1727 - # Assert to make sure the agreed upon token count is correct

1728 - # otherwise num_tokens_across_dp will no-longer be valid

1729 - assert batch_desc.num_tokens == num_tokens_padded

1730 - num_tokens_across_dp[dp_rank] = num_tokens_padded

1731 -

1732 - return cudagraph_mode, num_tokens_padded, num_tokens_across_dp

6 +from vllm.config import VllmConfig

7 +from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer

10 class EagleProposer(SpecDecodeBaseProposer):

@@ -1745,49 +20,3 @@ class EagleProposer(SpecDecodeBaseProposer):

20 pass_hidden_states_to_model=True,

21 runner=runner,

22 )

1748 -

1749 -

1750 -# NOTE(woosuk): Currently, the below code is not used and we always use argmax

1751 -# to sample the draft tokens. We will use this after we find a way to manage

1752 -# the draft prob tensor.

1753 -# Refer to https://github.com/vllm-project/vllm/pull/16899 for the details.

1754 -# FIXME(woosuk): The logic here is duplicated with the main sampling code.

1755 -# We should refactor this to reuse the same sampling implementation.

1756 -def compute_probs_and_sample_next_token(

1757 - logits: torch.Tensor,

1758 - sampling_metadata: SamplingMetadata,

1759 -) -> tuple[torch.Tensor, torch.Tensor]:

1760 - if sampling_metadata.all_greedy:

1761 - # For greedy requests, draft_probs is not used in rejection sampling.

1762 - # Therefore, we can just return the logits.

1763 - probs = logits

1764 - next_token_ids = logits.argmax(dim=-1)

1765 - return next_token_ids, probs

1766 -

1767 - assert sampling_metadata.temperature is not None

1768 -

1769 - # Use epsilon comparison to detect greedy sampling (temperature ~ 0.0)

1770 - # consistent with sampler.py's _SAMPLING_EPS threshold

1771 - temperature = sampling_metadata.temperature

1772 - # Avoid division by zero if there are greedy requests.

1773 - if not sampling_metadata.all_random:

1774 - is_greedy = temperature < _SAMPLING_EPS

1775 - temperature = torch.where(is_greedy, 1.0, temperature)

1776 - logits.div_(temperature.view(-1, 1))

1777 - probs = logits.softmax(dim=-1, dtype=torch.float32)

1778 -

1779 - # NOTE(woosuk): Currently, we ignore most of the sampling parameters in

1780 - # generating the draft tokens. We only use the temperature. While this

1781 - # could degrade the acceptance rate, it does not affect the distribution

1782 - # of the generated tokens after rejection sampling.

1783 -

1784 - # TODO(woosuk): Consider seeds.

1785 - q = torch.empty_like(probs)

1786 - q.exponential_()

1787 - # NOTE(woosuk): We shouldn't use `probs.div_(q)` because the draft_probs

1788 - # will be used later for rejection sampling.

1789 - next_token_ids = probs.div(q).argmax(dim=-1).view(-1)

1790 - if not sampling_metadata.all_random:

1791 - greedy_token_ids = probs.argmax(dim=-1)

1792 - next_token_ids = torch.where(is_greedy, greedy_token_ids, next_token_ids)

1793 - return next_token_ids, probs

vllm/v1/spec_decode/llm_base_proposer.py new

+1778

@@ -0,0 +1,1778 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +import ast

4 +from importlib.util import find_spec

5 +from typing import Any, cast

6 +

7 +import numpy as np

8 +import torch

9 +import torch.nn as nn

10 +

11 +from vllm.config import (

12 + CUDAGraphMode,

13 + VllmConfig,

14 + get_layers_from_vllm_config,

15 + replace,

16 +)

17 +from vllm.distributed.parallel_state import get_pp_group

18 +from vllm.forward_context import set_forward_context

19 +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_eagle3 import Eagle3DeepseekV2ForCausalLM

24 +from vllm.model_executor.models.interfaces import SupportsMultiModal

25 +from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

26 +from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM

27 +from vllm.multimodal import MULTIMODAL_REGISTRY

28 +from vllm.platforms import current_platform

29 +from vllm.utils.platform_utils import is_pin_memory_available

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

31 +from vllm.v1.attention.backends.registry import AttentionBackendEnum

32 +from vllm.v1.attention.backends.tree_attn import (

33 + TreeAttentionMetadata,

34 + TreeAttentionMetadataBuilder,

35 +)

36 +from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata

37 +from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

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

39 +from vllm.v1.sample.metadata import SamplingMetadata

40 +from vllm.v1.sample.sampler import _SAMPLING_EPS

41 +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

42 +from vllm.v1.spec_decode.utils import (

43 + PADDING_SLOT_ID,

44 + compute_new_slot_mapping,

45 + copy_and_expand_eagle_inputs_kernel,

46 + eagle_prepare_inputs_padded_kernel,

47 + eagle_prepare_next_token_padded_kernel,

48 + eagle_step_update_slot_mapping_and_metadata,

49 + extend_all_queries_by_N,

50 + next_power_of_2,

51 +)

52 +from vllm.v1.utils import CpuGpuBuffer

53 +from vllm.v1.worker.dp_utils import coordinate_batch_across_dp

54 +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

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

56 +

57 +logger = init_logger(__name__)

58 +

59 +

60 +class SpecDecodeBaseProposer:

61 + def __init__(

62 + self,

63 + vllm_config: VllmConfig,

64 + device: torch.device,

65 + pass_hidden_states_to_model: bool,

66 + runner=None,

67 + ):

68 + self.vllm_config = vllm_config

69 + assert vllm_config.speculative_config is not None

70 + self.speculative_config = vllm_config.speculative_config

71 + self.draft_model_config = self.speculative_config.draft_model_config

72 + self.method = self.speculative_config.method

73 + self.pass_hidden_states_to_model = pass_hidden_states_to_model

74 +

75 + self.device = device

76 + self.dtype = vllm_config.model_config.dtype

77 + self.max_model_len = vllm_config.model_config.max_model_len

78 + self.dp_rank = vllm_config.parallel_config.data_parallel_rank

79 + self.num_speculative_tokens = self.speculative_config.num_speculative_tokens

80 +

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

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

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

84 + self.hidden_size = self.draft_model_config.get_hidden_size()

85 + self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()

86 +

87 + # Unifying eagle, draft model, and parallel drafting support.

88 + # DFlash always uses parallel drafting (all tokens in one pass),

89 + # but has an additional slot for the next_token_id (does not shift like EAGLE)

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

91 + self.extra_slots_per_request = (

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

93 + )

94 + self.net_num_new_slots_per_request = self.extra_slots_per_request - (

95 + 1 if (self.pass_hidden_states_to_model and self.method != "dflash") else 0

96 + )

97 + self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0

98 +

99 + self.parallel_drafting_token_id: int = 0

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

101 + if self.parallel_drafting:

102 + self._init_parallel_drafting_params()

103 + self.use_local_argmax_reduction: bool = (

104 + self.speculative_config.use_local_argmax_reduction

105 + )

106 +

107 + self.max_batch_size = vllm_config.scheduler_config.max_num_seqs

108 + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

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

110 +

111 + # Can be specialized by methods like DFlash to reduce the limit

112 + self.max_query_tokens = self.max_num_tokens

113 + self.max_positions = self.max_num_tokens

114 +

115 + # Multi-modal data support

116 + self.mm_registry = MULTIMODAL_REGISTRY

117 + self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(

118 + vllm_config.model_config

119 + )

120 +

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

122 + self.kv_cache_gid: int = -1

123 + self.eagle3_use_aux_hidden_state: bool = (

124 + self._get_eagle3_use_aux_hidden_state_from_config()

125 + )

126 +

127 + self.compilation_config = self.vllm_config.compilation_config

128 +

129 + # Cudagraph dispatcher for PIECEWISE-only dispatching in eagle.

130 + # Keys are initialized later via initialize_cudagraph_keys() called from

131 + # gpu_model_runner._check_and_update_cudagraph_mode after

132 + # adjust_cudagraph_sizes_for_spec_decode is called.

133 + self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)

134 +

135 + # persistent buffers for cuda graph

136 + self.input_ids = torch.zeros(

137 + self.max_num_tokens, dtype=torch.int32, device=device

138 + )

139 + # Use draft model's M-RoPE setting, not target model's

140 + # Draft models may be text-only even if target is multimodal

141 + self.uses_mrope = self.draft_model_config.uses_mrope

142 + self.uses_xdrope_dim = self.vllm_config.model_config.uses_xdrope_dim

143 + self.draft_uses_xdrope_dim = self.draft_model_config.uses_xdrope_dim

144 + if self.uses_mrope:

145 + # NOTE: `mrope_positions` is implemented with one additional dummy

146 + # position on purpose to make it non-contiguous so that it can work

147 + # with torch compile.

148 + # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923

149 +

150 + # NOTE: When M-RoPE is enabled, position ids are 3D regardless of

151 + # the modality of inputs. For text-only inputs, each dimension has

152 + # identical position IDs, making M-RoPE functionally equivalent to

153 + # 1D-RoPE.

154 + # See page 5 of https://arxiv.org/abs/2409.12191

155 + self.mrope_positions = torch.zeros(

156 + (3, self.max_positions + 1), dtype=torch.int64, device=device

157 + )

158 + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

159 + self.xdrope_positions = torch.zeros(

160 + (self.uses_xdrope_dim, self.max_positions + 1),

161 + dtype=torch.int64,

162 + device=device,

163 + )

164 + else:

165 + # RoPE need (max_num_tokens,)

166 + self.positions = torch.zeros(

167 + self.max_positions,

168 + dtype=torch.int64,

169 + device=device,

170 + )

171 + self.hidden_states = torch.zeros(

172 + (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device

173 + )

174 +

175 + # Will be set when we initialize the attention backend

176 + self.block_size: int = -1

177 +

178 + # We need +1 here because the arange is used to set query_start_loc,

179 + # which has one more element than batch_size.

180 + max_num_slots_for_arange = max(self.max_batch_size + 1, self.max_num_tokens)

181 + self.arange = torch.arange(

182 + max_num_slots_for_arange, device=device, dtype=torch.int32

183 + )

184 +

185 + if self.needs_extra_input_slots:

186 + self._raise_if_padded_drafter_batch_disabled()

187 + self._raise_if_multimodal()

188 + self._raise_if_mrope()

189 +

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

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

192 + if self.needs_extra_input_slots:

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

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

195 + self.is_rejected_token_mask = torch.zeros(

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

197 + )

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

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

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

201 + self.is_masked_token_mask = torch.zeros(

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

203 + )

204 +

205 + self.inputs_embeds = torch.zeros(

206 + (self.max_num_tokens, self.inputs_embeds_size),

207 + dtype=self.dtype,

208 + device=device,

209 + )

210 +

211 + self.backup_next_token_ids = CpuGpuBuffer(

212 + self.max_batch_size,

213 + dtype=torch.int32,

214 + pin_memory=is_pin_memory_available(),

215 + device=device,

216 + with_numpy=True,

217 + )

218 +

219 + self._slot_mapping_buffer = torch.zeros(

220 + self.max_positions,

221 + dtype=torch.int64,

222 + device=device,

223 + )

224 +

225 + # Determine allowed attention backends once during initialization.

226 + self.allowed_attn_types: tuple | None = None

227 + if current_platform.is_rocm():

228 + from vllm.v1.attention.backends.mla.indexer import (

229 + DeepseekV32IndexerMetadata,

230 + )

231 + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (

232 + ROCMAiterMLASparseMetadata,

233 + )

234 + from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata

235 +

236 + rocm_types = [

237 + TritonAttentionMetadata,

238 + RocmAttentionMetadata,

239 + ROCMAiterMLASparseMetadata,

240 + DeepseekV32IndexerMetadata,

241 + ]

242 + # ROCM_AITER_FA is an optional backend

243 + # We check is_enabled() here to avoid importing the backend module during

244 + # auto-discovery when VLLM_ROCM_USE_AITER=0, which would trigger aiter

245 + # import and JIT compilation warnings. Explicit backend selection via

246 + # attention_config still works because the backend module is loaded

247 + # directly when selected, not through this auto-discovery path.

248 + # Check if backend module exists to allow explicit selection

249 + if find_spec(

250 + AttentionBackendEnum.ROCM_AITER_FA.get_path(include_classname=False)

251 + ):

252 + from vllm.v1.attention.backends.rocm_aiter_fa import (

253 + AiterFlashAttentionMetadata,

254 + )

255 +

256 + rocm_types.append(AiterFlashAttentionMetadata)

257 +

258 + # TRITON_MLA backend support for MLA models (e.g., DeepSeek)

259 + from vllm.model_executor.layers.attention.mla_attention import (

260 + MLACommonMetadata,

261 + )

262 +

263 + rocm_types.append(MLACommonMetadata)

264 +

265 + # FlexAttention backend support

266 + from vllm.v1.attention.backends.flex_attention import FlexAttentionMetadata

267 +

268 + rocm_types.append(FlexAttentionMetadata)

269 +

270 + self.allowed_attn_types = tuple(rocm_types)

271 +

272 + # Parse the speculative token tree.

273 + spec_token_tree = self.speculative_config.speculative_token_tree

274 + assert spec_token_tree is not None

275 + self.tree_choices: list[tuple[int, ...]] = ast.literal_eval(spec_token_tree)

276 + tree_depth = len(self.tree_choices[-1])

277 + # Precompute per-level properties of the tree.

278 + num_drafts_per_level = [0] * tree_depth

279 + for node in self.tree_choices:

280 + num_drafts_per_level[len(node) - 1] += 1

281 + self.cu_drafts_per_level = [num_drafts_per_level[0]]

282 + self.child_drafts_per_level = [num_drafts_per_level[0]]

283 + for level in range(1, tree_depth):

284 + self.cu_drafts_per_level.append(

285 + self.cu_drafts_per_level[-1] + num_drafts_per_level[level]

286 + )

287 + self.child_drafts_per_level.append(

288 + num_drafts_per_level[level] // num_drafts_per_level[level - 1]

289 + )

290 + # Precompute draft position offsets in flattened tree.

291 + self.tree_draft_pos_offsets = torch.arange(

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

293 + ).repeat(self.max_batch_size, 1)

294 +

295 + def _raise_if_padded_drafter_batch_disabled(self):

296 + if self.speculative_config.disable_padded_drafter_batch:

297 + raise NotImplementedError(

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

299 + "supports padded drafter batch. Please unset "

300 + "disable_padded_drafter_batch in the speculative_config."

301 + )

302 +

303 + def _raise_if_multimodal(self):

304 + if self.supports_mm_inputs:

305 + raise NotImplementedError(

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

307 + "does not support multimodal models yet"

308 + )

309 +

310 + def _raise_if_mrope(self):

311 + if self.draft_model_config.uses_mrope:

312 + raise NotImplementedError(

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

314 + "does not support M-RoPE yet"

315 + )

316 +

317 + def _init_parallel_drafting_params(self):

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

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

320 + # for those masked slots.

321 +

322 + model_hf_config = self.draft_model_config.hf_config

323 + # DFlash stores mask_token_id in dflash_config

324 + dflash_config = getattr(model_hf_config, "dflash_config", None)

325 + if dflash_config and "mask_token_id" in dflash_config:

326 + self.parallel_drafting_token_id = dflash_config["mask_token_id"]

327 + elif hasattr(model_hf_config, "pard_token"):

328 + self.parallel_drafting_token_id = model_hf_config.pard_token

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

330 + self.parallel_drafting_token_id = model_hf_config.ptd_token_id

331 + else:

332 + raise ValueError(

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

334 + "`pard_token`, `ptd_token_id`, or "

335 + "`dflash_config.mask_token_id` specified in its config.json."

336 + )

337 +

338 + if self.pass_hidden_states_to_model:

339 + self.parallel_drafting_hidden_state_tensor = torch.empty(

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

341 + )

342 +

343 + def _get_positions(self, num_tokens: int):

344 + if self.uses_mrope:

345 + return self.mrope_positions[:, :num_tokens]

346 + if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

347 + return self.xdrope_positions[:, :num_tokens]

348 + return self.positions[:num_tokens]

349 +

350 + def _set_positions(self, num_tokens: int, positions: torch.Tensor):

351 + if self.uses_mrope:

352 + self.mrope_positions[:, :num_tokens] = positions

353 + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

354 + self.xdrope_positions[:, :num_tokens] = positions

355 + else:

356 + # Convert M-RoPE positions if target model uses M-RoPE

357 + # but draft doesn't, For text inputs, all M-RoPE

358 + # dimensions are identical

359 + if self.vllm_config.model_config.uses_mrope:

360 + positions = positions[0]

361 + self.positions[:num_tokens] = positions

362 +

363 + def _get_slot_mapping(

364 + self,

365 + num_tokens: int,

366 + slot_mapping: torch.Tensor | None = None,

367 + ) -> dict[str, torch.Tensor]:

368 + """Return slot_mapping dict for EAGLE layers.

369 +

370 + If slot_mapping is provided, copies it into the buffer first.

371 + """

372 + if slot_mapping is not None:

373 + num_actual = slot_mapping.shape[0]

374 + self._slot_mapping_buffer[:num_actual].copy_(slot_mapping)

375 + if num_tokens > num_actual:

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

377 +

378 + view = self._slot_mapping_buffer[:num_tokens]

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

380 +

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

382 + """Initialize cudagraph dispatcher keys for eagle.

383 +

384 + Eagle only supports PIECEWISE cudagraphs (via mixed_mode).

385 + This should be called after adjust_cudagraph_sizes_for_spec_decode.

386 + """

387 + if (

388 + not self.speculative_config.enforce_eager

389 + and cudagraph_mode.mixed_mode()

390 + in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]

391 + ):

392 + eagle_cudagraph_mode = CUDAGraphMode.PIECEWISE

393 + else:

394 + eagle_cudagraph_mode = CUDAGraphMode.NONE

395 +

396 + self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode)

397 +

398 + def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor:

399 + """Greedy-sample draft tokens from hidden states."""

400 + if self.use_local_argmax_reduction:

401 + return self.model.get_top_tokens(hidden_states)

402 + return self.model.compute_logits(hidden_states).argmax(dim=-1)

403 +

404 + def propose(

405 + self,

406 + # [num_tokens]

407 + target_token_ids: torch.Tensor,

408 + # [num_tokens] or [3, num_tokens] when M-RoPE is enabled

409 + target_positions: torch.Tensor,

410 + # [num_tokens, hidden_size]

411 + target_hidden_states: torch.Tensor,

412 + # [batch_size]

413 + next_token_ids: torch.Tensor,

414 + token_indices_to_sample: torch.Tensor | None,

415 + common_attn_metadata: CommonAttentionMetadata,

416 + sampling_metadata: SamplingMetadata,

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

418 + num_rejected_tokens_gpu: torch.Tensor | None = None,

419 + slot_mappings: dict[str, torch.Tensor]

420 + | list[dict[str, torch.Tensor]]

421 + | None = None,

422 + ) -> torch.Tensor:

423 + batch_size = common_attn_metadata.batch_size()

424 +

425 + if self.method in ("eagle3", "dflash"):

426 + assert isinstance(

427 + self.model,

428 + (

429 + Eagle3LlamaForCausalLM,

430 + Eagle3DeepseekV2ForCausalLM,

431 + DFlashQwen3ForCausalLM,

432 + ),

433 + )

434 + target_hidden_states = self.model.combine_hidden_states(

435 + target_hidden_states

436 + )

437 + assert target_hidden_states.shape[-1] == self.hidden_size

438 +

439 + num_tokens, token_indices_to_sample, common_attn_metadata = (

440 + self.set_inputs_first_pass(

441 + target_token_ids=target_token_ids,

442 + next_token_ids=next_token_ids,

443 + target_positions=target_positions,

444 + target_hidden_states=target_hidden_states,

445 + token_indices_to_sample=token_indices_to_sample,

446 + cad=common_attn_metadata,

447 + num_rejected_tokens_gpu=num_rejected_tokens_gpu,

448 + )

449 + )

450 +

451 + per_group_attn_metadata, per_layer_attn_metadata = (

452 + self.build_per_group_and_layer_attn_metadata(common_attn_metadata)

453 + )

454 +

455 + cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

456 + self._determine_batch_execution_and_padding(num_tokens)

457 + )

458 +

459 + model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass(

460 + num_tokens, num_input_tokens, mm_embed_inputs

461 + )

462 +

463 + with set_forward_context(

464 + per_layer_attn_metadata,

465 + self.vllm_config,

466 + num_tokens=num_input_tokens,

467 + num_tokens_across_dp=num_tokens_across_dp,

468 + cudagraph_runtime_mode=cudagraph_runtime_mode,

469 + slot_mapping=self._get_slot_mapping(

470 + slot_mapping_size, common_attn_metadata.slot_mapping

471 + ),

472 + ):

473 + ret_hidden_states = self.model(**model_kwargs)

474 + if not self.model_returns_tuple():

475 + last_hidden_states = ret_hidden_states

476 + hidden_states = last_hidden_states

477 + else:

478 + last_hidden_states, hidden_states = ret_hidden_states

479 +

480 + sample_hidden_states = last_hidden_states[token_indices_to_sample]

481 +

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

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

484 + draft_token_ids = self._greedy_sample(sample_hidden_states)

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

486 +

487 + if self.uses_mrope:

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

489 + else:

490 + positions = self.positions[token_indices_to_sample]

491 + hidden_states = hidden_states[token_indices_to_sample]

492 +

493 + if any(isinstance(md, TreeAttentionMetadata) for md in per_group_attn_metadata):

494 + # Draft using tree attention - requires full logits for top-k

495 + logits = self.model.compute_logits(sample_hidden_states)

496 + draft_token_ids_list = self.propose_tree(

497 + batch_size=batch_size,

498 + logits=logits,

499 + positions=positions,

500 + hidden_states=hidden_states,

501 + common_attn_metadata=common_attn_metadata,

502 + slot_mappings=slot_mappings,

503 + )

504 + # [batch_size, num_tree_tokens]

505 + return torch.cat(draft_token_ids_list, dim=1)

506 +

507 + draft_token_ids = self._greedy_sample(sample_hidden_states)

508 +

509 + if self.allowed_attn_types is not None:

510 + for group_md in per_group_attn_metadata:

511 + if not isinstance(group_md, self.allowed_attn_types):

512 + raise ValueError(

513 + f"Unsupported attention metadata type for speculative "

514 + "decoding with num_speculative_tokens > 1: "

515 + f"{type(group_md)}. Supported types are: "

516 + f"{self.allowed_attn_types}"

517 + )

518 +

519 + # Generate the remaining draft tokens.

520 + draft_token_ids_list = [draft_token_ids]

521 +

522 + cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = (

523 + self._determine_batch_execution_and_padding(batch_size)

524 + )

525 +

526 + common_attn_metadata.num_actual_tokens = batch_size

527 + common_attn_metadata.max_query_len = 1

528 + common_attn_metadata.query_start_loc = self.arange[: batch_size + 1]

529 + common_attn_metadata.query_start_loc_cpu = torch.from_numpy(

530 + self.token_arange_np[: batch_size + 1]

531 + ).clone()

532 +

533 + # In padded drafter batch, we need to adjust the sequence lengths

534 + # to remove the "padding" (i.e. rejected tokens).

535 + # Only apply this adjustment when we have rejected tokens

536 + # (i.e., not the first proposal).

537 + if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None:

538 + common_attn_metadata.seq_lens -= num_rejected_tokens_gpu

539 + # Invalidate the CPU-side shadows to avoid H<>D sync.

540 + common_attn_metadata._seq_lens_cpu = None

541 + common_attn_metadata._num_computed_tokens_cpu = None

542 +

543 + block_size = self.block_size

544 + assert block_size > 0, "block_size has not been initialized."

545 + for token_index in range(self.num_speculative_tokens - 1):

546 + # Update the inputs.

547 + # cast to int32 is crucial when eagle model is compiled.

548 + # tensor.argmax() returns int64 by default.

549 + input_ids = draft_token_ids_list[-1].int()

550 + # Use fused kernel for slot mapping and metadata updates.

551 + # Write clamped positions directly into the positions buffer to

552 + # avoid an extra D2D copy for the common (non-mrope) case.

553 + positions_1d = positions[0] if self.uses_mrope else positions

554 + if self.uses_mrope:

555 + out_pos = self.mrope_positions[0, :batch_size]

556 + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

557 + out_pos = self.xdrope_positions[0, :batch_size]

558 + else:

559 + out_pos = self.positions[:batch_size]

560 + eagle_step_update_slot_mapping_and_metadata(

561 + positions_1d=positions_1d,

562 + block_table_tensor=common_attn_metadata.block_table_tensor,

563 + seq_lens=common_attn_metadata.seq_lens,

564 + block_size=block_size,

565 + max_model_len=self.max_model_len,

566 + out_clamped_positions=out_pos,

567 + out_slot_mapping=self._slot_mapping_buffer[:input_batch_size],

568 + input_batch_size=input_batch_size,

569 + )

570 + common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]

571 + if self.uses_mrope:

572 + self.mrope_positions[1:, :batch_size] = self.mrope_positions[

573 + 0, :batch_size

574 + ]

575 + positions = self.mrope_positions[:, :batch_size]

576 + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

577 + self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[

578 + 0, :batch_size

579 + ]

580 + positions = self.xdrope_positions[0, :batch_size]

581 + else:

582 + positions = self.positions[:batch_size]

583 + # Increment the maximum sequence length. We increment max_seq_len

584 + # unconditionally even though some seq_lens may have been capped above,

585 + # as max_seq_len serves as an upper bound for sequence lengths.

586 + common_attn_metadata.max_seq_len = min(

587 + common_attn_metadata.max_seq_len + 1, self.max_model_len

588 + )

589 +

590 + # Also update the CPU-side shadow; NOTE: this is hacky and should be

591 + # removed in when common_attn_metadata.seq_lens_cpu is deprecated.

592 + if common_attn_metadata._seq_lens_cpu is not None:

593 + common_attn_metadata._seq_lens_cpu += 1

594 + if common_attn_metadata._num_computed_tokens_cpu is not None:

595 + common_attn_metadata._num_computed_tokens_cpu += 1

596 +

597 + # Rebuild attention metadata

598 + _, per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata(

599 + common_attn_metadata, draft_index=token_index + 1

600 + )

601 +

602 + # copy inputs to buffer for cudagraph

603 + self.input_ids[:batch_size] = input_ids

604 + self.hidden_states[:batch_size] = hidden_states

605 + if self.supports_mm_inputs:

606 + self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids)

607 +

608 + input_ids = None

609 + inputs_embeds = self.inputs_embeds[:input_batch_size]

610 + else:

611 + input_ids = self.input_ids[:input_batch_size]

612 + inputs_embeds = None

613 +

614 + # Run the model.

615 + model_kwargs = {

616 + "input_ids": input_ids,

617 + "positions": self._get_positions(input_batch_size),

618 + "inputs_embeds": inputs_embeds,

619 + }

620 + if self.pass_hidden_states_to_model:

621 + model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size]

622 +

623 + with set_forward_context(

624 + per_layer_attn_metadata,

625 + self.vllm_config,

626 + num_tokens=input_batch_size,

627 + num_tokens_across_dp=batch_size_across_dp,

628 + cudagraph_runtime_mode=cudagraph_runtime_mode,

629 + slot_mapping=self._get_slot_mapping(input_batch_size),

630 + ):

631 + ret_hidden_states = self.model(**model_kwargs)

632 + if not self.model_returns_tuple():

633 + last_hidden_states = ret_hidden_states

634 + hidden_states = ret_hidden_states

635 + else:

636 + last_hidden_states, hidden_states = ret_hidden_states

637 +

638 + hidden_states = hidden_states[:batch_size]

639 + draft_token_ids = self._greedy_sample(last_hidden_states[:batch_size])

640 + draft_token_ids_list.append(draft_token_ids)

641 +

642 + # [batch_size, num_speculative_tokens]

643 + draft_token_ids = torch.stack(draft_token_ids_list, dim=1)

644 + return draft_token_ids

645 +

646 + def set_inputs_first_pass(

647 + self,

648 + target_token_ids: torch.Tensor,

649 + next_token_ids: torch.Tensor,

650 + target_positions: torch.Tensor,

651 + target_hidden_states: torch.Tensor,

652 + token_indices_to_sample: torch.Tensor | None,

653 + cad: CommonAttentionMetadata,

654 + num_rejected_tokens_gpu: torch.Tensor | None,

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

656 + if not self.needs_extra_input_slots:

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

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

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

660 + if token_indices_to_sample is None:

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

662 +

663 + num_tokens = target_token_ids.shape[0]

664 + # Shift the input ids by one token.

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

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

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

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

669 + self.input_ids[token_indices_to_sample] = next_token_ids

670 +

671 + # copy inputs to buffer for cudagraph

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

673 + target_positions = target_positions[0]

674 + self._set_positions(num_tokens, target_positions)

675 +

676 + self.hidden_states[:num_tokens] = target_hidden_states

677 +

678 + return num_tokens, token_indices_to_sample, cad

679 + else:

680 + assert self.is_rejected_token_mask is not None

681 + assert self.is_masked_token_mask is not None

682 + # 1.

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

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

685 + # self.positions.

686 + batch_size = cad.batch_size()

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

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

689 + max_num_tokens_per_request = (

690 + cad.max_query_len + self.net_num_new_slots_per_request

691 + )

692 + BLOCK_SIZE_TOKENS = min(256, next_power_of_2(max_num_tokens_per_request))

693 + num_blocks = (

694 + max_num_tokens_per_request + BLOCK_SIZE_TOKENS - 1

695 + ) // BLOCK_SIZE_TOKENS

696 + total_num_input_tokens = target_token_ids.shape[0]

697 + total_num_output_tokens = total_num_input_tokens + (

698 + self.net_num_new_slots_per_request * batch_size

699 + )

700 +

701 + token_indices_to_sample = torch.empty(

702 + batch_size * self.extra_slots_per_request,

703 + dtype=torch.int32,

704 + device=self.device,

705 + )

706 +

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

708 + out_hidden_state_mapping = torch.empty(

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

710 + )

711 +

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

713 + grid = (batch_size, num_blocks)

714 + query_start_loc = cad.query_start_loc

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

716 + if num_rejected_tokens_gpu is not None:

717 + query_end_loc = query_end_loc - num_rejected_tokens_gpu

718 +

719 + copy_and_expand_eagle_inputs_kernel[grid](

720 + # (Padded) Inputs from the target model

721 + target_token_ids_ptr=target_token_ids,

722 + target_positions_ptr=target_positions,

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

724 + # Outputs to the drafting buffers

725 + out_input_ids_ptr=self.input_ids,

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

727 + out_is_rejected_token_mask_ptr=self.is_rejected_token_mask,

728 + out_is_masked_token_mask_ptr=self.is_masked_token_mask,

729 + out_new_token_indices_ptr=token_indices_to_sample,

730 + out_hidden_state_mapping_ptr=out_hidden_state_mapping,

731 + # Input metadata

732 + query_start_loc_ptr=query_start_loc,

733 + query_end_loc_ptr=query_end_loc,

734 + padding_token_id=0,

735 + parallel_drafting_token_id=self.parallel_drafting_token_id,

736 + # Sizing info

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

738 + total_input_tokens=total_num_input_tokens,

739 + num_padding_slots_per_request=self.extra_slots_per_request,

740 + shift_input_ids=self.pass_hidden_states_to_model,

741 + BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,

742 + )

743 + if self.pass_hidden_states_to_model:

744 + assert self.parallel_drafting_hidden_state_tensor is not None

745 + self.hidden_states[out_hidden_state_mapping] = target_hidden_states

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

747 + mask = self.is_masked_token_mask[:total_num_output_tokens]

748 + torch.where(

749 + mask.unsqueeze(1),

750 + self.parallel_drafting_hidden_state_tensor,

751 + self.hidden_states[:total_num_output_tokens],

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

753 + )

754 +

755 + # 2.

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

757 + # rejection mask.

758 + assert self.block_size > 0, "block_size has not been initialized."

759 + new_slot_mapping = compute_new_slot_mapping(

760 + cad=cad,

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

762 + is_rejected_token_mask=self.is_rejected_token_mask[

763 + :total_num_output_tokens

764 + ],

765 + block_size=self.block_size,

766 + num_new_tokens=self.net_num_new_slots_per_request,

767 + max_model_len=self.max_model_len,

768 + )

769 +

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

771 + new_cad = extend_all_queries_by_N(

772 + cad,

773 + N=self.net_num_new_slots_per_request,

774 + arange=self.arange,

775 + new_slot_mapping=new_slot_mapping,

776 + )

777 +

778 + return total_num_output_tokens, token_indices_to_sample, new_cad

779 +

780 + def build_model_inputs_first_pass(

781 + self,

782 + num_tokens: int,

783 + num_input_tokens: int,

784 + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,

785 + ) -> tuple[dict[str, Any], int]:

786 + if self.supports_mm_inputs:

787 + mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

788 +

789 + self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(

790 + self.input_ids[:num_tokens],

791 + multimodal_embeddings=mm_embeds,

792 + is_multimodal=is_mm_embed,

793 + )

794 +

795 + input_ids = None

796 + inputs_embeds = self.inputs_embeds[:num_input_tokens]

797 + else:

798 + input_ids = self.input_ids[:num_input_tokens]

799 + inputs_embeds = None

800 +

801 + model_kwargs = {

802 + "input_ids": input_ids,

803 + "positions": self._get_positions(num_input_tokens),

804 + "inputs_embeds": inputs_embeds,

805 + }

806 + if self.pass_hidden_states_to_model:

807 + model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]

808 +

809 + return model_kwargs, num_input_tokens

810 +

811 + def build_per_group_and_layer_attn_metadata(

812 + self, common_attn_metadata: CommonAttentionMetadata, draft_index: int = 0

813 + ) -> tuple[list[object], dict[str, object]]:

814 + per_group_attn_metadata: list[object] = []

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

816 + for attn_group in self.draft_attn_groups:

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

818 + common_attn_metadata=common_attn_metadata, draft_index=draft_index

819 + )

820 + per_group_attn_metadata.append(attn_metadata)

821 + for layer_name in attn_group.layer_names:

822 + per_layer_attn_metadata[layer_name] = attn_metadata

823 + return per_group_attn_metadata, per_layer_attn_metadata

824 +

825 + def model_returns_tuple(self) -> bool:

826 + return self.method not in ("mtp", "draft_model", "dflash")

827 +

828 + def prepare_next_token_ids_cpu(

829 + self,

830 + sampled_token_ids: list[list[int]],

831 + requests: dict[str, CachedRequestState],

832 + gpu_input_batch: InputBatch,

833 + num_scheduled_tokens: dict[str, int],

834 + ) -> torch.Tensor:

835 + """

836 + This function is used to prepare the inputs for speculative decoding.

837 + It calculates the next token ids for each request based on the sampled

838 + token ids from the CPU. If a request has no sampled token ids (e.g.,

839 + during the initial decoding steps), it falls back to using the request

840 + state to get the next token id.

841 + """

842 + req_ids = gpu_input_batch.req_ids

843 + next_token_ids: list[int] = []

844 + for i, token_ids in enumerate(sampled_token_ids):

845 + if token_ids:

846 + # Common case.

847 + next_token_id = token_ids[-1]

848 + else:

849 + # Partial prefill (rare case).

850 + # Get the next token id from the request state.

851 + req_id = req_ids[i]

852 + req_state = requests[req_id]

853 + seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]

854 + next_token_id = req_state.get_token_id(seq_len)

855 + next_token_ids.append(next_token_id)

856 + next_token_ids = torch.tensor(

857 + next_token_ids, dtype=torch.int32, device=self.input_ids.device

858 + )

859 + return next_token_ids

860 +

861 + def prepare_next_token_ids_padded(

862 + self,

863 + sampled_token_ids: torch.Tensor,

864 + requests: dict[str, CachedRequestState],

865 + gpu_input_batch: InputBatch,

866 + discard_request_mask: torch.Tensor,

867 + ) -> tuple[torch.Tensor, torch.Tensor]:

868 + """

869 + This function is used to prepare the inputs for speculative decoding.

870 + It calculates the next token ids and the number of valid sampled tokens

871 + for each request, considering the "discarded" requests whose next token

872 + is not sampled and comes from `request.get_token_id()` instead. This is denoted

873 + the "backup" token id. It also counts rejected tokens via `sampled_token_ids`.

874 + """

875 + # Precompute get_token_id for when there is no valid next token

876 + num_reqs = gpu_input_batch.num_reqs

877 + seq_lens_list = (gpu_input_batch.num_tokens_no_spec[:num_reqs] - 1).tolist()

878 + self.backup_next_token_ids.np[:num_reqs] = np.array(

879 + [

880 + requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])

881 + for i in range(num_reqs)

882 + ],

883 + dtype=np.int32,

884 + )

885 + self.backup_next_token_ids.copy_to_gpu(num_reqs)

886 + backup_tokens_gpu = self.backup_next_token_ids.gpu

887 +

888 + batch_size, num_tokens = sampled_token_ids.shape

889 + device = sampled_token_ids.device

890 +

891 + assert discard_request_mask.dtype == torch.bool

892 + assert backup_tokens_gpu.dtype == torch.int32

893 +

894 + next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=device)

895 + valid_sampled_tokens_count = next_token_ids.new_empty(batch_size)

896 +

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

898 + grid = (batch_size,)

899 +

900 + # Find the next power of 2 for block sizes

901 + BLOCK_SIZE_TOKENS = next_power_of_2(num_tokens)

902 + eagle_prepare_next_token_padded_kernel[grid](

903 + sampled_token_ids,

904 + discard_request_mask,

905 + backup_tokens_gpu,

906 + next_token_ids,

907 + valid_sampled_tokens_count,

908 + gpu_input_batch.vocab_size,

909 + num_tokens,

910 + batch_size,

911 + sampled_token_ids.stride(0),

912 + BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,

913 + )

914 +

915 + return next_token_ids, valid_sampled_tokens_count

916 +

917 + def prepare_inputs_padded(

918 + self,

919 + common_attn_metadata: CommonAttentionMetadata,

920 + spec_decode_metadata: SpecDecodeMetadata,

921 + valid_sampled_tokens_count: torch.Tensor,

922 + ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:

923 + """

924 + This function is used to prepare the inputs for speculative decoding

925 + It updates the common_attn_metadata for speculative decoding,

926 + but does not consider the rejected tokens. Instead, all tokens

927 + are included as inputs to the speculator, with the rejected tokens

928 + used as padding and filtered out later by `token_indices_to_sample`.

929 + No blocking CPU operations should be introduced in this function.

930 + """

931 + num_reqs = common_attn_metadata.num_reqs

932 + device = valid_sampled_tokens_count.device

933 +

934 + token_indices_to_sample = torch.empty(

935 + (num_reqs,), dtype=torch.int32, device=device

936 + )

937 + num_rejected_tokens_gpu = torch.empty(

938 + (num_reqs,), dtype=torch.int32, device=device

939 + )

940 +

941 + grid = (num_reqs,)

942 + eagle_prepare_inputs_padded_kernel[grid](

943 + spec_decode_metadata.cu_num_draft_tokens,

944 + valid_sampled_tokens_count,

945 + common_attn_metadata.query_start_loc,

946 + token_indices_to_sample,

947 + num_rejected_tokens_gpu,

948 + num_reqs,

949 + )

950 +

951 + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

952 + new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

953 +

954 + total_num_tokens = query_start_loc_cpu[-1].item()

955 +

956 + spec_common_attn_metadata = CommonAttentionMetadata(

957 + query_start_loc=common_attn_metadata.query_start_loc,

958 + seq_lens=common_attn_metadata.seq_lens,

959 + query_start_loc_cpu=query_start_loc_cpu,

960 + _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,

961 + _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

962 + num_reqs=common_attn_metadata.num_reqs,

963 + num_actual_tokens=total_num_tokens,

964 + max_query_len=new_query_len_per_req.max().item(),

965 + max_seq_len=common_attn_metadata.max_seq_len,

966 + block_table_tensor=common_attn_metadata.block_table_tensor,

967 + slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens],

968 + causal=True,

969 + dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,

970 + )

971 +

972 + return (

973 + spec_common_attn_metadata,

974 + token_indices_to_sample,

975 + num_rejected_tokens_gpu,

976 + )

977 +

978 + def propose_tree(

979 + self,

980 + batch_size: int,

981 + # [num_tokens, vocab_size]

982 + logits: torch.Tensor,

983 + # [num_tokens]

984 + positions: torch.Tensor,

985 + # [num_tokens, hidden_size]

986 + hidden_states: torch.Tensor,

987 + common_attn_metadata: CommonAttentionMetadata,

988 + slot_mappings: dict[str, torch.Tensor]

989 + | list[dict[str, torch.Tensor]]

990 + | None = None,

991 + ) -> list[torch.Tensor]:

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

993 + assert isinstance(tree_attn_metadata_builder, TreeAttentionMetadataBuilder)

994 +

995 + total_num_drafts = self.cu_drafts_per_level[0]

996 + level_num_drafts = total_num_drafts

997 + # Sample a draft token for each child at the tree root level.

998 + num_children = self.child_drafts_per_level[0]

999 + if num_children == 1:

1000 + draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)

1001 + else:

1002 + draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(

1003 + batch_size, -1

1004 + )

1005 + draft_token_ids_list = [draft_token_ids]

1006 + draft_hidden_states = hidden_states.view(batch_size, 1, -1)

1007 +

1008 + # Initialize empty tensors for concatenation with the level outputs.

1009 + tree_input_ids = torch.empty(

1010 + 0, device=self.input_ids.device, dtype=self.input_ids.dtype

1011 + )

1012 + tree_positions = torch.empty(

1013 + 0, device=self.positions.device, dtype=self.positions.dtype

1014 + )

1015 + tree_hidden_states = torch.empty(

1016 + 0, device=self.hidden_states.device, dtype=self.hidden_states.dtype

1017 + )

1018 + # Precompute the draft token positions.

1019 + flattened_draft_positions = (

1020 + positions.view(batch_size, -1) + self.tree_draft_pos_offsets[:batch_size, :]

1021 + )

1022 + tree_depth = len(self.cu_drafts_per_level)

1023 + for level in range(tree_depth - 1):

1024 + # Get draft positions for RoPE.

1025 + draft_positions = positions + (level + 1)

1026 + exceeds_max_model_len = (positions + total_num_drafts) >= self.max_model_len

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

1028 + # Otherwise, we may get out-of-range error in RoPE.

1029 + draft_positions = torch.where(

1030 + exceeds_max_model_len,

1031 + 0,

1032 + draft_positions,

1033 + ).view(batch_size, -1)

1034 +

1035 + if level_num_drafts > 1:

1036 + # Repeat the positions for each draft at this level.

1037 + draft_positions = draft_positions.repeat_interleave(

1038 + level_num_drafts, dim=1

1039 + )

1040 +

1041 + if num_children > 1:

1042 + # Repeat draft hidden states for each child.

1043 + draft_hidden_states = draft_hidden_states.repeat_interleave(

1044 + num_children, dim=1

1045 + )

1046 +

1047 + # Concatenate the draft tokens, positions, and hidden states.

1048 + tree_input_ids = torch.cat([tree_input_ids, draft_token_ids], dim=1)

1049 + tree_positions = torch.cat([tree_positions, draft_positions], dim=1)

1050 + tree_hidden_states = torch.cat(

1051 + [tree_hidden_states, draft_hidden_states], dim=1

1052 + )

1053 +

1054 + # Build new attention metadata for the next level of drafts.

1055 + # This is necessary to support tree attention.

1056 + query_len = total_num_drafts

1057 + common_attn_metadata = replace(

1058 + common_attn_metadata,

1059 + query_start_loc=query_len * self.arange[: batch_size + 1],

1060 + seq_lens=common_attn_metadata.seq_lens + level_num_drafts,

1061 + num_actual_tokens=batch_size * query_len,

1062 + max_query_len=query_len,

1063 + )

1064 + attn_metadata = tree_attn_metadata_builder.build_for_drafting(

1065 + common_attn_metadata=common_attn_metadata, draft_index=level + 1

1066 + )

1067 +

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

1069 + per_layer_attn_metadata = {}

1070 + for attn_group in self.draft_attn_groups:

1071 + for layer_name in attn_group.layer_names:

1072 + per_layer_attn_metadata[layer_name] = attn_metadata

1073 +

1074 + # Consider max model length.

1075 + attn_metadata.max_seq_len = min(

1076 + attn_metadata.max_seq_len, self.max_model_len

1077 + )

1078 + # For the requests that exceed the max model length, we set the

1079 + # sequence length to 1 to minimize their overheads in attention.

1080 + attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)

1081 +

1082 + # Compute the slot mapping.

1083 + block_size = tree_attn_metadata_builder.kv_cache_spec.block_size

1084 + query_positions = flattened_draft_positions[:, level : level + query_len]

1085 + block_numbers = query_positions // block_size

1086 + block_ids = attn_metadata.block_table.gather(dim=1, index=block_numbers)

1087 + slot_mapping = block_ids * block_size + query_positions % block_size

1088 + # Mask out the slot mappings that exceed the max model length.

1089 + # Otherwise, the KV cache will be inadvertently updated with the

1090 + # padding tokens.

1091 + slot_mapping[exceeds_max_model_len] = PADDING_SLOT_ID

1092 + attn_metadata.slot_mapping = slot_mapping.view(-1)

1093 +

1094 + # Copy inputs to buffer for cudagraph.

1095 + num_tokens = attn_metadata.num_actual_tokens

1096 + input_ids = tree_input_ids.view(-1)

1097 + self.input_ids[:num_tokens] = input_ids

1098 + self.positions[:num_tokens] = tree_positions.view(-1)

1099 + self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1)

1100 +

1101 + cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

1102 + num_tokens

1103 + )

1104 + num_input_tokens = batch_desc.num_tokens

1105 + # Run the model.

1106 + with set_forward_context(

1107 + per_layer_attn_metadata,

1108 + self.vllm_config,

1109 + num_tokens=num_input_tokens,

1110 + cudagraph_runtime_mode=cudagraph_runtime_mode,

1111 + slot_mapping=self._get_slot_mapping(

1112 + num_input_tokens, attn_metadata.slot_mapping

1113 + ),

1114 + ):

1115 + last_hidden_states, hidden_states = self.model(

1116 + input_ids=self.input_ids[:num_input_tokens],

1117 + positions=self.positions[:num_input_tokens],

1118 + hidden_states=self.hidden_states[:num_input_tokens],

1119 + inputs_embeds=None,

1120 + )

1121 +

1122 + # Get the output hidden states for the draft tokens.

1123 + draft_hidden_states = hidden_states[:num_tokens].view(

1124 + batch_size, query_len, -1

1125 + )[:, -level_num_drafts:]

1126 + draft_last_hidden_states = last_hidden_states[:num_tokens].view(

1127 + batch_size, query_len, -1

1128 + )[:, -level_num_drafts:]

1129 +

1130 + # Get the output logits for the draft tokens.

1131 + logits = self.model.compute_logits(

1132 + draft_last_hidden_states.reshape(batch_size * level_num_drafts, -1)

1133 + )

1134 +

1135 + # Sample a draft token for each child at the next tree level.

1136 + num_children = self.child_drafts_per_level[level + 1]

1137 + if num_children == 1:

1138 + draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)

1139 + else:

1140 + draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(

1141 + batch_size, -1

1142 + )

1143 + draft_token_ids_list.append(draft_token_ids)

1144 +

1145 + # Update the # drafts counters for the next tree level.

1146 + level_num_drafts = self.cu_drafts_per_level[level + 1] - total_num_drafts

1147 + total_num_drafts = self.cu_drafts_per_level[level + 1]

1148 + return draft_token_ids_list

1149 +

1150 + def prepare_inputs(

1151 + self,

1152 + common_attn_metadata: CommonAttentionMetadata,

1153 + sampled_token_ids: list[list[int]],

1154 + num_draft_tokens: list[int],

1155 + ) -> tuple[CommonAttentionMetadata, torch.Tensor]:

1156 + """

1157 + This function is used to prepare the inputs for speculative decoding.

1158 + It updates to the common_attn_metadata to account for the rejected

1159 + tokens (and newly sampled tokens). It also returns the token indices

1160 + of the tokens that should be fed to the speculator.

1161 + """

1162 + # E.g.

1163 + # common_attn_metadata.query_start_loc{_cpu}:

1164 + # [0, q1, q1 + q2, q1 + q2 + q3]

1165 + # common_attn_metadata.seq_lens{_cpu}: [s1, s2, s3]

1166 + # num_rejected_tokens: [n1, n2, n3]

1167 + # This function computes the intermediate values:

1168 + # num_tokens_per_req: [q1 - n1, q2 - n2, q3 - n3]

1169 + # And returns:

1170 + # common_attn_metadata.query_start_loc{_cpu}:

1171 + # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]

1172 + # common_attn_metadata.seq_lens{_cpu}:

1173 + # [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1]

1174 + # token_indices: [0, 1, ..., q1 - n1 - 1,

1175 + # q1, q1 + 1, ..., q1 + q2 - n2 - 1,

1176 + # q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]

1177 +

1178 + num_rejected_tokens = [

1179 + n + 1 - len(sampled_token_ids[i]) if n > 0 else 0

1180 + for i, n in enumerate(num_draft_tokens)

1181 + ]

1182 + num_rejected_tokens = torch.tensor(num_rejected_tokens, dtype=torch.int32)

1183 +

1184 + device = common_attn_metadata.query_start_loc.device

1185 + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

1186 + new_seq_lens_cpu = common_attn_metadata.seq_lens_cpu - num_rejected_tokens

1187 +

1188 + # [0, q1, q1 + q2, q1 + q2 + q3] -> [q1, q2, q3]

1189 + new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

1190 + # [q1, q2, q3] -> [q1 - n1, q2 - n2, q3 - n3]

1191 + new_num_tokens_per_req = new_query_len_per_req - num_rejected_tokens

1192 + new_num_tokens_per_req_np = new_num_tokens_per_req.numpy()

1193 +

1194 + # [q1 - n1, q2 - n2, q3 - n3] ->

1195 + # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]

1196 + new_query_start_loc_cpu = torch.zeros(

1197 + query_start_loc_cpu.shape,

1198 + dtype=torch.int32,

1199 + pin_memory=is_pin_memory_available(),

1200 + )

1201 + new_query_start_loc_np = new_query_start_loc_cpu.numpy()

1202 + np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:])

1203 +

1204 + total_num_tokens = new_query_start_loc_np[-1]

1205 + # Example assuming num_tokens_per_req_np = [2, 4, 3]

1206 + # this implies that `new_query_start_locs` is:

1207 + # [0, 2, 6, 9] ->

1208 + # [0, 0, 2, 2, 2, 2, 6, 6, 6]

1209 + # _r1_ ____r2____ ___r3__

1210 + new_query_start_locs_expanded = np.repeat(

1211 + new_query_start_loc_np[:-1], new_num_tokens_per_req_np

1212 + )

1213 + # [0, 1, 2, 3, 4, 5, 6, 7, 8] ->

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

1215 + # _r1_ ____r2____ ___r3__

1216 + token_offsets = (

1217 + self.token_arange_np[:total_num_tokens] - new_query_start_locs_expanded

1218 + )

1219 +

1220 + # Expand starting positions to match token pattern

1221 + # [0, q1, q1 + q2] ->

1222 + # [0, 0, q1, q1, q1, q1, q1 + q2, q1 + q2, q1 + q2]

1223 + # _r1_ _____r2_______ ___________r3____________

1224 + old_query_start_locs_expanded = np.repeat(

1225 + query_start_loc_cpu[:-1].numpy(), new_num_tokens_per_req_np

1226 + )

1227 + # Final token indices are:

1228 + # [0, 1, // req 1

1229 + # q1 + 0, q1 + 1, q1 + 2, q1 + 3, // req 2

1230 + # q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3

1231 + token_indices_np = token_offsets + old_query_start_locs_expanded

1232 + token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True)

1233 +

1234 + spec_common_attn_metadata = CommonAttentionMetadata(

1235 + query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True),

1236 + seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),

1237 + query_start_loc_cpu=new_query_start_loc_cpu,

1238 + _seq_lens_cpu=new_seq_lens_cpu,

1239 + _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

1240 + num_reqs=common_attn_metadata.num_reqs,

1241 + num_actual_tokens=total_num_tokens,

1242 + max_query_len=new_query_len_per_req.max().item(),

1243 + max_seq_len=new_seq_lens_cpu.max().item(),

1244 + block_table_tensor=common_attn_metadata.block_table_tensor,

1245 + slot_mapping=common_attn_metadata.slot_mapping[token_indices],

1246 + causal=True,

1247 + dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,

1248 + )

1249 +

1250 + return spec_common_attn_metadata, token_indices

1251 +

1252 + def get_model_name(self, model: nn.Module) -> str:

1253 + if hasattr(model, "module"): # multi-GPU

1254 + model = model.module

1255 + return model.__class__.__name__

1256 +

1257 + def _create_draft_vllm_config(self) -> VllmConfig:

1258 + """Return a VllmConfig with kernel-level overrides for the proposer.

1259 + Subclasses may override to apply additional config changes.

1260 + """

1261 + spec_cfg = self.speculative_config

1262 + if spec_cfg.moe_backend is not None:

1263 + return replace(

1264 + self.vllm_config,

1265 + kernel_config=replace(

1266 + self.vllm_config.kernel_config,

1267 + moe_backend=spec_cfg.moe_backend,

1268 + ),

1269 + )

1270 + return self.vllm_config

1271 +

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

1273 + """

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

1275 + need to customize model loading.

1276 + """

1277 + from vllm.compilation.backends import set_model_tag

1278 +

1279 + draft_vllm_config = self._create_draft_vllm_config()

1280 + with set_model_tag("eagle_head"):

1281 + model = get_model(

1282 + vllm_config=draft_vllm_config,

1283 + model_config=self.speculative_config.draft_model_config,

1284 + load_config=self.speculative_config.draft_load_config,

1285 + )

1286 + return model

1287 +

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

1289 + target_attn_layer_names = set(

1290 + get_layers_from_vllm_config(

1291 + self.vllm_config,

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

1293 + ).keys()

1294 + )

1295 +

1296 + self.model = self._get_model()

1297 +

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

1299 + all_attn_layers = get_layers_from_vllm_config(

1300 + self.vllm_config,

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

1302 + )

1303 + self._draft_attn_layer_names = (

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

1305 + )

1306 +

1307 + if self.supports_mm_inputs:

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

1309 + # text-only draft models

1310 + try:

1311 + dummy_input_ids = torch.tensor([[1]], device=self.input_ids.device)

1312 + self.model.embed_input_ids(dummy_input_ids, multimodal_embeddings=None)

1313 + except (NotImplementedError, AttributeError, TypeError):

1314 + logger.warning(

1315 + "Draft model does not support multimodal inputs, "

1316 + "falling back to text-only mode"

1317 + )

1318 + self.supports_mm_inputs = False

1319 +

1320 + if supports_multimodal(target_model):

1321 + # handle multimodality

1322 + assert hasattr(target_model, "config")

1323 + if self.get_model_name(target_model) in [

1324 + "Exaone4_5_ForConditionalGeneration",

1325 + "GlmOcrForConditionalGeneration",

1326 + "HunYuanVLForConditionalGeneration",

1327 + "Qwen2_5_VLForConditionalGeneration",

1328 + "Qwen3_5ForConditionalGeneration",

1329 + "Qwen3_5MoeForConditionalGeneration",

1330 + "Qwen3VLForConditionalGeneration",

1331 + "Qwen3VLMoeForConditionalGeneration",

1332 + "Gemma4ForConditionalGeneration",

1333 + ]:

1334 + self.model.config.image_token_index = target_model.config.image_token_id

1335 + elif self.get_model_name(target_model) == "PixtralForConditionalGeneration":

1336 + self.model.config.image_token_index = (

1337 + target_model.config.vision_config.image_token_id

1338 + )

1339 + elif self.get_model_name(target_model) == "KimiK25ForConditionalGeneration":

1340 + self.model.config.image_token_index = (

1341 + target_model.config.media_placeholder_token_id

1342 + )

1343 + else:

1344 + self.model.config.image_token_index = (

1345 + target_model.config.image_token_index

1346 + )

1347 + target_language_model = cast(

1348 + SupportsMultiModal, target_model

1349 + ).get_language_model()

1350 + else:

1351 + target_language_model = target_model

1352 +

1353 + self._maybe_share_embeddings(target_language_model)

1354 + self._maybe_share_lm_head(target_language_model)

1355 +

1356 + if (

1357 + self.parallel_drafting

1358 + and self.pass_hidden_states_to_model

1359 + and self.parallel_drafting_hidden_state_tensor is not None

1360 + ):

1361 + flat_mask = self.model.mask_hidden.view(-1)

1362 + if self.eagle3_use_aux_hidden_state:

1363 + # EAGLE3: mask_hidden stores all aux hidden states,

1364 + # project through combine_hidden_states

1365 + self.parallel_drafting_hidden_state_tensor.copy_(

1366 + self.model.combine_hidden_states(flat_mask)

1367 + )

1368 + else:

1369 + self.parallel_drafting_hidden_state_tensor.copy_(flat_mask)

1370 +

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

1372 + """

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

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

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

1376 + memory.

1377 + """

1378 + if get_pp_group().world_size == 1:

1379 + inner_model = getattr(target_language_model, "model", None)

1380 + if inner_model is None:

1381 + raise AttributeError("Target model does not have 'model' attribute")

1382 + if hasattr(inner_model, "embed_tokens"):

1383 + target_embed_tokens = inner_model.embed_tokens

1384 + elif hasattr(inner_model, "embedding"):

1385 + target_embed_tokens = inner_model.embedding

1386 + else:

1387 + raise AttributeError(

1388 + "Target model does not have 'embed_tokens' or 'embedding' attribute"

1389 + )

1390 +

1391 + share_embeddings = False

1392 + if hasattr(self.model, "has_own_embed_tokens"):

1393 + # EAGLE model

1394 + if not self.model.has_own_embed_tokens:

1395 + share_embeddings = True

1396 + logger.info(

1397 + "Detected EAGLE model without its own embed_tokens in the"

1398 + " checkpoint. Sharing target model embedding weights with the"

1399 + " draft model."

1400 + )

1401 + elif (

1402 + isinstance(target_embed_tokens.weight, torch.Tensor)

1403 + and isinstance(self.model.model.embed_tokens.weight, torch.Tensor)

1404 + # TODO: Offload to CPU for comparison to avoid extra GPU memory

1405 + # usage in CI testing environments with limited GPU memory

1406 + and torch.equal(

1407 + target_embed_tokens.weight.cpu(),

1408 + self.model.model.embed_tokens.weight.cpu(),

1409 + )

1410 + ):

1411 + share_embeddings = True

1412 + logger.info(

1413 + "Detected EAGLE model with embed_tokens identical to the target"

1414 + " model. Sharing target model embedding weights with the draft"

1415 + " model."

1416 + )

1417 + else:

1418 + logger.info(

1419 + "Detected EAGLE model with distinct embed_tokens weights. "

1420 + "Keeping separate embedding weights from the target model."

1421 + )

1422 + else:

1423 + # MTP model

1424 + share_embeddings = True

1425 + logger.info(

1426 + "Detected MTP model. "

1427 + "Sharing target model embedding weights with the draft model."

1428 + )

1429 +

1430 + if share_embeddings:

1431 + if hasattr(self.model.model, "embed_tokens"):

1432 + del self.model.model.embed_tokens

1433 + self.model.model.embed_tokens = target_embed_tokens

1434 + else:

1435 + logger.info(

1436 + "The draft model's vocab embedding will be loaded separately"

1437 + " from the target model."

1438 + )

1439 +

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

1441 + """

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

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

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

1445 + """

1446 + share_lm_head = False

1447 + if hasattr(self.model, "has_own_lm_head"):

1448 + # EAGLE model

1449 + if not self.model.has_own_lm_head:

1450 + share_lm_head = True

1451 + logger.info(

1452 + "Detected EAGLE model without its own lm_head in the checkpoint. "

1453 + "Sharing target model lm_head weights with the draft model."

1454 + )

1455 + elif (

1456 + hasattr(target_language_model, "lm_head")

1457 + and hasattr(target_language_model.lm_head, "weight")

1458 + and hasattr(self.model.lm_head, "weight")

1459 + and isinstance(target_language_model.lm_head.weight, torch.Tensor)

1460 + and isinstance(self.model.lm_head.weight, torch.Tensor)

1461 + # TODO: Offload to CPU for comparison to avoid extra GPU memory

1462 + # usage in CI testing environments with limited GPU memory

1463 + and torch.equal(

1464 + target_language_model.lm_head.weight.cpu(),

1465 + self.model.lm_head.weight.cpu(),

1466 + )

1467 + ):

1468 + share_lm_head = True

1469 + logger.info(

1470 + "Detected EAGLE model with lm_head identical to the target model. "

1471 + "Sharing target model lm_head weights with the draft model."

1472 + )

1473 + else:

1474 + logger.info(

1475 + "Detected EAGLE model with distinct lm_head weights. "

1476 + "Keeping separate lm_head weights from the target model."

1477 + )

1478 + else:

1479 + # MTP model

1480 + share_lm_head = True

1481 + logger.info(

1482 + "Detected MTP model. "

1483 + "Sharing target model lm_head weights with the draft model."

1484 + )

1485 +

1486 + if share_lm_head and hasattr(target_language_model, "lm_head"):

1487 + if hasattr(self.model, "lm_head"):

1488 + del self.model.lm_head

1489 + self.model.lm_head = target_language_model.lm_head

1490 +

1491 + # MTP models call compute_logits via shared_head.head (a

1492 + # ParallelLMHead inside each MTP layer), not self.model.lm_head.

1493 + # If the checkpoint omits a copy of the lm_head weights at the

1494 + # MTP layer path, shared_head.head stays uninitialised and

1495 + # produces NaN logits. Always share it explicitly.

1496 + inner = getattr(self.model, "model", None)

1497 + layers = getattr(inner, "layers", None) if inner else None

1498 + if layers is not None:

1499 + items = layers.values() if isinstance(layers, nn.ModuleDict) else layers

1500 + for layer in items:

1501 + sh = getattr(layer, "shared_head", None)

1502 + if sh is not None and hasattr(sh, "head"):

1503 + del sh.head

1504 + sh.head = target_language_model.lm_head

1505 + logger.info(

1506 + "Shared target model lm_head with MTP shared_head.head."

1507 + )

1508 +

1509 + if self.use_local_argmax_reduction:

1510 + if not hasattr(self.model, "get_top_tokens"):

1511 + raise ValueError(

1512 + "use_local_argmax_reduction is enabled but draft model "

1513 + f"{self.model.__class__.__name__} does not implement "

1514 + "get_top_tokens()."

1515 + )

1516 + # Warn if draft model has vocab remapping, which forces fallback

1517 + # to the full-logits path (negating the optimization).

1518 + if (

1519 + hasattr(self.model, "draft_id_to_target_id")

1520 + and self.model.draft_id_to_target_id is not None

1521 + ):

1522 + logger.warning(

1523 + "use_local_argmax_reduction is enabled but draft model "

1524 + "uses draft_id_to_target_id vocab remapping. The "

1525 + "optimization will be bypassed (falling back to full "

1526 + "logits gather + argmax)."

1527 + )

1528 + else:

1529 + logger.info(

1530 + "Using local argmax reduction for draft token generation "

1531 + "(communication: O(2*tp_size) vs O(vocab_size))."

1532 + )

1533 +

1534 + @torch.inference_mode()

1535 + def dummy_run(

1536 + self,

1537 + num_tokens: int,

1538 + use_cudagraphs: bool = True,

1539 + is_graph_capturing: bool = False,

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

1541 + ) -> None:

1542 + # FIXME: when using tree-based specdec, adjust number of forward-passes

1543 + # according to the depth of the tree.

1544 + only_one_forward_pass = is_graph_capturing or self.parallel_drafting

1545 + for fwd_idx in range(

1546 + 1 if only_one_forward_pass else self.num_speculative_tokens

1547 + ):

1548 + if fwd_idx <= 1:

1549 + cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

1550 + self._determine_batch_execution_and_padding(

1551 + num_tokens, use_cudagraphs=use_cudagraphs

1552 + )

1553 + )

1554 +

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

1556 + if (

1557 + self._draft_attn_layer_names

1558 + and slot_mappings is not None

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

1560 + ):

1561 + slot_mapping_dict = self._get_slot_mapping(num_input_tokens)

1562 + else:

1563 + slot_mapping_dict = slot_mappings or {}

1564 +

1565 + with set_forward_context(

1566 + None,

1567 + self.vllm_config,

1568 + num_tokens=num_input_tokens,

1569 + num_tokens_across_dp=num_tokens_across_dp,

1570 + cudagraph_runtime_mode=cudagraph_runtime_mode,

1571 + slot_mapping=slot_mapping_dict,

1572 + ):

1573 + if self.supports_mm_inputs:

1574 + input_ids = None

1575 + inputs_embeds = self.inputs_embeds[:num_input_tokens]

1576 + else:

1577 + input_ids = self.input_ids[:num_input_tokens]

1578 + inputs_embeds = None

1579 +

1580 + kwargs = dict(

1581 + input_ids=input_ids,

1582 + positions=self._get_positions(num_input_tokens),

1583 + inputs_embeds=inputs_embeds,

1584 + )

1585 + if self.pass_hidden_states_to_model:

1586 + kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]

1587 + self.model(**kwargs)

1588 +

1589 + def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:

1590 + """

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

1592 + hidden states and directly uses the last layer output just like eagle1.

1593 + They might indicate this by setting "use_aux_hidden_state" to False

1594 + inside the "eagle_config" dict of their hf_config.

1595 + """

1596 + if self.method != "eagle3":

1597 + return False

1598 + # Assume that eagle3 heads use aux hidden states by default

1599 + use_aux_hidden_state = True

1600 + eagle_config = getattr(self.draft_model_config.hf_config, "eagle_config", None)

1601 + if eagle_config is not None:

1602 + use_aux_hidden_state = eagle_config.get("use_aux_hidden_state", True)

1603 + return use_aux_hidden_state

1604 +

1605 + def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:

1606 + """

1607 + Validate that all drafting layers belong to the same KVCacheGroup.

1608 + Need this assumption to ensure all drafting layers can use the

1609 + same AttentionMetadata.

1610 + May extend to multiple AttentionMetadata in the future.

1611 + """

1612 + kv_cache_groups: dict[str, int] = {}

1613 + for id, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):

1614 + for layer_name in kv_cache_group.layer_names:

1615 + kv_cache_groups[layer_name] = id

1616 + assert (

1617 + len(

1618 + set(

1619 + [

1620 + kv_cache_groups[layer_name]

1621 + for layer_name in self._draft_attn_layer_names

1622 + ]

1623 + )

1624 + )

1625 + == 1

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

1627 +

1628 + def initialize_attn_backend(

1629 + self,

1630 + kv_cache_config: KVCacheConfig,

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

1632 + ) -> None:

1633 + """

1634 + Initialize AttentionGroups for draft layers using kv_cache_config.

1635 + Called from the model runner's initialize_metadata_builders.

1636 + """

1637 + all_attn_layers = get_layers_from_vllm_config(

1638 + self.vllm_config,

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

1640 + )

1641 +

1642 + # Find which kv_cache_group the draft layers belong to

1643 + self.validate_same_kv_cache_group(kv_cache_config)

1644 + kv_cache_spec = None

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

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

1647 + self.kv_cache_gid = gid

1648 + kv_cache_spec = group.kv_cache_spec

1649 + break

1650 +

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

1652 + if kv_cache_spec is not None:

1653 + for layer_name in self._draft_attn_layer_names:

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

1655 + backend_key = attn_backend.full_cls_name()

1656 + if backend_key not in attention_groups:

1657 + layer_kv_cache_spec = kv_cache_spec

1658 + if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):

1659 + layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[

1660 + layer_name

1661 + ]

1662 +

1663 + kernel_block_size = (

1664 + kernel_block_sizes[self.kv_cache_gid]

1665 + if kernel_block_sizes is not None

1666 + and self.kv_cache_gid < len(kernel_block_sizes)

1667 + else None

1668 + )

1669 + attn_group = AttentionGroup(

1670 + backend=attn_backend,

1671 + layer_names=[layer_name],

1672 + kv_cache_spec=layer_kv_cache_spec,

1673 + kv_cache_group_id=self.kv_cache_gid,

1674 + )

1675 + attn_group.create_metadata_builders(

1676 + self.vllm_config,

1677 + self.device,

1678 + kernel_block_size=kernel_block_size,

1679 + )

1680 + attention_groups[backend_key] = attn_group

1681 + else:

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

1683 +

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

1685 + self.block_size = (

1686 + self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size

1687 + )

1688 + logger.debug("Using block size %d for drafting layers", self.block_size)

1689 +

1690 + def _determine_batch_execution_and_padding(

1691 + self,

1692 + num_tokens: int,

1693 + use_cudagraphs: bool = True,

1694 + ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:

1695 + cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

1696 + num_tokens,

1697 + valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None),

1698 + )

1699 + num_tokens_padded = batch_desc.num_tokens

1700 +

1701 + # Extra coordination when running data-parallel since we need to

1702 + # coordinate across ranks

1703 + # TODO(Flechman): support DBO ubatching

1704 + should_ubatch, num_tokens_across_dp = False, None

1705 + if self.vllm_config.parallel_config.data_parallel_size > 1:

1706 + should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = (

1707 + coordinate_batch_across_dp(

1708 + num_tokens_unpadded=num_tokens,

1709 + parallel_config=self.vllm_config.parallel_config,

1710 + allow_microbatching=False,

1711 + num_tokens_padded=num_tokens_padded,

1712 + cudagraph_mode=cudagraph_mode.value,

1713 + )

1714 + )

1715 + assert not should_ubatch, "DBO ubatching not implemented for EAGLE"

1716 +

1717 + # Extract DP-synced values

1718 + if num_tokens_across_dp is not None:

1719 + dp_rank = self.dp_rank

1720 + num_tokens_padded = int(num_tokens_across_dp[dp_rank].item())

1721 + # Re-dispatch with DP padding so we have the correct

1722 + # batch_descriptor

1723 + cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

1724 + num_tokens_padded,

1725 + valid_modes={CUDAGraphMode(synced_cudagraph_mode)},

1726 + )

1727 + # Assert to make sure the agreed upon token count is correct

1728 + # otherwise num_tokens_across_dp will no-longer be valid

1729 + assert batch_desc.num_tokens == num_tokens_padded

1730 + num_tokens_across_dp[dp_rank] = num_tokens_padded

1731 +

1732 + return cudagraph_mode, num_tokens_padded, num_tokens_across_dp

1733 +

1734 +

1735 +# NOTE(woosuk): Currently, the below code is not used and we always use argmax

1736 +# to sample the draft tokens. We will use this after we find a way to manage

1737 +# the draft prob tensor.

1738 +# Refer to https://github.com/vllm-project/vllm/pull/16899 for the details.

1739 +# FIXME(woosuk): The logic here is duplicated with the main sampling code.

1740 +# We should refactor this to reuse the same sampling implementation.

1741 +def compute_probs_and_sample_next_token(

1742 + logits: torch.Tensor,

1743 + sampling_metadata: SamplingMetadata,

1744 +) -> tuple[torch.Tensor, torch.Tensor]:

1745 + if sampling_metadata.all_greedy:

1746 + # For greedy requests, draft_probs is not used in rejection sampling.

1747 + # Therefore, we can just return the logits.

1748 + probs = logits

1749 + next_token_ids = logits.argmax(dim=-1)

1750 + return next_token_ids, probs

1751 +

1752 + assert sampling_metadata.temperature is not None

1753 +

1754 + # Use epsilon comparison to detect greedy sampling (temperature ~ 0.0)

1755 + # consistent with sampler.py's _SAMPLING_EPS threshold

1756 + temperature = sampling_metadata.temperature

1757 + # Avoid division by zero if there are greedy requests.

1758 + if not sampling_metadata.all_random:

1759 + is_greedy = temperature < _SAMPLING_EPS

1760 + temperature = torch.where(is_greedy, 1.0, temperature)

1761 + logits.div_(temperature.view(-1, 1))

1762 + probs = logits.softmax(dim=-1, dtype=torch.float32)

1763 +

1764 + # NOTE(woosuk): Currently, we ignore most of the sampling parameters in

1765 + # generating the draft tokens. We only use the temperature. While this

1766 + # could degrade the acceptance rate, it does not affect the distribution

1767 + # of the generated tokens after rejection sampling.

1768 +

1769 + # TODO(woosuk): Consider seeds.

1770 + q = torch.empty_like(probs)

1771 + q.exponential_()

1772 + # NOTE(woosuk): We shouldn't use `probs.div_(q)` because the draft_probs

1773 + # will be used later for rejection sampling.

1774 + next_token_ids = probs.div(q).argmax(dim=-1).view(-1)

1775 + if not sampling_metadata.all_random:

1776 + greedy_token_ids = probs.argmax(dim=-1)

1777 + next_token_ids = torch.where(is_greedy, greedy_token_ids, next_token_ids)

1778 + return next_token_ids, probs

vllm/v1/worker/cpu_model_runner.py

+4 -4

@@ -65,16 +65,16 @@ class CPUModelRunner(GPUModelRunner):

66 # Speculative decoding fallbacks

67 import vllm.v1.sample.rejection_sampler

68 - import vllm.v1.spec_decode.eagle

68 + import vllm.v1.spec_decode.llm_base_proposer

69 import vllm.v1.spec_decode.utils

71 - vllm.v1.spec_decode.eagle.eagle_prepare_inputs_padded_kernel = (

71 + vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_inputs_padded_kernel = (

72 cpu_tl.eagle_prepare_inputs_padded_kernel

73 )

74 - vllm.v1.spec_decode.eagle.eagle_prepare_next_token_padded_kernel = (

74 + vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_next_token_padded_kernel = (

75 cpu_tl.eagle_prepare_next_token_padded_kernel

76 )

77 - vllm.v1.spec_decode.eagle.copy_and_expand_eagle_inputs_kernel = (

77 + vllm.v1.spec_decode.llm_base_proposer.copy_and_expand_eagle_inputs_kernel = (

78 cpu_tl.copy_and_expand_eagle_inputs_kernel

79 )

80 vllm.v1.spec_decode.utils.eagle_step_slot_mapping_metadata_kernel = (