eagle.py 49.7 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

11
from vllm.config import (
12
    CompilationMode,
13
14
15
16
    CUDAGraphMode,
    VllmConfig,
    get_layers_from_vllm_config,
)
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_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.utils.platform_utils import is_pin_memory_available
28
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
29
30
31
32
from vllm.v1.attention.backends.tree_attn import (
    TreeAttentionMetadata,
    TreeAttentionMetadataBuilder,
)
33
from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
34
35
36
37
from vllm.v1.attention.backends.utils import (
    AttentionMetadataBuilder,
    CommonAttentionMetadata,
)
38
from vllm.v1.kv_cache_interface import KVCacheConfig
39
from vllm.v1.sample.metadata import SamplingMetadata
40
from vllm.v1.sample.sampler import _SAMPLING_EPS
41
42
43
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
from vllm.v1.utils import CpuGpuBuffer
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
44

45
46
logger = init_logger(__name__)

47
48
PADDING_SLOT_ID = -1

49
50
51
52
53
54

class EagleProposer:
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
Jiayi Yao's avatar
Jiayi Yao committed
55
        runner=None,
56
57
    ):
        self.vllm_config = vllm_config
58
        self.speculative_config = vllm_config.speculative_config
59
        assert self.speculative_config is not None
60
61
        self.draft_model_config = self.speculative_config.draft_model_config
        self.method = self.speculative_config.method
62

Jiayi Yao's avatar
Jiayi Yao committed
63
        self.runner = runner
64
        self.device = device
65
        self.dtype = vllm_config.model_config.dtype
66
67
        self.max_model_len = vllm_config.model_config.max_model_len
        self.block_size = vllm_config.cache_config.block_size
68
69
        self.num_speculative_tokens = self.speculative_config.num_speculative_tokens
        self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
70
        self.token_arange_np = np.arange(self.max_num_tokens)
71
72
73
74
        # 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()
75

76
77
78
        # Multi-modal data support
        self.mm_registry = MULTIMODAL_REGISTRY
        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
79
80
            vllm_config.model_config
        )
81

82
83
        self.attn_metadata_builder: AttentionMetadataBuilder | None = None
        self.draft_indexer_metadata_builder: AttentionMetadataBuilder | None = None
84
85
        self.attn_layer_names: list[str] = []
        self.indexer_layer_names: list[str] = []
86

87
88
89
        self.use_cuda_graph = False

        compilation_config = self.vllm_config.compilation_config
90
        if compilation_config.mode == CompilationMode.VLLM_COMPILE:
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
            cudagraph_mode = compilation_config.cudagraph_mode
            if cudagraph_mode != CUDAGraphMode.NONE and not cudagraph_mode.has_mode(
                CUDAGraphMode.PIECEWISE
            ):
                logger.warning(
                    "Currently the eagle proposer only supports cudagraph_mode "
                    "PIECEWISE, if you want the drafter to use cuda graphs, "
                    "please set compilation_config.cudagraph_mode to PIECEWISE "
                    "or FULL_AND_PIECEWISE"
                )
            self.use_cuda_graph = (
                cudagraph_mode.has_mode(CUDAGraphMode.PIECEWISE)
                and not self.speculative_config.enforce_eager
            )

106
        self.cudagraph_batch_sizes = (
107
            (sorted(self.vllm_config.compilation_config.cudagraph_capture_sizes))
108
109
110
            if self.use_cuda_graph
            else []
        )
111
112

        # persistent buffers for cuda graph
113
114
115
        self.input_ids = torch.zeros(
            self.max_num_tokens, dtype=torch.int32, device=device
        )
116
117
118
        self.uses_mrope = self.vllm_config.model_config.uses_mrope
        if self.uses_mrope:
            # M-RoPE need (3, max_num_tokens)
119
120
121
            self.mrope_positions = torch.zeros(
                (3, self.max_num_tokens), dtype=torch.int64, device=device
            )
122
123
        else:
            # RoPE need (max_num_tokens,)
124
125
126
            self.positions = torch.zeros(
                self.max_num_tokens, dtype=torch.int64, device=device
            )
