eagle.py 70.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import ast
from dataclasses import replace
5
from importlib.util import find_spec
6

7
import numpy as np
8
9
10
import torch
import torch.nn as nn

王敏's avatar
王敏 committed
11
import vllm.envs as envs
12
13
14
15
16
from vllm.config import (
    CUDAGraphMode,
    VllmConfig,
    get_layers_from_vllm_config,
)
王敏's avatar
王敏 committed
17
from vllm.distributed.parallel_state import get_pp_group, get_tensor_model_parallel_rank
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_v2 import DeepseekV32IndexerCache
24
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
25
from vllm.multimodal import MULTIMODAL_REGISTRY
26
from vllm.platforms import current_platform
27
from vllm.triton_utils import triton
28
from vllm.utils.platform_utils import is_pin_memory_available
29
30
31
from vllm.v1.attention.backend import (
    AttentionMetadataBuilder,
    CommonAttentionMetadata,
王敏's avatar
王敏 committed
32
    CpCommonAttentionMetadata,
33
)
34
from vllm.v1.attention.backends.registry import AttentionBackendEnum
35
36
37
38
from vllm.v1.attention.backends.tree_attn import (
    TreeAttentionMetadata,
    TreeAttentionMetadataBuilder,
)
39
from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
40
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
41
from vllm.v1.kv_cache_interface import KVCacheConfig
42
from vllm.v1.sample.metadata import SamplingMetadata
43
from vllm.v1.sample.sampler import _SAMPLING_EPS
44
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
45
46
47
48
from vllm.v1.spec_decode.utils import (
    eagle_prepare_inputs_padded_kernel,
    eagle_prepare_next_token_padded_kernel,
)
49
from vllm.v1.utils import CpuGpuBuffer
Rémi Delacourt's avatar
Rémi Delacourt committed
50
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
51
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
王敏's avatar
王敏 committed
52
from vllm.utils.math_utils import cdiv, round_up
53

54
55
logger = init_logger(__name__)

56
57
PADDING_SLOT_ID = -1

58

59
class SpecDecodeBaseProposer:
60
61
62
63
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
64
        pass_hidden_states_to_model: bool,
Jiayi Yao's avatar
Jiayi Yao committed
65
        runner=None,
66
67
    ):
        self.vllm_config = vllm_config
68
        self.speculative_config = vllm_config.speculative_config
69
        assert self.speculative_config is not None
70
71
        self.draft_model_config = self.speculative_config.draft_model_config
        self.method = self.speculative_config.method
72
        self.pass_hidden_states_to_model = pass_hidden_states_to_model
73

Jiayi Yao's avatar
Jiayi Yao committed
74
        self.runner = runner
75
        self.device = device
76
        self.dtype = vllm_config.model_config.dtype
77
        self.max_model_len = vllm_config.model_config.max_model_len
Rémi Delacourt's avatar
Rémi Delacourt committed
78
        self.dp_rank = vllm_config.parallel_config.data_parallel_rank
79
        self.num_speculative_tokens = self.speculative_config.num_speculative_tokens
80
        # The drafter can get longer sequences than the target model.
81
82
        max_batch_size = vllm_config.scheduler_config.max_num_seqs if not \
            vllm_config.parallel_config.enable_lightly_cplb \
王敏's avatar
王敏 committed
83
            else vllm_config.scheduler_config.max_num_seqs * 2
84
85
86
        self.max_num_tokens = (
            vllm_config.scheduler_config.max_num_batched_tokens + max_batch_size
        )
87
        self.token_arange_np = np.arange(self.max_num_tokens)
88
89
90
91
        # We need to get the hidden size from the draft model config because
        # the draft model's hidden size can be different from the target model's
        # hidden size (e.g., Llama 3.3 70B).
        self.hidden_size = self.draft_model_config.get_hidden_size()
92
        self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()
93

94
95
96
        # Multi-modal data support
        self.mm_registry = MULTIMODAL_REGISTRY
        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
97
98
            vllm_config.model_config
        )
99

100
101
        self.attn_metadata_builder: AttentionMetadataBuilder | None = None
        self.draft_indexer_metadata_builder: AttentionMetadataBuilder | None = None
102
103
        self.attn_layer_names: list[str] = []
        self.indexer_layer_names: list[str] = []
104
105
106
        self.eagle3_use_aux_hidden_state: bool = (
            self._get_eagle3_use_aux_hidden_state_from_config()
        )
107

108
        self.compilation_config = self.vllm_config.compilation_config
109
110
111
112
113
114

        # Cudagraph dispatcher for PIECEWISE-only dispatching in eagle.
        # Keys are initialized later via initialize_cudagraph_keys() called from
        # gpu_model_runner._check_and_update_cudagraph_mode after
        # adjust_cudagraph_sizes_for_spec_decode is called.
        self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)
115
116

        # persistent buffers for cuda graph
117
118
119
        self.input_ids = torch.zeros(
            self.max_num_tokens, dtype=torch.int32, device=device
        )
120
121
122
        # Use draft model's M-RoPE setting, not target model's
        # Draft models may be text-only even if target is multimodal
        self.uses_mrope = self.draft_model_config.uses_mrope
123
        if self.uses_mrope:
124
125
126
127
128
129
130
131
132
133
            # NOTE: `mrope_positions` is implemented with one additional dummy
            # position on purpose to make it non-contiguous so that it can work
            # with torch compile.
            # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923

            # NOTE: When M-RoPE is enabled, position ids are 3D regardless of
            # the modality of inputs. For text-only inputs, each dimension has
            # identical position IDs, making M-RoPE functionally equivalent to
            # 1D-RoPE.
            # See page 5 of https://arxiv.org/abs/2409.12191
134
            self.mrope_positions = torch.zeros(
135
                (3, self.max_num_tokens + 1), dtype=torch.int64, device=device
136
            )
137
138
        else:
            # RoPE need (max_num_tokens,)
139
140
141
            self.positions = torch.zeros(
                self.max_num_tokens, dtype=torch.int64, device=device
            )