127
        self.hidden_states = torch.zeros(
128
129
            (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device
        )
130

131
132
        # We need +1 here because the arange is used to set query_start_loc,
        # which has one more element than batch_size.
133
        max_batch_size = vllm_config.scheduler_config.max_num_seqs
134
        max_num_slots_for_arange = max(max_batch_size + 1, self.max_num_tokens)
135
136
137
        self.arange = torch.arange(
            max_num_slots_for_arange, device=device, dtype=torch.int32
        )
138

139
        self.inputs_embeds = torch.zeros(
140
141
            (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device
        )
142

143
144
145
146
147
        self.backup_next_token_ids = CpuGpuBuffer(
            max_batch_size,
            dtype=torch.int32,
            pin_memory=is_pin_memory_available(),
            device=device,
148
149
            with_numpy=True,
        )
150

151
        # Determine allowed attention backends once during initialization.
152
        self.allowed_attn_types: tuple | None = None
153
154
155
156
157
        if current_platform.is_rocm():
            rocm_types = [TritonAttentionMetadata, FlashAttentionMetadata]
            # vllm.v1.attention.backends.rocm_aiter_fa is an optional backend
            if find_spec("vllm.v1.attention.backends.rocm_aiter_fa"):
                from vllm.v1.attention.backends.rocm_aiter_fa import (
158
159
160
                    AiterFlashAttentionMetadata,
                )

161
162
163
                rocm_types.append(AiterFlashAttentionMetadata)
            self.allowed_attn_types = tuple(rocm_types)

164
165
        # Parse the speculative token tree.
        spec_token_tree = self.speculative_config.speculative_token_tree
166
        self.tree_choices: list[tuple[int, ...]] = ast.literal_eval(spec_token_tree)
167
168
169
170
171
172
173
174
        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):
175
176
177
178
179
180
            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]
            )
181
182
183
184
185
186
187
188
        # Precompute draft position offsets in flattened tree.
        self.tree_draft_pos_offsets = torch.arange(
            1,
            len(self.tree_choices) + 1,
            device=device,
            dtype=torch.int32,
        ).repeat(max_batch_size, 1)

189
190
191
192
193
194
195
196
197
198
199
    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:
            self.positions[:num_tokens] = positions

200
201
202
203
    def propose(
        self,
        # [num_tokens]
        target_token_ids: torch.Tensor,
204
        # [num_tokens] or [3, num_tokens] when M-RoPE is enabled
205
206
207
208
209
        target_positions: torch.Tensor,
        # [num_tokens, hidden_size]
        target_hidden_states: torch.Tensor,
        # [batch_size]
        next_token_ids: torch.Tensor,
210
        last_token_indices: torch.Tensor | None,
211
        common_attn_metadata: CommonAttentionMetadata,
212
        sampling_metadata: SamplingMetadata,
213
        mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
214
    ) -> torch.Tensor:
215
216
        num_tokens = target_token_ids.shape[0]
        batch_size = next_token_ids.shape[0]
217
218
219

        if last_token_indices is None:
            last_token_indices = common_attn_metadata.query_start_loc[1:] - 1
220

221
222
223
        if self.method == "eagle3":
            assert isinstance(self.model, Eagle3LlamaForCausalLM)
            target_hidden_states = self.model.combine_hidden_states(
224
225
                target_hidden_states
            )
226
            assert target_hidden_states.shape[-1] == self.hidden_size
227
228
        # Shift the input ids by one token.
        # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
229
        self.input_ids[: num_tokens - 1] = target_token_ids[1:]
230
231
        # Replace the last token with the next token.
        # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]
232
        self.input_ids[last_token_indices] = next_token_ids
233

234
        assert self.runner is not None
Jiayi Yao's avatar
Jiayi Yao committed
235

236
237
238
239
240
        if self.attn_metadata_builder is None:
            attn_metadata_builder = self._get_attention_metadata_builder()
        else:
            attn_metadata_builder = self.attn_metadata_builder

241
        attn_metadata = attn_metadata_builder.build_for_drafting(
242
243
            common_attn_metadata=common_attn_metadata, draft_index=0
        )
244
245
246
247
248
249
        # 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,
250
251
                )
            )
252
253
        else:
            draft_indexer_metadata = None
254
255
256
257
258
        # 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
259

260
261
262
263
        for layer_name in self.indexer_layer_names:
            assert draft_indexer_metadata is not None
            per_layer_attn_metadata[layer_name] = draft_indexer_metadata

264
        cudagraph_runtime_mode = CUDAGraphMode.NONE
265
        if self.use_cuda_graph and num_tokens <= self.cudagraph_batch_sizes[-1]:
266
            num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
267
            cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
268
269
270
        else:
            num_input_tokens = num_tokens
        # copy inputs to buffer for cudagraph
271
        self._set_positions(num_tokens, target_positions)
272
        self.hidden_states[:num_tokens] = target_hidden_states
273
274
275
276
277
278
279
280

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

            self.inputs_embeds[:num_tokens] = self.model.get_input_embeddings(
                self.input_ids[:num_tokens],
                multimodal_embeddings=mm_embeds,
                is_multimodal=is_mm_embed,
281
            )
282

283
            input_ids = None