142
        self.hidden_states = torch.zeros(
143
144
            (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device
        )
145

146
147
148
        # We need +1 here because the arange is used to set query_start_loc,
        # which has one more element than batch_size.
        max_num_slots_for_arange = max(max_batch_size + 1, self.max_num_tokens)
149
150
151
        self.arange = torch.arange(
            max_num_slots_for_arange, device=device, dtype=torch.int32
        )
152

153
        self.inputs_embeds = torch.zeros(
154
            (self.max_num_tokens, self.inputs_embeds_size),
155
            dtype=self.dtype,
156
            device=device,
157
        )
158

159
160
161
162
163
        self.backup_next_token_ids = CpuGpuBuffer(
            max_batch_size,
            dtype=torch.int32,
            pin_memory=is_pin_memory_available(),
            device=device,
164
165
            with_numpy=True,
        )
166

167
168
169
170
        self._slot_mapping_buffer = torch.zeros(
            self.max_num_tokens, dtype=torch.int64, device=device
        )

171
        # Determine allowed attention backends once during initialization.
172
        self.allowed_attn_types: tuple | None = None
zhuwenwen's avatar
zhuwenwen committed
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
        # if current_platform.is_rocm():
        #     from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata

        #     rocm_types = [
        #         TritonAttentionMetadata,
        #         RocmAttentionMetadata,
        #     ]
        #     # ROCM_AITER_FA is an optional backend
        #     if find_spec(
        #         AttentionBackendEnum.ROCM_AITER_FA.get_path(include_classname=False)
        #     ):
        #         from vllm.v1.attention.backends.rocm_aiter_fa import (
        #             AiterFlashAttentionMetadata,
        #         )

        #         rocm_types.append(AiterFlashAttentionMetadata)

        #     # TRITON_MLA backend support for MLA models (e.g., DeepSeek)
        #     from vllm.model_executor.layers.attention.mla_attention import (
        #         MLACommonMetadata,
        #     )
194

zhuwenwen's avatar
zhuwenwen committed
195
        #     rocm_types.append(MLACommonMetadata)
196

zhuwenwen's avatar
zhuwenwen committed
197
198
        #     # FlexAttention backend support
        #     from vllm.v1.attention.backends.flex_attention import FlexAttentionMetadata
199

zhuwenwen's avatar
zhuwenwen committed
200
        #     rocm_types.append(FlexAttentionMetadata)
201

zhuwenwen's avatar
zhuwenwen committed
202
        #     self.allowed_attn_types = tuple(rocm_types)
203

204
205
        # Parse the speculative token tree.
        spec_token_tree = self.speculative_config.speculative_token_tree
206
        self.tree_choices: list[tuple[int, ...]] = ast.literal_eval(spec_token_tree)
207
208
209
210
211
212
213
214
        tree_depth = len(self.tree_choices[-1])
        # Precompute per-level properties of the tree.
        num_drafts_per_level = [0] * tree_depth
        for node in self.tree_choices:
            num_drafts_per_level[len(node) - 1] += 1
        self.cu_drafts_per_level = [num_drafts_per_level[0]]
        self.child_drafts_per_level = [num_drafts_per_level[0]]
        for level in range(1, tree_depth):
215
216
217
218
219
220
            self.cu_drafts_per_level.append(
                self.cu_drafts_per_level[-1] + num_drafts_per_level[level]
            )
            self.child_drafts_per_level.append(
                num_drafts_per_level[level] // num_drafts_per_level[level - 1]
            )
221
222
        # Precompute draft position offsets in flattened tree.
        self.tree_draft_pos_offsets = torch.arange(
223
            1, len(self.tree_choices) + 1, device=device, dtype=torch.int32
224
        ).repeat(max_batch_size, 1)
225

226
227
        self.scatter_indexes_tensor = None
        self.gather_indexes_tensor = None
228
229
230
231

        self.enable_lightly_cp = vllm_config.parallel_config.enable_lightly_cp
        self.enable_lightly_cplb = self.enable_lightly_cp and vllm_config.parallel_config.enable_lightly_cplb
        if self.enable_lightly_cp:
王敏's avatar
王敏 committed
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
            self.query_start_loc = CpuGpuBuffer(
                max_batch_size + 1,
                dtype=torch.int32,
                pin_memory=is_pin_memory_available(),
                device=device,
                with_numpy=True,
            )

            self.seq_lens = CpuGpuBuffer(
                max_batch_size,
                dtype=torch.int32,
                pin_memory=is_pin_memory_available(),
                device=device,
                with_numpy=True,
            )

248
249
250
251
252
253
254
255
256
    def _get_positions(self, num_tokens: int):
        if self.uses_mrope:
            return self.mrope_positions[:, :num_tokens]
        return self.positions[:num_tokens]

    def _set_positions(self, num_tokens: int, positions: torch.Tensor):
        if self.uses_mrope:
            self.mrope_positions[:, :num_tokens] = positions
        else:
257
258
259
260
261
            # Convert M-RoPE positions if target model uses M-RoPE
            # but draft doesn't, For text inputs, all M-RoPE
            # dimensions are identical
            if self.vllm_config.model_config.uses_mrope:
                positions = positions[0]
262
263
            self.positions[:num_tokens] = positions

264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
    def _get_slot_mapping(
        self,
        num_tokens: int,
        slot_mapping: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """Return slot_mapping dict for EAGLE layers.

        If slot_mapping is provided, copies it into the buffer first.
        """
        if slot_mapping is not None:
            num_actual = slot_mapping.shape[0]
            self._slot_mapping_buffer[:num_actual].copy_(slot_mapping)
            if num_tokens > num_actual:
                self._slot_mapping_buffer[num_actual:num_tokens].fill_(PADDING_SLOT_ID)

        view = self._slot_mapping_buffer[:num_tokens]
        return {name: view for name in self.attn_layer_names + self.indexer_layer_names}

282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
    def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:
        """Initialize cudagraph dispatcher keys for eagle.

        Eagle only supports PIECEWISE cudagraphs (via mixed_mode).
        This should be called after adjust_cudagraph_sizes_for_spec_decode.
        """
        if (
            not self.speculative_config.enforce_eager
            and cudagraph_mode.mixed_mode()
            in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]
        ):
            eagle_cudagraph_mode = CUDAGraphMode.PIECEWISE
        else:
            eagle_cudagraph_mode = CUDAGraphMode.NONE

        self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode)

王敏's avatar
王敏 committed
299
300
301
302
    def _pad_for_mla_cp(self, num_scheduled_tokens: int) -> int:
        tp_size = self.vllm_config.parallel_config.tensor_parallel_size
        return round_up(num_scheduled_tokens, tp_size)

303
304
305
306
    def propose(
        self,
        # [num_tokens]
        target_token_ids: torch.Tensor,
307
        # [num_tokens] or [3, num_tokens] when M-RoPE is enabled
308
309
310
311
312
        target_positions: torch.Tensor,
        # [num_tokens, hidden_size]
        target_hidden_states: torch.Tensor,
        # [batch_size]
        next_token_ids: torch.Tensor,
313
        last_token_indices: torch.Tensor | None,
314
        common_attn_metadata: CommonAttentionMetadata,
315
        sampling_metadata: SamplingMetadata,
316
        mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
317
        num_rejected_tokens_gpu: torch.Tensor | None = None,
318
319
320
        slot_mappings: dict[str, torch.Tensor]
        | list[dict[str, torch.Tensor]]
        | None = None,
321
    ) -> torch.Tensor:
322
        batch_size = common_attn_metadata.batch_size()
323

324
325
326
        if self.method == "eagle3":
            assert isinstance(self.model, Eagle3LlamaForCausalLM)
            target_hidden_states = self.model.combine_hidden_states(
327
328
                target_hidden_states
            )
329
            assert target_hidden_states.shape[-1] == self.hidden_size
330
331
332
333
334
335
336
337
338
339
340

        num_tokens, last_token_indices, common_attn_metadata = (
            self.set_inputs_first_pass(
                target_token_ids=target_token_ids,
                next_token_ids=next_token_ids,
                target_positions=target_positions,
                last_token_indices=last_token_indices,
                cad=common_attn_metadata,
                num_rejected_tokens_gpu=num_rejected_tokens_gpu,
            )
        )
341

王敏's avatar
王敏 committed
342
343
344
345
        num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(
            num_tokens_unpadded=num_tokens, num_tokens_padded=num_tokens
        )

346
347
        enable_lightly_cp = self.enable_lightly_cp and num_tokens > self.runner.lightly_cp_threshould
        if enable_lightly_cp:
王敏's avatar
王敏 committed
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
            num_tokens_dp_padded = self._pad_for_mla_cp(num_tokens_dp_padded)

            common_attn_metadata = self._prepare_cp_metadata(
                num_reqs_padded=common_attn_metadata.num_reqs,
                max_query_len=common_attn_metadata.max_query_len,
                max_seq_len=common_attn_metadata.seq_lens_cpu.max().item(),
                num_tokens=num_tokens,
                block_table_gid_0=common_attn_metadata.block_table_tensor,
                slot_mapping_gid_0=common_attn_metadata.slot_mapping,
                query_start_loc=common_attn_metadata.query_start_loc,
                query_start_loc_cpu=common_attn_metadata.query_start_loc_cpu,
                seq_lens=common_attn_metadata.seq_lens,
                seq_lens_cpu=common_attn_metadata.seq_lens_cpu,
                num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
            )
            self.scatter_indexes_tensor = common_attn_metadata.scatter_indexes_tensor
            self.gather_indexes_tensor = common_attn_metadata.gather_indexes_tensor


367
        assert self.runner is not None
Jiayi Yao's avatar
Jiayi Yao committed
368

369
370
371
        if self.attn_metadata_builder is None:
            attn_metadata_builder = self._get_attention_metadata_builder()
        else:
王敏's avatar
王敏 committed
372
            attn_metadata_builder = self.attn_metadata_builder            
373

374
        attn_metadata = attn_metadata_builder.build_for_drafting(
375
376
            common_attn_metadata=common_attn_metadata, draft_index=0
        )
377
378
379
380
381
382
        # FIXME: support hybrid kv for draft model (remove separate indexer)
        if self.draft_indexer_metadata_builder:
            draft_indexer_metadata = (
                self.draft_indexer_metadata_builder.build_for_drafting(
                    common_attn_metadata=common_attn_metadata,
                    draft_index=0,
383
384
                )
            )
385
386
        else:
            draft_indexer_metadata = None
387
388
389
390
391
        # At this moment, we assume all eagle layers belong to the same KV
        # cache group, thus using the same attention metadata.
        per_layer_attn_metadata = {}
        for layer_name in self.attn_layer_names:
            per_layer_attn_metadata[layer_name] = attn_metadata
392

393
394
395
396
        for layer_name in self.indexer_layer_names:
            assert draft_indexer_metadata is not None
            per_layer_attn_metadata[layer_name] = draft_indexer_metadata

397
398
399
400
        cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
            num_tokens_dp_padded
        )
        num_input_tokens = batch_desc.num_tokens
Rémi Delacourt's avatar
Rémi Delacourt committed
401
402
403
        if num_tokens_across_dp is not None:
            num_tokens_across_dp[self.dp_rank] = num_input_tokens

404
405
406
407
        if self.pass_hidden_states_to_model:
            # target_hidden_states and self.hidden_states can have different
            # hidden dims. E.g. large target model and small draft model.
            self.hidden_states[:num_tokens] = target_hidden_states
408
409
410
411

        if self.supports_mm_inputs:
            mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

412
            self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(
413
414
415
                self.input_ids[:num_tokens],
                multimodal_embeddings=mm_embeds,
                is_multimodal=is_mm_embed,
416
            )
417

418
            input_ids = None
419
            inputs_embeds = self.inputs_embeds[:num_input_tokens]
420
421
        else:
            input_ids = self.input_ids[:num_input_tokens]
422
            inputs_embeds = None
423

424
425
426
427
428
429
430
431
        model_kwargs = {
            "input_ids": input_ids,
            "positions": self._get_positions(num_input_tokens),
            "inputs_embeds": inputs_embeds,
        }
        if self.pass_hidden_states_to_model:
            model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]

432
        with set_forward_context(
433
434
435
            per_layer_attn_metadata,
            self.vllm_config,
            num_tokens=num_input_tokens,
Rémi Delacourt's avatar
Rémi Delacourt committed
436
            num_tokens_across_dp=num_tokens_across_dp,
437
            cudagraph_runtime_mode=cudagraph_runtime_mode,
438
439
440
            slot_mapping=self._get_slot_mapping(
                num_input_tokens, common_attn_metadata.slot_mapping
            ),
王敏's avatar
王敏 committed
441
442
            scatter_indexes_tensor=self.scatter_indexes_tensor,
            gather_indexes_tensor=self.gather_indexes_tensor,
443
444
            enable_lightly_cp=self.enable_lightly_cp and num_tokens > self.runner.lightly_cp_threshould,
            enable_lightly_cplb=self.enable_lightly_cplb
445
        ):
446
447
            ret_hidden_states = self.model(**model_kwargs)
            if not self.model_returns_tuple():
Jiayi Yao's avatar
Jiayi Yao committed
448
                last_hidden_states = ret_hidden_states
449
                hidden_states = last_hidden_states
Jiayi Yao's avatar
Jiayi Yao committed
450
451
            else:
                last_hidden_states, hidden_states = ret_hidden_states
452

453
        sample_hidden_states = last_hidden_states[last_token_indices]
454
        logits = self.model.compute_logits(sample_hidden_states)
455

王敏's avatar
王敏 committed
456
457
458
        if envs.VLLM_REJECT_SAMPLE_OPT:
            draft_prob = logits.softmax(dim=-1, dtype=torch.float32)

459
460
461
        # Early exit if there is only one draft token to be generated.
        if self.num_speculative_tokens == 1:
            draft_token_ids = logits.argmax(dim=-1)
王敏's avatar
王敏 committed
462
463
464
465

            if envs.VLLM_REJECT_SAMPLE_OPT:
                return draft_token_ids.view(-1, 1), draft_prob.view(-1, 1, logits.shape[-1])

466
            return draft_token_ids.view(-1, 1)
467

468
        if self.uses_mrope:
469
            positions = self.mrope_positions[:, last_token_indices]
470
        else:
471
            positions = self.positions[last_token_indices]
472
473
474
475
476
477
        if self.method in (
            "deepseek_mtp",
            "ernie_mtp",
            "longcat_flash_mtp",
            "pangu_ultra_moe_mtp",
        ):
zhuwenwen's avatar
zhuwenwen committed
478
479
480
            hidden_states = self.hidden_states[last_token_indices]
        else:
            hidden_states = hidden_states[last_token_indices]
481
482
483

        if isinstance(attn_metadata, TreeAttentionMetadata):
            # Draft using tree attention.
484
485
486
487
488
489
            draft_token_ids_list = self.propose_tree(
                batch_size=batch_size,
                logits=logits,
                positions=positions,
                hidden_states=hidden_states,
                common_attn_metadata=common_attn_metadata,
490
                slot_mappings=slot_mappings,
491
492
493
494
            )
            # [batch_size, num_tree_tokens]
            return torch.cat(draft_token_ids_list, dim=1)

495
        draft_token_ids = logits.argmax(dim=-1)
496

497
498
499
        if self.allowed_attn_types is not None and not isinstance(
            attn_metadata, self.allowed_attn_types
        ):
500
501
502
503
            raise ValueError(
                f"Unsupported attention metadata type for speculative "
                "decoding with num_speculative_tokens > 1: "
                f"{type(attn_metadata)}. Supported types are: "
504
505
                f"{self.allowed_attn_types}"
            )
506

507
508
        # Generate the remaining draft tokens.
        draft_token_ids_list = [draft_token_ids]
509

Rémi Delacourt's avatar
Rémi Delacourt committed
510
        batch_size_dp_padded, batch_size_across_dp = self._pad_batch_across_dp(
511
            num_tokens_unpadded=batch_size, num_tokens_padded=batch_size
Rémi Delacourt's avatar
Rémi Delacourt committed
512
513
        )

514
515
516
517
        cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
            batch_size_dp_padded
        )
        input_batch_size = batch_desc.num_tokens
Rémi Delacourt's avatar
Rémi Delacourt committed
518
519
        if batch_size_across_dp is not None:
            batch_size_across_dp[self.dp_rank] = input_batch_size
520

521
        if enable_lightly_cp:
王敏's avatar
王敏 committed
522
523
            common_attn_metadata = common_attn_metadata.cp_common_metadata

524
525
        common_attn_metadata.num_actual_tokens = batch_size
        common_attn_metadata.max_query_len = 1
526
        common_attn_metadata.query_start_loc = self.arange[: batch_size + 1]
527
        common_attn_metadata.query_start_loc_cpu = torch.from_numpy(
528
529
            self.token_arange_np[: batch_size + 1]
        ).clone()
530
531
532
533
534
535
536
537
538
539
540

        # In padded drafter batch, we need to adjust the sequence lengths
        # to remove the "padding" (i.e. rejected tokens).
        # Only apply this adjustment when we have rejected tokens
        # (i.e., not the first proposal).
        if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None:
            common_attn_metadata.seq_lens -= num_rejected_tokens_gpu
            # Invalidate the CPU-side shadows to avoid H<>D sync.
            common_attn_metadata._seq_lens_cpu = None
            common_attn_metadata._num_computed_tokens_cpu = None

王敏's avatar
王敏 committed
541
542
543
        if envs.VLLM_REJECT_SAMPLE_OPT:
            draft_probs_list = [draft_prob]

544
        for token_index in range(self.num_speculative_tokens - 1):
545
            # Update the inputs.
546
547
548
            # cast to int32 is crucial when eagle model is compiled.
            # tensor.argmax() returns int64 by default.
            input_ids = draft_token_ids_list[-1].int()
549
550
551
552
553
554
555
556
557
558
559
560
561
            if self.uses_mrope:
                positions += 1
                # NOTE(woosuk): We should handle the case where the draft model
                # generates tokens beyond the max model length.
                # Since it is complex to remove such requests from the batch,
                # we keep them in the batch but adjust the position ids
                # and slot mappings to avoid the
                # out-of-range access during the model execution.
                # The draft tokens generated with this adjustment
                # should be ignored.
                exceeds_max_model_len = positions[0] >= self.max_model_len
                # Mask out the position ids that exceed the max model length.
                # Otherwise, we may get out-of-range error in RoPE.
562
563
564
565
566
                clamped_positions = torch.where(
                    exceeds_max_model_len.unsqueeze(0),
                    torch.zeros_like(positions),
                    positions,
                )
567
            else:
568
569
                positions += 1
                exceeds_max_model_len = positions >= self.max_model_len
570
                clamped_positions = torch.where(exceeds_max_model_len, 0, positions)
571
572
573
            # For data integrity when async scheduling, we shouldn't use in place
            # operations in case they are modified in next step's `prepare_input`
            # of main model.
574
            # Increment the sequence lengths.
575
            common_attn_metadata.seq_lens += 1
576
577
            # For the requests that exceed the max model length, we set the
            # sequence length to 1 to minimize their overheads in attention.
578
            common_attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)
579
580
581
582
583
584
            # Increment the maximum sequence length. We increment max_seq_len
            # unconditionally even though some seq_lens may have been capped above,
            # as max_seq_len serves as an upper bound for sequence lengths.
            common_attn_metadata.max_seq_len = min(
                common_attn_metadata.max_seq_len + 1, self.max_model_len
            )
585

586
587
588
589
590
591
            # Also update the CPU-side shadow; NOTE: this is hacky and should be
            # removed in when common_attn_metadata.seq_lens_cpu is deprecated.
            if common_attn_metadata._seq_lens_cpu is not None:
                common_attn_metadata._seq_lens_cpu += 1
            if common_attn_metadata._num_computed_tokens_cpu is not None:
                common_attn_metadata._num_computed_tokens_cpu += 1
592

593
            # Compute the slot mapping.
594
            block_size = attn_metadata_builder.kv_cache_spec.block_size
595
596
            if self.uses_mrope:
                # all dimensions of positions are the same
597
                block_numbers = clamped_positions[0] // block_size
598
            else:
599
                block_numbers = clamped_positions // block_size
600
            block_ids = common_attn_metadata.block_table_tensor.gather(
601
602
                dim=1, index=block_numbers.view(-1, 1)
            )
603
            block_ids = block_ids.view(-1)
604
605
            if self.uses_mrope:
                common_attn_metadata.slot_mapping = (
606
                    block_ids * block_size + clamped_positions[0] % block_size
607
                )
608
609
            else:
                common_attn_metadata.slot_mapping = (
610
                    block_ids * block_size + clamped_positions % block_size
611
                )