284
            inputs_embeds = self.inputs_embeds[:num_input_tokens]
285
286
        else:
            input_ids = self.input_ids[:num_input_tokens]
287
            inputs_embeds = None
288

289
        with set_forward_context(
290
291
292
293
            per_layer_attn_metadata,
            self.vllm_config,
            num_tokens=num_input_tokens,
            cudagraph_runtime_mode=cudagraph_runtime_mode,
294
        ):
Jiayi Yao's avatar
Jiayi Yao committed
295
            ret_hidden_states = self.model(
296
                input_ids=input_ids,
297
                positions=self._get_positions(num_input_tokens),
298
299
                hidden_states=self.hidden_states[:num_input_tokens],
                inputs_embeds=inputs_embeds,
300
            )
301
            if self.method == "mtp":
Jiayi Yao's avatar
Jiayi Yao committed
302
                last_hidden_states = ret_hidden_states
303
                hidden_states = last_hidden_states
Jiayi Yao's avatar
Jiayi Yao committed
304
305
            else:
                last_hidden_states, hidden_states = ret_hidden_states
306
        sample_hidden_states = last_hidden_states[last_token_indices]
307
        logits = self.model.compute_logits(sample_hidden_states)
308
309
310
311
312
313

        # 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)
            return draft_token_ids.view(-1, 1)

314
315
316
317
        if self.uses_mrope:
            positions = target_positions[:, last_token_indices]
        else:
            positions = target_positions[last_token_indices]
XuruiYang's avatar
XuruiYang committed
318
319
320
321
        if self.method in ("deepseek_mtp", "ernie_mtp", "longcat_flash_mtp"):
            hidden_states = self.hidden_states[last_token_indices]
        else:
            hidden_states = hidden_states[last_token_indices]
322
323
324

        if isinstance(attn_metadata, TreeAttentionMetadata):
            # Draft using tree attention.
325
326
327
328
329
330
331
332
333
334
            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,
            )
            # [batch_size, num_tree_tokens]
            return torch.cat(draft_token_ids_list, dim=1)

335
        draft_token_ids = logits.argmax(dim=-1)
336

337
338
339
        if self.allowed_attn_types is not None and not isinstance(
            attn_metadata, self.allowed_attn_types
        ):
340
341
342
343
            raise ValueError(
                f"Unsupported attention metadata type for speculative "
                "decoding with num_speculative_tokens > 1: "
                f"{type(attn_metadata)}. Supported types are: "
344
345
                f"{self.allowed_attn_types}"
            )
346

347
348
349
        # Generate the remaining draft tokens.
        draft_token_ids_list = [draft_token_ids]

350
        if self.use_cuda_graph and batch_size <= self.cudagraph_batch_sizes[-1]:
351
            input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
352
            cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
353
354
        else:
            input_batch_size = batch_size
355
            cudagraph_runtime_mode = CUDAGraphMode.NONE
356
357
358

        common_attn_metadata.num_actual_tokens = batch_size
        common_attn_metadata.max_query_len = 1
359
        common_attn_metadata.query_start_loc = self.arange[: batch_size + 1]
360
        common_attn_metadata.query_start_loc_cpu = torch.from_numpy(
361
362
            self.token_arange_np[: batch_size + 1]
        ).clone()
363
        for token_index in range(self.num_speculative_tokens - 1):
364
            # Update the inputs.
365
366
367
            # cast to int32 is crucial when eagle model is compiled.
            # tensor.argmax() returns int64 by default.
            input_ids = draft_token_ids_list[-1].int()
368
369
370
371
372
373
374
375
376
377
378
379
380
            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.
381
382
383
384
385
                clamped_positions = torch.where(
                    exceeds_max_model_len.unsqueeze(0),
                    torch.zeros_like(positions),
                    positions,
                )
386
387
388
            else:
                positions += 1
                exceeds_max_model_len = positions >= self.max_model_len
389
                clamped_positions = torch.where(exceeds_max_model_len, 0, positions)
390
391

            # Increment the sequence lengths.
392
393
            common_attn_metadata.seq_lens += 1
            common_attn_metadata.seq_lens_cpu += 1
394
395
            # For the requests that exceed the max model length, we set the
            # sequence length to 1 to minimize their overheads in attention.
396

397
            common_attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)
398

399
            common_attn_metadata.num_computed_tokens_cpu = (
400
                common_attn_metadata.seq_lens_cpu - 1
401
            )
402

403
            # Compute the slot mapping.
404
405
406
407
408
            if self.uses_mrope:
                # all dimensions of positions are the same
                block_numbers = clamped_positions[0] // self.block_size
            else:
                block_numbers = clamped_positions // self.block_size