612
613
614
            # Mask out the slot mappings that exceed the max model length.
            # Otherwise, the KV cache will be inadvertently updated with the
            # padding tokens.
615
            common_attn_metadata.slot_mapping.masked_fill_(
616
617
                exceeds_max_model_len, PADDING_SLOT_ID
            )
618
619

            # Rebuild attention metadata
620
            attn_metadata = attn_metadata_builder.build_for_drafting(  # type: ignore
621
622
                common_attn_metadata=common_attn_metadata, draft_index=token_index + 1
            )
王敏's avatar
王敏 committed
623
624
625
626
627
628
629
630
631
632
633

            if self.draft_indexer_metadata_builder:
                draft_indexer_metadata = (
                    self.draft_indexer_metadata_builder.build_for_drafting(
                        common_attn_metadata=common_attn_metadata,
                        draft_index=token_index + 1,
                    )
                )
            else:
                draft_indexer_metadata = None

634
635
            for layer_name in self.attn_layer_names:
                per_layer_attn_metadata[layer_name] = attn_metadata
636

王敏's avatar
王敏 committed
637
638
639
            for layer_name in self.indexer_layer_names:
                per_layer_attn_metadata[layer_name] = draft_indexer_metadata

640
641
            # copy inputs to buffer for cudagraph
            self.input_ids[:batch_size] = input_ids
642
            self._set_positions(batch_size, clamped_positions)
643
            self.hidden_states[:batch_size] = hidden_states
644
            if self.supports_mm_inputs:
645
                self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids)
646

647
                input_ids = None
648
                inputs_embeds = self.inputs_embeds[:input_batch_size]
649
650
            else:
                input_ids = self.input_ids[:input_batch_size]
651
                inputs_embeds = None
652

653
            # Run the model.
654
655
656
657
658
659
660
661
            model_kwargs = {
                "input_ids": input_ids,
                "positions": self._get_positions(input_batch_size),
                "inputs_embeds": inputs_embeds,
            }
            if self.pass_hidden_states_to_model:
                model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size]

662
            with set_forward_context(
663
664
665
                per_layer_attn_metadata,
                self.vllm_config,
                num_tokens=input_batch_size,
Rémi Delacourt's avatar
Rémi Delacourt committed
666
                num_tokens_across_dp=batch_size_across_dp,
667
                cudagraph_runtime_mode=cudagraph_runtime_mode,
668
669
670
                slot_mapping=self._get_slot_mapping(
                    input_batch_size, common_attn_metadata.slot_mapping
                ),
671
            ):
672
673
                ret_hidden_states = self.model(**model_kwargs)
                if not self.model_returns_tuple():
674
                    last_hidden_states = ret_hidden_states
675
                    hidden_states = ret_hidden_states
676
677
                else:
                    last_hidden_states, hidden_states = ret_hidden_states
678

679
            hidden_states = hidden_states[:batch_size]
680
            logits = self.model.compute_logits(last_hidden_states[:batch_size])
681
            draft_token_ids = logits.argmax(dim=-1)
682
683
            draft_token_ids_list.append(draft_token_ids)

王敏's avatar
王敏 committed
684
685
686
687
            if envs.VLLM_REJECT_SAMPLE_OPT:
                draft_prob = logits.softmax(dim=-1, dtype=torch.float32)
                draft_probs_list.append(draft_prob)

688
689
        # [batch_size, num_speculative_tokens]
        draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
王敏's avatar
王敏 committed
690
691
692
693
694

        if envs.VLLM_REJECT_SAMPLE_OPT:
            draft_probs = torch.stack(draft_probs_list, dim=1).contiguous()
            return draft_token_ids, draft_probs

695
        return draft_token_ids
696

697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
    def set_inputs_first_pass(
        self,
        target_token_ids: torch.Tensor,
        next_token_ids: torch.Tensor,
        target_positions: torch.Tensor,
        last_token_indices: torch.Tensor | None,
        cad: CommonAttentionMetadata,
        num_rejected_tokens_gpu: torch.Tensor | None,
    ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:
        if last_token_indices is None:
            last_token_indices = cad.query_start_loc[1:] - 1

        num_tokens = target_token_ids.shape[0]
        # Shift the input ids by one token.
        # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
        self.input_ids[: num_tokens - 1] = target_token_ids[1:]
        # Replace the last token with the next token.
        # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]
        self.input_ids[last_token_indices] = next_token_ids

        # copy inputs to buffer for cudagraph
        self._set_positions(num_tokens, target_positions)

        return num_tokens, last_token_indices, cad

    def model_returns_tuple(self) -> bool:
        return self.method not in ("mtp", "draft_model")

725
    def prepare_next_token_ids_cpu(
726
        self,
727
        sampled_token_ids: list[list[int]],
728
729
730
731
        requests: dict[str, CachedRequestState],
        gpu_input_batch: InputBatch,
        num_scheduled_tokens: dict[str, int],
    ) -> torch.Tensor:
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
        """
        This function is used to prepare the inputs for speculative decoding.
        It calculates the next token ids for each request based on the sampled
        token ids from the CPU. If a request has no sampled token ids (e.g.,
        during the initial decoding steps), it falls back to using the request
        state to get the next token id.
        """
        req_ids = gpu_input_batch.req_ids
        next_token_ids: list[int] = []
        for i, token_ids in enumerate(sampled_token_ids):
            if token_ids:
                # Common case.
                next_token_id = token_ids[-1]
            else:
                # Partial prefill (rare case).
                # Get the next token id from the request state.
                req_id = req_ids[i]
                req_state = requests[req_id]
750
                seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]
751
752
                next_token_id = req_state.get_token_id(seq_len)
            next_token_ids.append(next_token_id)
753
        next_token_ids = torch.tensor(
754
755
            next_token_ids, dtype=torch.int32, device=self.input_ids.device
        )
756
757
        return next_token_ids

758
759
760
761
762
763
    def prepare_next_token_ids_padded(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        sampled_token_ids: torch.Tensor,
        requests: dict[str, CachedRequestState],
        gpu_input_batch: InputBatch,
764
        discard_request_mask: torch.Tensor,
765
    ) -> tuple[torch.Tensor, torch.Tensor]:
766
767
768
769
        """
        This function is used to prepare the inputs for speculative decoding.
        It calculates the next token ids and the number of valid sampled tokens
        for each request, considering the "discarded" requests whose next token
770
771
        is not sampled and comes from `request.get_token_id()` instead. This is denoted
        the "backup" token id. It also counts rejected tokens via `sampled_token_ids`.
772
773
774
        """
        # Precompute get_token_id for when there is no valid next token
        num_reqs = gpu_input_batch.num_reqs
775
776
777
778
779
780
        self.backup_next_token_ids.np[:num_reqs] = np.array(
            [
                requests[gpu_input_batch.req_ids[i]].get_token_id(
                    common_attn_metadata.seq_lens_cpu[i].item()
                )
                for i in range(num_reqs)
781
782
            ],
            dtype=np.int32,
783
        )
784
        self.backup_next_token_ids.copy_to_gpu(num_reqs)
785
        backup_tokens_gpu = self.backup_next_token_ids.gpu
786

787
788
        batch_size, num_tokens = sampled_token_ids.shape
        device = sampled_token_ids.device
789

790
791
        assert discard_request_mask.dtype == torch.bool
        assert backup_tokens_gpu.dtype == torch.int32
792

793
794
        next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=device)
        valid_sampled_tokens_count = next_token_ids.new_empty(batch_size)
795

796
797
        # Kernel grid: one program per request (row)
        grid = (batch_size,)
798

799
800
801
802
803
804
805
806
807
808
809
810
811
        # Find the next power of 2 for block sizes
        BLOCK_SIZE_TOKENS = triton.next_power_of_2(num_tokens)
        eagle_prepare_next_token_padded_kernel[grid](
            sampled_token_ids,
            discard_request_mask,
            backup_tokens_gpu,
            next_token_ids,
            valid_sampled_tokens_count,
            gpu_input_batch.vocab_size,
            num_tokens,
            batch_size,
            sampled_token_ids.stride(0),
            BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,
812
        )
813
814
815

        return next_token_ids, valid_sampled_tokens_count

816
817
818
819
820
    def prepare_inputs_padded(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        spec_decode_metadata: SpecDecodeMetadata,
        valid_sampled_tokens_count: torch.Tensor,
821
    ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:
822
823
824
825
826
827
        """
        This function is used to prepare the inputs for speculative decoding
        It updates the common_attn_metadata for speculative decoding,
        but does not consider the rejected tokens. Instead, all tokens
        are included as inputs to the speculator, with the rejected tokens
        used as padding and filtered out later by `token_indices_to_sample`.
828
        No blocking CPU operations should be introduced in this function.
829
        """
830
831
        num_reqs = common_attn_metadata.num_reqs
        device = valid_sampled_tokens_count.device
832

833
834
        token_indices_to_sample = torch.empty(
            (num_reqs,), dtype=torch.int32, device=device
835
        )
836
837
838
        num_rejected_tokens_gpu = torch.empty(
            (num_reqs,), dtype=torch.int32, device=device
        )
839

840
841
842
843
844
845
        grid = (num_reqs,)
        eagle_prepare_inputs_padded_kernel[grid](
            spec_decode_metadata.cu_num_draft_tokens,
            valid_sampled_tokens_count,
            common_attn_metadata.query_start_loc,
            token_indices_to_sample,
846
            num_rejected_tokens_gpu,
847
            num_reqs,
848
        )
849
850

        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
851
        new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
852
853
854
855
856
857
858

        total_num_tokens = query_start_loc_cpu[-1].item()

        spec_common_attn_metadata = CommonAttentionMetadata(
            query_start_loc=common_attn_metadata.query_start_loc,
            seq_lens=common_attn_metadata.seq_lens,
            query_start_loc_cpu=query_start_loc_cpu,
859
860
            _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,
            _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
861
862
            num_reqs=common_attn_metadata.num_reqs,
            num_actual_tokens=total_num_tokens,
王敏's avatar
王敏 committed
863
            num_kv_actual_tokens=total_num_tokens,
864
865
866
            max_query_len=new_query_len_per_req.max().item(),
            max_seq_len=common_attn_metadata.seq_lens_cpu.max().item(),
            block_table_tensor=common_attn_metadata.block_table_tensor,
867
            slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens],
868
            causal=True,
869
            dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
870
871
        )

872
873
874
875
876
        return (
            spec_common_attn_metadata,
            token_indices_to_sample,
            num_rejected_tokens_gpu,
        )
877

878
879
880
881
882
883
884
885
886
887
    def propose_tree(
        self,
        batch_size: int,
        # [num_tokens, vocab_size]
        logits: torch.Tensor,
        # [num_tokens]
        positions: torch.Tensor,
        # [num_tokens, hidden_size]
        hidden_states: torch.Tensor,
        common_attn_metadata: CommonAttentionMetadata,
888
889
890
        slot_mappings: dict[str, torch.Tensor]
        | list[dict[str, torch.Tensor]]
        | None = None,
891
    ) -> list[torch.Tensor]:
892
893
894
895
        tree_attn_metadata_builder = self.runner.attn_groups[0][
            0
        ].get_metadata_builder()
        assert isinstance(tree_attn_metadata_builder, TreeAttentionMetadataBuilder)
896

897
        total_num_drafts = self.cu_drafts_per_level[0]
898
899
        level_num_drafts = total_num_drafts
        # Sample a draft token for each child at the tree root level.
900
        num_children = self.child_drafts_per_level[0]
901
902
903
        if num_children == 1:
            draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)
        else:
904
905
906
            draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(
                batch_size, -1
            )
907
908
909
910
        draft_token_ids_list = [draft_token_ids]
        draft_hidden_states = hidden_states.view(batch_size, 1, -1)

        # Initialize empty tensors for concatenation with the level outputs.
911
912
913
914
915
916
917
918
919
        tree_input_ids = torch.empty(
            0, device=self.input_ids.device, dtype=self.input_ids.dtype
        )
        tree_positions = torch.empty(
            0, device=self.positions.device, dtype=self.positions.dtype
        )
        tree_hidden_states = torch.empty(
            0, device=self.hidden_states.device, dtype=self.hidden_states.dtype
        )
920
921
        # Precompute the draft token positions.
        flattened_draft_positions = (
922
923
            positions.view(batch_size, -1) + self.tree_draft_pos_offsets[:batch_size, :]
        )
924
        tree_depth = len(self.cu_drafts_per_level)
925
        for level in range(tree_depth - 1):
926
927
            # Get draft positions for RoPE.
            draft_positions = positions + (level + 1)
928
            exceeds_max_model_len = (positions + total_num_drafts) >= self.max_model_len
929
930
            # Mask out the position ids that exceed the max model length.
            # Otherwise, we may get out-of-range error in RoPE.
931
            draft_positions = torch.where(
932
933
934
                exceeds_max_model_len,
                0,
                draft_positions,
935
936
            ).view(batch_size, -1)

937
938
            if level_num_drafts > 1:
                # Repeat the positions for each draft at this level.
939
                draft_positions = draft_positions.repeat_interleave(
940
941
                    level_num_drafts, dim=1
                )
942
943
944
945

            if num_children > 1:
                # Repeat draft hidden states for each child.
                draft_hidden_states = draft_hidden_states.repeat_interleave(
946
947
                    num_children, dim=1
                )
948
949

            # Concatenate the draft tokens, positions, and hidden states.
950
951
            tree_input_ids = torch.cat([tree_input_ids, draft_token_ids], dim=1)
            tree_positions = torch.cat([tree_positions, draft_positions], dim=1)
952
            tree_hidden_states = torch.cat(
953
954
                [tree_hidden_states, draft_hidden_states], dim=1
            )
955
956
957

            # Build new attention metadata for the next level of drafts.
            # This is necessary to support tree attention.
958
            query_len = total_num_drafts
959
960
            common_attn_metadata = replace(
                common_attn_metadata,
961
                query_start_loc=query_len * self.arange[: batch_size + 1],
962
963
964
965
966
                seq_lens=common_attn_metadata.seq_lens + level_num_drafts,
                num_actual_tokens=batch_size * query_len,
                max_query_len=query_len,
            )
            attn_metadata = tree_attn_metadata_builder.build_for_drafting(
967
                common_attn_metadata=common_attn_metadata, draft_index=level + 1
968
969
970
971
972
973
974
975
            )

            # Apply new attention metadata to all layers.
            per_layer_attn_metadata = {}
            for layer_name in self.attn_layer_names:
                per_layer_attn_metadata[layer_name] = attn_metadata

            # Consider max model length.
976
977
978
            attn_metadata.max_seq_len = min(
                attn_metadata.max_seq_len, self.max_model_len
            )
979
980
981
982
983
            # For the requests that exceed the max model length, we set the
            # sequence length to 1 to minimize their overheads in attention.
            attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)

            # Compute the slot mapping.
984
            block_size = tree_attn_metadata_builder.kv_cache_spec.block_size
985
            query_positions = flattened_draft_positions[:, level : level + query_len]
986
            block_numbers = query_positions // block_size
987
            block_ids = attn_metadata.block_table.gather(dim=1, index=block_numbers)
988
            slot_mapping = block_ids * block_size + query_positions % block_size
989
990
991
992
993
994
995
996
997
998
999
            # Mask out the slot mappings that exceed the max model length.
            # Otherwise, the KV cache will be inadvertently updated with the
            # padding tokens.
            slot_mapping[exceeds_max_model_len] = PADDING_SLOT_ID
            attn_metadata.slot_mapping = slot_mapping.view(-1)

            # Copy inputs to buffer for cudagraph.
            num_tokens = attn_metadata.num_actual_tokens
            input_ids = tree_input_ids.view(-1)
            self.input_ids[:num_tokens] = input_ids
            self.positions[:num_tokens] = tree_positions.view(-1)
1000
            self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1)
1001

1002
1003
1004
1005
            cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
                num_tokens
            )
            num_input_tokens = batch_desc.num_tokens
1006
            # Run the model.
1007
            with set_forward_context(
1008
1009
1010
1011
                per_layer_attn_metadata,
                self.vllm_config,
                num_tokens=num_input_tokens,
                cudagraph_runtime_mode=cudagraph_runtime_mode,
1012
1013
1014
                slot_mapping=self._get_slot_mapping(
                    num_input_tokens, attn_metadata.slot_mapping
                ),
1015
            ):
1016
1017
1018
1019
1020
1021
1022
1023
1024
                last_hidden_states, hidden_states = self.model(
                    input_ids=self.input_ids[:num_input_tokens],
                    positions=self.positions[:num_input_tokens],
                    hidden_states=self.hidden_states[:num_input_tokens],
                    inputs_embeds=None,
                )

            # Get the output hidden states for the draft tokens.
            draft_hidden_states = hidden_states[:num_tokens].view(
1025
1026
                batch_size, query_len, -1
            )[:, -level_num_drafts:]
1027
            draft_last_hidden_states = last_hidden_states[:num_tokens].view(
1028
1029
                batch_size, query_len, -1
            )[:, -level_num_drafts:]
1030
1031
1032

            # Get the output logits for the draft tokens.
            logits = self.model.compute_logits(
1033
1034
                draft_last_hidden_states.reshape(batch_size * level_num_drafts, -1)
            )
1035
1036
1037
1038
1039
1040

            # Sample a draft token for each child at the next tree level.
            num_children = self.child_drafts_per_level[level + 1]
            if num_children == 1:
                draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)
            else:
1041
1042
1043
                draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(
                    batch_size, -1
                )
1044
1045
1046
            draft_token_ids_list.append(draft_token_ids)

            # Update the # drafts counters for the next tree level.