409
            block_ids = common_attn_metadata.block_table_tensor.gather(
410
411
                dim=1, index=block_numbers.view(-1, 1)
            )
412
            block_ids = block_ids.view(-1)
413
414
            if self.uses_mrope:
                common_attn_metadata.slot_mapping = (
415
416
                    block_ids * self.block_size + clamped_positions[0] % self.block_size
                )
417
418
            else:
                common_attn_metadata.slot_mapping = (
419
420
                    block_ids * self.block_size + clamped_positions % self.block_size
                )
421
422
423
            # Mask out the slot mappings that exceed the max model length.
            # Otherwise, the KV cache will be inadvertently updated with the
            # padding tokens.
424
            common_attn_metadata.slot_mapping.masked_fill_(
425
426
                exceeds_max_model_len, PADDING_SLOT_ID
            )
427
428

            # Rebuild attention metadata
429
            attn_metadata = attn_metadata_builder.build_for_drafting(  # type: ignore
430
431
                common_attn_metadata=common_attn_metadata, draft_index=token_index + 1
            )
432
433
            for layer_name in self.attn_layer_names:
                per_layer_attn_metadata[layer_name] = attn_metadata
434

435
436
            # copy inputs to buffer for cudagraph
            self.input_ids[:batch_size] = input_ids
437
            self._set_positions(batch_size, clamped_positions)
438
            self.hidden_states[:batch_size] = hidden_states
439
            if self.supports_mm_inputs:
440
441
442
                self.inputs_embeds[:batch_size] = self.model.get_input_embeddings(
                    input_ids
                )
443

444
                input_ids = None
445
                inputs_embeds = self.inputs_embeds[:input_batch_size]
446
447
            else:
                input_ids = self.input_ids[:input_batch_size]
448
                inputs_embeds = None
449

450
            # Run the model.
451
            with set_forward_context(
452
453
454
455
                per_layer_attn_metadata,
                self.vllm_config,
                num_tokens=input_batch_size,
                cudagraph_runtime_mode=cudagraph_runtime_mode,
456
            ):
457
                ret_hidden_states = self.model(
458
                    input_ids=input_ids,
459
                    positions=self._get_positions(input_batch_size),
460
461
                    hidden_states=self.hidden_states[:input_batch_size],
                    inputs_embeds=inputs_embeds,
462
                )
463
                if self.method == "mtp":
464
465
466
467
                    last_hidden_states = ret_hidden_states
                    hidden_states = ret_hidden_states
                else:
                    last_hidden_states, hidden_states = ret_hidden_states
468
            hidden_states = hidden_states[:batch_size]
469
            logits = self.model.compute_logits(last_hidden_states[:batch_size])
470
            draft_token_ids = logits.argmax(dim=-1)
471
472
473
474
            draft_token_ids_list.append(draft_token_ids)

        # [batch_size, num_speculative_tokens]
        draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
475
        return draft_token_ids
476