1047
            level_num_drafts = self.cu_drafts_per_level[level + 1] - total_num_drafts
1048
1049
            total_num_drafts = self.cu_drafts_per_level[level + 1]
        return draft_token_ids_list
王敏's avatar
王敏 committed
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
    
    def _prepare_cp_metadata(
        self,
        num_reqs_padded,
        max_query_len,
        max_seq_len,
        num_tokens,
        block_table_gid_0,
        slot_mapping_gid_0,
        query_start_loc,
        query_start_loc_cpu,
        seq_lens,
        seq_lens_cpu,
        num_computed_tokens_cpu,
    ):
        tp_size = self.vllm_config.parallel_config.tensor_parallel_size
        tp_rank = get_tensor_model_parallel_rank()

        cp_common_metadata = CpCommonAttentionMetadata(
            query_start_loc=query_start_loc.clone(),
            query_start_loc_cpu=query_start_loc_cpu.clone(),
            seq_lens=seq_lens.clone(),
            _seq_lens_cpu=seq_lens_cpu.clone(),
            max_query_len=max_query_len,
            max_seq_len=max_seq_len,
            num_reqs=num_reqs_padded,
            req_ids=self.runner.input_batch.req_ids,
            num_actual_tokens=num_tokens,
            num_kv_actual_tokens=num_tokens,
            block_table_tensor=block_table_gid_0,
            slot_mapping=slot_mapping_gid_0,
            _num_computed_tokens_cpu=num_computed_tokens_cpu
        )

        q_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
        kv_lens_cpu = seq_lens_cpu
        total_q_len = num_tokens
        total_kv_len = num_tokens

        (
            total_q_len,
            q_lens_cpu,
            seq_count,
            kv_lens_cpu,
            local_req_ids,
            scatter_indexes_tensor,
            gather_indexes_tensor,
            seq_indexes_list,
        ) = self.runner._distribute_tokens_to_cp_ranks(
            total_q_len,
            q_lens_cpu,
            kv_lens_cpu,
            tp_rank,
            tp_size,
            self.runner.input_batch.req_ids,
        )

        num_reqs = seq_count

        cu_num_tokens = np.cumsum(q_lens_cpu)
        self.query_start_loc.np[0] = 0
        self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens
        self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1])
        self.query_start_loc.copy_to_gpu()
        q_acc_lens = self.query_start_loc.gpu[: num_reqs + 1]
        q_acc_lens_cpu = self.query_start_loc.cpu[: num_reqs + 1]
        max_q_len = max(q_acc_lens_cpu)

        self.seq_lens.np[:num_reqs] = kv_lens_cpu
        self.seq_lens.np[num_reqs:].fill(0)
        self.seq_lens.copy_to_gpu()
        kv_lens = self.seq_lens.gpu[:num_reqs]
        kv_lens_cpu = self.seq_lens.cpu[:num_reqs]
        max_kv_len = max(kv_lens_cpu)

        num_computed_tokens_cpu = kv_lens_cpu - q_acc_lens_cpu[1:]
        blk_table_tensor = block_table_gid_0[seq_indexes_list]

        cm_base = CommonAttentionMetadata(
            query_start_loc=q_acc_lens,
            query_start_loc_cpu=q_acc_lens_cpu,
            seq_lens=kv_lens,
            _seq_lens_cpu=kv_lens_cpu,
            _num_computed_tokens_cpu=num_computed_tokens_cpu,
            num_reqs=num_reqs,
            num_actual_tokens=total_q_len,
            max_query_len=max_q_len,
            max_seq_len=max_kv_len,
            block_table_tensor=blk_table_tensor,
            slot_mapping=slot_mapping_gid_0,
            causal=True,
            num_kv_actual_tokens=total_kv_len,
            seq_indexes_list=seq_indexes_list,
            cp_common_metadata=cp_common_metadata,
            scatter_indexes_tensor=scatter_indexes_tensor,
            gather_indexes_tensor=gather_indexes_tensor,
        )
        return cm_base
1148

1149
    def prepare_inputs(
1150
1151
        self,
        common_attn_metadata: CommonAttentionMetadata,
1152
1153
        sampled_token_ids: list[list[int]],
        num_draft_tokens: list[int],
1154
1155
    ) -> tuple[CommonAttentionMetadata, torch.Tensor]:
        """
1156
        This function is used to prepare the inputs for speculative decoding.
1157
1158
1159
1160
1161
1162
        It updates to the common_attn_metadata to account for the rejected
        tokens (and newly sampled tokens). It also returns the token indices
        of the tokens that should be fed to the speculator.
        """
        # E.g.
        #  common_attn_metadata.query_start_loc{_cpu}:
1163
        #       [0, q1, q1 + q2, q1 + q2 + q3]
1164
1165
1166
1167
1168
1169
        #  common_attn_metadata.seq_lens{_cpu}: [s1, s2, s3]
        #  num_rejected_tokens: [n1, n2, n3]
        # This function computes the intermediate values:
        #  num_tokens_per_req: [q1 - n1, q2 - n2, q3 - n3]
        # And returns:
        #  common_attn_metadata.query_start_loc{_cpu}:
1170
        #       [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
1171
        #  common_attn_metadata.seq_lens{_cpu}:
1172
        #       [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1]
1173
        #  token_indices: [0, 1, ..., q1 - n1 - 1,
1174
1175
        #                 q1, q1 + 1, ..., q1 + q2 - n2 - 1,
        #                 q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]
1176

1177
1178
1179
1180
        num_rejected_tokens = [
            n + 1 - len(sampled_token_ids[i]) if n > 0 else 0
            for i, n in enumerate(num_draft_tokens)
        ]
1181
        num_rejected_tokens = torch.tensor(num_rejected_tokens, dtype=torch.int32)
1182

1183
1184
        device = common_attn_metadata.query_start_loc.device
        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
1185
        new_seq_lens_cpu = common_attn_metadata.seq_lens_cpu - num_rejected_tokens
1186
1187

        # [0, q1, q1 + q2, q1 + q2 + q3] -> [q1, q2, q3]
1188
        new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
1189
1190
1191
1192
1193
1194
1195
1196
        # [q1, q2, q3] -> [q1 - n1, q2 - n2, q3 - n3]
        new_num_tokens_per_req = new_query_len_per_req - num_rejected_tokens
        new_num_tokens_per_req_np = new_num_tokens_per_req.numpy()

        # [q1 - n1, q2 - n2, q3 - n3] ->
        # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
        new_query_start_loc_cpu = torch.zeros(
            query_start_loc_cpu.shape,
1197
            dtype=torch.int32,
1198
1199
            pin_memory=is_pin_memory_available(),
        )
1200
1201
1202
1203
1204
1205
1206
1207
1208
        new_query_start_loc_np = new_query_start_loc_cpu.numpy()
        np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:])

        total_num_tokens = new_query_start_loc_np[-1]
        # Example assuming num_tokens_per_req_np = [2, 4, 3]
        # this implies that `new_query_start_locs` is:
        # [0, 2, 6, 9] ->
        # [0, 0, 2, 2, 2, 2, 6, 6, 6]
        #  _r1_  ____r2____  ___r3__
1209
1210
1211
        new_query_start_locs_expanded = np.repeat(
            new_query_start_loc_np[:-1], new_num_tokens_per_req_np
        )
1212
1213
1214
        # [0, 1, 2, 3, 4, 5, 6, 7, 8] ->
        # [0, 1, 0, 1, 2, 3, 0, 1, 2]
        #  _r1_  ____r2____  ___r3__
1215
1216
1217
        token_offests = (
            self.token_arange_np[:total_num_tokens] - new_query_start_locs_expanded
        )
1218
1219
1220
1221
1222
1223

        # Expand starting positions to match token pattern
        # [0, q1, q1 + q2] ->
        # [0, 0, q1, q1, q1, q1, q1 + q2, q1 + q2, q1 + q2]
        #  _r1_  _____r2_______  ___________r3____________
        old_query_start_locs_expanded = np.repeat(
1224
1225
            query_start_loc_cpu[:-1].numpy(), new_num_tokens_per_req_np
        )
1226
        # Final token indices are:
1227
1228
1229
        # [0, 1,                                // req 1
        #  q1 + 0, q1 + 1, q1 + 2, q1 + 3,       // req 2
        #  q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3
1230
        token_indices_np = token_offests + old_query_start_locs_expanded
1231
        token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True)
1232
1233

        spec_common_attn_metadata = CommonAttentionMetadata(
1234
            query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True),
1235
1236
            seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),
            query_start_loc_cpu=new_query_start_loc_cpu,
1237
1238
            _seq_lens_cpu=new_seq_lens_cpu,
            _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
1239
1240
1241
            num_reqs=common_attn_metadata.num_reqs,
            num_actual_tokens=total_num_tokens,
            max_query_len=new_query_len_per_req.max().item(),
1242
            max_seq_len=new_seq_lens_cpu.max().item(),
1243
1244
            block_table_tensor=common_attn_metadata.block_table_tensor,
            slot_mapping=common_attn_metadata.slot_mapping[token_indices],