477
    def prepare_next_token_ids_cpu(
478
479
480
481
482
483
        self,
        sampled_token_ids: list[list[int]],
        requests: dict[str, CachedRequestState],
        gpu_input_batch: InputBatch,
        num_scheduled_tokens: dict[str, int],
    ) -> torch.Tensor:
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
        """
        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]
502
                seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]
503
504
                next_token_id = req_state.get_token_id(seq_len)
            next_token_ids.append(next_token_id)
505
506
507
        next_token_ids = torch.tensor(
            next_token_ids, dtype=torch.int32, device=self.input_ids.device
        )
508
509
        return next_token_ids

510
511
512
513
514
515
516
517
518
    def prepare_next_token_ids_padded(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        sampled_token_ids: torch.Tensor,
        requests: dict[str, CachedRequestState],
        gpu_input_batch: InputBatch,
        discard_request_indices: torch.Tensor,
        num_discarded_requests: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
519
520
521
522
523
524
525
526
527
528
529
530
531
        """
        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
        is not sampled and comes from `request.get_token_id()` instead.
        It also accounts for the rejected tokens in `sampled_token_ids`.
        This function must use device functions to operate on the inputs, and
        should not introduce any blocking CPU-GPU synchronization.
        """
        # TODO(Ben): Combine this into a custom fused kernel

        # Precompute get_token_id for when there is no valid next token
        num_reqs = gpu_input_batch.num_reqs
532
533
534
535
536
537
538
539
        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)
            ]
        )
540
541
542
        self.backup_next_token_ids.copy_to_gpu(num_reqs)

        # Mask out the sampled tokens indices that should not be sampled.
543
544
545
        discard_sampled_tokens_req_indices = discard_request_indices[
            :num_discarded_requests
        ]
546
547
548

        valid_sampled_token_ids_gpu = sampled_token_ids.clone()
        valid_sampled_token_ids_gpu.index_fill_(
549
550
            0, discard_sampled_tokens_req_indices, -1
        )
551
552

        # Generate a mask for all valid tokens within those requests
553
554
555
        valid_mask = (valid_sampled_token_ids_gpu != -1) & (
            valid_sampled_token_ids_gpu < gpu_input_batch.vocab_size
        )
556
557
558
559
560
561
562
563
564
565
566

        # Count the number of valid tokens in each request
        valid_sampled_tokens_count = valid_mask.sum(dim=1)

        # Get the rightmost valid index per row
        last_valid_indices = valid_sampled_tokens_count - 1
        last_valid_indices_safe = torch.clamp(last_valid_indices, min=0)

        # Get last valid token from each row
        # (assume undefined state where there is no valid token)
        selected_tokens = torch.gather(
567
568
            valid_sampled_token_ids_gpu, 1, last_valid_indices_safe.unsqueeze(1)
        ).squeeze(1)
569
570
571
572

        # Use last token if valid, pre-computed backup if not
        batch_size = valid_sampled_token_ids_gpu.shape[0]
        next_token_ids = torch.where(
573
574
575
576
            last_valid_indices != -1,
            selected_tokens,
            self.backup_next_token_ids.gpu[:batch_size],
        )
577
578
579

        return next_token_ids, valid_sampled_tokens_count

580
581
582
583
584
585
    def prepare_inputs_padded(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        spec_decode_metadata: SpecDecodeMetadata,
        valid_sampled_tokens_count: torch.Tensor,
    ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:
586
587
588
589
590
591
592
593
        """
        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`.
        No blocking CPU operations should be introduced in this function.
        """
594
595
596
597
598
599
600
        num_draft_tokens_gpu = torch.cat(
            [
                spec_decode_metadata.cu_num_draft_tokens[0:1],
                spec_decode_metadata.cu_num_draft_tokens[1:]
                - spec_decode_metadata.cu_num_draft_tokens[:-1],
            ]
        )
601
602
603
604

        num_rejected_tokens_gpu = torch.where(
            num_draft_tokens_gpu > 0,
            num_draft_tokens_gpu + 1 - valid_sampled_tokens_count,
605
606
            torch.zeros_like(num_draft_tokens_gpu),
        )
607
608
609

        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu

610
        new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
611
612
613
614
615
616
617
618
619

        total_num_tokens = query_start_loc_cpu[-1].item()
        token_indices = self.arange[:total_num_tokens]

        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,
            seq_lens_cpu=common_attn_metadata.seq_lens_cpu,
620
            num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu,
621
622
623
624
625
626
627
            num_reqs=common_attn_metadata.num_reqs,
            num_actual_tokens=total_num_tokens,
            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,
            slot_mapping=common_attn_metadata.slot_mapping[token_indices],
            causal=True,
628
            dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
629
630
        )

631
632
633
        token_indices_to_sample = (
            common_attn_metadata.query_start_loc[1:] - 1 - num_rejected_tokens_gpu
        )
634
635
636

        return spec_common_attn_metadata, token_indices, token_indices_to_sample

637
638
639
640
641
642
643
644
645
646
647
    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,
    ) -> list[torch.Tensor]:
648
649
650
651
        tree_attn_metadata_builder = self.runner.attn_groups[0][
            0
        ].get_metadata_builder()
        assert isinstance(tree_attn_metadata_builder, TreeAttentionMetadataBuilder)
652

653
        total_num_drafts = self.cu_drafts_per_level[0]
654
655
        level_num_drafts = total_num_drafts
        # Sample a draft token for each child at the tree root level.
656
        num_children = self.child_drafts_per_level[0]
657
658
659
        if num_children == 1:
            draft_token_ids = logits.argmax(dim=-1).view(batch_size, -1)
        else:
660
661
662
            draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(
                batch_size, -1
            )
663
664
665
666
        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.
667
668
669
670
671
672
673
674
675
        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
        )
676
677
        # Precompute the draft token positions.
        flattened_draft_positions = (
678
679
            positions.view(batch_size, -1) + self.tree_draft_pos_offsets[:batch_size, :]
        )
680
        tree_depth = len(self.cu_drafts_per_level)
681
        for level in range(tree_depth - 1):
682
683
            # Get draft positions for RoPE.
            draft_positions = positions + (level + 1)
684
            exceeds_max_model_len = (positions + total_num_drafts) >= self.max_model_len
685
686
            # Mask out the position ids that exceed the max model length.
            # Otherwise, we may get out-of-range error in RoPE.
687
            draft_positions = torch.where(
688
689
690
                exceeds_max_model_len,
                0,
                draft_positions,
691
692
            ).view(batch_size, -1)

693
694
            if level_num_drafts > 1:
                # Repeat the positions for each draft at this level.
695
                draft_positions = draft_positions.repeat_interleave(
696
697
                    level_num_drafts, dim=1
                )
698
699
700
701

            if num_children > 1:
                # Repeat draft hidden states for each child.
                draft_hidden_states = draft_hidden_states.repeat_interleave(
702
703
                    num_children, dim=1
                )
704
705

            # Concatenate the draft tokens, positions, and hidden states.
706
707
            tree_input_ids = torch.cat([tree_input_ids, draft_token_ids], dim=1)
            tree_positions = torch.cat([tree_positions, draft_positions], dim=1)
708
            tree_hidden_states = torch.cat(
709
710
                [tree_hidden_states, draft_hidden_states], dim=1
            )
711
712
713

            # Build new attention metadata for the next level of drafts.
            # This is necessary to support tree attention.
714
            query_len = total_num_drafts
715
716
            common_attn_metadata = replace(
                common_attn_metadata,
717
                query_start_loc=query_len * self.arange[: batch_size + 1],
718
719
720
721
722
723
                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(
                common_attn_metadata=common_attn_metadata,
724
                draft_index=level + 1,
725
726
727
728
729
730
731
732
            )

            # 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.
733
734
735
            attn_metadata.max_seq_len = min(
                attn_metadata.max_seq_len, self.max_model_len
            )
736
737
738
739
740
            # 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.
741
            query_positions = flattened_draft_positions[:, level : level + query_len]
742
            block_numbers = query_positions // self.block_size
743
744
745
746
            block_ids = attn_metadata.block_table.gather(dim=1, index=block_numbers)
            slot_mapping = (
                block_ids * self.block_size + query_positions % self.block_size
            )
747
748
749
750
751
752
753
754
755
756
757
            # 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)
758
            self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1)
759

760
761
            if self.use_cuda_graph and num_tokens <= self.cudagraph_batch_sizes[-1]:
                num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
762
                cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
763
764
            else:
                num_input_tokens = num_tokens
765
                cudagraph_runtime_mode = CUDAGraphMode.NONE
766
            # Run the model.
767
            with set_forward_context(
768
769
770
771
                per_layer_attn_metadata,
                self.vllm_config,
                num_tokens=num_input_tokens,
                cudagraph_runtime_mode=cudagraph_runtime_mode,
772
            ):
773
774
775
776
777
778
779
780
781
                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(
782
783
                batch_size, query_len, -1
            )[:, -level_num_drafts:]
784
            draft_last_hidden_states = last_hidden_states[:num_tokens].view(
785
786
                batch_size, query_len, -1
            )[:, -level_num_drafts:]
787
788
789

            # Get the output logits for the draft tokens.
            logits = self.model.compute_logits(
790
791
                draft_last_hidden_states.reshape(batch_size * level_num_drafts, -1)
            )
792
793
794
795
796
797

            # 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:
798
799
800
                draft_token_ids = torch.topk(logits, num_children, dim=-1).indices.view(
                    batch_size, -1
                )
801
802
803
            draft_token_ids_list.append(draft_token_ids)

            # Update the # drafts counters for the next tree level.
804
            level_num_drafts = self.cu_drafts_per_level[level + 1] - total_num_drafts
805
806
807
            total_num_drafts = self.cu_drafts_per_level[level + 1]
        return draft_token_ids_list

808
    def prepare_inputs(
809
810
        self,
        common_attn_metadata: CommonAttentionMetadata,
811
812
        sampled_token_ids: list[list[int]],
        num_draft_tokens: list[int],
813
814
    ) -> tuple[CommonAttentionMetadata, torch.Tensor]:
        """
815
        This function is used to prepare the inputs for speculative decoding.
816
817
818
819
820
821
        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}:
822
        #       [0, q1, q1 + q2, q1 + q2 + q3]
823
824
825
826
827
828
        #  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}:
829
        #       [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
830
        #  common_attn_metadata.seq_lens{_cpu}:
831
        #       [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1]
832
        #  token_indices: [0, 1, ..., q1 - n1 - 1,
833
834
        #                 q1, q1 + 1, ..., q1 + q2 - n2 - 1,
        #                 q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]
835

836
837
838
839
        num_rejected_tokens = [
            n + 1 - len(sampled_token_ids[i]) if n > 0 else 0
            for i, n in enumerate(num_draft_tokens)
        ]
840
        num_rejected_tokens = torch.tensor(num_rejected_tokens, dtype=torch.int32)
841

842
843
        device = common_attn_metadata.query_start_loc.device
        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
844
        new_seq_lens_cpu = common_attn_metadata.seq_lens_cpu - num_rejected_tokens
845
846

        # [0, q1, q1 + q2, q1 + q2 + q3] -> [q1, q2, q3]
847
        new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
848
849
850
851
852
853
854
855
        # [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,
856
            dtype=torch.int32,
857
858
            pin_memory=is_pin_memory_available(),
        )
859
860
861
862
863
864
865
866
867
        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__
868
869
870
        new_query_start_locs_expanded = np.repeat(
            new_query_start_loc_np[:-1], new_num_tokens_per_req_np
        )
871
872
873
        # [0, 1, 2, 3, 4, 5, 6, 7, 8] ->
        # [0, 1, 0, 1, 2, 3, 0, 1, 2]
        #  _r1_  ____r2____  ___r3__
874
875
876
        token_offests = (
            self.token_arange_np[:total_num_tokens] - new_query_start_locs_expanded
        )
877
878
879
880
881
882

        # 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(
883
884
            query_start_loc_cpu[:-1].numpy(), new_num_tokens_per_req_np
        )
885
        # Final token indices are:
886
887
888
        # [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
889
        token_indices_np = token_offests + old_query_start_locs_expanded
890
        token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True)
891
892

        spec_common_attn_metadata = CommonAttentionMetadata(
893
            query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True),
894
895
896
            seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),
            query_start_loc_cpu=new_query_start_loc_cpu,
            seq_lens_cpu=new_seq_lens_cpu,
897
            num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu,
898
899
900
            num_reqs=common_attn_metadata.num_reqs,
            num_actual_tokens=total_num_tokens,
            max_query_len=new_query_len_per_req.max().item(),
901
            max_seq_len=new_seq_lens_cpu.max().item(),
902
903
            block_table_tensor=common_attn_metadata.block_table_tensor,
            slot_mapping=common_attn_metadata.slot_mapping[token_indices],
904
            causal=True,
905
            dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
906
        )
907
908

        return spec_common_attn_metadata, token_indices
909

910
    def get_model_name(self, model: nn.Module) -> str:
911
        if hasattr(model, "module"):  # multi-GPU
912
913
914
            model = model.module
        return model.__class__.__name__

915
    def load_model(self, target_model: nn.Module) -> None:
916
        draft_model_config = self.vllm_config.speculative_config.draft_model_config
917
        target_attn_layer_names = set(
918
            get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()
919
        )
920
921
        # FIXME: support hybrid kv for draft model
        target_indexer_layer_names = set(
922
923
924
925
            get_layers_from_vllm_config(
                self.vllm_config, DeepseekV32IndexerCache
            ).keys()
        )
926

927
        from vllm.compilation.backends import set_model_tag
928

929
        with set_model_tag("eagle_head"):
930
931
932
            self.model = get_model(
                vllm_config=self.vllm_config, model_config=draft_model_config
            )
933

934
        draft_attn_layer_names = (
935
            get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys()
936
937
938
939
940
941
            - 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
942
        self.attn_layer_names = list(draft_attn_layer_names)
943
944
945
946
947
        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 = (
948
949
950
                indexer_layers[first_layer]
                .get_attn_backend()
                .get_builder_cls()(
951
                    indexer_layers[first_layer].get_kv_cache_spec(self.vllm_config),
952
953
954
                    self.indexer_layer_names,
                    self.vllm_config,
                    self.device,
955
956
                )
            )
957
958
        else:
            self.draft_indexer_metadata_builder = None
959

960
        if self.supports_mm_inputs:
961
962
963
            # Even if the target model is multimodal, we can also use
            # text-only draft models
            try:
964
965
966
967
                dummy_input_ids = torch.tensor([[1]], device=self.input_ids.device)
                self.model.get_input_embeddings(
                    dummy_input_ids, multimodal_embeddings=None
                )
968
969
970
            except (NotImplementedError, AttributeError, TypeError):
                logger.warning(
                    "Draft model does not support multimodal inputs, "
971
972
                    "falling back to text-only mode"
                )
973
                self.supports_mm_inputs = False
974

975
976
        if supports_multimodal(target_model):
            # handle multimodality
977
978
979
980
981
            if (
                self.get_model_name(target_model)
                == "Qwen2_5_VLForConditionalGeneration"
            ):
                self.model.config.image_token_index = target_model.config.image_token_id
982
983
            else:
                self.model.config.image_token_index = (
984
985
                    target_model.config.image_token_index
                )
986
987
988
            target_language_model = target_model.get_language_model()
        else:
            target_language_model = target_model
989
        # share embed_tokens with the target model if needed
990
        if get_pp_group().world_size == 1:
991
            if hasattr(target_language_model.model, "embed_tokens"):
992
                target_embed_tokens = target_language_model.model.embed_tokens
993
            elif hasattr(target_language_model.model, "embedding"):
994
995
996
                target_embed_tokens = target_language_model.model.embedding
            else:
                raise AttributeError(
997
998
                    "Target model does not have 'embed_tokens' or 'embedding' attribute"
                )
999
1000
1001
1002
1003
1004
1005

            # Check if shapes match and we found the embedding
            eagle_shape = self.model.model.embed_tokens.weight.shape
            target_shape = target_embed_tokens.weight.shape
            if eagle_shape == target_shape:
                logger.info(
                    "Assuming the EAGLE head shares the same vocab embedding"
1006
1007
                    " with the target model."
                )
1008
1009
1010
1011
1012
                del self.model.model.embed_tokens
                self.model.model.embed_tokens = target_embed_tokens
            else:
                logger.info(
                    "The EAGLE head's vocab embedding will be loaded separately"
1013
1014
                    " from the target model."
                )
1015
        else:
1016
            logger.info(
1017
                "The EAGLE head's vocab embedding will be loaded separately"
1018
1019
                " from the target model."
            )
1020
1021
1022
1023

        # share lm_head with the target model if needed
        # some model definition do not define lm_head explicitly
        # and reuse embed_tokens for lm_head, e.g., CohereForCausalLM
1024
1025
        if self.vllm_config.speculative_config.method != "eagle3":
            if hasattr(target_language_model, "lm_head"):
1026
                logger.info("Loading EAGLE LM head weights from the target model.")
1027
1028
                self.model.lm_head = target_language_model.lm_head
        else:
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
            if (
                hasattr(self.model, "lm_head")
                and hasattr(target_language_model, "lm_head")
                and self.model.lm_head.weight.shape
                == target_language_model.lm_head.weight.shape
            ):
                logger.info(
                    "Assuming the EAGLE head shares the same lm_head"
                    " with the target model."
                )
1039
1040
1041
1042
1043
                del self.model.lm_head
                self.model.lm_head = target_language_model.lm_head
            else:
                logger.info(
                    "The EAGLE head's lm_head will be loaded separately"
1044
1045
                    " from the target model."
                )
1046

1047
1048
1049
1050
    @torch.inference_mode()
    def dummy_run(
        self,
        num_tokens: int,
1051
        use_cudagraphs=True,
1052
    ) -> None:
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
        if use_cudagraphs and num_tokens <= self.cudagraph_batch_sizes[-1]:
            num_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)

        with set_forward_context(
            None,
            self.vllm_config,
            num_tokens=num_tokens,
            cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE
            if use_cudagraphs
            else CUDAGraphMode.NONE,
        ):
1064
            if self.supports_mm_inputs:
1065
1066
1067
1068
1069
1070
                input_ids = None
                inputs_embeds = self.inputs_embeds[:num_tokens]
            else:
                input_ids = self.input_ids[:num_tokens]
                inputs_embeds = None

1071
            self.model(
1072
                input_ids=input_ids,
1073
                positions=self._get_positions(num_tokens),
1074
1075
                hidden_states=self.hidden_states[:num_tokens],
                inputs_embeds=inputs_embeds,
1076
            )
1077

1078
    def _get_attention_metadata_builder(self) -> AttentionMetadataBuilder:
1079
        """Find and return the attention metadata builders for EAGLE layers.
1080

1081
1082
        Returns:
            The metadata builders for EAGLE layers.
1083

1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
        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, (
1099
1100
            "Failed to find attention metadata builder for EAGLE layers."
        )
1101
1102
        return builder

1103
    def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
        """
        Validate that all eagle layers belong to the same KVCacheGroup.
        Need this assumption to ensure all eagle layers can use the
        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
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
        assert (
            len(
                set(
                    [
                        kv_cache_groups[layer_name]
                        for layer_name in self.attn_layer_names
                    ]
                )
            )
            == 1
        ), "All eagle layers should belong to the same kv cache group"
1125

1126

1127
1128
1129
1130
# 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.
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
# 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

1144
1145
1146
1147
1148
1149
1150
1151
1152
    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)
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
    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_()
1164
1165
1166
    # 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)
1167
1168
1169
1170
1171
1172
1173
1174
    if not sampling_metadata.all_random:
        greedy_token_ids = probs.argmax(dim=-1)
        next_token_ids = torch.where(
            is_greedy,
            greedy_token_ids,
            next_token_ids,
        )
    return next_token_ids, probs