1245
            causal=True,
1246
            dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
1247
        )
1248
1249

        return spec_common_attn_metadata, token_indices
1250

1251
    def get_model_name(self, model: nn.Module) -> str:
1252
        if hasattr(model, "module"):  # multi-GPU
1253
1254
1255
            model = model.module
        return model.__class__.__name__

1256
    def load_model(self, target_model: nn.Module) -> None:
1257
        draft_model_config = self.vllm_config.speculative_config.draft_model_config
1258
        target_attn_layer_names = set(
1259
            get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()
1260
        )
1261
1262
        # FIXME: support hybrid kv for draft model
        target_indexer_layer_names = set(
1263
1264
1265
1266
            get_layers_from_vllm_config(
                self.vllm_config, DeepseekV32IndexerCache
            ).keys()
        )
1267

1268
        from vllm.compilation.backends import set_model_tag
1269

1270
        with set_model_tag("eagle_head"):
1271
1272
1273
            self.model = get_model(
                vllm_config=self.vllm_config, model_config=draft_model_config
            )
1274

1275
        draft_attn_layer_names = (
1276
            get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()
1277
1278
1279
1280
1281
1282
            - target_attn_layer_names
        )
        indexer_layers = get_layers_from_vllm_config(
            self.vllm_config, DeepseekV32IndexerCache
        )
        draft_indexer_layer_names = indexer_layers.keys() - target_indexer_layer_names
1283
        self.attn_layer_names = list(draft_attn_layer_names - draft_indexer_layer_names)
1284
1285
1286
1287
1288
        self.indexer_layer_names = list(draft_indexer_layer_names)

        if self.indexer_layer_names:
            first_layer = self.indexer_layer_names[0]
            self.draft_indexer_metadata_builder = (
1289
1290
1291
                indexer_layers[first_layer]
                .get_attn_backend()
                .get_builder_cls()(
1292
                    indexer_layers[first_layer].get_kv_cache_spec(self.vllm_config),
1293
1294
1295
                    self.indexer_layer_names,
                    self.vllm_config,
                    self.device,
1296
1297
                )
            )
1298
1299
        else:
            self.draft_indexer_metadata_builder = None
1300

1301
        if self.supports_mm_inputs:
1302
1303
1304
            # Even if the target model is multimodal, we can also use
            # text-only draft models
            try:
1305
                dummy_input_ids = torch.tensor([[1]], device=self.input_ids.device)
1306
                self.model.embed_input_ids(dummy_input_ids, multimodal_embeddings=None)
1307
1308
1309
            except (NotImplementedError, AttributeError, TypeError):
                logger.warning(
                    "Draft model does not support multimodal inputs, "
1310
1311
                    "falling back to text-only mode"
                )
1312
                self.supports_mm_inputs = False
1313

1314
1315
        if supports_multimodal(target_model):
            # handle multimodality
1316
1317
1318
            if self.get_model_name(target_model) in [
                "Qwen2_5_VLForConditionalGeneration",
                "Qwen3VLForConditionalGeneration",
1319
                "Qwen3VLMoeForConditionalGeneration",
1320
                "GlmOcrForConditionalGeneration",
Rayyyyy's avatar
Rayyyyy committed
1321
1322
                "Qwen3_5ForConditionalGeneration",
                "Qwen3_5MoeForConditionalGeneration",
1323
            ]:
1324
                self.model.config.image_token_index = target_model.config.image_token_id
1325
1326
1327
1328
            elif self.get_model_name(target_model) == "PixtralForConditionalGeneration":
                self.model.config.image_token_index = (
                    target_model.config.vision_config.image_token_id
                )
1329
1330
            else:
                self.model.config.image_token_index = (
1331
1332
                    target_model.config.image_token_index
                )
1333
1334
1335
            target_language_model = target_model.get_language_model()
        else:
            target_language_model = target_model
1336

1337
        # share embed_tokens with the target model if needed
1338
        if get_pp_group().world_size == 1:
1339
            if hasattr(target_language_model.model, "embed_tokens"):
1340
                target_embed_tokens = target_language_model.model.embed_tokens
1341
            elif hasattr(target_language_model.model, "embedding"):
1342
1343
1344
                target_embed_tokens = target_language_model.model.embedding
            else:
                raise AttributeError(
1345
1346
                    "Target model does not have 'embed_tokens' or 'embedding' attribute"
                )
1347

1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
            share_embeddings = False
            if hasattr(self.model, "has_own_embed_tokens"):
                # EAGLE model
                if not self.model.has_own_embed_tokens:
                    share_embeddings = True
                    logger.info(
                        "Detected EAGLE model without its own embed_tokens in the"
                        " checkpoint. Sharing target model embedding weights with the"
                        " draft model."
                    )
                elif (
                    isinstance(target_embed_tokens.weight, torch.Tensor)
                    and isinstance(self.model.model.embed_tokens.weight, torch.Tensor)
1361
1362
1363
                    # TODO: Offload to CPU for comparison to avoid extra GPU memory
                    # usage in CI testing environments with limited GPU memory
                    and torch.equal(
1364
1365
                        target_embed_tokens.weight.cpu(),
                        self.model.model.embed_tokens.weight.cpu(),
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
                    )
                ):
                    share_embeddings = True
                    logger.info(
                        "Detected EAGLE model with embed_tokens identical to the target"
                        " model. Sharing target model embedding weights with the draft"
                        " model."
                    )
                else:
                    logger.info(
                        "Detected EAGLE model with distinct embed_tokens weights. "
                        "Keeping separate embedding weights from the target model."
                    )
1379
            else:
1380
1381
                # MTP model
                share_embeddings = True
1382
                logger.info(
1383
1384
                    "Detected MTP model. "
                    "Sharing target model embedding weights with the draft model."
1385
                )
1386
1387
1388
1389
1390

            if share_embeddings:
                if hasattr(self.model.model, "embed_tokens"):
                    del self.model.model.embed_tokens
                self.model.model.embed_tokens = target_embed_tokens
1391
        else:
1392
            logger.info(
1393
                "The draft model's vocab embedding will be loaded separately"
1394
1395
                " from the target model."
            )
1396
1397

        # share lm_head with the target model if needed
1398
1399
1400
1401
1402
        share_lm_head = False
        if hasattr(self.model, "has_own_lm_head"):
            # EAGLE model
            if not self.model.has_own_lm_head:
                share_lm_head = True
1403
                logger.info(
1404
1405
1406
1407
1408
1409
1410
                    "Detected EAGLE model without its own lm_head in the checkpoint. "
                    "Sharing target model lm_head weights with the draft model."
                )
            elif (
                hasattr(target_language_model, "lm_head")
                and isinstance(target_language_model.lm_head.weight, torch.Tensor)
                and isinstance(self.model.lm_head.weight, torch.Tensor)
1411
1412
                # TODO: Offload to CPU for comparison to avoid extra GPU memory
                # usage in CI testing environments with limited GPU memory
1413
                and torch.equal(
1414
1415
                    target_language_model.lm_head.weight.cpu(),
                    self.model.lm_head.weight.cpu(),
1416
                )
1417
            ):
1418
                share_lm_head = True
1419
                logger.info(
1420
1421
                    "Detected EAGLE model with lm_head identical to the target model. "
                    "Sharing target model lm_head weights with the draft model."
1422
                )
1423
1424
            else:
                logger.info(
1425
1426
                    "Detected EAGLE model with distinct lm_head weights. "
                    "Keeping separate lm_head weights from the target model."
1427
                )
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
        else:
            # MTP model
            share_lm_head = True
            logger.info(
                "Detected MTP model. "
                "Sharing target model lm_head weights with the draft model."
            )

        if share_lm_head and hasattr(target_language_model, "lm_head"):
            if hasattr(self.model, "lm_head"):
                del self.model.lm_head
            self.model.lm_head = target_language_model.lm_head
1440

王敏's avatar
王敏 committed
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
            # MTP models call compute_logits via shared_head.head (a
            # ParallelLMHead inside each MTP layer), not self.model.lm_head.
            # If the checkpoint omits a copy of the lm_head weights at the
            # MTP layer path, shared_head.head stays uninitialised and
            # produces NaN logits. Always share it explicitly.
            inner = getattr(self.model, "model", None)
            layers = getattr(inner, "layers", None) if inner else None
            if layers is not None:
                items = layers.values() if isinstance(layers, nn.ModuleDict) else layers
                for layer in items:
                    sh = getattr(layer, "shared_head", None)
                    if sh is not None and hasattr(sh, "head"):
                        del sh.head
                        sh.head = target_language_model.lm_head
                        logger.info(
                            "Shared target model lm_head with MTP shared_head.head."
                        )

1459
1460
1461
1462
    @torch.inference_mode()
    def dummy_run(
        self,
        num_tokens: int,
1463
1464
        use_cudagraphs: bool = True,
        is_graph_capturing: bool = False,
1465
        slot_mappings: dict[str, torch.Tensor] | None = None,
1466
    ) -> None:
Rémi Delacourt's avatar
Rémi Delacourt committed
1467
1468
1469
1470
        # FIXME: when using tree-based specdec, adjust number of forward-passes
        # according to the depth of the tree.
        for fwd_idx in range(
            self.num_speculative_tokens if not is_graph_capturing else 1
1471
        ):
Rémi Delacourt's avatar
Rémi Delacourt committed
1472
1473
            if fwd_idx <= 1:
                num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(
1474
                    num_tokens_unpadded=num_tokens, num_tokens_padded=num_tokens
Rémi Delacourt's avatar
Rémi Delacourt committed
1475
                )
1476
1477
1478
                if use_cudagraphs:
                    cudagraph_runtime_mode, batch_desc = (
                        self.cudagraph_dispatcher.dispatch(num_tokens_dp_padded)
Rémi Delacourt's avatar
Rémi Delacourt committed
1479
                    )
1480
                    num_input_tokens = batch_desc.num_tokens
Rémi Delacourt's avatar
Rémi Delacourt committed
1481
                else:
1482
                    cudagraph_runtime_mode = CUDAGraphMode.NONE
Rémi Delacourt's avatar
Rémi Delacourt committed
1483
1484
1485
                    num_input_tokens = num_tokens_dp_padded
                if num_tokens_across_dp is not None:
                    num_tokens_across_dp[self.dp_rank] = num_input_tokens
1486

1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
            # Make sure to use EAGLE's own buffer during cudagraph capture.
            if (
                self.attn_layer_names
                and slot_mappings is not None
                and self.attn_layer_names[0] in slot_mappings
            ):
                slot_mapping_dict = self._get_slot_mapping(num_input_tokens)
            else:
                slot_mapping_dict = slot_mappings or {}

Rémi Delacourt's avatar
Rémi Delacourt committed
1497
1498
1499
1500
1501
            with set_forward_context(
                None,
                self.vllm_config,
                num_tokens=num_input_tokens,
                num_tokens_across_dp=num_tokens_across_dp,
1502
                cudagraph_runtime_mode=cudagraph_runtime_mode,
1503
                slot_mapping=slot_mapping_dict,
Rémi Delacourt's avatar
Rémi Delacourt committed
1504
1505
1506
1507
1508
1509
1510
            ):
                if self.supports_mm_inputs:
                    input_ids = None
                    inputs_embeds = self.inputs_embeds[:num_input_tokens]
                else:
                    input_ids = self.input_ids[:num_input_tokens]
                    inputs_embeds = None
1511

1512
                kwargs = dict(
Rémi Delacourt's avatar
Rémi Delacourt committed
1513
1514
1515
1516
                    input_ids=input_ids,
                    positions=self._get_positions(num_input_tokens),
                    inputs_embeds=inputs_embeds,
                )
1517
1518
1519
                if self.pass_hidden_states_to_model:
                    kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]
                self.model(**kwargs)
1520

1521
    def _get_attention_metadata_builder(self) -> AttentionMetadataBuilder:
1522
        """Find and return the attention metadata builders for EAGLE layers.
1523

1524
1525
        Returns:
            The metadata builders for EAGLE layers.
1526

1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
        Raises:
            AssertionError: If no metadata builders are found for EAGLE layers.
        """
        builder = None
        chosen_layer = self.attn_layer_names[0]

        for kv_cache_group in self.runner.attn_groups:
            for attn_group in kv_cache_group:
                if chosen_layer in attn_group.layer_names:
                    builder = attn_group.get_metadata_builder()
                    break
            if builder is not None:
                break

        assert builder is not None, (
1542
1543
            "Failed to find attention metadata builder for EAGLE layers."
        )
1544
1545
        return builder

1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
    def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:
        """
        Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary
        hidden states and directly uses the last layer output just like eagle1.
        They might indicate this by setting "use_aux_hidden_state" to False
        inside the "eagle_config" dict of their hf_config.
        """
        if self.method != "eagle3":
            return False
        # Assume that eagle3 heads use aux hidden states by default
        use_aux_hidden_state = True
        eagle_config = getattr(self.draft_model_config.hf_config, "eagle_config", None)
        if eagle_config is not None:
            use_aux_hidden_state = eagle_config.get("use_aux_hidden_state", True)
        return use_aux_hidden_state

1562
    def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
1563
        """
1564
1565
        Validate that all drafting layers belong to the same KVCacheGroup.
        Need this assumption to ensure all drafting layers can use the
1566
1567
1568
1569
1570
1571
1572
        same AttentionMetadata.
        May extend to multiple AttentionMetadata in the future.
        """
        kv_cache_groups: dict[str, int] = {}
        for id, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
            for layer_name in kv_cache_group.layer_names:
                kv_cache_groups[layer_name] = id
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
        assert (
            len(
                set(
                    [
                        kv_cache_groups[layer_name]
                        for layer_name in self.attn_layer_names
                    ]
                )
            )
            == 1
1583
        ), "All drafting layers should belong to the same kv cache group"
1584

Rémi Delacourt's avatar
Rémi Delacourt committed
1585
1586
1587
1588
1589
1590
    def _pad_batch_across_dp(
        self,
        num_tokens_unpadded: int,
        num_tokens_padded: int,
    ) -> tuple[int, torch.Tensor]:
        # TODO(Flechman): support DBO ubatching
1591
        should_ubatch, num_toks_across_dp, _ = coordinate_batch_across_dp(
Rémi Delacourt's avatar
Rémi Delacourt committed
1592
1593
1594
            num_tokens_unpadded=num_tokens_unpadded,
            parallel_config=self.vllm_config.parallel_config,
            allow_microbatching=False,
1595
1596
            allow_dp_padding=self.cudagraph_dispatcher.cudagraph_mode
            != CUDAGraphMode.NONE,
Rémi Delacourt's avatar
Rémi Delacourt committed
1597
1598
1599
1600
            num_tokens_padded=num_tokens_padded,
            uniform_decode=None,
            num_scheduled_tokens_per_request=None,
        )
1601
        assert not should_ubatch, "DBO ubatching not implemented for EAGLE"
Rémi Delacourt's avatar
Rémi Delacourt committed
1602
1603
1604
1605
1606

        num_tokens_dp_padded = num_tokens_padded
        if num_toks_across_dp is not None:
            num_tokens_dp_padded = int(num_toks_across_dp[self.dp_rank].item())
        return num_tokens_dp_padded, num_toks_across_dp
1607
1608


1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
class EagleProposer(SpecDecodeBaseProposer):
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
        runner=None,
    ):
        super().__init__(
            vllm_config,
            device,
            pass_hidden_states_to_model=True,
            runner=runner,
        )


1624
1625
1626
1627
# NOTE(woosuk): Currently, the below code is not used and we always use argmax
# to sample the draft tokens. We will use this after we find a way to manage
# the draft prob tensor.
# Refer to https://github.com/vllm-project/vllm/pull/16899 for the details.
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
# FIXME(woosuk): The logic here is duplicated with the main sampling code.
# We should refactor this to reuse the same sampling implementation.
def compute_probs_and_sample_next_token(
    logits: torch.Tensor,
    sampling_metadata: SamplingMetadata,
) -> tuple[torch.Tensor, torch.Tensor]:
    if sampling_metadata.all_greedy:
        # For greedy requests, draft_probs is not used in rejection sampling.
        # Therefore, we can just return the logits.
        probs = logits
        next_token_ids = logits.argmax(dim=-1)
        return next_token_ids, probs

1641
1642
1643
1644
1645
1646
1647
1648
1649
    assert sampling_metadata.temperature is not None

    # Use epsilon comparison to detect greedy sampling (temperature ~ 0.0)
    # consistent with sampler.py's _SAMPLING_EPS threshold
    temperature = sampling_metadata.temperature
    # Avoid division by zero if there are greedy requests.
    if not sampling_metadata.all_random:
        is_greedy = temperature < _SAMPLING_EPS
        temperature = torch.where(is_greedy, 1.0, temperature)
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
    logits.div_(temperature.view(-1, 1))
    probs = logits.softmax(dim=-1, dtype=torch.float32)

    # NOTE(woosuk): Currently, we ignore most of the sampling parameters in
    # generating the draft tokens. We only use the temperature. While this
    # could degrade the acceptance rate, it does not affect the distribution
    # of the generated tokens after rejection sampling.

    # TODO(woosuk): Consider seeds.
    q = torch.empty_like(probs)
    q.exponential_()
1661
1662
1663
    # NOTE(woosuk): We shouldn't use `probs.div_(q)` because the draft_probs
    # will be used later for rejection sampling.
    next_token_ids = probs.div(q).argmax(dim=-1).view(-1)
1664
1665
    if not sampling_metadata.all_random:
        greedy_token_ids = probs.argmax(dim=-1)
1666
        next_token_ids = torch.where(is_greedy, greedy_token_ids, next_token_ids)
1667
    return next_token_ids, probs