gpu_model_runner.py 187 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

王敏's avatar
王敏 committed
4
import os
Robert Shaw's avatar
Robert Shaw committed
5
import copy
6
import gc
7
import time
8
import weakref
9
from contextlib import contextmanager
10
from typing import TYPE_CHECKING, Any, Optional, Union
11
12
13
14
15

import numpy as np
import torch
import torch.distributed
import torch.nn as nn
16
from tqdm import tqdm
17

18
import vllm.envs as envs
19
from vllm.attention import AttentionType, get_attn_backend
20
from vllm.attention.backends.abstract import AttentionBackend
21
from vllm.attention.layer import Attention
22
from vllm.compilation.counter import compilation_counter
23
24
from vllm.config import (CompilationLevel, VllmConfig,
                         get_layers_from_vllm_config)
25
from vllm.distributed.eplb.eplb_state import EplbState
26
27
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
                                          has_kv_transfer_group)
Robert Shaw's avatar
Robert Shaw committed
28
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
29
from vllm.distributed.parallel_state import (
30
    get_pp_group, get_tp_group, graph_capture, is_global_first_rank,
王敏's avatar
王敏 committed
31
32
    prepare_communication_buffer_for_model,
    get_tensor_model_parallel_world_size)
33
from vllm.forward_context import (DPMetadata, get_forward_context,
34
                                  set_forward_context, set_profilling)
35
from vllm.logger import init_logger
Chen Zhang's avatar
Chen Zhang committed
36
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
37
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
38
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
39
40
from vllm.model_executor.models.interfaces import (has_step_pooler,
                                                   is_mixture_of_experts)
41
42
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
43
from vllm.multimodal.utils import group_mm_inputs_by_modality
44
from vllm.pooling_params import PoolingParams
45
from vllm.sampling_params import SamplingType
46
from vllm.sequence import IntermediateTensors
47
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
48
                        GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,
Chen Zhang's avatar
Chen Zhang committed
49
                        check_use_alibi, get_dtype_size,
50
                        is_pin_memory_available, round_up, round_down)
Chen Zhang's avatar
Chen Zhang committed
51
from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
52
53
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
                                              CommonAttentionMetadata)
54
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
55
from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
Chen Zhang's avatar
Chen Zhang committed
56
                                        KVCacheConfig, KVCacheSpec, MambaSpec,
57
                                        SlidingWindowSpec)
58
59
from vllm.v1.kv_compression.budget import (compute_prompt_topk_keep_total,
                                           compute_topk_budget_step)
60
61
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
                             ModelRunnerOutput)
62
from vllm.v1.pool.metadata import PoolingMetadata
63
from vllm.v1.sample.metadata import SamplingMetadata
64
from vllm.v1.sample.rejection_sampler import RejectionSampler
王敏's avatar
王敏 committed
65
from vllm.v1.sample.rejection_sampler_opt import OptRejectionSampler
66
from vllm.v1.sample.sampler import Sampler
67
from vllm.v1.spec_decode.eagle import EagleProposer
68
from vllm.v1.spec_decode.medusa import MedusaProposer
69
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
70
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
71
from vllm.v1.utils import bind_kv_cache
72
from vllm.v1.worker.block_table import BlockTable
73
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
74
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
zhuwenwen's avatar
zhuwenwen committed
75
from vllm.platforms import current_platform
王敏's avatar
王敏 committed
76
from vllm.two_batch_overlap.v1.model_input_split_v1 import tbo_split_and_execute_model
maxiao1's avatar
maxiao1 committed
77
from vllm.profiler.prof import profile
78
from ..sample.logits_processor import LogitsProcessorManager
79
80
from .utils import (gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,
                    sanity_check_mm_encoder_outputs, scatter_mm_placeholders)
81
from vllm.zero_overhead.v1.eagle import V1ZeroEagleProposer
王敏's avatar
王敏 committed
82
from vllm.v1.spec_decode.utils import DraftProbs
83

84
if TYPE_CHECKING:
85
    import xgrammar as xgr
86
    import xgrammar.kernels.apply_token_bitmask_inplace_torch_compile as xgr_torch_compile  # noqa: E501
87

88
    from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
89
    from vllm.v1.core.sched.output import SchedulerOutput
90
91
else:
    xgr = LazyLoader("xgr", globals(), "xgrammar")
92
93
94
    xgr_torch_compile = LazyLoader(
        "xgr_torch_compile", globals(),
        "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile")
95
96
97
98

logger = init_logger(__name__)


99
class GPUModelRunnerBase(LoRAModelRunnerMixin):
100
101
102

    def __init__(
        self,
103
        vllm_config: VllmConfig,
104
        device: torch.device,
105
    ):
106
107
108
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
        self.cache_config = vllm_config.cache_config
109
        self.compilation_config = vllm_config.compilation_config
110
111
112
113
114
115
116
        self.lora_config = vllm_config.lora_config
        self.load_config = vllm_config.load_config
        self.parallel_config = vllm_config.parallel_config
        self.scheduler_config = vllm_config.scheduler_config
        self.speculative_config = vllm_config.speculative_config
        self.prompt_adapter_config = vllm_config.prompt_adapter_config
        self.observability_config = vllm_config.observability_config
maxiao1's avatar
maxiao1 committed
117
118
119
        if envs.VLLM_P2P_ASYNC:
            self.p2p_event = torch.cuda.Event(enable_timing=False)
            self.p2p_stream = torch.cuda.Stream()
120

121
122
123
124
        from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
        set_cpu_offload_max_bytes(
            int(self.cache_config.cpu_offload_gb * 1024**3))

125
126
127
128
        model_config = self.model_config
        cache_config = self.cache_config
        scheduler_config = self.scheduler_config
        parallel_config = self.parallel_config
129
        self.device = device
130
131
132
133
134
135
136
137
        self.pin_memory = is_pin_memory_available()
        self.dtype = self.model_config.dtype
        if cache_config.cache_dtype == "auto":
            self.kv_cache_dtype = self.dtype
        else:
            self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
                cache_config.cache_dtype]

138
        self.is_multimodal_model = model_config.is_multimodal_model
139
        self.is_pooling_model = model_config.pooler_config is not None
140
141
        self.max_model_len = model_config.max_model_len
        self.max_num_tokens = scheduler_config.max_num_batched_tokens
142
        self.max_num_reqs = scheduler_config.max_num_seqs
143
144

        # Model-related.
145
146
        self.num_query_heads = model_config.get_num_attention_heads(
            parallel_config)
147
        self.hidden_size = model_config.get_hidden_size()
148
        self.attention_chunk_size = model_config.attention_chunk_size
149

150
        self.cascade_attn_enabled = not self.model_config.disable_cascade_attn
laibao's avatar
laibao committed
151
152
153
154
155
156
157
        if envs.VLLM_ENABLE_KV_COMPRESSION:
            # KV compression changes the effective KV sequence layout and
            # invalidates cascade attention assumptions (common-prefix blocks).
            self.cascade_attn_enabled = False
        # Whether the current step needs KV compaction work (score/topk/dst).
        # This is set per-step in `_prepare_inputs`.
        self.kv_compression_needs_compaction: bool = False
158

159
        # Multi-modal data support
160
        self.mm_registry = MULTIMODAL_REGISTRY
161
        self.uses_mrope = model_config.uses_mrope
162

163
164
165
        encoder_compute_budget, encoder_cache_size = compute_encoder_budget(
            model_config=model_config,
            scheduler_config=scheduler_config,
166
            mm_registry=self.mm_registry,
167
168
169
        )
        self.max_num_encoder_input_tokens = encoder_compute_budget
        self.encoder_cache_size = encoder_cache_size
170

171
172
173
        # Sampler
        self.sampler = Sampler()

174
175
176
177
178
179
180
        self.eplb_state: Optional[EplbState] = None
        """
        State of the expert parallelism load balancer.

        Will be lazily initialized when the model is loaded.
        """

181
        # Lazy initializations
182
        # self.model: nn.Module  # Set after load_model
183
        # Initialize in initialize_kv_cache
184
        self.kv_caches: list[torch.Tensor] = []
185
186
        self.attn_metadata_builders: list[AttentionMetadataBuilder] = []
        self.attn_backends: list[type[AttentionBackend]] = []
187
188
        # self.kv_cache_config: KVCacheConfig

189
        # req_id -> (input_id -> encoder_output)
190
        self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}
191

192
        self.use_aux_hidden_state_outputs = False
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
        # Set up speculative decoding.
        # NOTE(Jiayi): currently we put the entire draft model on
        # the last PP rank. This is not ideal if there are many
        # layers in the draft model.
        if self.speculative_config and get_pp_group().is_last_rank:
            if self.speculative_config.method == "ngram":
                self.drafter = NgramProposer(self.vllm_config)
            elif self.speculative_config.use_eagle():
                self.drafter = EagleProposer(self.vllm_config, self.device,
                                             self)  # type: ignore
                if self.speculative_config.method == "eagle3":
                    self.use_aux_hidden_state_outputs = True
            elif self.speculative_config.method == "medusa":
                self.drafter = MedusaProposer(
                    vllm_config=self.vllm_config,
                    device=self.device)  # type: ignore
            else:
                raise ValueError("Unknown speculative decoding method: "
                                 f"{self.speculative_config.method}")
王敏's avatar
王敏 committed
212
213
214
215
            if not envs.VLLM_REJECT_SAMPLE_OPT:
                self.rejection_sampler = RejectionSampler()
            else:
                self.rejection_sampler = OptRejectionSampler()
216

217
        # Request states.
218
        self.requests: dict[str, CachedRequestState] = {}
219

220
221
222
223
224
225
226
227
228
        # Input Batch
        # NOTE(Chen): Ideally, we should initialize the input batch inside
        # `initialize_kv_cache` based on the kv cache config. However, as in
        # https://github.com/vllm-project/vllm/pull/18298, due to some unknown
        # reasons, we have to initialize the input batch before `load_model`,
        # quantization + weight offloading will fail otherwise. As a temporary
        # solution, we initialize the input batch here, and re-initialize it
        # in `initialize_kv_cache` if the block_sizes here is different from
        # the block_sizes in the kv cache config.
229
230
231
232
233
234
        self.input_batch = InputBatch(
            max_num_reqs=self.max_num_reqs,
            max_model_len=self.max_model_len,
            max_num_batched_tokens=self.max_num_tokens,
            device=self.device,
            pin_memory=self.pin_memory,
235
            vocab_size=self.model_config.get_vocab_size(),
236
            block_sizes=[self.cache_config.block_size],
237
            is_spec_decode=bool(self.vllm_config.speculative_config),
238
        )
239

240
241
242
243
244
        self.use_cuda_graph = (
            self.vllm_config.compilation_config.level
            == CompilationLevel.PIECEWISE
            and self.vllm_config.compilation_config.use_cudagraph
            and not self.model_config.enforce_eager)
245
        # TODO(woosuk): Provide an option to tune the max cudagraph batch size.
246
247
248
249
        # The convention is different.
        # self.cudagraph_batch_sizes sorts in ascending order.
        # The batch sizes in the config are in descending order.
        self.cudagraph_batch_sizes = list(
250
251
252
            reversed(self.compilation_config.cudagraph_capture_sizes))

        self.full_cuda_graph = self.compilation_config.full_cuda_graph
253

254
        # Cache the device properties.
255
        self._init_device_properties()
256

257
258
259
260
        # Persistent buffers for CUDA graphs.
        self.input_ids = torch.zeros(self.max_num_tokens,
                                     dtype=torch.int32,
                                     device=self.device)
261
262
263
        self.positions = torch.zeros(self.max_num_tokens,
                                     dtype=torch.int64,
                                     device=self.device)
264
265
266
267
268
269
270
271
272
273
        self.query_start_loc = torch.zeros(self.max_num_reqs + 1,
                                           dtype=torch.int32,
                                           device=self.device)
        self.seq_lens = torch.zeros(self.max_num_reqs,
                                    dtype=torch.int32,
                                    device=self.device)
        self.slot_mapping = torch.zeros(self.max_num_tokens,
                                        dtype=torch.int64,
                                        device=self.device)

274
275
        # None in the first PP rank. The rest are set after load_model.
        self.intermediate_tensors: Optional[IntermediateTensors] = None
276
277

        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
278
        if self.uses_mrope:
Roger Wang's avatar
Roger Wang committed
279
280
281
282
            # 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
283
284
285
286
287
288

            # 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
Roger Wang's avatar
Roger Wang committed
289
            self.mrope_positions = torch.zeros((3, self.max_num_tokens + 1),
290
291
                                               dtype=torch.int64,
                                               device=self.device)
Roger Wang's avatar
Roger Wang committed
292
293
294
295
296
            self.mrope_positions_cpu = torch.zeros(
                (3, self.max_num_tokens + 1),
                dtype=torch.int64,
                device="cpu",
                pin_memory=self.pin_memory)
297
            self.mrope_positions_np = self.mrope_positions_cpu.numpy()
298

299
300
301
        # Only relevant for models using ALiBi (e.g, MPT)
        self.use_alibi = check_use_alibi(model_config)

302
303
304
305
        self.inputs_embeds = torch.zeros(
            (self.max_num_tokens, self.hidden_size),
            dtype=self.dtype,
            device=self.device)
306

307
        # OPTIMIZATION: Cache the tensors rather than creating them every step.
308
        # Keep in int64 to avoid overflow with long context
309
        self.arange_np = np.arange(max(self.max_num_reqs + 1,
310
311
                                       self.max_model_len,
                                       self.max_num_tokens),
312
                                   dtype=np.int64)
313
314
315
        # NOTE(woosuk): These tensors are "stateless", i.e., they are literally
        # a faster version of creating a new tensor every time. Thus, we should
        # not make any assumptions about the values in these tensors.
316
317
318
319
320
321
322
323
324
        self.input_ids_cpu = torch.zeros(self.max_num_tokens,
                                         dtype=torch.int32,
                                         device="cpu",
                                         pin_memory=self.pin_memory)
        self.positions_cpu = torch.zeros(self.max_num_tokens,
                                         dtype=torch.int64,
                                         device="cpu",
                                         pin_memory=self.pin_memory)
        self.positions_np = self.positions_cpu.numpy()
laibao's avatar
laibao committed
325
326
327
328
329
330
331
332
        # KV positions are decoupled from logical positions when KV compression
        # is enabled. We keep a separate buffer to avoid recomputing or
        # overwriting `positions_np` (used for RoPE / input token lookup).
        self.kv_positions_cpu = torch.zeros(self.max_num_tokens,
                                            dtype=torch.int64,
                                            device="cpu",
                                            pin_memory=self.pin_memory)
        self.kv_positions_np = self.kv_positions_cpu.numpy()
333
334
335
336
337
        self.query_start_loc_cpu = torch.zeros(self.max_num_reqs + 1,
                                               dtype=torch.int32,
                                               device="cpu",
                                               pin_memory=self.pin_memory)
        self.query_start_loc_np = self.query_start_loc_cpu.numpy()
338
339
340
341
342
        self.seq_lens_cpu = torch.zeros(self.max_num_reqs,
                                        dtype=torch.int32,
                                        device="cpu",
                                        pin_memory=self.pin_memory)
        self.seq_lens_np = self.seq_lens_cpu.numpy()
laibao's avatar
laibao committed
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
        # KV compression metadata buffers (used by the "topk" policy).
        # Per-token: whether this scheduled token must be kept in KV cache.
        self.kv_compression_must_keep_cpu = torch.zeros(
            self.max_num_tokens,
            dtype=torch.bool,
            device="cpu",
            pin_memory=self.pin_memory,
        )
        self.kv_compression_must_keep_np = self.kv_compression_must_keep_cpu.numpy()
        self.kv_compression_must_keep = torch.zeros(
            self.max_num_tokens,
            dtype=torch.bool,
            device=self.device,
        )
        # Per-request: how many additional prompt tokens to keep among
        # non-protected candidates (budget from env; selection uses scores).
        self.kv_compression_topk_budget_cpu = torch.zeros(
            self.max_num_reqs,
            dtype=torch.int32,
            device="cpu",
            pin_memory=self.pin_memory,
        )
        self.kv_compression_topk_budget_np = self.kv_compression_topk_budget_cpu.numpy()
        self.kv_compression_topk_budget = torch.zeros(
            self.max_num_reqs,
            dtype=torch.int32,
            device=self.device,
        )
371

372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
        # Chunked-prefill prompt-end KV compression metadata (scheme 3).
        # Per-request: whether this step finishes the prompt and should compute
        # global prompt indices (score/topk) for a one-shot compaction.
        self.kv_compression_prompt_end_cpu = torch.zeros(
            self.max_num_reqs,
            dtype=torch.bool,
            device="cpu",
            pin_memory=self.pin_memory,
        )
        self.kv_compression_prompt_end_np = self.kv_compression_prompt_end_cpu.numpy()
        self.kv_compression_prompt_end = torch.zeros(
            self.max_num_reqs,
            dtype=torch.bool,
            device=self.device,
        )
        # Per-request: prompt length (tokens) and Top-K keep count among prompt
        # candidates (excluding protected prefix/suffix).
        self.kv_compression_prompt_lens_cpu = torch.zeros(
            self.max_num_reqs,
            dtype=torch.int32,
            device="cpu",
            pin_memory=self.pin_memory,
        )
        self.kv_compression_prompt_lens_np = self.kv_compression_prompt_lens_cpu.numpy()
        self.kv_compression_prompt_lens = torch.zeros(
            self.max_num_reqs,
            dtype=torch.int32,
            device=self.device,
        )
        self.kv_compression_prompt_topk_keep_cpu = torch.zeros(
            self.max_num_reqs,
            dtype=torch.int32,
            device="cpu",
            pin_memory=self.pin_memory,
        )
        self.kv_compression_prompt_topk_keep_np = self.kv_compression_prompt_topk_keep_cpu.numpy()
        self.kv_compression_prompt_topk_keep = torch.zeros(
            self.max_num_reqs,
            dtype=torch.int32,
            device=self.device,
        )
        self.kv_compression_prompt_topk_keep_max: Optional[int] = None

415
416
417
418
419
420
        # Layer pairings for cross-layer KV sharing.
        # If an Attention layer `layer_name` is in the keys of this dict, it
        # means this layer will perform attention using the keys and values
        # from the KV cache of `shared_kv_cache_layers[layer_name]`.
        self.shared_kv_cache_layers: dict[str, str] = {}

王敏's avatar
王敏 committed
421
422
        self.draft_probs : Optional[DraftProbs] = None

423
424
425
426
427
428
429
        self.ep_sp = False
        self.dp_size = self.parallel_config.data_parallel_size
        self.tp_size = self.parallel_config.tensor_parallel_size
        self.enable_expert_parallel = self.parallel_config.enable_expert_parallel
        if self.enable_expert_parallel and self.dp_size > 1 and self.tp_size > 1:
            self.ep_sp = True

430
    def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
431
432
        """
        Update the order of requests in the batch based on the attention
433
        backend's needs. For example, some attention backends (namely MLA) may
434
435
436
437
438
439
        want to separate requests based on if the attention computation will be
        compute-bound or memory-bound.

        Args:
            scheduler_output: The scheduler output.
        """
440
441
        self.attn_metadata_builders[0].reorder_batch(self.input_batch,
                                                     scheduler_output)
442
443
444
445
446

        # For models with multiple KV cache groups, the groups should agree on
        # the same order of requests. We ensure this by only allowing the first
        # group to reorder the batch and asserting that all other groups do not
        # reorder the batch.
447
448
449
        # TODO(tdoublep): make this more flexible so that any group can
        # re-order the batch (not only the first).
        # TODO(tdoublep): verify this during engine init instead of at runtime
450
        for i in range(1, len(self.kv_cache_config.kv_cache_groups)):
451
            batch_reordered = self.attn_metadata_builders[i].reorder_batch(
452
                self.input_batch, scheduler_output)
453
            assert not batch_reordered
454

455
456
457
458
459
460
461
462
463
464
465
    # Note: used for model runner override.
    def _init_device_properties(self) -> None:
        """Initialize attributes from torch.cuda.get_device_properties
        """
        self.device_properties = torch.cuda.get_device_properties(self.device)
        self.num_sms = self.device_properties.multi_processor_count

    # Note: used for model runner override.
    def _sync_device(self) -> None:
        torch.cuda.synchronize()

466
    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
467
468
469
470
471
472
        """Update the cached states and the persistent batch with the scheduler
        output.

        The updated states are used by the `_prepare_inputs` function to create
        the input GPU tensors for the model.

473
474
        The SamplingMetadata is updated and copied to the GPU if there is a
        new/resumed/paused/finished request in the batch.
475
476
        """
        # Remove finished requests from the cached states.
477
478
        for req_id in scheduler_output.finished_req_ids:
            self.requests.pop(req_id, None)
479
            self.encoder_cache.pop(req_id, None)
480
481
482
483
484
485
486
        # Remove the finished requests from the persistent batch.
        # NOTE(woosuk): There could be an edge case where finished_req_ids and
        # scheduled_req_ids overlap. This happens when a request is aborted and
        # then resubmitted with the same ID. In this case, we treat them as two
        # distinct requests - clearing the cached states for the first request
        # and handling the second as a new request.
        for req_id in scheduler_output.finished_req_ids:
487
            self.input_batch.remove_request(req_id)
488

王敏's avatar
王敏 committed
489
490
491
492
        # prune draft probs of finished requests
        if envs.VLLM_REJECT_SAMPLE_OPT and self.draft_probs is not None and len(scheduler_output.finished_req_ids) > 0:
            self.draft_probs.prune(list(scheduler_output.finished_req_ids))

493
494
495
496
497
498
499
        # Free the cached encoder outputs.
        for req_id, input_id in scheduler_output.free_encoder_input_ids:
            encoder_outputs = self.encoder_cache.get(req_id)
            if encoder_outputs is not None:
                encoder_outputs.pop(input_id, None)
                if not encoder_outputs:
                    self.encoder_cache.pop(req_id, None)
500

501
502
503
504
505
506
507
508
509
510
511
512
513
        # Remove the unscheduled requests from the persistent batch.
        # NOTE(woosuk): The unscheduled requests are either preempted requests
        # or running requests that are not scheduled in this step. We remove
        # them from the persistent batch but keep their cached states since
        # they will be scheduled again sometime in the future.
        scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
        cached_req_ids = self.input_batch.req_id_to_index.keys()
        unscheduled_req_ids = cached_req_ids - scheduled_req_ids
        # NOTE(woosuk): The persistent batch optimization assumes that
        # consecutive batches contain mostly the same requests. If batches
        # have low request overlap (e.g., alternating between two distinct
        # sets of requests), this optimization becomes very inefficient.
        for req_id in unscheduled_req_ids:
514
            self.input_batch.remove_request(req_id)
515

516
        req_ids_to_add: list[str] = []
517
        # Add new requests to the cached states.
518
519
520
        for new_req_data in scheduler_output.scheduled_new_reqs:
            req_id = new_req_data.req_id
            sampling_params = new_req_data.sampling_params
521
522
523
            pooling_params = new_req_data.pooling_params
            if sampling_params and \
                sampling_params.sampling_type == SamplingType.RANDOM_SEED:
524
525
526
527
528
                generator = torch.Generator(device=self.device)
                generator.manual_seed(sampling_params.seed)
            else:
                generator = None

529
530
            self.requests[req_id] = CachedRequestState(
                req_id=req_id,
531
532
533
                prompt_token_ids=new_req_data.prompt_token_ids,
                mm_inputs=new_req_data.mm_inputs,
                mm_positions=new_req_data.mm_positions,
534
                sampling_params=sampling_params,
535
                pooling_params=pooling_params,
536
                generator=generator,
537
538
                block_ids=new_req_data.block_ids,
                num_computed_tokens=new_req_data.num_computed_tokens,
laibao's avatar
laibao committed
539
                num_kv_tokens=new_req_data.num_kv_tokens,
540
                output_token_ids=[],
541
                lora_request=new_req_data.lora_request,
542
            )
543
544

            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
545
            if self.uses_mrope:
546
547
                image_grid_thw = []
                video_grid_thw = []
Roger Wang's avatar
Roger Wang committed
548
                second_per_grid_ts = []
549
550
                audio_feature_lengths = []
                use_audio_in_video = False
551
552
553
554
555
556
557
                for mm_input in self.requests[req_id].mm_inputs:
                    if mm_input.get("image_grid_thw") is not None:
                        image_grid_thw.extend(
                            mm_input["image_grid_thw"].tolist())
                    if mm_input.get("video_grid_thw") is not None:
                        video_grid_thw.extend(
                            mm_input["video_grid_thw"].tolist())
Roger Wang's avatar
Roger Wang committed
558
559
560
                    if mm_input.get("second_per_grid_ts") is not None:
                        second_per_grid_ts.extend(
                            mm_input["second_per_grid_ts"])
561
562
563
564
565
                    if mm_input.get("audio_feature_lengths") is not None:
                        audio_feature_lengths.extend(
                            mm_input["audio_feature_lengths"])
                    if mm_input.get("use_audio_in_video") is True:
                        use_audio_in_video = True
566
567
568
569
570
571
572

                hf_config = self.model_config.hf_config

                self.requests[req_id].mrope_positions, \
                    self.requests[req_id].mrope_position_delta = \
                    MRotaryEmbedding.get_input_positions_tensor(
                        self.requests[req_id].prompt_token_ids,
Roger Wang's avatar
Roger Wang committed
573
                        hf_config=hf_config,
574
575
                        image_grid_thw=image_grid_thw,
                        video_grid_thw=video_grid_thw,
Roger Wang's avatar
Roger Wang committed
576
                        second_per_grid_ts=second_per_grid_ts,
577
578
                        audio_feature_lengths=audio_feature_lengths,
                        use_audio_in_video=use_audio_in_video,
579
580
                    )

581
582
            req_ids_to_add.append(req_id)

583
        # Update the states of the running/resumed requests.
584
        is_last_rank = get_pp_group().is_last_rank
585
586
        req_data = scheduler_output.scheduled_cached_reqs
        for i, req_id in enumerate(req_data.req_ids):
587
            req_state = self.requests[req_id]
588
            num_computed_tokens = req_data.num_computed_tokens[i]
laibao's avatar
laibao committed
589
            num_kv_tokens = req_data.num_kv_tokens[i]
590
591
            new_block_ids = req_data.new_block_ids[i]
            resumed_from_preemption = req_data.resumed_from_preemption[i]
592

593
            # Update the cached states.
594
            req_state.num_computed_tokens = num_computed_tokens
laibao's avatar
laibao committed
595
            req_state.num_kv_tokens = num_kv_tokens
lizhigong's avatar
lizhigong committed
596
597
            spec_token_ids = (
                scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613

            if not is_last_rank:
                # When using PP, the scheduler sends the sampled tokens back,
                # because there's no direct communication between the first-
                # stage worker and the last-stage worker.
                new_token_ids = req_data.new_token_ids[i]
                # Add the sampled token(s) from the previous step (if any).
                # This doesn't include "unverified" tokens like spec tokens.
                num_new_tokens = (num_computed_tokens + len(new_token_ids) -
                                  req_state.num_tokens)
                if num_new_tokens == 1:
                    # Avoid slicing list in most common case.
                    req_state.output_token_ids.append(new_token_ids[-1])
                elif num_new_tokens > 0:
                    req_state.output_token_ids.extend(
                        new_token_ids[-num_new_tokens:])
lizhigong's avatar
lizhigong committed
614
615
            if len(spec_token_ids) > 0:
                req_state.spec_token_ids = spec_token_ids
616

617
            # Update the block IDs.
618
            if not resumed_from_preemption:
619
                # Append the new blocks to the existing block IDs.
620
621
622
                for block_ids, new_ids in zip(req_state.block_ids,
                                              new_block_ids):
                    block_ids.extend(new_ids)
623
624
625
            else:
                # The request is resumed from preemption.
                # Replace the existing block IDs with the new ones.
626
                req_state.block_ids = new_block_ids
627
628
629
630
631
632
633
634
635
636
637

            req_index = self.input_batch.req_id_to_index.get(req_id)
            if req_index is None:
                # The request is not in the persistent batch.
                # The request was either preempted and resumed later, or was not
                # scheduled in the previous step and needs to be added again.
                req_ids_to_add.append(req_id)
                continue

            # Update the persistent batch.
            self.input_batch.num_computed_tokens_cpu[req_index] = (
638
                num_computed_tokens)
laibao's avatar
laibao committed
639
640
641
642
643
            self.input_batch.num_kv_tokens_cpu[req_index] = num_kv_tokens
            if resumed_from_preemption:
                self.input_batch.block_table.add_row(new_block_ids, req_index)
            else:
                self.input_batch.block_table.append_row(new_block_ids, req_index)
644
645
646
647
648
649

            # For the last rank, we don't need to update the token_ids_cpu
            # because the sampled tokens are already cached.
            if not is_last_rank:
                # Add new_token_ids to token_ids_cpu.
                start_token_index = num_computed_tokens
lizhigong's avatar
lizhigong committed
650
                end_token_index = num_computed_tokens + 1
651
                self.input_batch.token_ids_cpu[
652
                    req_index,
lizhigong's avatar
lizhigong committed
653
                    start_token_index:end_token_index] = new_token_ids[-1]
654
655
656
                self.input_batch.num_tokens_no_spec[
                    req_index] = end_token_index
                self.input_batch.num_tokens[req_index] = end_token_index
657

658
659
            # Add spec_token_ids to token_ids_cpu.
            if spec_token_ids:
660
661
662
                num_spec_tokens = len(spec_token_ids)
                start_index = self.input_batch.num_tokens_no_spec[req_index]
                end_token_index = start_index + num_spec_tokens
663
664
                self.input_batch.token_ids_cpu[
                    req_index, start_index:end_token_index] = spec_token_ids
665
666
                # NOTE(woosuk): `num_tokens` here may include spec tokens.
                self.input_batch.num_tokens[req_index] += num_spec_tokens
667

668
669
670
671
        # Add the new or resumed requests to the persistent batch.
        # The smaller empty indices are filled first.
        for req_id in req_ids_to_add:
            req_state = self.requests[req_id]
672
            self.input_batch.add_request(req_state)
673

674
675
676
677
        # Condense the batched states if there are gaps left by removed requests
        self.input_batch.condense()
        # Allow attention backend to reorder the batch, potentially
        self._may_reorder_batch(scheduler_output)
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
        # Refresh batch metadata with any pending updates. If we are in spec
        # decode + reject mode, also expand sampling metadata to token shape
        # using per-request repeat counts.
        repeat_counts: Optional[torch.Tensor] = None
        if envs.VLLM_REJECT_SAMPLE_OPT and \
                scheduler_output.scheduled_spec_decode_tokens:
            num_reqs = self.input_batch.num_reqs
            num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
            for req_id, draft_token_ids in (
                    scheduler_output.scheduled_spec_decode_tokens.items()):
                req_idx = self.input_batch.req_id_to_index.get(req_id)
                if req_idx is not None:
                    num_draft_tokens[req_idx] = len(draft_token_ids)
            repeat_counts = torch.from_numpy(num_draft_tokens).add_(1)
        self.input_batch.refresh_metadata(repeat_counts)
693

694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
    def _get_cumsum_and_arange(
        self,
        num_tokens: np.ndarray,
        cumsum_dtype: Optional[np.dtype] = None,
    ) -> tuple[np.ndarray, np.ndarray]:
        """Get the cumulative sum and batched arange of the given array.
        # E.g., [2, 5, 3] -> ([2, 7, 10], [0, 1, 0, 1, 2, 3, 4, 0, 1, 2])
        # Equivalent to but faster than:
        # np.concatenate([np.arange(n) for n in num_tokens])
        """
        # Step 1. [2, 5, 3] -> [2, 7, 10]
        cu_num_tokens = np.cumsum(num_tokens, dtype=cumsum_dtype)
        total_num_tokens = cu_num_tokens[-1]
        # Step 2. [2, 7, 10] -> [0, 0, 2, 2, 2, 2, 2, 7, 7, 7]
        cumsums_offsets = np.repeat(cu_num_tokens - num_tokens, num_tokens)
        # Step 3. [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
        arange = self.arange_np[:total_num_tokens] - cumsums_offsets

        return cu_num_tokens, arange

714
    def _prepare_inputs(
715
716
        self,
        scheduler_output: "SchedulerOutput",
717
    ) -> tuple[dict[str, Any], bool, torch.Tensor,
718
               Optional[SpecDecodeMetadata], np.ndarray]:
719
720
721
722
723
724
725
        """
        :return: tuple[
            attn_metadata: layer-to-attention_metadata mapping,
            attention_cuda_graphs: whether attention can run in cudagraph
            logits_indices, spec_decode_metadata
        ]
        """
726
727
728
729
730
731
732
        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
        assert total_num_scheduled_tokens > 0
        num_reqs = self.input_batch.num_reqs
        assert num_reqs > 0

        # OPTIMIZATION: Start copying the block table first.
        # This way, we can overlap the copy with the following CPU operations.
733
        self.input_batch.block_table.commit(num_reqs)
734
735

        # Get the number of scheduled tokens for each request.
736
737
738
739
        req_ids = self.input_batch.req_ids
        tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
        num_scheduled_tokens = np.array(tokens, dtype=np.int32)
        max_num_scheduled_tokens = max(tokens)
740
741
742

        # Get request indices.
        # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
743
744
        req_indices = np.repeat(self.arange_np[:num_reqs],
                                num_scheduled_tokens)
745

746
747
748
749
        # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
        # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
        cu_num_tokens, arange = self._get_cumsum_and_arange(
            num_scheduled_tokens)
750
751

        # Get positions.
752
        positions_np = self.positions_np[:total_num_scheduled_tokens]
753
754
755
        np.add(self.input_batch.num_computed_tokens_cpu[req_indices],
               arange,
               out=positions_np)
laibao's avatar
laibao committed
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
        # KV positions (where the KV for each scheduled token is temporarily
        # written). When KV compression is enabled, KV positions are decoupled
        # from logical positions.
        use_kv_compression = envs.VLLM_ENABLE_KV_COMPRESSION
        if use_kv_compression:
            kv_positions_np = self.kv_positions_np[:total_num_scheduled_tokens]
            np.add(self.input_batch.num_kv_tokens_cpu[req_indices],
                   arange,
                   out=kv_positions_np)
        else:
            kv_positions_np = None

        if use_kv_compression:
            prompt_ratio = envs.VLLM_KV_COMPRESSION_PROMPT_RATIO
            prompt_budget = envs.VLLM_KV_COMPRESSION_PROMPT_BUDGET
            protected_prefix = envs.VLLM_KV_COMPRESSION_PROTECTED_PREFIX
            protected_suffix = envs.VLLM_KV_COMPRESSION_PROTECTED_SUFFIX
            keep_last = envs.VLLM_KV_COMPRESSION_KEEP_LAST_TOKEN

775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
            if self.scheduler_config.chunked_prefill_enabled:
                # Scheme 3: with chunked prefill, defer compaction until after
                # the full prompt is ingested. Otherwise, the next prefill chunk
                # would attend to a truncated history and quality can collapse.
                prompt_end_np = self.kv_compression_prompt_end_np[:num_reqs]
                prompt_end_np.fill(False)
                prompt_lens_np = self.kv_compression_prompt_lens_np[:num_reqs]
                prompt_lens_np.fill(0)
                topk_keep_np = self.kv_compression_prompt_topk_keep_np[:num_reqs]
                topk_keep_np.fill(0)

                for req_idx in range(num_reqs):
                    qlen = int(num_scheduled_tokens[req_idx])
                    if qlen <= 0:
                        continue
                    base_pos = int(self.input_batch.num_computed_tokens_cpu[req_idx])
                    prompt_len = int(self.input_batch.num_prompt_tokens[req_idx])
                    end_pos = base_pos + qlen
                    ends_prompt = (base_pos < prompt_len) and (end_pos >= prompt_len)
                    if not ends_prompt:
                        continue

                    prompt_end_np[req_idx] = True
                    prompt_lens_np[req_idx] = prompt_len
                    topk_keep_np[req_idx] = compute_prompt_topk_keep_total(
laibao's avatar
laibao committed
800
801
802
803
804
805
806
807
                        prompt_len=prompt_len,
                        protected_prefix=protected_prefix,
                        protected_suffix=protected_suffix,
                        keep_last_token=keep_last,
                        prompt_ratio=prompt_ratio,
                        prompt_budget=prompt_budget,
                    )

808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
                self.kv_compression_prompt_topk_keep_max = int(
                    topk_keep_np[:num_reqs].max()) if num_reqs > 0 else 0
                self.kv_compression_needs_compaction = False
            else:
                must_keep_np = self.kv_compression_must_keep_np[
                    :total_num_scheduled_tokens]
                must_keep_np.fill(False)
                topk_budget_np = self.kv_compression_topk_budget_np[:num_reqs]
                topk_budget_np.fill(0)

                for req_idx in range(num_reqs):
                    qlen = int(num_scheduled_tokens[req_idx])
                    if qlen <= 0:
                        continue
                    start = 0 if req_idx == 0 else int(
                        cu_num_tokens[req_idx - 1])
                    end = int(cu_num_tokens[req_idx])
                    assert end - start == qlen

                    base_pos = int(self.input_batch.num_computed_tokens_cpu[req_idx])
                    prompt_len = int(self.input_batch.num_prompt_tokens[req_idx])
                    end_pos = base_pos + qlen
                    pos = base_pos + np.arange(qlen, dtype=np.int64)

                    prompt_mask = pos < prompt_len
                    # Decode tokens are always kept.
                    must_keep = ~prompt_mask

                    if np.any(prompt_mask):
                        suffix_start = max(prompt_len - protected_suffix, 0)
                        must_keep |= prompt_mask & (pos < protected_prefix)
                        must_keep |= prompt_mask & (pos >= suffix_start)
                        if keep_last:
                            last = prompt_len - 1
                            if base_pos <= last < end_pos:
                                must_keep[last - base_pos] = True

                        topk_budget_np[req_idx] = compute_topk_budget_step(
                            prompt_len=prompt_len,
                            start_pos=base_pos,
                            end_pos=end_pos,
                            protected_prefix=protected_prefix,
                            protected_suffix=protected_suffix,
                            keep_last_token=keep_last,
                            prompt_ratio=prompt_ratio,
                            prompt_budget=prompt_budget,
                        )

                    must_keep_np[start:end] = must_keep
laibao's avatar
laibao committed
857

858
859
860
861
862
863
                # Decode-only fast path: if all scheduled tokens are
                # unconditionally kept and there is no Top-K budget, KV
                # compaction is a no-op and we can skip score/topk/dst entirely
                # in the attention backend.
                self.kv_compression_needs_compaction = (not must_keep_np.all(
                )) or (topk_budget_np > 0).any()
laibao's avatar
laibao committed
864
865
        else:
            self.kv_compression_needs_compaction = False
866

867
868
        # Calculate M-RoPE positions.
        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
869
        if self.uses_mrope:
870
871
            self._calc_mrope_positions(scheduler_output)

872
873
874
875
        # Get token indices.
        # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
        # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
        # where M is the max_model_len.
876
877
        token_indices = (positions_np +
                         req_indices * self.input_batch.token_ids_cpu.shape[1])
lizhigong's avatar
lizhigong committed
878
        
879
880
881
882
        # NOTE(woosuk): We use torch.index_select instead of np.take here
        # because torch.index_select is much faster than np.take for large
        # tensors.
        torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(),
883
                           0,
884
885
                           torch.from_numpy(token_indices),
                           out=self.input_ids_cpu[:total_num_scheduled_tokens])
886

887
888
889
890
891
892
        # Calculate the slot mapping for each KV cache group.
        for kv_cache_group_id, kv_cache_group_spec in enumerate(
                self.kv_cache_config.kv_cache_groups):
            block_size = kv_cache_group_spec.kv_cache_spec.block_size
            block_table: BlockTable = self.input_batch.block_table[
                kv_cache_group_id]
laibao's avatar
laibao committed
893
894
            slot_positions_np = (kv_positions_np
                                 if use_kv_compression else positions_np)
895
896
897
898
899
900
901
902
            # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
            # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
            # where K is the max_num_blocks_per_req and the block size is 2.
            # NOTE(woosuk): We can't simply use `token_indices // block_size`
            # here because M (max_model_len) is not necessarily divisible by
            # block_size.
            block_table_indices = (
                req_indices * block_table.max_num_blocks_per_req +
laibao's avatar
laibao committed
903
                slot_positions_np // block_size)
904
905
906
            block_table_cpu = block_table.get_cpu_tensor()
            block_numbers = block_table_cpu.flatten(
            )[block_table_indices].numpy()
laibao's avatar
laibao committed
907
            block_offsets = slot_positions_np % block_size
908
909
910
911
            np.add(
                block_numbers * block_size,
                block_offsets,
                out=block_table.slot_mapping_np[:total_num_scheduled_tokens])
912
913

        # Prepare the attention metadata.
914
        self.query_start_loc_np[0] = 0
915
        self.query_start_loc_np[1:num_reqs + 1] = cu_num_tokens
916

laibao's avatar
laibao committed
917
918
919
920
921
922
923
924
        if use_kv_compression:
            self.seq_lens_np[:num_reqs] = (
                self.input_batch.num_kv_tokens_cpu[:num_reqs] +
                num_scheduled_tokens)
        else:
            self.seq_lens_np[:num_reqs] = (
                self.input_batch.num_computed_tokens_cpu[:num_reqs] +
                num_scheduled_tokens)
925
926
927
928

        # Copy the tensors to the GPU.
        self.input_ids[:total_num_scheduled_tokens].copy_(
            self.input_ids_cpu[:total_num_scheduled_tokens], non_blocking=True)
929
        if self.uses_mrope:
930
931
932
933
934
935
936
937
938
            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
            self.mrope_positions[:, :total_num_scheduled_tokens].copy_(
                self.mrope_positions_cpu[:, :total_num_scheduled_tokens],
                non_blocking=True)
        else:
            # Common case (1D positions)
            self.positions[:total_num_scheduled_tokens].copy_(
                self.positions_cpu[:total_num_scheduled_tokens],
                non_blocking=True)
939

940
941
942
943
        self.query_start_loc[:num_reqs + 1].copy_(
            self.query_start_loc_cpu[:num_reqs + 1], non_blocking=True)
        self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
                                       non_blocking=True)
laibao's avatar
laibao committed
944
        if use_kv_compression:
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
            if self.scheduler_config.chunked_prefill_enabled:
                self.kv_compression_prompt_end[:num_reqs].copy_(
                    self.kv_compression_prompt_end_cpu[:num_reqs],
                    non_blocking=True,
                )
                self.kv_compression_prompt_lens[:num_reqs].copy_(
                    self.kv_compression_prompt_lens_cpu[:num_reqs],
                    non_blocking=True,
                )
                self.kv_compression_prompt_topk_keep[:num_reqs].copy_(
                    self.kv_compression_prompt_topk_keep_cpu[:num_reqs],
                    non_blocking=True,
                )
            elif self.kv_compression_needs_compaction:
                self.kv_compression_must_keep[:total_num_scheduled_tokens].copy_(
                    self.kv_compression_must_keep_cpu[:total_num_scheduled_tokens],
                    non_blocking=True,
                )
                self.kv_compression_topk_budget[:num_reqs].copy_(
                    self.kv_compression_topk_budget_cpu[:num_reqs],
                    non_blocking=True,
                )
967
968
969

        # Fill unused with -1. Needed for reshape_and_cache
        self.seq_lens[num_reqs:].fill_(0)
970
971
972
973
        # Note: pad query_start_loc to be non-decreasing, as kernels
        # like FlashAttention requires that
        self.query_start_loc[num_reqs + 1:].fill_(
            self.query_start_loc_cpu[num_reqs].item())
974
975
976
977

        query_start_loc = self.query_start_loc[:num_reqs + 1]
        seq_lens = self.seq_lens[:num_reqs]

978
        common_attn_metadata = CommonAttentionMetadata(
979
980
            query_start_loc=query_start_loc,
            seq_lens=seq_lens,
981
            # seq_lens_tensor=seq_lens_tensor,
982
983
984
985
            num_reqs=num_reqs,
            num_actual_tokens=total_num_scheduled_tokens,
            max_query_len=max_num_scheduled_tokens,
        )
986

987
        attn_metadata: dict[str, Any] = {}
988
989
990
991
992
993
994
        # Prepare the attention metadata for each KV cache group and make layers
        # in the same group share the same metadata.
        for kv_cache_group_id, kv_cache_group_spec in enumerate(
                self.kv_cache_config.kv_cache_groups):

            # Prepare for cascade attention if enabled & beneficial.
            common_prefix_len = 0
995
            builder = self.attn_metadata_builders[kv_cache_group_id]
996
997
998
            if self.cascade_attn_enabled:
                common_prefix_len = self._compute_cascade_attn_prefix_len(
                    num_scheduled_tokens,
999
1000
1001
                    scheduler_output.
                    num_common_prefix_blocks[kv_cache_group_id],
                    kv_cache_group_spec.kv_cache_spec,
1002
                    builder,
1003
                )
1004

1005
1006
1007
1008
1009
            attn_metadata_i = (builder.build(
                common_prefix_len=common_prefix_len,
                common_attn_metadata=common_attn_metadata,
            ))

1010
1011
            for layer_name in kv_cache_group_spec.layer_names:
                attn_metadata[layer_name] = attn_metadata_i
1012

1013
1014
1015
1016
        attention_cuda_graphs = all(
            b.can_run_in_cudagraph(common_attn_metadata)
            for b in self.attn_metadata_builders)

1017
1018
        use_spec_decode = len(
            scheduler_output.scheduled_spec_decode_tokens) > 0
1019
        if not use_spec_decode:
1020
1021
1022
1023
1024
            # NOTE(woosuk): Due to chunked prefills, the batch may contain
            # partial requests. While we should not sample any token
            # from these partial requests, we do so for simplicity.
            # We will ignore the sampled tokens from the partial requests.
            # TODO: Support prompt logprobs.
1025
            logits_indices = query_start_loc[1:] - 1
1026
1027
1028
1029
1030
1031
            spec_decode_metadata = None
        else:
            # Get the number of draft tokens for each request.
            # Iterate over the dictionary rather than all requests since not all
            # requests have draft tokens.
            num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
王敏's avatar
王敏 committed
1032

1033
1034
1035
1036
1037
            for req_id, draft_token_ids in (
                    scheduler_output.scheduled_spec_decode_tokens.items()):
                req_idx = self.input_batch.req_id_to_index[req_id]
                num_draft_tokens[req_idx] = len(draft_token_ids)

王敏's avatar
王敏 committed
1038
1039
1040
1041
            spec_decode_ids = None
            if envs.VLLM_REJECT_SAMPLE_OPT:
                spec_decode_ids = scheduler_output.scheduled_spec_decode_tokens.keys()

1042
            spec_decode_metadata = self._calc_spec_decode_metadata(
王敏's avatar
王敏 committed
1043
                num_draft_tokens, cu_num_tokens, spec_decode_ids)
1044
            logits_indices = spec_decode_metadata.logits_indices
1045

1046
1047
1048
1049
        # Hot-Swap lora model
        if self.lora_config:
            self.set_active_loras(self.input_batch, num_scheduled_tokens)

1050
        return (attn_metadata, attention_cuda_graphs, logits_indices,
1051
                spec_decode_metadata, num_scheduled_tokens)
1052

1053
1054
1055
1056
    def _compute_cascade_attn_prefix_len(
        self,
        num_scheduled_tokens: np.ndarray,
        num_common_prefix_blocks: int,
1057
1058
        kv_cache_spec: KVCacheSpec,
        attn_metadata_builder: AttentionMetadataBuilder,
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
    ) -> int:
        """Compute the length of the common prefix for cascade attention.

        NOTE(woosuk): The common prefix length returned by this function
        represents the length used specifically for cascade attention, not the
        actual number of tokens shared between requests. When cascade attention
        is disabled (use_cascade=False), this function returns 0 even if
        requests share common tokens. Additionally, the common prefix length is
        truncated to a multiple of the block size and may be further truncated
        due to implementation details explained below.

        Args:
            num_scheduled_tokens: Number of tokens scheduled per request.
            num_common_prefix_blocks: Number of shared KV cache blocks.

        Returns:
            int: Length of common prefix in tokens.
        """
1077
        common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size
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
        if common_prefix_len == 0:
            # Common case.
            return 0

        # NOTE(woosuk): Cascade attention uses two attention kernels: one
        # for the common prefix and the other for the rest. For the first
        # kernel, we concatenate all the query tokens (possibly from
        # different requests) and treat them as if they are from the same
        # request. Then, we use bi-directional attention to process the
        # common prefix in the KV cache. Importantly, this means that the
        # first kernel does not do any masking.

        # Consider the following example:
        # Request 1's input query: [D, E, X]
        # Request 1's kv cache: [A, B, C, D, E, X]
        # Request 1's num_computed_tokens: 3 (i.e., [A, B, C])
        # Request 2's input query: [E, Y]
        # Request 2's kv cache: [A, B, C, D, E, Y]
        # Request 2's num_computed_tokens: 4 (i.e., [A, B, C, D])

        # If we use [A, B, C, D, E] as the common prefix, then the
        # first kernel will compute the bi-directional attention between
        # input query [D, E, X, E, Y] and common prefix [A, B, C, D, E].
        # However, this is wrong because D in Request 1 should not attend to
        # E in the common prefix (i.e., we need masking).
        # To avoid this, [A, B, C, D] should be the common prefix.
        # That is, the common prefix should be capped by the minimum
        # num_computed_tokens among the requests, and plus one to include
        # the first token of the query.

        # In practice, we use [A, B, C] as the common prefix, instead of
        # [A, B, C, D] (i.e., the common prefix is capped by the minimum
        # num_computed_tokens, without plus one).
        # This is because of an implementation detail: We want to always
        # use two kernels for cascade attention. Let's imagine:
        # Request 3's input query: [D]
        # Request 3's kv cache: [A, B, C, D]
1115
        # Request 3's num_computed_tokens: 3 (i.e., [A, B, C])
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
        # If we use [A, B, C, D] as the common prefix for Request 1-3,
        # then Request 3 will be processed only by the first kernel,
        # and the second kernel will get an empty input. While this is not
        # a fundamental problem, our current implementation does not support
        # this case.
        num_reqs = len(num_scheduled_tokens)
        common_prefix_len = min(
            common_prefix_len,
            self.input_batch.num_computed_tokens_cpu[:num_reqs].min())
        # common_prefix_len should be a multiple of the block size.
1126
1127
1128
1129
1130
1131
1132
        common_prefix_len = (common_prefix_len // kv_cache_spec.block_size *
                             kv_cache_spec.block_size)
        use_sliding_window = (isinstance(kv_cache_spec, SlidingWindowSpec) or
                              (isinstance(kv_cache_spec, FullAttentionSpec)
                               and kv_cache_spec.sliding_window is not None))
        assert isinstance(kv_cache_spec, AttentionSpec)
        use_cascade = attn_metadata_builder.use_cascade_attention(
1133
1134
1135
            common_prefix_len=common_prefix_len,
            query_lens=num_scheduled_tokens,
            num_query_heads=self.num_query_heads,
1136
            num_kv_heads=kv_cache_spec.num_kv_heads,
1137
            use_alibi=self.use_alibi,
1138
            use_sliding_window=use_sliding_window,
1139
1140
1141
1142
            num_sms=self.num_sms,
        )
        return common_prefix_len if use_cascade else 0

1143
1144
    def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"):
        mrope_pos_ptr = 0
1145
        for index, req_id in enumerate(self.input_batch.req_ids):
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
            req = self.requests[req_id]
            assert req.mrope_positions is not None

            num_computed_tokens = \
                self.input_batch.num_computed_tokens_cpu[index]
            num_scheduled_tokens = \
                scheduler_output.num_scheduled_tokens[req_id]
            num_prompt_tokens = len(req.prompt_token_ids)

            if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens:
                prompt_part_len = max(0,
                                      num_prompt_tokens - num_computed_tokens)
                completion_part_len = max(
                    0, num_scheduled_tokens - prompt_part_len)
            else:
                prompt_part_len = num_scheduled_tokens
                completion_part_len = 0

            assert num_scheduled_tokens == prompt_part_len + completion_part_len

            if prompt_part_len > 0:
                # prompt's mrope_positions are pre-computed
                dst_start = mrope_pos_ptr
                dst_end = mrope_pos_ptr + prompt_part_len
                src_start = num_computed_tokens
                src_end = num_computed_tokens + prompt_part_len

                self.mrope_positions_cpu[:, dst_start:dst_end] = \
                    req.mrope_positions[:,src_start:src_end]

                mrope_pos_ptr += prompt_part_len

            if completion_part_len > 0:
                # compute completion's mrope_positions on-the-fly
                dst_start = mrope_pos_ptr
                dst_end = mrope_pos_ptr + completion_part_len

1183
1184
1185
1186
1187
1188
1189
                MRotaryEmbedding.get_next_input_positions_tensor(
                    out=self.mrope_positions_np,
                    out_offset=dst_start,
                    mrope_position_delta=req.mrope_position_delta,
                    context_len=num_computed_tokens + prompt_part_len,
                    num_new_tokens=completion_part_len,
                )
1190
1191
1192

                mrope_pos_ptr += completion_part_len

1193
1194
    def _calc_spec_decode_metadata(
        self,
1195
1196
        num_draft_tokens: np.ndarray,
        cu_num_scheduled_tokens: np.ndarray,
王敏's avatar
王敏 committed
1197
        spec_decode_ids: Optional[list[str]] = None
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
    ) -> SpecDecodeMetadata:
        # Inputs:
        # cu_num_scheduled_tokens:  [  4, 104, 107, 207, 209]
        # num_draft_tokens:         [  3,   0,   2,   0,   1]
        # Outputs:
        # cu_num_draft_tokens:      [  3,   3,   5,   5,   6]
        # logits_indices:           [  0,   1,   2,   3, 103, 104, 105, 106,
        #                            206, 207, 208]
        # target_logits_indices:    [  0,   1,   2,   5,   6,   9]
        # bonus_logits_indices:     [  3,   4,   7,   8,  10]

        # Compute the logits indices.
        # [4, 1, 3, 1, 2]
        num_sampled_tokens = num_draft_tokens + 1
1212
1213
1214
1215
1216
1217

        # Step 1. cu_num_sampled_tokens: [4, 5, 8, 9, 11]
        # arange: [0, 1, 2, 3, 0, 0, 1, 2, 0, 0, 1]
        cu_num_sampled_tokens, arange = self._get_cumsum_and_arange(
            num_sampled_tokens, cumsum_dtype=np.int32)
        # Step 2. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]
1218
1219
        logits_indices = np.repeat(
            cu_num_scheduled_tokens - num_sampled_tokens, num_sampled_tokens)
1220
        # Step 3. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
1221
1222
1223
1224
1225
1226
        logits_indices += arange

        # Compute the bonus logits indices.
        bonus_logits_indices = cu_num_sampled_tokens - 1

        # Compute the draft logits indices.
1227
1228
1229
1230
        # cu_num_draft_tokens: [3, 3, 5, 5, 6]
        # arange: [0, 1, 2, 0, 1, 0]
        cu_num_draft_tokens, arange = self._get_cumsum_and_arange(
            num_draft_tokens, cumsum_dtype=np.int32)
1231
1232
1233
1234
1235
1236
        # [0, 0, 0, 5, 5, 9]
        target_logits_indices = np.repeat(
            cu_num_sampled_tokens - num_sampled_tokens, num_draft_tokens)
        # [0, 1, 2, 5, 6, 9]
        target_logits_indices += arange

1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
        if envs.VLLM_ZERO_OVERHEAD:
            cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).pin_memory().to(
                self.device, non_blocking=True)
            logits_indices = torch.from_numpy(logits_indices).pin_memory().to(self.device,
                                                                non_blocking=True)
            target_logits_indices = torch.from_numpy(target_logits_indices).pin_memory().to(
                self.device, non_blocking=True)
            bonus_logits_indices = torch.from_numpy(bonus_logits_indices).pin_memory().to(
                self.device, non_blocking=True)
        else:
            # TODO: Optimize the CPU -> GPU copy.
            cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(
                self.device, non_blocking=True)
            logits_indices = torch.from_numpy(logits_indices).to(self.device,
                                                                non_blocking=True)
            target_logits_indices = torch.from_numpy(target_logits_indices).to(
                self.device, non_blocking=True)
            bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to(
                self.device, non_blocking=True)
1256

1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
        # Compute the draft token ids.
        # draft_token_indices:      [  1,   2,   3, 105, 106, 208]
        draft_token_ids = self.input_ids[logits_indices]
        draft_token_ids = draft_token_ids[target_logits_indices + 1]

        metadata = SpecDecodeMetadata(
            draft_token_ids=draft_token_ids,
            num_draft_tokens=num_draft_tokens.tolist(),
            cu_num_draft_tokens=cu_num_draft_tokens,
            target_logits_indices=target_logits_indices,
            bonus_logits_indices=bonus_logits_indices,
            logits_indices=logits_indices,
王敏's avatar
王敏 committed
1269
            spec_decode_ids=spec_decode_ids,
1270
1271
1272
        )
        return metadata

1273
    def _execute_mm_encoder(self, scheduler_output: "SchedulerOutput"):
1274
1275
1276
1277
1278
        scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
        if not scheduled_encoder_inputs:
            return

        # Batch the multi-modal inputs.
1279
1280
        mm_inputs = list[MultiModalKwargs]()
        req_ids_pos = list[tuple[str, int, PlaceholderRange]]()
1281
1282
        for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
            req_state = self.requests[req_id]
1283
1284
1285
1286
1287

            for mm_input_id in encoder_input_ids:
                mm_inputs.append(req_state.mm_inputs[mm_input_id])
                req_ids_pos.append(
                    (req_id, mm_input_id, req_state.mm_positions[mm_input_id]))
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299

        # Batch mm inputs as much as we can: if a request in the batch has
        # multiple modalities or a different modality than the previous one,
        # we process it separately to preserve item order.
        # FIXME(ywang96): This is a hacky way to deal with multiple modalities
        # in the same batch while still being able to benefit from batching
        # multimodal inputs. The proper solution should be reordering the
        # encoder outputs.
        grouped_mm_inputs_list = group_mm_inputs_by_modality(mm_inputs)

        encoder_outputs = []
        for grouped_mm_inputs in grouped_mm_inputs_list:
1300
1301
            batched_mm_inputs = MultiModalKwargs.batch(
                grouped_mm_inputs, pin_memory=self.pin_memory)
1302
1303
1304
1305
            batched_mm_inputs = MultiModalKwargs.as_kwargs(
                batched_mm_inputs,
                device=self.device,
            )
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316

            # Run the encoder.
            # `curr_group_outputs` is either of the following:
            # 1. A tensor of shape (num_items, feature_size, hidden_size)
            # in case feature_size is fixed across all multimodal items.
            # 2. A list or tuple (length: num_items) of tensors, each of shape
            # (feature_size, hidden_size) in case the feature size is dynamic
            # depending on the input multimodal items.
            curr_group_outputs = self.model.get_multimodal_embeddings(
                **batched_mm_inputs)

1317
1318
1319
1320
1321
            sanity_check_mm_encoder_outputs(
                curr_group_outputs,
                expected_num_items=len(grouped_mm_inputs),
            )

1322
1323
            for output in curr_group_outputs:
                encoder_outputs.append(output)
1324
1325

        # Cache the encoder outputs.
1326
1327
1328
1329
        for (req_id, input_id, pos_info), output in zip(
                req_ids_pos,
                encoder_outputs,
        ):
1330
1331
1332
            if req_id not in self.encoder_cache:
                self.encoder_cache[req_id] = {}

1333
1334
1335
1336
1337
1338
            self.encoder_cache[req_id][input_id] = scatter_mm_placeholders(
                output,
                is_embed=pos_info.is_embed,
            )

    def _gather_mm_embeddings(
1339
1340
        self,
        scheduler_output: "SchedulerOutput",
1341
    ) -> list[torch.Tensor]:
1342
        mm_embeds: list[torch.Tensor] = []
1343
        for req_id in self.input_batch.req_ids:
1344
1345
1346
1347
1348
1349
            num_scheduled_tokens = scheduler_output.num_scheduled_tokens[
                req_id]
            req_state = self.requests[req_id]
            num_computed_tokens = req_state.num_computed_tokens
            mm_positions = req_state.mm_positions
            for i, pos_info in enumerate(mm_positions):
1350
1351
                start_pos = pos_info.offset
                num_encoder_tokens = pos_info.length
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372

                # The encoder output is needed if the two ranges overlap:
                # [num_computed_tokens,
                #  num_computed_tokens + num_scheduled_tokens) and
                # [start_pos, start_pos + num_encoder_tokens)
                if start_pos >= num_computed_tokens + num_scheduled_tokens:
                    # The encoder output is not needed in this step.
                    break
                if start_pos + num_encoder_tokens <= num_computed_tokens:
                    # The encoder output is already processed and stored
                    # in the decoder's KV cache.
                    continue

                start_idx = max(num_computed_tokens - start_pos, 0)
                end_idx = min(
                    num_computed_tokens - start_pos + num_scheduled_tokens,
                    num_encoder_tokens)
                assert start_idx < end_idx
                assert req_id in self.encoder_cache
                assert i in self.encoder_cache[req_id]
                encoder_output = self.encoder_cache[req_id][i]
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382

                if (is_embed := pos_info.is_embed) is not None:
                    is_embed = is_embed[start_idx:end_idx]

                mm_embeds_item = gather_mm_placeholders(
                    encoder_output[start_idx:end_idx],
                    is_embed=is_embed,
                )
                mm_embeds.append(mm_embeds_item)
        return mm_embeds
1383

1384
1385
1386
    def get_model(self) -> nn.Module:
        return self.model

1387
1388
1389
1390
1391
1392
1393
1394
1395
    def apply_grammar_bitmask(
        self,
        scheduler_output: "SchedulerOutput",
        logits: torch.Tensor,
    ):
        grammar_bitmask = scheduler_output.grammar_bitmask
        if grammar_bitmask is None:
            return

1396
1397
1398
1399
1400
1401
1402
1403
1404
        # We receive the structured output bitmask from the scheduler,
        # compacted to contain bitmasks only for structured output requests.
        # The order of the requests in the bitmask is not guaranteed to be the
        # same as the order of the requests in the gpu runner's batch. We need
        # to sort the bitmask to match the order of the requests used here.

        # Get the batch indices of the structured output requests.
        # Keep track of the number of speculative tokens scheduled for every
        # request in the batch, as the logit indices are offset by this amount.
1405
        struct_out_req_batch_indices: dict[str, int] = {}
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
        cumulative_offset = 0
        seq = sorted(self.input_batch.req_id_to_index.items(),
                     key=lambda x: x[1])
        for req_id, batch_index in seq:
            logit_index = batch_index + cumulative_offset
            cumulative_offset += len(
                scheduler_output.scheduled_spec_decode_tokens.get(req_id, []))
            if req_id in scheduler_output.structured_output_request_ids:
                struct_out_req_batch_indices[req_id] = logit_index

        out_indices = []

        # Reorder the bitmask to match the order of the requests in the batch.
        sorted_bitmask = np.zeros_like(grammar_bitmask,
                                       shape=(logits.shape[0],
                                              grammar_bitmask.shape[1]))
        cumulative_index = 0
        seq = sorted(scheduler_output.structured_output_request_ids.items(),
                     key=lambda x: x[1])
        for req_id, _ in seq:
            logit_index = struct_out_req_batch_indices[req_id]
            num_spec_tokens = len(
                scheduler_output.scheduled_spec_decode_tokens.get(req_id, []))
            for i in range(1 + num_spec_tokens):
                sorted_bitmask[logit_index + i] = \
                    grammar_bitmask[cumulative_index + i]
                out_indices.append(logit_index + i)
            cumulative_index += 1 + num_spec_tokens
        grammar_bitmask = sorted_bitmask
1435

1436
1437
        # Serialization of np.ndarray is much more efficient than a tensor,
        # so we receive it in that format.
1438
1439
        grammar_bitmask = torch.from_numpy(grammar_bitmask)

1440
1441
1442
1443
        # Force use of the torch.compile implementation from xgrammar to work
        # around issues with the Triton kernel in concurrent structured output
        # scenarios. See PR #19565 and issues #19493, #18376 for details.
        xgr_torch_compile.apply_token_bitmask_inplace_torch_compile(
1444
1445
            logits,
            grammar_bitmask.to(self.device, non_blocking=True),
1446
            indices=out_indices,
1447
1448
        )

1449
1450
1451
1452
1453
1454
1455
    def sync_and_slice_intermediate_tensors(
            self, num_tokens: int, intermediate_tensors: IntermediateTensors,
            sync_self: bool) -> IntermediateTensors:

        assert self.intermediate_tensors is not None

        tp = self.vllm_config.parallel_config.tensor_parallel_size
1456
        enabled_sp = self.compilation_config.pass_config. \
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
            enable_sequence_parallelism
        if enabled_sp:
            # When sequence parallelism is enabled, we always pad num_tokens
            # to be a multiple of tensor_parallel_size (tp) earlier
            assert num_tokens % tp == 0
        is_residual_scattered = tp > 1 and enabled_sp \
            and num_tokens % tp == 0

        # When sequence parallelism is enabled, the "residual" tensor is sharded
        # across tensor parallel ranks, so each rank only needs its own slice.
        if sync_self:
            assert intermediate_tensors is not None
            for k, v in intermediate_tensors.items():
                is_scattered = "residual" and is_residual_scattered
                copy_len = num_tokens // tp if is_scattered else \
                    num_tokens
                self.intermediate_tensors[k][:copy_len].copy_(
                    v[:copy_len], non_blocking=True)

        return IntermediateTensors({
            k:
            v[:num_tokens // tp]
            if k == "residual" and is_residual_scattered else v[:num_tokens]
            for k, v in self.intermediate_tensors.items()
        })

1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
    def eplb_step(self,
                  is_dummy: bool = False,
                  is_profile: bool = False) -> None:
        """
        Step for the EPLB (Expert Parallelism Load Balancing) state.
        """
        if not self.parallel_config.enable_eplb:
            return

        assert self.eplb_state is not None
        assert is_mixture_of_experts(self.model)
        self.eplb_state.step(
            self.model,
            is_dummy,
            is_profile,
            log_stats=self.parallel_config.eplb_log_balancedness,
        )

1501
1502
    def get_dp_padding(self,
                       num_tokens: int) -> tuple[int, Optional[torch.Tensor]]:
1503
1504
        dp_size = self.vllm_config.parallel_config.data_parallel_size
        dp_rank = self.vllm_config.parallel_config.data_parallel_rank
1505
1506
1507
1508
1509
1510
1511
1512

        # For DP: Don't pad when setting enforce_eager.
        # This lets us set enforce_eager on the prefiller in a P/D setup and
        # still use CUDA graphs (enabled by this padding) on the decoder.
        #
        # TODO(tms) : There are many cases where padding is enabled for
        # prefills, causing unnecessary and excessive padding of activations.

yangql's avatar
yangql committed
1513
        if dp_size == 1 or self.vllm_config.model_config.enforce_eager:
yangql's avatar
yangql committed
1514
1515
1516
1517
            # auto 
            if not envs.VLLM_ALL2ALL_BACKEND == "deepep_auto":
                # Early exit.
                return 0, None
1518
1519
1520
1521
1522


        num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
            num_tokens, dp_size, dp_rank)
        max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp).item()
1523
1524
1525
1526
1527
        num_tokens_after_padding = torch.tensor([max_tokens_across_dp_cpu] *
                                                dp_size,
                                                device="cpu",
                                                dtype=torch.int32)
        return max_tokens_across_dp_cpu - num_tokens, num_tokens_after_padding
1528

1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
    def _pool(
        self,
        hidden_states: torch.Tensor,
        num_scheduled_tokens: int,
        num_scheduled_tokens_np: np.ndarray,
        finished_sending: Optional[set[str]],
        finished_recving: Optional[set[str]],
    ) -> ModelRunnerOutput:
        assert self.input_batch.num_reqs ==\
            len(self.input_batch.pooling_params), \
        "Either all or none of the requests in" \
        " a batch must be pooling request"

        extracted_hidden_states = list(
            torch.split(hidden_states[:num_scheduled_tokens],
                        num_scheduled_tokens_np.tolist()))

        pooling_metadata = self.input_batch.pooling_metadata

        raw_pooler_output = self.model.pooler(
            hidden_states=extracted_hidden_states,
            pooling_metadata=pooling_metadata)

        pooler_output: list[Optional[torch.Tensor]] = []
        seq_lens = self.seq_lens[:self.input_batch.num_reqs]
        for raw_output, seq_len, prompt_len in zip(
                raw_pooler_output, seq_lens, pooling_metadata.prompt_lens):

            if seq_len == prompt_len:
                pooler_output.append(raw_output.data.cpu())
            else:
                pooler_output.append(None)

        return ModelRunnerOutput(
            req_ids=self.input_batch.req_ids,
            req_id_to_index=self.input_batch.req_id_to_index,
            sampled_token_ids=[],
            spec_token_ids=None,
            logprobs=None,
            prompt_logprobs_dict={},
            pooler_output=pooler_output,
            finished_sending=finished_sending,
            finished_recving=finished_recving,
        )

1574
1575
1576
1577
    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: "SchedulerOutput",
1578
        intermediate_tensors: Optional[IntermediateTensors] = None,
1579
    ) -> Union[ModelRunnerOutput, IntermediateTensors]:
1580
        self._update_states(scheduler_output)
1581
        if not scheduler_output.total_num_scheduled_tokens:
Robert Shaw's avatar
Robert Shaw committed
1582
1583
1584
1585
1586
            if not has_kv_transfer_group():
                # Return empty ModelRunnerOutput if there's no work to do.
                return EMPTY_MODEL_RUNNER_OUTPUT

            return self.kv_connector_no_forward(scheduler_output)
1587
1588

        # Prepare the decoder inputs.
1589
        (attn_metadata, attention_cuda_graphs, logits_indices,
1590
1591
         spec_decode_metadata,
         num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
1592
        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
        
        # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim.
        if self.ep_sp:
            num_input_tokens = round_up(num_scheduled_tokens, self.tp_size)
            if (self.use_cuda_graph
                    and num_input_tokens <= self.cudagraph_batch_sizes[-1]):
                # Use piecewise CUDA graphs.
                # Add padding to the batch size.
                num_input_tokens = self.vllm_config.pad_for_cudagraph(
                    num_input_tokens)
1603
        else:
1604
1605
1606
1607
1608
1609
            if (self.use_cuda_graph
                    and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
                # Use piecewise CUDA graphs.
                # Add padding to the batch size.
                num_input_tokens = self.vllm_config.pad_for_cudagraph(
                    num_scheduled_tokens)
1610
            else:
1611
1612
1613
1614
1615
1616
1617
1618
1619
                # Eager mode.
                # Pad tokens to multiple of tensor_parallel_size when
                # enabled collective fusion for SP
                tp_size = self.vllm_config.parallel_config.tensor_parallel_size
                if self.compilation_config.pass_config. \
                        enable_sequence_parallelism and tp_size > 1:
                    num_input_tokens = round_up(num_scheduled_tokens, tp_size)
                else:
                    num_input_tokens = num_scheduled_tokens
王敏's avatar
王敏 committed
1620

1621
        # Padding for DP
1622
1623
        num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
        num_input_tokens += num_pad
1624

1625
1626
1627
1628
1629
1630
1631
1632
1633
        # _prepare_inputs may reorder the batch, so we must gather multi
        # modal outputs after that to ensure the correct order
        if self.is_multimodal_model:
            # Run the multimodal encoder if any.
            self._execute_mm_encoder(scheduler_output)
            mm_embeds = self._gather_mm_embeddings(scheduler_output)
        else:
            mm_embeds = []

1634
        if self.is_multimodal_model and get_pp_group().is_first_rank:
1635
1636
1637
1638
            # NOTE(woosuk): To unify token ids and soft tokens (vision
            # embeddings), we always use embeddings (rather than token ids)
            # as input to the multimodal model, even when the input is text.
            input_ids = self.input_ids[:num_scheduled_tokens]
1639
            if mm_embeds:
1640
                inputs_embeds = self.model.get_input_embeddings(
1641
                    input_ids, mm_embeds)
1642
1643
1644
1645
1646
1647
            else:
                inputs_embeds = self.model.get_input_embeddings(input_ids)
            # TODO(woosuk): Avoid the copy. Optimize.
            self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds)
            inputs_embeds = self.inputs_embeds[:num_input_tokens]
            input_ids = None
1648
        else:
1649
1650
1651
1652
1653
1654
            # For text-only models, we use token ids as input.
            # While it is possible to use embeddings as input just like the
            # multimodal models, it is not desirable for performance since
            # then the embedding layer is not included in the CUDA graph.
            input_ids = self.input_ids[:num_input_tokens]
            inputs_embeds = None
1655
1656
1657
1658
        if self.uses_mrope:
            positions = self.mrope_positions[:, :num_input_tokens]
        else:
            positions = self.positions[:num_input_tokens]
1659

1660
1661
1662
        if get_pp_group().is_first_rank:
            intermediate_tensors = None
        else:
1663
1664
            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
                num_input_tokens, intermediate_tensors, True)
1665

1666
1667
1668
1669
        # Some attention backends only support CUDA Graphs in pure decode.
        # If attention doesn't support CUDA Graphs for this batch, but we
        # compiled with full CUDA graphs, we have to skip them entirely.
        skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
maxiao1's avatar
maxiao1 committed
1670
        if envs.VLLM_ENABLE_TBO and scheduler_output.total_num_scheduled_tokens >= envs.VLLM_TBO_MIN_TOKENS: 
王敏's avatar
王敏 committed
1671
1672
1673
            model_output, finished_sending, finished_recving = \
                 tbo_split_and_execute_model(self, attn_metadata, num_input_tokens,
                                             num_tokens_across_dp, input_ids, positions,
1674
1675
                                             inputs_embeds, scheduler_output, intermediate_tensors,
                                             skip_cuda_graphs)
maxiao1's avatar
maxiao1 committed
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
            
        elif envs.VLLM_P2P_ASYNC:
            self.p2p_event.record()
            current_stream = torch.cuda.current_stream()
            with torch.cuda.stream(self.p2p_stream):
                self.p2p_stream.wait_event(self.p2p_event)
                with set_forward_context(
                    attn_metadata,
                    self.vllm_config,
                    num_tokens=num_input_tokens,
                    num_tokens_across_dp=num_tokens_across_dp,
                    skip_cuda_graphs=skip_cuda_graphs,
                ):
                    self.maybe_setup_kv_connector(scheduler_output)

                    model_output = self.model(
                        input_ids=input_ids,
                        positions=positions,
                        intermediate_tensors=intermediate_tensors,
                        inputs_embeds=inputs_embeds,
                    )
                    
                    self.maybe_wait_for_kv_save()
                    finished_sending, finished_recving = (
                        self.get_finished_kv_transfers(scheduler_output))
                self.p2p_event.record()
            current_stream.wait_event(self.p2p_event)
                                    
王敏's avatar
王敏 committed
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
        else:
            # Run the model.
            # Use persistent buffers for CUDA graphs.
            with set_forward_context(
                    attn_metadata,
                    self.vllm_config,
                    num_tokens=num_input_tokens,
                    num_tokens_across_dp=num_tokens_across_dp,
                    skip_cuda_graphs=skip_cuda_graphs,
            ):
                self.maybe_setup_kv_connector(scheduler_output)

                model_output = self.model(
                    input_ids=input_ids,
                    positions=positions,
                    intermediate_tensors=intermediate_tensors,
                    inputs_embeds=inputs_embeds,
                )
1722

王敏's avatar
王敏 committed
1723
1724
1725
                self.maybe_wait_for_kv_save()
                finished_sending, finished_recving = (
                    self.get_finished_kv_transfers(scheduler_output))
Robert Shaw's avatar
Robert Shaw committed
1726

1727
        if self.use_aux_hidden_state_outputs:
Robert Shaw's avatar
Robert Shaw committed
1728
            hidden_states, aux_hidden_states = model_output
1729
        else:
Robert Shaw's avatar
Robert Shaw committed
1730
            hidden_states = model_output
1731
1732
            aux_hidden_states = None

1733
1734
1735
1736
1737
1738
1739
        # Broadcast PP output for external_launcher (torchrun)
        # to make sure we are synced across pp ranks
        # TODO: Support overlapping mirco-batches
        # https://github.com/vllm-project/vllm/issues/18019
        broadcast_pp_output = \
            self.parallel_config.distributed_executor_backend \
            == "external_launcher" and len(get_pp_group().ranks) > 0
1740
        if not get_pp_group().is_last_rank:
1741
            # For mid-pipeline stages, return the hidden states.
1742
1743
1744
1745
1746
1747
1748
            if not broadcast_pp_output:
                return hidden_states
            assert isinstance(hidden_states, IntermediateTensors)
            get_pp_group().send_tensor_dict(hidden_states.tensors,
                                            all_gather_group=get_tp_group())
            logits = None
        else:
1749
1750
1751
1752
1753
            if self.input_batch.pooling_params:
                return self._pool(hidden_states, num_scheduled_tokens,
                                  num_scheduled_tokens_np, finished_sending,
                                  finished_recving)

1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
            sample_hidden_states = hidden_states[logits_indices]
            logits = self.model.compute_logits(sample_hidden_states, None)
        if broadcast_pp_output:
            model_output_broadcast_data = {
                "logits": logits.contiguous(),
            } if logits is not None else {}
            model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
            assert model_output_broadcast_data is not None
            logits = model_output_broadcast_data["logits"]
1764

1765
1766
1767
1768
        # Apply structured output bitmasks if present
        if scheduler_output.grammar_bitmask is not None:
            self.apply_grammar_bitmask(scheduler_output, logits)

1769
        # Sample the next token and get logprobs if needed.
1770
        sampling_metadata = self.input_batch.sampling_metadata
1771
        if spec_decode_metadata is None:
1772
            sampler_output = self.sampler(
1773
1774
1775
1776
                logits=logits,
                sampling_metadata=sampling_metadata,
            )
        else:
1777
1778
1779
1780
            # When indexing with a tensor (bonus_logits_indices), PyTorch
            # creates a new tensor with separate storage from the original
            # logits tensor. This means any in-place operations on bonus_logits
            # won't affect the original logits tensor.
1781
            assert logits is not None
1782

王敏's avatar
王敏 committed
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
            if not envs.VLLM_REJECT_SAMPLE_OPT:
                bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
                sampler_output = self.sampler(
                    logits=bonus_logits,
                    sampling_metadata=sampling_metadata,
                )
                bonus_token_ids = sampler_output.sampled_token_ids

                # Just like `bonus_logits`, `target_logits` is a new tensor with
                # separate storage from the original `logits` tensor. Therefore,
                # it is safe to update `target_logits` in place.
                target_logits = logits[spec_decode_metadata.target_logits_indices]
                output_token_ids = self.rejection_sampler(
                    spec_decode_metadata,
                    None,  # draft_probs
                    target_logits,
                    bonus_token_ids,
                    sampling_metadata,
                )
                sampler_output.sampled_token_ids = output_token_ids
            else:
1804
1805
                sampling_metadata.all_greedy = True
                sampling_metadata.all_random = False
王敏's avatar
王敏 committed
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
                sampler_output = self.sampler(
                    logits=logits,
                    sampling_metadata=sampling_metadata,
                )
                target_token_ids = sampler_output.sampled_token_ids[spec_decode_metadata.target_logits_indices]
                target_logits = logits[spec_decode_metadata.target_logits_indices]

                bonus_token_ids = sampler_output.sampled_token_ids[spec_decode_metadata.bonus_logits_indices]

                output_token_ids = self.rejection_sampler(
                    spec_decode_metadata,
                    self.draft_probs.get_probs(spec_decode_metadata.spec_decode_ids),
                    target_logits,
                    target_token_ids,
                    bonus_token_ids,
                    sampling_metadata,
                )
                sampler_output.sampled_token_ids = output_token_ids               

1825

1826
1827
1828
1829
        num_nans_in_logits = {}
        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
            num_nans_in_logits = self._get_nans_in_logits(logits)

1830
1831
        # TODO(woosuk): The following loop can be slow since it iterates over
        # the requests one by one. Optimize.
1832
1833
        discard_sampled_tokens_req_indices = []
        for i, req_id in enumerate(self.input_batch.req_ids):
1834
1835
1836
            req_state = self.requests[req_id]
            seq_len = (req_state.num_computed_tokens +
                       scheduler_output.num_scheduled_tokens[req_id])
1837
            if seq_len < req_state.num_tokens:
1838
                # Ignore the sampled token for partial prefills.
1839
                # Rewind the generator state as if the token was not sampled.
1840
                # This relies on cuda-specific torch-internal impl details
1841
1842
1843
1844
1845
1846
                generator = self.input_batch.generators.get(i)
                if generator is not None:
                    generator.set_offset(generator.get_offset() - 4)
                # Record the index of the request that should not be sampled,
                # so that we could clear the sampled tokens before returning.
                discard_sampled_tokens_req_indices.append(i)
1847

1848
1849
        # NOTE: GPU -> CPU Sync happens here.
        # Move as many CPU operations as possible before this sync point.
1850
1851
1852
1853
1854
1855
        logprobs_tensors = sampler_output.logprobs_tensors
        logprobs_lists = logprobs_tensors.tolists() \
            if logprobs_tensors is not None else None

        # Compute prompt logprobs if needed.
        prompt_logprobs_dict = self._get_prompt_logprobs_dict(
1856
            hidden_states[:num_scheduled_tokens],
1857
1858
1859
            scheduler_output,
        )

1860
        # Get the valid generated tokens.
1861
1862
        sampled_token_ids = sampler_output.sampled_token_ids
        max_gen_len = sampled_token_ids.shape[-1]
lizhigong's avatar
lizhigong committed
1863

1864
        if max_gen_len == 1:
1865
            # No spec decode tokens.
1866
1867
            valid_sampled_token_ids = sampled_token_ids.tolist()
        else:
1868
            # Includes spec decode tokens.
1869
            valid_sampled_token_ids = self.rejection_sampler.parse_output(
1870
1871
1872
                sampled_token_ids,
                self.input_batch.vocab_size,
            )
1873
1874
1875
        # Mask out the sampled tokens that should not be sampled.
        for i in discard_sampled_tokens_req_indices:
            valid_sampled_token_ids[i].clear()
1876
1877
1878
        if spec_token_ids is not None:
            for i in discard_sampled_tokens_req_indices:
                spec_token_ids[i].clear()
1879

1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
        # Cache the sampled tokens in the model runner, so that the scheduler
        # doesn't need to send them back.
        # NOTE(woosuk): As an exception, when using PP, the scheduler sends
        # the sampled tokens back, because there's no direct communication
        # between the first-stage worker and the last-stage worker.
        for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
            if not sampled_ids:
                continue

            start_idx = self.input_batch.num_tokens_no_spec[req_idx]
            end_idx = start_idx + len(sampled_ids)
            assert end_idx <= self.max_model_len, (
                "Sampled token IDs exceed the max model length. "
                f"Total number of tokens: {end_idx} > max_model_len: "
                f"{self.max_model_len}")

            self.input_batch.token_ids_cpu[req_idx,
                                           start_idx:end_idx] = sampled_ids
            self.input_batch.num_tokens_no_spec[req_idx] = end_idx
            self.input_batch.num_tokens[req_idx] = end_idx
            req_id = self.input_batch.req_ids[req_idx]
            req_state = self.requests[req_id]
            req_state.output_token_ids.extend(sampled_ids)

1904
        if not self.speculative_config:
1905
            # Speculative decoding is not enabled.
1906
            spec_token_ids = None
1907
        else:
王敏's avatar
王敏 committed
1908
            spec_token_ids = self.propose_draft_token_ids(
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
                scheduler_output,
                valid_sampled_token_ids,
                sampling_metadata,
                hidden_states,
                sample_hidden_states,
                aux_hidden_states,
                spec_decode_metadata,
                attn_metadata,
            )

        # Clear KVConnector state after all KVs are generated.
        if has_kv_transfer_group():
            get_kv_transfer_group().clear_connector_metadata()

        self.eplb_step()

        return ModelRunnerOutput(
            req_ids=self.input_batch.req_ids,
            req_id_to_index=self.input_batch.req_id_to_index,
            sampled_token_ids=valid_sampled_token_ids,
            spec_token_ids=spec_token_ids,
            logprobs=logprobs_lists,
            prompt_logprobs_dict=prompt_logprobs_dict,
            pooler_output=[],
            finished_sending=finished_sending,
            finished_recving=finished_recving,
王敏's avatar
王敏 committed
1935
            num_nans_in_logits=num_nans_in_logits,
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
        )

    def propose_draft_token_ids(
        self,
        scheduler_output: "SchedulerOutput",
        sampled_token_ids: list[list[int]],
        sampling_metadata: SamplingMetadata,
        hidden_states: torch.Tensor,
        sample_hidden_states: torch.Tensor,
        aux_hidden_states: Optional[torch.Tensor],
        spec_decode_metadata: Optional[SpecDecodeMetadata],
        attn_metadata: dict[str, Any],
王敏's avatar
王敏 committed
1948
    ) -> list[list[int]]:
1949
1950
        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
        if self.speculative_config.method == "ngram":
1951
            assert isinstance(self.drafter, NgramProposer)
1952
1953
            spec_token_ids = self.propose_ngram_draft_token_ids(
                sampled_token_ids)
1954
1955
        elif self.speculative_config.method == "medusa":
            assert isinstance(self.drafter, MedusaProposer)
1956
1957
            if sample_hidden_states.shape[0] == len(sampled_token_ids):
                # The input to the target model does not include draft tokens.
1958
1959
1960
1961
1962
1963
                hidden_states = sample_hidden_states
            else:
                indices = []
                offset = 0
                for num_draft, tokens in zip(
                        spec_decode_metadata.num_draft_tokens,
1964
                        sampled_token_ids):
1965
1966
                    indices.append(offset + len(tokens) - 1)
                    offset += num_draft + 1
1967
                indices = torch.tensor(indices, device=self.device)
1968
1969
1970
1971
1972
1973
                hidden_states = sample_hidden_states[indices]

            spec_token_ids = self.drafter.propose(
                target_hidden_states=hidden_states,
                sampling_metadata=sampling_metadata,
            )
1974
        elif self.speculative_config.use_eagle():
1975
1976
1977
            assert isinstance(self.drafter, EagleProposer)
            # TODO(woosuk): Refactor the loop.
            next_token_ids: list[int] = []
1978
            for i, token_ids in enumerate(sampled_token_ids):
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
                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 = self.input_batch.req_ids[i]
                    req_state = self.requests[req_id]
                    seq_len = (req_state.num_computed_tokens +
                               scheduler_output.num_scheduled_tokens[req_id])
                    next_token_id = req_state.get_token_id(seq_len)
                next_token_ids.append(next_token_id)
1991
1992
1993
1994
1995
1996
1997
            next_token_ids = torch.tensor(next_token_ids,
                                          dtype=torch.int32,
                                          device=self.device)
            # At this moment, we assume all eagle layers belong to the same KV
            # cache group, thus using the same attention metadata.
            eagle_attn_metadata = attn_metadata[
                self.drafter.attn_layer_names[0]]
1998

Jiayi Yao's avatar
Jiayi Yao committed
1999
2000
2001
2002
2003
2004
            # NOTE: deepseek_mtp uses MLA which does not have `block_table`
            if hasattr(eagle_attn_metadata, "block_table"):
                block_table = eagle_attn_metadata.block_table
            else:
                block_table = None

2005
2006
2007
            if spec_decode_metadata is None:
                # input_ids can be None for multimodal models.
                target_token_ids = self.input_ids[:num_scheduled_tokens]
2008
2009
                # TODO(woosuk): Support M-RoPE.
                target_positions = self.positions[:num_scheduled_tokens]
2010
                if self.use_aux_hidden_state_outputs:
2011
2012
2013
                    target_hidden_states = torch.cat(
                        [h[:num_scheduled_tokens] for h in aux_hidden_states],
                        dim=-1)
2014
2015
                else:
                    target_hidden_states = hidden_states[:num_scheduled_tokens]
2016
2017
                target_slot_mapping = eagle_attn_metadata.slot_mapping
                cu_num_tokens = eagle_attn_metadata.query_start_loc
2018
2019
            else:
                # TODO(woosuk): Refactor this.
2020
2021
2022
                num_accepted_tokens = [len(s) - 1 for s in sampled_token_ids]
                num_accepted_tokens_tensor = async_tensor_h2d(
                    num_accepted_tokens,
2023
                    dtype=torch.int32,
2024
2025
                    target_device=self.device,
                    pin_memory=True)
2026
                cu_num_tokens, token_indices = self.drafter.prepare_inputs(
2027
                    eagle_attn_metadata.query_start_loc,
2028
                    num_accepted_tokens_tensor,
2029
2030
                )
                target_token_ids = self.input_ids[token_indices]
2031
2032
                # TODO(woosuk): Support M-RoPE.
                target_positions = self.positions[token_indices]
2033
                if self.use_aux_hidden_state_outputs:
2034
2035
                    target_hidden_states = torch.cat(
                        [h[token_indices] for h in aux_hidden_states], dim=-1)
2036
2037
                else:
                    target_hidden_states = hidden_states[token_indices]
2038
2039
                target_slot_mapping = eagle_attn_metadata.slot_mapping[
                    token_indices]
王敏's avatar
王敏 committed
2040
            draft_result = self.drafter.propose(
2041
2042
2043
2044
2045
2046
                target_token_ids=target_token_ids,
                target_positions=target_positions,
                target_hidden_states=target_hidden_states,
                target_slot_mapping=target_slot_mapping,
                next_token_ids=next_token_ids,
                cu_num_tokens=cu_num_tokens,
Jiayi Yao's avatar
Jiayi Yao committed
2047
                block_table=block_table,
2048
                sampling_metadata=sampling_metadata,
2049
                decoding=spec_decode_metadata is not None
2050
            )
王敏's avatar
王敏 committed
2051
2052
2053

            if not envs.VLLM_REJECT_SAMPLE_OPT:
                draft_token_ids = draft_result
2054
            else:
王敏's avatar
王敏 committed
2055
                draft_token_ids, draft_probs = draft_result
王敏's avatar
王敏 committed
2056
2057
2058
2059
            spec_token_ids = draft_token_ids.tolist()

            if envs.VLLM_REJECT_SAMPLE_OPT:
                draft_req_ids = list(scheduler_output.num_scheduled_tokens.keys())
王敏's avatar
王敏 committed
2060
2061
2062
2063
2064
                if self.draft_probs is None:
                    self.draft_probs = DraftProbs(
                        draft_probs, draft_req_ids)
                else:
                    self.draft_probs.update(draft_probs, draft_req_ids)
王敏's avatar
王敏 committed
2065

王敏's avatar
王敏 committed
2066
        return spec_token_ids
2067

Robert Shaw's avatar
Robert Shaw committed
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
    def kv_connector_no_forward(
            self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
        # KV send/recv even if no work to do.
        with set_forward_context(None, self.vllm_config):
            self.maybe_setup_kv_connector(scheduler_output)
            finished_sending, finished_recving = (
                self.get_finished_kv_transfers(scheduler_output))

        if not finished_sending and not finished_recving:
            return EMPTY_MODEL_RUNNER_OUTPUT

        output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
        output.finished_sending = finished_sending
        output.finished_recving = finished_recving
        return output

    @staticmethod
    def maybe_setup_kv_connector(scheduler_output: "SchedulerOutput"):
        # Update KVConnector with the KVConnector metadata forward().
        if has_kv_transfer_group():
            kv_connector = get_kv_transfer_group()
            assert isinstance(kv_connector, KVConnectorBase_V1)
            assert scheduler_output.kv_connector_metadata is not None
            kv_connector.bind_connector_metadata(
                scheduler_output.kv_connector_metadata)

            # Background KV cache transfers happen here.
            # These transfers are designed to be async and the requests
            # involved may be disjoint from the running requests.
            # Do this here to save a collective_rpc.
            kv_connector.start_load_kv(get_forward_context())

    @staticmethod
    def maybe_wait_for_kv_save() -> None:
        if has_kv_transfer_group():
            get_kv_transfer_group().wait_for_save()

    @staticmethod
    def get_finished_kv_transfers(
        scheduler_output: "SchedulerOutput",
    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
        if has_kv_transfer_group():
            return get_kv_transfer_group().get_finished(
                scheduler_output.finished_req_ids)
        return None, None

2114
    def propose_ngram_draft_token_ids(
2115
        self,
2116
2117
        sampled_token_ids: list[list[int]],
    ) -> list[list[int]]:
2118
        # TODO(woosuk): Optimize.
2119
        draft_token_ids: list[list[int]] = []
2120
2121
2122
        for i, sampled_ids in enumerate(sampled_token_ids):
            num_sampled_ids = len(sampled_ids)
            if not num_sampled_ids:
2123
2124
2125
2126
                # Skip speculative decoding.
                draft_token_ids.append([])
                continue

2127
2128
            # Skip requests that require sampling parameters that are not
            # supported with speculative decoding.
2129
            req_id = self.input_batch.req_ids[i]
2130
            if req_id in self.input_batch.spec_decode_unsupported_reqs:
2131
2132
2133
                draft_token_ids.append([])
                continue

2134
2135
            num_tokens = self.input_batch.num_tokens_no_spec[i]
            if num_tokens >= self.max_model_len:
2136
2137
2138
2139
                # Skip requests that have already reached the max model length.
                draft_token_ids.append([])
                continue

2140
            drafter_output = self.drafter.propose(
2141
                self.input_batch.token_ids_cpu[i, :num_tokens])
2142
2143
2144
2145
2146
2147
            if drafter_output is None or len(drafter_output) == 0:
                draft_token_ids.append([])
            else:
                draft_token_ids.append(drafter_output.tolist())
        return draft_token_ids

2148
2149
2150
    def load_model(self) -> None:
        logger.info("Starting to load model %s...", self.model_config.model)
        with DeviceMemoryProfiler() as m:  # noqa: SIM117
2151
            time_before_load = time.perf_counter()
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
            model_loader = get_model_loader(self.load_config)
            if not hasattr(self, "model"):
                logger.info("Loading model from scratch...")
                self.model = model_loader.load_model(
                    vllm_config=self.vllm_config,
                    model_config=self.model_config)
            else:
                logger.info(
                    "Model was already initialized. Loading weights inplace..."
                )
                model_loader.load_weights(self.model,
                                          model_config=self.model_config)
2164
2165
            if has_step_pooler(self.model):
                self.input_batch.logits_processing_needs_token_ids = True
2166
2167
2168
2169
2170
2171
            if self.lora_config:
                self.model = self.load_lora_model(self.model,
                                                  self.model_config,
                                                  self.scheduler_config,
                                                  self.lora_config,
                                                  self.device)
2172
2173
2174
            if hasattr(self, "drafter"):
                logger.info("Loading drafter model...")
                self.drafter.load_model(self.model)
2175
2176
2177
            if self.use_aux_hidden_state_outputs:
                self.model.set_aux_hidden_state_layers(
                    self.model.get_eagle3_aux_hidden_state_layers())
2178
            time_after_load = time.perf_counter()
2179
        self.model_memory_usage = m.consumed_memory
2180
2181
        logger.info("Model loading took %.4f GiB and %.6f seconds",
                    self.model_memory_usage / GiB_bytes,
2182
                    time_after_load - time_before_load)
2183
        prepare_communication_buffer_for_model(self.model)
2184

王敏's avatar
王敏 committed
2185
2186
2187
        if hasattr(self, "drafter"):
            prepare_communication_buffer_for_model(self.drafter.model)

2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
        if is_mixture_of_experts(
                self.model) and self.parallel_config.enable_eplb:
            logger.info("EPLB is enabled for model %s.",
                        self.model_config.model)
            self.eplb_state = EplbState.build(
                self.model,
                self.device,
                self.parallel_config,
            )

2198
2199
2200
2201
2202
2203
2204
2205
2206
    def save_tensorized_model(
        self,
        tensorizer_config: "TensorizerConfig",
    ) -> None:
        TensorizerLoader.save_model(
            self.model,
            tensorizer_config=tensorizer_config,
        )

2207
2208
2209
2210
    def _get_prompt_logprobs_dict(
        self,
        hidden_states: torch.Tensor,
        scheduler_output: "SchedulerOutput",
2211
    ) -> dict[str, Optional[LogprobsTensors]]:
2212
2213
2214
2215
        num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs
        if not num_prompt_logprobs_dict:
            return {}

2216
        in_progress_dict = self.input_batch.in_progress_prompt_logprobs_cpu
2217
        prompt_logprobs_dict: dict[str, Optional[LogprobsTensors]] = {}
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231

        # Since prompt logprobs are a rare feature, prioritize simple,
        # maintainable loop over optimal performance.
        completed_prefill_reqs = []
        for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items():

            num_tokens = scheduler_output.num_scheduled_tokens[req_id]

            # Get metadata for this request.
            request = self.requests[req_id]
            num_prompt_tokens = len(request.prompt_token_ids)
            prompt_token_ids = torch.tensor(request.prompt_token_ids).to(
                self.device, non_blocking=True)

2232
2233
2234
2235
2236
2237
2238
2239
2240
            # Set up target LogprobsTensors object.
            logprobs_tensors = in_progress_dict.get(req_id)
            if not logprobs_tensors:
                # Create empty logprobs CPU tensors for the entire prompt.
                # If chunked, we'll copy in slice by slice.
                logprobs_tensors = LogprobsTensors.empty_cpu(
                    num_prompt_tokens - 1, num_prompt_logprobs + 1)
                in_progress_dict[req_id] = logprobs_tensors

2241
            # Determine number of logits to retrieve.
2242
2243
            start_idx = request.num_computed_tokens
            start_tok = start_idx + 1
2244
            num_remaining_tokens = num_prompt_tokens - start_tok
2245
            if num_tokens <= num_remaining_tokens:
2246
                # This is a chunk, more tokens remain.
2247
2248
2249
                # In the == case, there are no more prompt logprobs to produce
                # but we want to defer returning them to the next step where we
                # have new generated tokens to return.
2250
2251
2252
2253
2254
                num_logits = num_tokens
            else:
                # This is the last chunk of prompt tokens to return.
                num_logits = num_remaining_tokens
                completed_prefill_reqs.append(req_id)
2255
2256
2257
2258
2259
2260
2261
                prompt_logprobs_dict[req_id] = logprobs_tensors

            if num_logits <= 0:
                # This can happen for the final chunk if we prefilled exactly
                # (num_prompt_tokens - 1) tokens for this request in the prior
                # step. There are no more prompt logprobs to produce.
                continue
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276

            # Get the logits corresponding to this req's prompt tokens.
            # If this is a partial request (i.e. chunked prefill),
            # then there is prompt logprob generated for each index.
            req_idx = self.input_batch.req_id_to_index[req_id]
            offset = self.query_start_loc_np[req_idx].item()
            prompt_hidden_states = hidden_states[offset:offset + num_logits]
            logits = self.model.compute_logits(prompt_hidden_states, None)

            # Get the "target" tokens for each index. For prompt at index i,
            # the token at prompt index i+1 is the "sampled" token we want
            # to gather the logprob for.
            tgt_token_ids = prompt_token_ids[start_tok:start_tok + num_logits]

            # Compute prompt logprobs.
2277
2278
            logprobs = self.sampler.compute_logprobs(logits)
            token_ids, logprobs, ranks = self.sampler.gather_logprobs(
2279
2280
2281
                logprobs, num_prompt_logprobs, tgt_token_ids)

            # Transfer GPU->CPU async.
2282
2283
2284
2285
2286
2287
2288
            chunk_slice = slice(start_idx, start_idx + num_logits)
            logprobs_tensors.logprob_token_ids[chunk_slice].copy_(
                token_ids, non_blocking=True)
            logprobs_tensors.logprobs[chunk_slice].copy_(logprobs,
                                                         non_blocking=True)
            logprobs_tensors.selected_token_ranks[chunk_slice].copy_(
                ranks, non_blocking=True)
2289
2290
2291
2292
2293

        # Remove requests that have completed prefill from the batch
        # num_prompt_logprobs_dict.
        for req_id in completed_prefill_reqs:
            del num_prompt_logprobs_dict[req_id]
2294
            del in_progress_dict[req_id]
2295
2296

        # Must synchronize the non-blocking GPU->CPU transfers.
2297
        if prompt_logprobs_dict:
2298
            self._sync_device()
2299
2300
2301

        return prompt_logprobs_dict

2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
    def _get_nans_in_logits(
        self,
        logits: Optional[torch.Tensor],
    ) -> dict[str, int]:
        try:
            if logits is None:
                return {req_id: 0 for req_id in self.input_batch.req_ids}

            num_nans_in_logits = {}
            num_nans_for_index = logits.isnan().sum(dim=-1).cpu().numpy()
            for req_id in self.input_batch.req_ids:
                req_index = self.input_batch.req_id_to_index[req_id]
                num_nans_in_logits[req_id] = (
                    int(num_nans_for_index[req_index])
                    if num_nans_for_index is not None
                    and req_index < logits.shape[0] else 0)
            return num_nans_in_logits
        except IndexError:
            return {}

2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
    @contextmanager
    def maybe_randomize_inputs(self, input_ids: torch.Tensor):
        """
        Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set.
        This is to help balance expert-selection
         - during profile_run
         - during DP rank dummy run 
        """
        dp_size = self.vllm_config.parallel_config.data_parallel_size
        randomize_inputs = envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1
        if not randomize_inputs:
            yield
        else:
            import functools

            @functools.cache
            def rand_input_ids() -> torch.Tensor:
                return torch.randint_like(
                    self.input_ids,
                    low=0,
                    high=self.model_config.get_vocab_size(),
                    dtype=input_ids.dtype)

            logger.debug("Randomizing dummy data for DP Rank")
            input_ids.copy_(rand_input_ids()[:input_ids.size(0)],
                            non_blocking=True)
            yield
            input_ids.fill_(0)

2351
2352
2353
2354
    @torch.inference_mode()
    def _dummy_run(
        self,
        num_tokens: int,
2355
        capture_attn_cudagraph: bool = False,
2356
2357
        skip_eplb: bool = False,
        is_profile: bool = False,
2358
    ) -> tuple[torch.Tensor, torch.Tensor]:
2359
        
王敏's avatar
王敏 committed
2360
        # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim.
2361
2362
2363
        if self.ep_sp:
            if num_tokens < self.tp_size:
                num_tokens = self.tp_size
王敏's avatar
王敏 committed
2364

2365
        # Padding for DP
2366
2367
        num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
        num_tokens += num_pad
2368

2369
2370
2371
        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
        # for dummy run with LoRA so that the num_reqs collectively
        # has num_tokens in total.
2372
        
2373
2374
        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
        max_num_reqs = self.scheduler_config.max_num_seqs
2375
        num_reqs = min(num_tokens, max_num_reqs)
2376
        min_tokens_per_req = num_tokens // num_reqs
2377
        num_actual_tokens = num_tokens
王敏's avatar
王敏 committed
2378

王敏's avatar
王敏 committed
2379
2380
2381
2382
        if not is_profile and self.speculative_config is not None \
            and self.speculative_config.num_lookahead_slots > 0 \
            and num_tokens >= 1 + self.speculative_config.num_lookahead_slots:

王敏's avatar
王敏 committed
2383
2384
            min_tokens_per_req = (1 + self.speculative_config.num_lookahead_slots)
            num_reqs = num_tokens // min_tokens_per_req
2385
2386
2387
2388
2389

            if self.ep_sp:
                num_actual_tokens = round_down(num_tokens, 1 + self.speculative_config.num_lookahead_slots)
                num_reqs = num_actual_tokens // min_tokens_per_req
        
2390
        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
2391
2392
2393
2394

        if not self.ep_sp:
            num_scheduled_tokens_list[-1] += num_tokens % num_reqs
        else:
2395
2396
2397
2398
            if self.speculative_config is not None:
                num_scheduled_tokens_list[-1] += num_tokens % min_tokens_per_req
            else:
                num_scheduled_tokens_list[-1] += num_tokens % num_reqs
2399
2400
2401
2402
        assert sum(num_scheduled_tokens_list) == num_tokens
        assert len(num_scheduled_tokens_list) == num_reqs
        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
                                        dtype=np.int32)
2403

2404
2405
2406
2407
        attn_metadata: Optional[dict[str, Any]] = None
        if capture_attn_cudagraph:
            attn_metadata = {}

2408
            query_start_loc = self.query_start_loc[:num_reqs + 1]
2409
2410
2411
2412
2413
            # Make sure max_model_len is used at the graph capture time.
            self.seq_lens_np[:num_reqs] = self.max_model_len
            self.seq_lens_np[num_reqs:] = 0
            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
                                           non_blocking=True)
2414
2415
            seq_lens = self.seq_lens[:num_reqs]

王敏's avatar
王敏 committed
2416
            num_speculative_tokens = 0 if self.speculative_config is None else self.speculative_config.num_lookahead_slots
2417
            common_attn_metadata = CommonAttentionMetadata(
2418
2419
                query_start_loc=query_start_loc,
                seq_lens=seq_lens,
2420
                # seq_lens_tensor=seq_lens_tensor,
2421
                num_reqs=num_reqs,
2422
                num_actual_tokens=num_actual_tokens,
2423
                max_query_len=num_tokens,
王敏's avatar
王敏 committed
2424
                num_speculative_tokens=num_speculative_tokens,
2425
            )
2426

2427
2428
            for kv_cache_group_id, kv_cache_group_spec in enumerate(
                    self.kv_cache_config.kv_cache_groups):
2429
2430
2431
2432

                attn_metadata_i = self.attn_metadata_builders[
                    kv_cache_group_id].build_for_cudagraph_capture(
                        common_attn_metadata)
2433
2434
                for layer_name in kv_cache_group_spec.layer_names:
                    attn_metadata[layer_name] = attn_metadata_i
2435

2436
2437
2438
2439
2440
2441
2442
        with self.maybe_dummy_run_with_lora(self.lora_config,
                                            num_scheduled_tokens):
            model = self.model
            if self.is_multimodal_model:
                input_ids = None
                inputs_embeds = self.inputs_embeds[:num_tokens]
            else:
王敏's avatar
王敏 committed
2443
2444
                self.input_ids[:num_tokens] = torch.randint(0, self.model_config.get_vocab_size(), (num_tokens,),
                                                            dtype=torch.int32)
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
                input_ids = self.input_ids[:num_tokens]
                inputs_embeds = None
            if self.uses_mrope:
                positions = self.mrope_positions[:, :num_tokens]
            else:
                positions = self.positions[:num_tokens]

            if get_pp_group().is_first_rank:
                intermediate_tensors = None
            else:
                if self.intermediate_tensors is None:
                    self.intermediate_tensors = (
                        self.model.make_empty_intermediate_tensors(
                            batch_size=self.max_num_tokens,
                            dtype=self.model_config.dtype,
                            device=self.device))
2461
2462
2463

                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
                    num_tokens, None, False)
2464

2465
            with self.maybe_randomize_inputs(input_ids), set_forward_context(
2466
2467
2468
2469
                    attn_metadata,
                    self.vllm_config,
                    num_tokens=num_tokens,
                    num_tokens_across_dp=num_tokens_across_dp):
2470
                outputs = model(
2471
2472
2473
2474
2475
                    input_ids=input_ids,
                    positions=positions,
                    intermediate_tensors=intermediate_tensors,
                    inputs_embeds=inputs_embeds,
                )
2476
2477
2478
2479
            if self.use_aux_hidden_state_outputs:
                hidden_states, _ = outputs
            else:
                hidden_states = outputs
2480

2481
            if self.speculative_config and self.speculative_config.use_eagle() and not is_profile:
lizhigong's avatar
lizhigong committed
2482
2483
2484
                #assert isinstance(self.drafter, EagleProposer)
                if hasattr(self, 'drafter') and isinstance(self.drafter, EagleProposer):
                    self.drafter.dummy_run(num_tokens, attn_metadata)
2485

2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
        # This is necessary to avoid blocking DP.
        # For dummy runs, we typically skip EPLB since we don't have any real
        # requests to process.
        # However, in DP settings, there may be cases when some DP ranks do
        # not have any requests to process, so they're executing dummy batches.
        # In such cases, we still have to trigger EPLB to make sure
        # ranks execute the rearrangement in synchronization.
        if not skip_eplb:
            self.eplb_step(is_dummy=True, is_profile=is_profile)

2496
        logit_indices = np.cumsum(num_scheduled_tokens) - 1
2497
        return hidden_states, hidden_states[logit_indices]
2498
2499
2500
2501
2502
2503

    @torch.inference_mode()
    def _dummy_sampler_run(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
2504
2505
2506
2507
        # The dummy hidden states may contain special values,
        # like `inf` or `nan`.
        # To avoid breaking the sampler, we use a random tensor here instead.
        hidden_states = torch.rand_like(hidden_states)
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530

        logits = self.model.compute_logits(hidden_states, None)
        num_reqs = logits.size(0)

        dummy_tensors = lambda v: torch.full(
            (num_reqs, ), v, device=self.device)

        dummy_metadata = SamplingMetadata(
            temperature=dummy_tensors(0.5),
            all_greedy=False,
            all_random=False,
            top_p=dummy_tensors(0.9),
            top_k=dummy_tensors(logits.size(1) - 1),
            generators={},
            max_num_logprobs=None,
            no_penalties=True,
            prompt_token_ids=None,
            frequency_penalties=dummy_tensors(0.1),
            presence_penalties=dummy_tensors(0.1),
            repetition_penalties=dummy_tensors(0.1),
            output_token_ids=[[] for _ in range(num_reqs)],
            allowed_token_ids_mask=None,
            bad_words_token_ids={},
2531
            logitsprocs=LogitsProcessorManager(),
2532
        )
2533
        try:
2534
2535
            sampler_output = self.sampler(logits=logits,
                                          sampling_metadata=dummy_metadata)
2536
2537
2538
2539
2540
2541
2542
2543
2544
        except RuntimeError as e:
            if 'out of memory' in str(e):
                raise RuntimeError(
                    "CUDA out of memory occurred when warming up sampler with "
                    f"{num_reqs} dummy requests. Please try lowering "
                    "`max_num_seqs` or `gpu_memory_utilization` when "
                    "initializing the engine.") from e
            else:
                raise e
2545
        if self.speculative_config:
王敏's avatar
王敏 committed
2546
            draft_token_ids = [[0]*self.speculative_config.num_lookahead_slots for _ in range(num_reqs)]
2547
2548
2549
2550
            dummy_spec_decode_metadata = SpecDecodeMetadata.make_dummy(
                draft_token_ids, self.device)

            num_tokens = sum(len(ids) for ids in draft_token_ids)
王敏's avatar
王敏 committed
2551
2552
2553
2554
2555
2556
2557
2558
2559
            
            if not envs.VLLM_REJECT_SAMPLE_OPT:
                draft_probs = None
            else:
                draft_probs = torch.randn(
                    num_reqs, self.speculative_config.num_lookahead_slots, logits.shape[-1], device=self.device,
                    dtype=logits.dtype)
                target_token_ids = torch.zeros(num_tokens, device=self.device,
                    dtype=torch.int32)
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
            target_logits = torch.randn(num_tokens,
                                        logits.shape[-1],
                                        device=self.device,
                                        dtype=logits.dtype)
            # NOTE(woosuk): Here, we should use int32 because the sampler uses
            # int32 for bonus_token_ids. If the dtype mismatches, re-compilation
            # will occur at runtime.
            bonus_token_ids = torch.zeros(num_reqs,
                                          device=self.device,
                                          dtype=torch.int32)
王敏's avatar
王敏 committed
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
            if not envs.VLLM_REJECT_SAMPLE_OPT:
                self.rejection_sampler(
                    dummy_spec_decode_metadata,
                    draft_probs,
                    target_logits,
                    bonus_token_ids,
                    dummy_metadata,
                )
            else:
                self.rejection_sampler(
                    dummy_spec_decode_metadata,
                    draft_probs,
                    target_logits,
                    target_token_ids,
                    bonus_token_ids,
                    dummy_metadata,
                )
2587
        return sampler_output
2588

2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
    @torch.inference_mode()
    def _dummy_pooler_run(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:

        num_tokens = hidden_states.shape[0]
        max_num_reqs = self.scheduler_config.max_num_seqs
        num_reqs = min(num_tokens, max_num_reqs)
        min_tokens_per_req = num_tokens // num_reqs
        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
        assert sum(num_scheduled_tokens_list) == num_tokens
        assert len(num_scheduled_tokens_list) == num_reqs

        hidden_states_list = list(
            torch.split(hidden_states, num_scheduled_tokens_list))

        req_num_tokens = num_tokens // num_reqs

        dummy_metadata = PoolingMetadata(
            prompt_lens=torch.tensor([h.shape[0] for h in hidden_states_list],
                                     device=self.device),
            prompt_token_ids=torch.zeros((num_reqs, req_num_tokens),
                                         dtype=torch.int32,
                                         device=self.device),
            pooling_params=[PoolingParams()] * num_reqs)

        try:
            pooler_output = self.model.pooler(hidden_states=hidden_states_list,
                                              pooling_metadata=dummy_metadata)
        except RuntimeError as e:
            if 'out of memory' in str(e):
                raise RuntimeError(
                    "CUDA out of memory occurred when warming up pooler with "
                    f"{num_reqs} dummy requests. Please try lowering "
                    "`max_num_seqs` or `gpu_memory_utilization` when "
                    "initializing the engine.") from e
            else:
                raise e
        return pooler_output

2631
    def profile_run(self) -> None:
2632
        # set profiling flag to avoid torch compile
2633
2634
        #set_profilling(True)
        #self._sync_device()
2635

2636
        # Profile with multimodal encoder & encoder cache.
2637
2638
2639
        # TODO: handle encoder-decoder models once we support them.
        if (self.is_multimodal_model and self.max_num_encoder_input_tokens > 0
                and self.encoder_cache_size > 0):
2640

2641
            # NOTE: Currently model is profiled with a single non-text
2642
2643
            # modality with the max possible input tokens even when
            # it supports multiple.
2644
2645
            max_tokens_by_modality_dict = self.mm_registry \
                .get_max_tokens_per_item_by_nonzero_modality(self.model_config)
2646
2647
2648
2649
            dummy_data_modality, max_tokens_per_mm_item = max(
                max_tokens_by_modality_dict.items(), key=lambda item: item[1])

            # Check how many items of this modality can be supported by
2650
2651
2652
2653
2654
2655
            # the encoder budget.
            encoder_budget = min(self.max_num_encoder_input_tokens,
                                 self.encoder_cache_size)

            max_num_mm_items_encoder_budget = cdiv(encoder_budget,
                                                   max_tokens_per_mm_item)
2656
2657
2658

            # Check how many items of this modality can be supported by
            # the decoder budget.
2659
2660
            max_mm_items_per_req = self.mm_registry.get_mm_limits_per_prompt(
                self.model_config)[dummy_data_modality]
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670

            # NOTE: We do not consider max_num_batched_tokens on purpose
            # because the multimodal embeddings can be generated in advance
            # and chunked prefilled.
            max_num_mm_items_decoder_budget = self.max_num_reqs * \
                max_mm_items_per_req

            max_num_mm_items = min(max_num_mm_items_encoder_budget,
                                   max_num_mm_items_decoder_budget)

2671
2672
2673
2674
2675
2676
            logger.info(
                "Encoder cache will be initialized with a budget of %s tokens,"
                " and profiled with %s %s items of the maximum feature size.",
                encoder_budget, max_num_mm_items, dummy_data_modality)

            # Create dummy batch of multimodal inputs.
2677
            dummy_mm_kwargs = self.mm_registry.get_decoder_dummy_data(
2678
2679
                model_config=self.model_config,
                seq_len=self.max_num_tokens,
2680
2681
2682
2683
                mm_counts={
                    dummy_data_modality: 1
                },
            ).multi_modal_data
2684

2685
            batched_dummy_mm_inputs = MultiModalKwargs.batch(
2686
2687
                [dummy_mm_kwargs] * max_num_mm_items,
                pin_memory=self.pin_memory)
2688
            batched_dummy_mm_inputs = MultiModalKwargs.as_kwargs(
2689
2690
2691
                batched_dummy_mm_inputs,
                device=self.device,
            )
2692
2693
2694
2695

            # Run multimodal encoder.
            dummy_encoder_outputs = self.model.get_multimodal_embeddings(
                **batched_dummy_mm_inputs)
2696
2697
2698
2699
2700

            sanity_check_mm_encoder_outputs(
                dummy_encoder_outputs,
                expected_num_items=max_num_mm_items,
            )
2701
2702
2703
2704

            # Cache the dummy encoder outputs.
            self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))

2705
        # Add `is_profile` here to pre-allocate communication buffers
2706
        hidden_states, last_hidden_states \
2707
            = self._dummy_run(self.max_num_tokens, is_profile=True)
2708
        if get_pp_group().is_last_rank:
2709
2710
2711
2712
            if self.is_pooling_model:
                output = self._dummy_pooler_run(hidden_states)
            else:
                output = self._dummy_sampler_run(last_hidden_states)
2713
        else:
2714
            output = None
2715
        self._sync_device()
2716
        del hidden_states, output
2717
        self.encoder_cache.clear()
2718
        gc.collect()
2719
        #set_profilling(False)
2720
2721

    def capture_model(self) -> None:
2722
2723
        if not self.use_cuda_graph:
            logger.warning(
2724
2725
2726
                "Skipping CUDA graph capture. To turn on CUDA graph capture, "
                "set -O %s and ensure `use_cudagraph` was not manually set to "
                "False", CompilationLevel.PIECEWISE)
2727
2728
            return

2729
2730
        compilation_counter.num_gpu_runner_capture_triggers += 1

2731
2732
2733
        start_time = time.perf_counter()
        start_free_gpu_memory = torch.cuda.mem_get_info()[0]

2734
2735
2736
        # Trigger CUDA graph capture for specific shapes.
        # Capture the large shapes first so that the smaller shapes
        # can reuse the memory pool allocated for the large shapes.
2737
        with graph_capture(device=self.device):
2738
            full_cg = self.full_cuda_graph
2739
2740
2741
2742
2743
2744
            # Only rank 0 should print progress bar during capture
            compilation_cases = reversed(self.cudagraph_batch_sizes)
            if is_global_first_rank():
                compilation_cases = tqdm(list(compilation_cases),
                                         desc="Capturing CUDA graph shapes")
            for num_tokens in compilation_cases:
2745
                # We skip EPLB here since we don't want to record dummy metrics
2746
2747
                for _ in range(
                        self.compilation_config.cudagraph_num_of_warmups):
2748
2749
2750
2751
2752
2753
                    self._dummy_run(num_tokens,
                                    capture_attn_cudagraph=full_cg,
                                    skip_eplb=True)
                self._dummy_run(num_tokens,
                                capture_attn_cudagraph=full_cg,
                                skip_eplb=True)
2754
2755
2756
2757
2758
2759
2760
2761

        end_time = time.perf_counter()
        end_free_gpu_memory = torch.cuda.mem_get_info()[0]
        elapsed_time = end_time - start_time
        cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
        # This usually takes 5~20 seconds.
        logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
                    elapsed_time, cuda_graph_size / (1 << 30))
2762

2763
2764
2765
2766
2767
2768
2769
    def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
        """
        Initialize the attention backends and attention metadata builders.
        """
        assert len(self.attn_backends) == 0 and len(
            self.attn_metadata_builders
        ) == 0, "Attention backends are already initialized"
laibao's avatar
laibao committed
2770
2771
2772
2773
        if envs.VLLM_ENABLE_KV_COMPRESSION and self.full_cuda_graph:
            raise ValueError(
                "KV compression is currently incompatible with full CUDA "
                "graph mode.")
2774
2775
2776
        for i, kv_cache_group_spec in enumerate(
                kv_cache_config.kv_cache_groups):
            kv_cache_spec = kv_cache_group_spec.kv_cache_spec
Chen Zhang's avatar
Chen Zhang committed
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
            if isinstance(kv_cache_spec, AttentionSpec):
                attn_backend_i = get_attn_backend(
                    kv_cache_spec.head_size,
                    self.dtype,
                    kv_cache_spec.dtype,
                    kv_cache_spec.block_size,
                    self.model_config.is_attention_free,
                    use_mla=kv_cache_spec.use_mla,
                )
                if attn_backend_i is None:
                    error_msg = (f"Error with get_attn_backend: "
                                 f"{kv_cache_spec.head_size=}, "
                                 f"{self.dtype=}, {kv_cache_spec.dtype=}, "
                                 f"{kv_cache_spec.block_size=}, "
                                 f"{self.model_config.is_attention_free=}, "
                                 f"{kv_cache_spec.use_mla=}")
                    logger.error(error_msg)
                    raise NotImplementedError(
                        "Non-Attention backend is not supported by V1 "
                        "GPUModelRunner.")
laibao's avatar
laibao committed
2797
2798
2799
2800
2801
                if (envs.VLLM_ENABLE_KV_COMPRESSION
                        and attn_backend_i.get_name() != "FLASH_ATTN_VLLM_V1"):
                    raise ValueError(
                        "KV compression currently requires "
                        "VLLM_ATTENTION_BACKEND=FLASH_ATTN_VLLM_V1.")
Chen Zhang's avatar
Chen Zhang committed
2802
            elif isinstance(kv_cache_spec, MambaSpec):
laibao's avatar
laibao committed
2803
2804
2805
2806
                if envs.VLLM_ENABLE_KV_COMPRESSION:
                    raise ValueError(
                        "KV compression is currently only supported for "
                        "Transformer attention layers.")
Chen Zhang's avatar
Chen Zhang committed
2807
2808
2809
2810
                attn_backend_i = Mamba2AttentionBackend
            else:
                raise ValueError(
                    f"Unknown KV cache spec type: {type(kv_cache_spec)}")
2811
2812
2813

            block_table_i = self.input_batch.block_table[i]
            attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
2814
2815
2816
2817
2818
                weakref.proxy(self),
                kv_cache_spec,
                block_table_i,
            )

zhuwenwen's avatar
zhuwenwen committed
2819
            if (self.full_cuda_graph
2820
2821
2822
2823
2824
2825
                    and not attn_metadata_builder_i.full_cudagraph_supported):
                raise ValueError(
                    f"Full CUDAGraph not supported for "
                    f"{attn_backend_i.__name__}. Turn off CompilationConfig."
                    f"full_cuda_graph or use a different attention backend.")

2826
2827
2828
            self.attn_backends.append(attn_backend_i)
            self.attn_metadata_builders.append(attn_metadata_builder_i)

2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
    def may_reinitialize_input_batch(self,
                                     kv_cache_config: KVCacheConfig) -> None:
        """
        Re-initialize the input batch if the block sizes are different from
        `[self.cache_config.block_size]`. This usually happens when there
        are multiple KV cache groups.

        Args:
            kv_cache_config: The KV cache configuration.
        """
        block_sizes = [
            kv_cache_group.kv_cache_spec.block_size
            for kv_cache_group in kv_cache_config.kv_cache_groups
        ]
        if block_sizes != [self.cache_config.block_size]:
            assert self.cache_config.cpu_offload_gb == 0, (
                "Cannot re-initialize the input batch when CPU weight "
                "offloading is enabled. See https://github.com/vllm-project/vllm/pull/18298 "  # noqa: E501
                "for more details.")
            self.input_batch = InputBatch(
                max_num_reqs=self.max_num_reqs,
                max_model_len=self.max_model_len,
                max_num_batched_tokens=self.max_num_tokens,
                device=self.device,
                pin_memory=self.pin_memory,
                vocab_size=self.model_config.get_vocab_size(),
                block_sizes=block_sizes,
2856
                is_spec_decode=bool(self.vllm_config.speculative_config),
2857
2858
            )

2859
2860
    def _allocate_kv_cache_tensors(
            self, kv_cache_config: KVCacheConfig) -> dict[str, torch.Tensor]:
2861
        """
2862
2863
2864
        Initializes the KV cache buffer with the correct size. The buffer needs
        to be reshaped to the desired shape before being used by the models.

2865
        Args:
2866
            kv_cache_config: The KV cache config
2867
        Returns:
2868
            dict[str, torch.Tensor]: A map between layer names to their
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
            corresponding memory buffer for KV cache.
         """
        kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
        for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
            tensor = torch.zeros(kv_cache_tensor.size,
                                 dtype=torch.int8,
                                 device=self.device)
            for layer_name in kv_cache_tensor.shared_by:
                kv_cache_raw_tensors[layer_name] = tensor

        layer_names = set()
        for group in kv_cache_config.kv_cache_groups:
            layer_names.update(group.layer_names)
        assert layer_names == set(kv_cache_raw_tensors.keys(
        )), "Some layers are not correctly initialized"
        return kv_cache_raw_tensors

    def _reshape_kv_cache_tensors(
        self,
        kv_cache_config: KVCacheConfig,
        kv_cache_raw_tensors: dict[str, torch.Tensor],
    ) -> dict[str, torch.Tensor]:
2891
        """
2892
        Reshape the KV cache tensors to the desired shape and dtype.
2893

2894
        Args:
2895
2896
            kv_cache_config: The KV cache config
            kv_cache_raw_tensors: The KV cache buffer of each layer, with
2897
2898
            correct size but uninitialized shape.
        Returns:
2899
            Dict[str, torch.Tensor]: A map between layer names to their
2900
2901
            corresponding memory buffer for KV cache.
        """
2902
        kv_caches: dict[str, torch.Tensor] = {}
2903
        has_attn, has_mamba = False, False
2904
2905
2906
2907
2908
2909
2910
2911
        for i, kv_cache_group_spec in enumerate(
                kv_cache_config.kv_cache_groups):
            kv_cache_spec = kv_cache_group_spec.kv_cache_spec
            for layer_name in kv_cache_group_spec.layer_names:
                raw_tensor = kv_cache_raw_tensors[layer_name]
                assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
                num_blocks = (raw_tensor.numel() //
                              kv_cache_spec.page_size_bytes)
2912
                if isinstance(kv_cache_spec, AttentionSpec):
2913
                    has_attn = True
zhuwenwen's avatar
zhuwenwen committed
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
                    if envs.VLLM_USE_FLASH_ATTN_PA and not kv_cache_spec.use_mla:
                        key_cache_shape, value_cache_shape = self.attn_backends[i].get_kv_cache_shape(
                            num_blocks, kv_cache_spec.block_size,
                            kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
                        dtype = kv_cache_spec.dtype
                        try:
                            key_stride_order, value_stride_order = self.attn_backends[
                                i].get_kv_cache_stride_order()
                            assert len(key_stride_order) == len(
                                key_cache_shape)
                            assert len(value_stride_order) == len(
                                value_cache_shape)
                        except (AttributeError, NotImplementedError):
                            key_stride_order = tuple(
                                range(len(key_cache_shape)))
                            value_stride_order = tuple(
                                range(len(value_cache_shape)))
                        # The allocation respects the backend-defined stride order
                        # to ensure the semantic remains consistent for each
                        # backend. We first obtain the generic kv cache shape and
                        # then permute it according to the stride order which could
                        # result in a non-contiguous tensor.
                                                
                        key_cache_shape = tuple(key_cache_shape[i]
                                            for i in key_stride_order)
                        value_cache_shape = tuple(value_cache_shape[i]
                                            for i in value_stride_order)
                        # Maintain original KV shape view.
                        inv_key_order = [
                            key_stride_order.index(i)
                            for i in range(len(key_stride_order))
                        ]
                        inv_value_order = [
                            value_stride_order.index(i)
                            for i in range(len(value_stride_order))
                        ]
                        
                        raw_tensor = kv_cache_raw_tensors[layer_name].view(dtype)
                        total_elements = raw_tensor.numel()
                        key_elements = (key_cache_shape[0] * key_cache_shape[1] * 
                                        key_cache_shape[2] * key_cache_shape[3])
                        value_elements = (value_cache_shape[0] * value_cache_shape[1] *
                                        value_cache_shape[2] * value_cache_shape[3])

                        assert total_elements == key_elements + value_elements

                        key_cache = raw_tensor[:key_elements].view(key_cache_shape).permute(
                            *inv_key_order)
                        value_cache = raw_tensor[key_elements:].view(value_cache_shape).permute(
                            *inv_value_order)
                        
                        kv_caches[layer_name] = (key_cache, value_cache)
                    else:
                        kv_cache_shape = self.attn_backends[i].get_kv_cache_shape(
                            num_blocks, kv_cache_spec.block_size,
                            kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
                        dtype = kv_cache_spec.dtype
                        try:
                            kv_cache_stride_order = self.attn_backends[
                                i].get_kv_cache_stride_order()
                            assert len(kv_cache_stride_order) == len(
                                kv_cache_shape)
                        except (AttributeError, NotImplementedError):
                            kv_cache_stride_order = tuple(
                                range(len(kv_cache_shape)))
                        # The allocation respects the backend-defined stride order
                        # to ensure the semantic remains consistent for each
                        # backend. We first obtain the generic kv cache shape and
                        # then permute it according to the stride order which could
                        # result in a non-contiguous tensor.
                        kv_cache_shape = tuple(kv_cache_shape[i]
                                            for i in kv_cache_stride_order)
                        # Maintain original KV shape view.
                        inv_order = [
                            kv_cache_stride_order.index(i)
                            for i in range(len(kv_cache_stride_order))
                        ]
                        kv_caches[layer_name] = kv_cache_raw_tensors[
                            layer_name].view(dtype).view(kv_cache_shape).permute(
                                *inv_order)
Chen Zhang's avatar
Chen Zhang committed
2994
                elif isinstance(kv_cache_spec, MambaSpec):
2995
                    has_mamba = True
Chen Zhang's avatar
Chen Zhang committed
2996
2997
                    raw_tensor = kv_cache_raw_tensors[layer_name]
                    dtype = kv_cache_spec.dtype
2998
2999
                    num_element_per_page = (kv_cache_spec.page_size_bytes //
                                            get_dtype_size(dtype))
Chen Zhang's avatar
Chen Zhang committed
3000
                    state_tensors = []
3001
                    storage_offset = 0
Chen Zhang's avatar
Chen Zhang committed
3002
3003
                    for shape in kv_cache_spec.shapes:
                        target_shape = (num_blocks, *shape)
3004
3005
3006
3007
3008
3009
3010
3011
                        stride = torch.empty(target_shape).stride()
                        target_stride = (num_element_per_page, *stride[1:])
                        tensor = torch.as_strided(
                            raw_tensor.view(dtype),
                            size=target_shape,
                            stride=target_stride,
                            storage_offset=storage_offset,
                        )
Chen Zhang's avatar
Chen Zhang committed
3012
                        state_tensors.append(tensor)
3013
3014
3015
                        storage_offset += stride[0]

                    kv_caches[layer_name] = state_tensors
3016
                else:
3017
                    raise NotImplementedError
3018
3019
3020
3021
3022

        if has_attn and has_mamba:
            self._verify_hybrid_attention_mamba_layout(kv_cache_config,
                                                       kv_cache_raw_tensors)

3023
3024
        return kv_caches

3025
3026
3027
3028
3029
3030
3031
3032
3033
3034
3035
3036
3037
3038
3039
3040
3041
3042
3043
3044
3045
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
    def _verify_hybrid_attention_mamba_layout(
            self, kv_cache_config: KVCacheConfig,
            kv_cache_raw_tensors: dict[str, torch.Tensor]) -> None:
        """
        Verify that the KV cache memory layout is compatible for
        models with both attention and mamba KV cache groups.

        Args:
            kv_cache_config: The KV cache config
            kv_cache_raw_tensors: The KV cache buffer of each layer.
        """

        for i, kv_cache_group_spec in enumerate(
                kv_cache_config.kv_cache_groups):
            kv_cache_spec = kv_cache_group_spec.kv_cache_spec
            for layer_name in kv_cache_group_spec.layer_names:
                raw_tensor = kv_cache_raw_tensors[layer_name]
                num_blocks = (raw_tensor.numel() //
                              kv_cache_spec.page_size_bytes)
                if isinstance(kv_cache_spec, AttentionSpec):
                    kv_cache_shape = self.attn_backends[i].get_kv_cache_shape(
                        num_blocks, kv_cache_spec.block_size,
                        kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
                    if kv_cache_shape[0] != num_blocks or kv_cache_shape[
                            1] != 2:
                        raise ValueError(
                            "Hybrid models in V1 require an attention "
                            "backend with kv_cache_shape="
                            "(num_blocks, 2, ...). Please try setting "
                            "VLLM_ATTENTION_BACKEND=FLASHINFER")

3056
3057
3058
3059
3060
3061
3062
3063
    def initialize_kv_cache_tensors(
            self, kv_cache_config: KVCacheConfig) -> dict[str, torch.Tensor]:
        """
        Initialize the memory buffer for KV cache.

        Args:
            kv_cache_config: The KV cache config
        Returns:
3064
            Dict[str, torch.Tensor]: A map between layer names to their
3065
3066
3067
3068
3069
3070
3071
            corresponding memory buffer for KV cache.
        """
        # Initialize the memory buffer for KV cache
        kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
        # Change the memory buffer to the desired shape
        kv_caches = self._reshape_kv_cache_tensors(kv_cache_config,
                                                   kv_cache_raw_tensors)
3072

3073
3074
3075
3076
3077
3078
3079
3080
3081
        # Setup `kv_cache_config` and `kv_caches` for models
        # with cross-layer KV sharing
        if self.shared_kv_cache_layers:
            initialize_kv_cache_for_kv_sharing(
                self.shared_kv_cache_layers,
                kv_cache_config.kv_cache_groups,
                kv_caches,
            )

3082
3083
3084
        bind_kv_cache(kv_caches,
                      self.compilation_config.static_forward_context,
                      self.kv_caches)
3085
3086
3087
3088
3089
3090
3091
3092
3093
3094
3095
3096
3097
3098
        return kv_caches

    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
        """
        Initialize KV cache based on `kv_cache_config`.
        Args:
            kv_cache_config: Configuration for the KV cache, including the KV
            cache size of each layer
        """
        self.kv_cache_config = kv_cache_config
        self.may_reinitialize_input_batch(kv_cache_config)
        self.initialize_attn_backend(kv_cache_config)
        kv_caches = self.initialize_kv_cache_tensors(kv_cache_config)

3099
        if self.speculative_config and self.speculative_config.use_eagle():
lizhigong's avatar
lizhigong committed
3100
            #assert isinstance(self.drafter, EagleProposer)
3101
3102
            # validate all draft model layers belong to the same kv cache
            # group
lizhigong's avatar
lizhigong committed
3103
3104
            if hasattr(self, 'drafter') and isinstance(self.drafter, EagleProposer):
                self.drafter.validate_same_kv_cache_group(kv_cache_config)
3105

Robert Shaw's avatar
Robert Shaw committed
3106
3107
3108
        if has_kv_transfer_group():
            get_kv_transfer_group().register_kv_caches(kv_caches)

3109
    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
3110
        """
3111
        Generates the KVCacheSpec by parsing the kv cache format from each
3112
3113
        Attention module in the static forward context.
        Returns:
3114
            KVCacheSpec: A dictionary mapping layer names to their KV cache
3115
3116
3117
3118
            format. Layers that do not need KV cache are not included.
        """

        block_size = self.vllm_config.cache_config.block_size
3119
        use_mla = self.vllm_config.model_config.use_mla
3120
        kv_cache_spec: dict[str, KVCacheSpec] = {}
Chen Zhang's avatar
Chen Zhang committed
3121
3122
        attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
        for layer_name, attn_module in attn_layers.items():
3123
3124
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
            if (kv_tgt_layer :=
                    attn_module.kv_sharing_target_layer_name) is not None:
                # The layer doesn't need its own KV cache and will use that of
                # the target layer. We skip creating a KVCacheSpec for it, so
                # that KV cache management logic will act as this layer does
                # not exist, and doesn't allocate KV cache for the layer. This
                # enables the memory saving of cross-layer kv sharing, allowing
                # a given amount of memory to accommodate longer context lengths
                # or enable more requests to be processed simultaneously.
                self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
                continue

3135
            # TODO: Support other attention modules, e.g., cross-attention
3136
            if attn_module.attn_type == AttentionType.DECODER:
3137
3138
3139
3140
3141
3142
3143
3144
3145
3146
3147
3148
3149
3150
3151
                if attn_module.sliding_window is not None:
                    kv_cache_spec[layer_name] = SlidingWindowSpec(
                        block_size=block_size,
                        num_kv_heads=attn_module.num_kv_heads,
                        head_size=attn_module.head_size,
                        dtype=self.kv_cache_dtype,
                        sliding_window=attn_module.sliding_window,
                        use_mla=use_mla)
                else:
                    kv_cache_spec[layer_name] = FullAttentionSpec(
                        block_size=block_size,
                        num_kv_heads=attn_module.num_kv_heads,
                        head_size=attn_module.head_size,
                        dtype=self.kv_cache_dtype,
                        use_mla=use_mla)
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
            elif attn_module.attn_type in (AttentionType.ENCODER,
                                           AttentionType.ENCODER_ONLY):
                # encoder-only attention does not need KV cache.
                continue
            elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
                raise NotImplementedError
            else:
                raise ValueError(
                    f"Unknown attention type: {attn_module.attn_type}")

Chen Zhang's avatar
Chen Zhang committed
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
        mamba_layers = get_layers_from_vllm_config(self.vllm_config,
                                                   MambaMixer2)
        if len(mamba_layers) > 0:
            if self.vllm_config.speculative_config is not None:
                raise NotImplementedError(
                    "Mamba with speculative decoding is not supported yet.")
            if not self.vllm_config.model_config.enforce_eager:
                raise NotImplementedError(
                    "Mamba with cuda graph is not supported yet.")
            if self.vllm_config.cache_config.enable_prefix_caching:
                raise NotImplementedError(
                    "Prefix caching is not supported for Mamba yet.")
            max_model_len = self.vllm_config.model_config.max_model_len
3175
3176
3177
3178
3179

            page_size_padded = self._maybe_pad_mamba_page_size(
                attn_layers, mamba_layers, kv_cache_spec, max_model_len,
                block_size)

Chen Zhang's avatar
Chen Zhang committed
3180
3181
3182
3183
3184
3185
            # Set block_size to max_model_len, so that mamba model will always
            # have only one block in the KV cache.
            for layer_name, mamba_module in mamba_layers.items():
                kv_cache_spec[layer_name] = MambaSpec(
                    shapes=mamba_module.get_state_shape(),
                    dtype=self.kv_cache_dtype,
3186
3187
3188
                    block_size=max_model_len,
                    page_size_padded=page_size_padded)

3189
        return kv_cache_spec
3190
3191
3192
3193
3194
3195
3196
3197
3198
3199
3200
3201
3202
3203
3204
3205
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
3218
3219
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230
3231
3232
3233
3234
3235
3236
3237
3238
3239
3240

    def _maybe_pad_mamba_page_size(
        self,
        attn_layers: dict[str, Attention],
        mamba_layers: dict[str, MambaMixer2],
        kv_cache_spec: dict[str, KVCacheSpec],
        max_model_len: int,
        block_size: int,
    ) -> Optional[int]:
        """
        Ensure that page size of attention KV cache groups is greater than or
        equal to the mamba KV cache groups. If not, we suggest to the user
        how to set the attention block size to ensure that it is.

        If the attention page size is strictly greater than the mamba page size,
        we pad the mamba page size to make them equal.

        Args:
            attn_layers: Attention layers
            mamba_layers: Mamba layers
            kv_cache_spec: KV cache spec (populated with attention layers)

        Returns:
            Optional[int]: Mamba page size with padding (None if no padding).
        """

        if len(attn_layers) == 0:
            return None

        attn_layer_name = next(iter(attn_layers))
        attn_page_size = kv_cache_spec[attn_layer_name].page_size_bytes
        mamba_layer_name = next(iter(mamba_layers))
        mamba_page_size = MambaSpec(
            shapes=mamba_layers[mamba_layer_name].get_state_shape(),
            dtype=self.kv_cache_dtype,
            block_size=max_model_len).page_size_bytes
        if attn_page_size < mamba_page_size:
            # attention page size (for 16 tokens)
            attn_page_size_16 = 16 * attn_page_size // block_size
            # some attention backends (e.g. FA) only support setting
            # block size to multiple of 16, so let's suggest a value
            # that would work (note: FA is currently not compatible
            # with mamba layers, use FlashInfer instead).
            suggest_attn_block_size = 16 * cdiv(mamba_page_size,
                                                attn_page_size_16)
            raise ValueError(
                "Attention block size should be increased to at least "
                f"{suggest_attn_block_size} in order to match "
                "the mamba page size")

        return attn_page_size
3241
3242
3243
3244
3245
3246
3247
3248
3249
3250
3251
3252
3253
3254
3255
3256
3257
3258
3259
3260
3261
3262
3263
3264
3265
3266
3267
3268
3269
3270
3271
3272
3273
3274
3275
3276
3277
3278
3279
3280
3281
3282
3283
3284
3285
3286
3287
3288
3289
3290
3291
3292
3293
3294
3295
3296
3297
3298
3299
3300
3301
3302
3303
3304
3305
3306
3307
3308
3309
3310
3311
3312
3313
3314
3315
3316
3317
3318
3319
3320
3321
3322
3323
3324
3325
3326
3327
3328
3329
3330
3331
3332
3333
3334
3335
3336
3337
3338
3339
3340
3341
3342
3343
3344
3345
3346
3347
3348
3349
3350
3351
3352
3353
3354
3355
3356
3357
3358
3359
3360
3361
3362
3363
3364
3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
3376
3377
3378
3379
3380
3381
3382
3383
3384
3385
3386
3387
3388
3389
3390
3391
3392
3393
3394
3395
3396
3397
3398
3399
3400
3401
3402
3403
3404
3405
3406
3407
3408
3409
3410
3411
3412
3413
3414
3415
3416
3417
3418
3419
3420
3421
3422
3423
3424
3425
3426
3427
3428
3429
3430
3431
3432
3433

class GPUModelRunnerMTP(GPUModelRunnerBase):
    def __init__(self, vllm_config, device):
        super().__init__(vllm_config, device)
        if hasattr(self, 'drafter') and isinstance(self.drafter, EagleProposer):
            self.drafter = V1ZeroEagleProposer(self.vllm_config, self.device,self)
        self.spec_sampler_event = torch.cuda.Event(enable_timing=False)
        self.spec_scheduler_max_num_tokens = 0


    def _prepare_inputs(
            self,
            scheduler_output: "SchedulerOutput",
    ) -> tuple[dict[str, Any], bool, torch.Tensor,
    Optional[SpecDecodeMetadata], np.ndarray]:
        """
        :return: tuple[
            attn_metadata: layer-to-attention_metadata mapping,
            attention_cuda_graphs: whether attention can run in cudagraph
            logits_indices, spec_decode_metadata
        ]
        """
        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
        assert total_num_scheduled_tokens > 0
        num_reqs = self.input_batch.num_reqs
        assert num_reqs > 0

        # OPTIMIZATION: Start copying the block table first.
        # This way, we can overlap the copy with the following CPU operations.
        self.input_batch.block_table.commit(num_reqs)

        # Get the number of scheduled tokens for each request.
        req_ids = self.input_batch.req_ids
        tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
        num_scheduled_tokens = np.array(tokens, dtype=np.int32)
        max_num_scheduled_tokens = max(tokens)
        self.spec_scheduler_max_num_tokens = max_num_scheduled_tokens
        # Get request indices.
        # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
        req_indices = np.repeat(self.arange_np[:num_reqs],
                                num_scheduled_tokens)

        # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
        # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
        cu_num_tokens, arange = self._get_cumsum_and_arange(
            num_scheduled_tokens)

        # Get positions.
        positions_np = self.positions_np[:total_num_scheduled_tokens]
        np.add(self.input_batch.num_computed_tokens_cpu[req_indices],
               arange,
               out=positions_np)

        # Calculate M-RoPE positions.
        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
        if self.uses_mrope:
            self._calc_mrope_positions(scheduler_output)

        # Get token indices.
        # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
        # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
        # where M is the max_model_len.
        token_indices = (positions_np +
                         req_indices * self.input_batch.token_ids_cpu.shape[1])

        # NOTE(woosuk): We use torch.index_select instead of np.take here
        # because torch.index_select is much faster than np.take for large
        # tensors.
        torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(),
                           0,
                           torch.from_numpy(token_indices),
                           out=self.input_ids_cpu[:total_num_scheduled_tokens])

        # Calculate the slot mapping for each KV cache group.
        for kv_cache_group_id, kv_cache_group_spec in enumerate(
                self.kv_cache_config.kv_cache_groups):
            block_size = kv_cache_group_spec.kv_cache_spec.block_size
            block_table: BlockTable = self.input_batch.block_table[
                kv_cache_group_id]
            # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
            # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
            # where K is the max_num_blocks_per_req and the block size is 2.
            # NOTE(woosuk): We can't simply use `token_indices // block_size`
            # here because M (max_model_len) is not necessarily divisible by
            # block_size.
            block_table_indices = (
                    req_indices * block_table.max_num_blocks_per_req +
                    positions_np // block_size)
            block_table_cpu = block_table.get_cpu_tensor()
            block_numbers = block_table_cpu.flatten(
            )[block_table_indices].numpy()
            block_offsets = positions_np % block_size
            np.add(
                block_numbers * block_size,
                block_offsets,
                out=block_table.slot_mapping_np[:total_num_scheduled_tokens])

        # Prepare the attention metadata.
        self.query_start_loc_np[0] = 0
        self.query_start_loc_np[1:num_reqs + 1] = cu_num_tokens

        self.seq_lens_np[:num_reqs] = (
                self.input_batch.num_computed_tokens_cpu[:num_reqs] +
                num_scheduled_tokens)

        # Copy the tensors to the GPU.
        self.input_ids[:total_num_scheduled_tokens].copy_(
            self.input_ids_cpu[:total_num_scheduled_tokens], non_blocking=True)
        if self.uses_mrope:
            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
            self.mrope_positions[:, :total_num_scheduled_tokens].copy_(
                self.mrope_positions_cpu[:, :total_num_scheduled_tokens],
                non_blocking=True)
        else:
            # Common case (1D positions)
            self.positions[:total_num_scheduled_tokens].copy_(
                self.positions_cpu[:total_num_scheduled_tokens],
                non_blocking=True)

        self.query_start_loc[:num_reqs + 1].copy_(
            self.query_start_loc_cpu[:num_reqs + 1], non_blocking=True)
        self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
                                       non_blocking=True)

        # Fill unused with -1. Needed for reshape_and_cache
        self.seq_lens[num_reqs:].fill_(0)
        # Note: pad query_start_loc to be non-decreasing, as kernels
        # like FlashAttention requires that
        self.query_start_loc[num_reqs + 1:].fill_(
            self.query_start_loc_cpu[num_reqs].item())

        query_start_loc = self.query_start_loc[:num_reqs + 1]
        seq_lens = self.seq_lens[:num_reqs]

        common_attn_metadata = CommonAttentionMetadata(
            query_start_loc=query_start_loc,
            seq_lens=seq_lens,
            # seq_lens_tensor=seq_lens_tensor,
            num_reqs=num_reqs,
            num_actual_tokens=total_num_scheduled_tokens,
            max_query_len=max_num_scheduled_tokens,
        )

        attn_metadata: dict[str, Any] = {}
        # Prepare the attention metadata for each KV cache group and make layers
        # in the same group share the same metadata.
        for kv_cache_group_id, kv_cache_group_spec in enumerate(
                self.kv_cache_config.kv_cache_groups):

            # Prepare for cascade attention if enabled & beneficial.
            common_prefix_len = 0
            builder = self.attn_metadata_builders[kv_cache_group_id]
            if self.cascade_attn_enabled:
                common_prefix_len = self._compute_cascade_attn_prefix_len(
                    num_scheduled_tokens,
                    scheduler_output.
                    num_common_prefix_blocks[kv_cache_group_id],
                    kv_cache_group_spec.kv_cache_spec,
                    builder,
                    )

            attn_metadata_i = (builder.build(
                common_prefix_len=common_prefix_len,
                common_attn_metadata=common_attn_metadata,
            ))

            for layer_name in kv_cache_group_spec.layer_names:
                attn_metadata[layer_name] = attn_metadata_i

        attention_cuda_graphs = all(
            b.can_run_in_cudagraph(common_attn_metadata)
            for b in self.attn_metadata_builders)

        use_spec_decode = len(
            scheduler_output.scheduled_spec_decode_tokens) > 0
        if not use_spec_decode:
            # NOTE(woosuk): Due to chunked prefills, the batch may contain
            # partial requests. While we should not sample any token
            # from these partial requests, we do so for simplicity.
            # We will ignore the sampled tokens from the partial requests.
            # TODO: Support prompt logprobs.
            logits_indices = query_start_loc[1:] - 1
            spec_decode_metadata = None
        else:
            # Get the number of draft tokens for each request.
            # Iterate over the dictionary rather than all requests since not all
            # requests have draft tokens.
            num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
            for req_id, draft_token_ids in (
                    scheduler_output.scheduled_spec_decode_tokens.items()):
                req_idx = self.input_batch.req_id_to_index[req_id]
                num_draft_tokens[req_idx] = len(draft_token_ids)

王敏's avatar
王敏 committed
3434
3435
3436
3437
            spec_decode_ids = None
            if envs.VLLM_REJECT_SAMPLE_OPT:
                spec_decode_ids = scheduler_output.scheduled_spec_decode_tokens.keys()

3438
            spec_decode_metadata = self._calc_spec_decode_metadata(
王敏's avatar
王敏 committed
3439
                num_draft_tokens, cu_num_tokens, spec_decode_ids)
3440
3441
3442
3443
3444
3445
3446
3447
3448
            logits_indices = spec_decode_metadata.logits_indices

        # Hot-Swap lora model
        if self.lora_config:
            self.set_active_loras(self.input_batch, num_scheduled_tokens)

        return (attn_metadata, attention_cuda_graphs, logits_indices,
                spec_decode_metadata, num_scheduled_tokens)

3449
3450
3451
3452
3453
3454
3455
3456
3457
3458
3459
3460
3461
3462
3463
3464
3465
3466
3467
3468
3469
3470
3471
3472
3473
3474
3475
3476
3477
3478
3479
3480
3481
3482
3483
3484
3485
3486
3487
3488
3489
3490
3491
3492
3493
3494
3495
3496
3497
3498
3499
3500
3501
3502
3503
3504
3505
3506
3507
3508
3509
3510
3511
3512
3513
3514
3515
3516
3517
3518
3519
3520
3521
3522
3523
3524
3525
3526
3527
3528
3529
3530
3531
3532
3533
3534
3535
3536
3537
3538
3539
3540
3541
3542
3543
3544
3545
3546
3547
3548
3549
3550
3551
3552
3553
3554
3555
3556
3557
3558
3559
3560
3561
3562
3563
3564
3565
3566
3567
3568
3569
3570
3571
3572
3573
3574
3575
3576
3577
3578
3579
3580
3581
3582
3583
3584
3585
3586
3587
3588
3589
3590
3591
3592
3593
3594
3595
3596
3597
    def _stash_kv_compression_prompt_payload(self) -> None:
        """Persist prompt-end compaction indices from the forward context."""
        if (not envs.VLLM_ENABLE_KV_COMPRESSION
                or not self.scheduler_config.chunked_prefill_enabled):
            return
        forward_context = get_forward_context()
        payload = getattr(forward_context, "_kv_compression_prompt_payload",
                          None)
        if payload is None:
            return

        req_indices = payload.get("req_indices")
        idx_sorted = payload.get("idx_sorted")
        keep_len = payload.get("keep_len")
        prompt_lens = payload.get("prompt_lens")
        if (req_indices is None or idx_sorted is None or keep_len is None
                or prompt_lens is None):
            return

        req_indices_cpu = req_indices.to(device="cpu",
                                         dtype=torch.int64).tolist()
        keep_cpu = keep_len.to(device="cpu", dtype=torch.int64).tolist()
        prompt_cpu = prompt_lens.to(device="cpu", dtype=torch.int64).tolist()
        for i, b in enumerate(req_indices_cpu):
            if b < 0 or b >= len(self.input_batch.req_ids):
                continue
            req_id = self.input_batch.req_ids[b]
            if req_id is None:
                continue
            rs = self.requests.get(req_id)
            if rs is None:
                continue
            rs.kv_compression_prompt_idx_sorted = idx_sorted[i]
            rs.kv_compression_prompt_keep_len = int(keep_cpu[i])
            rs.kv_compression_prompt_prompt_len = int(prompt_cpu[i])

    def _maybe_apply_kv_compression_prompt_compaction(self) -> None:
        """Apply one-shot prompt KV compaction before the first decode step."""
        if (not envs.VLLM_ENABLE_KV_COMPRESSION
                or not self.scheduler_config.chunked_prefill_enabled):
            return

        pending_req_ids: list[str] = []
        for req_id in self.input_batch.req_ids:
            if req_id is None:
                continue
            rs = self.requests.get(req_id)
            if rs is None:
                continue
            if rs.kv_compression_prompt_idx_sorted is None:
                continue
            # Only apply once the prompt is fully ingested (decode stage).
            if rs.num_computed_tokens < rs.num_prompt_tokens:
                continue
            pending_req_ids.append(req_id)

        if not pending_req_ids:
            return

        device = self.device
        pending_states: list[tuple[str, torch.Tensor, int]] = []
        for req_id in pending_req_ids:
            rs = self.requests[req_id]
            keep = rs.kv_compression_prompt_keep_len
            idx = rs.kv_compression_prompt_idx_sorted
            if keep is None or idx is None:
                continue
            keep_i = int(keep)
            if keep_i <= 0:
                # No prompt tokens kept; clear and skip.
                rs.kv_compression_prompt_idx_sorted = None
                rs.kv_compression_prompt_keep_len = None
                rs.kv_compression_prompt_prompt_len = None
                continue
            pending_states.append((req_id, idx, keep_i))

        if not pending_states:
            return

        B = len(pending_states)
        keep_list = [k for _, _, k in pending_states]
        K_max = max(keep_list)
        idx_batch = torch.zeros((B, K_max), device=device, dtype=torch.int32)
        for i, (_, row, k) in enumerate(pending_states):
            idx_batch[i, :k] = row[:k].to(device=device, dtype=torch.int32)
        keep_tensor = torch.tensor(keep_list,
                                   device=device,
                                   dtype=torch.int32)

        from vllm.v1.attention.kv_compression.kv_cache_triton import (
            front_compact_inplace_fa_triton, make_fa_cache_view)

        # Apply compaction to every attention layer's KV cache in-place.
        for group_id, kv_cache_group_spec in enumerate(
                self.kv_cache_config.kv_cache_groups):
            max_blocks = 0
            for req_id, _, _ in pending_states:
                rs = self.requests[req_id]
                if group_id >= len(rs.block_ids):
                    continue
                max_blocks = max(max_blocks, len(rs.block_ids[group_id]))
            if max_blocks == 0:
                continue

            block_table_cpu = torch.zeros((B, max_blocks),
                                          dtype=torch.int32,
                                          device="cpu")
            for i, (req_id, _, _) in enumerate(pending_states):
                rs = self.requests[req_id]
                if group_id >= len(rs.block_ids):
                    continue
                ids = rs.block_ids[group_id]
                if ids:
                    block_table_cpu[i, :len(ids)] = torch.tensor(
                        ids, dtype=torch.int32, device="cpu")
            block_table = block_table_cpu.to(device=device, non_blocking=True)

            for layer_name in kv_cache_group_spec.layer_names:
                layer_index = self._extract_layer_index(layer_name)
                if layer_index >= len(self.kv_caches):
                    continue
                kv_cache = self.kv_caches[layer_index]
                if not current_platform.is_rocm():
                    if not isinstance(kv_cache, torch.Tensor):
                        continue
                    key_cache, value_cache = kv_cache.unbind(0)
                else:
                    if not isinstance(kv_cache, (tuple, list)) or len(kv_cache) != 2:
                        continue
                    key_cache, value_cache = kv_cache
                k_view, v_view = make_fa_cache_view(key_cache=key_cache,
                                                    value_cache=value_cache)
                front_compact_inplace_fa_triton(
                    k_view,
                    v_view,
                    block_table,
                    idx_batch,
                    keep_tensor,
                )

        # Clear pending state after successful compaction.
        for req_id, _, _ in pending_states:
            rs = self.requests.get(req_id)
            if rs is None:
                continue
            rs.kv_compression_prompt_idx_sorted = None
            rs.kv_compression_prompt_keep_len = None
            rs.kv_compression_prompt_prompt_len = None

3598
3599
3600
3601
3602
3603
3604
3605
3606
3607
3608
3609
3610
3611
3612
3613
3614
3615
3616
3617
    @torch.inference_mode()
    def execute_model(
            self,
            scheduler_output: "SchedulerOutput",
            intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> Union[ModelRunnerOutput, IntermediateTensors]:
        self._update_states(scheduler_output)
        if not scheduler_output.total_num_scheduled_tokens:
            if not has_kv_transfer_group():
                # Return empty ModelRunnerOutput if there's no work to do.
                return EMPTY_MODEL_RUNNER_OUTPUT

            return self.kv_connector_no_forward(scheduler_output)

        # Prepare the decoder inputs.
        (attn_metadata, attention_cuda_graphs, logits_indices,
         spec_decode_metadata,
         num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

王敏's avatar
王敏 committed
3618
        # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim.
3619
3620
3621
3622
3623
3624
3625
3626
3627
3628
3629
3630
3631
3632
3633
3634
3635
3636
3637
3638
3639
3640
3641
3642
3643
        if self.ep_sp:
            num_input_tokens = round_up(num_scheduled_tokens, self.tp_size)
            if (self.use_cuda_graph
                    and num_input_tokens <= self.cudagraph_batch_sizes[-1]):
                # Use piecewise CUDA graphs.
                # Add padding to the batch size.
                num_input_tokens = self.vllm_config.pad_for_cudagraph(
                    num_input_tokens)
        else:
            if (self.use_cuda_graph
                    and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
                # Use piecewise CUDA graphs.
                # Add padding to the batch size.
                num_input_tokens = self.vllm_config.pad_for_cudagraph(
                    num_scheduled_tokens)
            else:
                # Eager mode.
                # Pad tokens to multiple of tensor_parallel_size when
                # enabled collective fusion for SP
                tp_size = self.vllm_config.parallel_config.tensor_parallel_size
                if self.compilation_config.pass_config. \
                        enable_sequence_parallelism and tp_size > 1:
                    num_input_tokens = round_up(num_scheduled_tokens, tp_size)
                else:
                    num_input_tokens = num_scheduled_tokens
王敏's avatar
王敏 committed
3644

3645
3646
3647
3648
3649
3650
3651
3652
3653
3654
3655
3656
3657
3658
3659
3660
3661
3662
3663
3664
3665
3666
3667
3668
3669
3670
3671
3672
3673
3674
3675
3676
3677
3678
3679
3680
3681
3682
3683
3684
3685
3686
3687
3688
3689
3690
3691
3692
3693
        # Padding for DP
        num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
        num_input_tokens += num_pad

        # _prepare_inputs may reorder the batch, so we must gather multi
        # modal outputs after that to ensure the correct order
        if self.is_multimodal_model:
            # Run the multimodal encoder if any.
            self._execute_mm_encoder(scheduler_output)
            mm_embeds = self._gather_mm_embeddings(scheduler_output)
        else:
            mm_embeds = []

        if self.is_multimodal_model and get_pp_group().is_first_rank:
            # NOTE(woosuk): To unify token ids and soft tokens (vision
            # embeddings), we always use embeddings (rather than token ids)
            # as input to the multimodal model, even when the input is text.
            input_ids = self.input_ids[:num_scheduled_tokens]
            if mm_embeds:
                inputs_embeds = self.model.get_input_embeddings(
                    input_ids, mm_embeds)
            else:
                inputs_embeds = self.model.get_input_embeddings(input_ids)
            # TODO(woosuk): Avoid the copy. Optimize.
            self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds)
            inputs_embeds = self.inputs_embeds[:num_input_tokens]
            input_ids = None
        else:
            # For text-only models, we use token ids as input.
            # While it is possible to use embeddings as input just like the
            # multimodal models, it is not desirable for performance since
            # then the embedding layer is not included in the CUDA graph.
            input_ids = self.input_ids[:num_input_tokens]
            inputs_embeds = None
        if self.uses_mrope:
            positions = self.mrope_positions[:, :num_input_tokens]
        else:
            positions = self.positions[:num_input_tokens]

        if get_pp_group().is_first_rank:
            intermediate_tensors = None
        else:
            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
                num_input_tokens, intermediate_tensors, True)

        # Some attention backends only support CUDA Graphs in pure decode.
        # If attention doesn't support CUDA Graphs for this batch, but we
        # compiled with full CUDA graphs, we have to skip them entirely.
        skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
3694
3695
3696
3697

        # Chunked prefill (scheme 3): apply one-shot prompt KV compaction before
        # the first decode step writes/reads KV at the compressed positions.
        self._maybe_apply_kv_compression_prompt_compaction()
3698
3699
3700
3701
3702
3703
3704
3705
3706
3707
3708
3709
3710
3711
3712
3713
3714
3715
3716
3717
3718
3719
3720
3721
3722
3723
3724
        if envs.VLLM_ENABLE_TBO and scheduler_output.total_num_scheduled_tokens >= envs.VLLM_TBO_MIN_TOKENS:
            model_output, finished_sending, finished_recving = \
                tbo_split_and_execute_model(self, attn_metadata, num_input_tokens,
                                            num_tokens_across_dp, input_ids, positions,
                                            inputs_embeds, scheduler_output, intermediate_tensors,
                                            skip_cuda_graphs)

        elif envs.VLLM_P2P_ASYNC:
            self.p2p_event.record()
            current_stream = torch.cuda.current_stream()
            with torch.cuda.stream(self.p2p_stream):
                self.p2p_stream.wait_event(self.p2p_event)
                with set_forward_context(
                        attn_metadata,
                        self.vllm_config,
                        num_tokens=num_input_tokens,
                        num_tokens_across_dp=num_tokens_across_dp,
                        skip_cuda_graphs=skip_cuda_graphs,
                ):
                    self.maybe_setup_kv_connector(scheduler_output)

                    model_output = self.model(
                        input_ids=input_ids,
                        positions=positions,
                        intermediate_tensors=intermediate_tensors,
                        inputs_embeds=inputs_embeds,
                    )
3725
                    self._stash_kv_compression_prompt_payload()
3726
3727
3728
3729
3730
3731
3732
3733
3734
3735
3736
3737
3738
3739
3740
3741
3742
3743
3744
3745
3746
3747
3748
3749
3750

                    self.maybe_wait_for_kv_save()
                    finished_sending, finished_recving = (
                        self.get_finished_kv_transfers(scheduler_output))
                self.p2p_event.record()
            current_stream.wait_event(self.p2p_event)

        else:
            # Run the model.
            # Use persistent buffers for CUDA graphs.
            with set_forward_context(
                    attn_metadata,
                    self.vllm_config,
                    num_tokens=num_input_tokens,
                    num_tokens_across_dp=num_tokens_across_dp,
                    skip_cuda_graphs=skip_cuda_graphs,
            ):
                self.maybe_setup_kv_connector(scheduler_output)

                model_output = self.model(
                    input_ids=input_ids,
                    positions=positions,
                    intermediate_tensors=intermediate_tensors,
                    inputs_embeds=inputs_embeds,
                )
3751
                self._stash_kv_compression_prompt_payload()
3752
3753
3754
3755
3756
3757
3758
3759
3760
3761
3762
3763
3764
3765
3766
3767
3768
3769
3770
3771
3772
3773
3774
3775
3776
3777
3778
3779
3780
3781
3782
3783
3784
3785
3786
3787
3788
3789
3790
3791
3792
3793
3794
3795
3796
3797
3798
3799
3800
3801
3802
3803
3804
3805
3806
3807
3808
3809
3810
3811
3812

                self.maybe_wait_for_kv_save()
                finished_sending, finished_recving = (
                    self.get_finished_kv_transfers(scheduler_output))

        if self.use_aux_hidden_state_outputs:
            hidden_states, aux_hidden_states = model_output
        else:
            hidden_states = model_output
            aux_hidden_states = None

        # Broadcast PP output for external_launcher (torchrun)
        # to make sure we are synced across pp ranks
        # TODO: Support overlapping mirco-batches
        # https://github.com/vllm-project/vllm/issues/18019
        broadcast_pp_output = \
            self.parallel_config.distributed_executor_backend \
            == "external_launcher" and len(get_pp_group().ranks) > 0
        if not get_pp_group().is_last_rank:
            # For mid-pipeline stages, return the hidden states.
            if not broadcast_pp_output:
                return hidden_states
            assert isinstance(hidden_states, IntermediateTensors)
            get_pp_group().send_tensor_dict(hidden_states.tensors,
                                            all_gather_group=get_tp_group())
            logits = None
        else:
            if self.input_batch.pooling_params:
                return self._pool(hidden_states, num_scheduled_tokens,
                                  num_scheduled_tokens_np, finished_sending,
                                  finished_recving)

            sample_hidden_states = hidden_states[logits_indices]
            logits = self.model.compute_logits(sample_hidden_states, None)
        if broadcast_pp_output:
            model_output_broadcast_data = {
                "logits": logits.contiguous(),
            } if logits is not None else {}
            model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
            assert model_output_broadcast_data is not None
            logits = model_output_broadcast_data["logits"]

        # Apply structured output bitmasks if present
        if scheduler_output.grammar_bitmask is not None:
            self.apply_grammar_bitmask(scheduler_output, logits)

        # Sample the next token and get logprobs if needed.
        sampling_metadata = self.input_batch.sampling_metadata
        if spec_decode_metadata is None:
            sampler_output = self.sampler(
                logits=logits,
                sampling_metadata=sampling_metadata,
            )
        else:
            # When indexing with a tensor (bonus_logits_indices), PyTorch
            # creates a new tensor with separate storage from the original
            # logits tensor. This means any in-place operations on bonus_logits
            # won't affect the original logits tensor.
            assert logits is not None

王敏's avatar
王敏 committed
3813
3814
3815
3816
3817
3818
3819
3820
3821
3822
3823
3824
3825
3826
3827
3828
3829
3830
3831
3832
3833
            if not envs.VLLM_REJECT_SAMPLE_OPT:
                bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
                sampler_output = self.sampler(
                    logits=bonus_logits,
                    sampling_metadata=sampling_metadata,
                )
                bonus_token_ids = sampler_output.sampled_token_ids

                # Just like `bonus_logits`, `target_logits` is a new tensor with
                # separate storage from the original `logits` tensor. Therefore,
                # it is safe to update `target_logits` in place.
                target_logits = logits[spec_decode_metadata.target_logits_indices]
                output_token_ids = self.rejection_sampler(
                    spec_decode_metadata,
                    None,  # draft_probs
                    target_logits,
                    bonus_token_ids,
                    sampling_metadata,
                )
                sampler_output.sampled_token_ids = output_token_ids
            else:
3834
3835
                # sampling_metadata.all_greedy = True
                # sampling_metadata.all_random = False
王敏's avatar
王敏 committed
3836
3837
3838
3839
3840
3841
3842
3843
3844
3845
3846
3847
3848
3849
3850
3851
3852
3853
                sampler_output = self.sampler(
                    logits=logits,
                    sampling_metadata=sampling_metadata,
                )
                target_token_ids = sampler_output.sampled_token_ids[spec_decode_metadata.target_logits_indices]
                target_logits = logits[spec_decode_metadata.target_logits_indices]

                bonus_token_ids = sampler_output.sampled_token_ids[spec_decode_metadata.bonus_logits_indices]

                output_token_ids = self.rejection_sampler(
                    spec_decode_metadata,
                    self.draft_probs.get_probs(spec_decode_metadata.spec_decode_ids),
                    target_logits,
                    target_token_ids,
                    bonus_token_ids,
                    sampling_metadata,
                )
                sampler_output.sampled_token_ids = output_token_ids
3854
3855
3856
3857
3858
3859
3860
3861
3862
3863
3864
3865
3866
3867
3868
3869
3870
3871
3872
3873
3874
3875
3876
3877
3878
3879
3880
3881
3882
3883
3884
3885
3886
3887
3888
3889
3890
3891
3892
3893
3894
3895

        num_nans_in_logits = {}
        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
            num_nans_in_logits = self._get_nans_in_logits(logits)

        # TODO(woosuk): The following loop can be slow since it iterates over
        # the requests one by one. Optimize.
        discard_sampled_tokens_req_indices = []
        for i, req_id in enumerate(self.input_batch.req_ids):
            req_state = self.requests[req_id]
            seq_len = (req_state.num_computed_tokens +
                       scheduler_output.num_scheduled_tokens[req_id])
            if seq_len < req_state.num_tokens:
                # Ignore the sampled token for partial prefills.
                # Rewind the generator state as if the token was not sampled.
                # This relies on cuda-specific torch-internal impl details
                generator = self.input_batch.generators.get(i)
                if generator is not None:
                    generator.set_offset(generator.get_offset() - 4)
                # Record the index of the request that should not be sampled,
                # so that we could clear the sampled tokens before returning.
                discard_sampled_tokens_req_indices.append(i)

        # NOTE: GPU -> CPU Sync happens here.
        # Move as many CPU operations as possible before this sync point.
        logprobs_tensors = sampler_output.logprobs_tensors
        logprobs_lists = logprobs_tensors.tolists() \
            if logprobs_tensors is not None else None

        # Compute prompt logprobs if needed.
        prompt_logprobs_dict = self._get_prompt_logprobs_dict(
            hidden_states[:num_scheduled_tokens],
            scheduler_output,
        )
        #-----------------------------------
        # Get the valid generated tokens.
        sampled_token_ids = sampler_output.sampled_token_ids
        max_gen_len = sampled_token_ids.shape[-1]
        if not self.speculative_config:
            # Speculative decoding is not enabled.
            spec_token_ids = None
        else:
3896
            sampled_token_ids_cpu = sampled_token_ids.to('cpu', non_blocking=True)
3897
3898
3899
3900
3901
            self.spec_sampler_event.record()
            mask = (sampled_token_ids == -1)
            mask_int = mask.int()
            first_neg_one_indices = torch.argmax(mask_int, dim=1)
            num_accepted_tokens_tensor = torch.where(torch.any(mask, dim=1), first_neg_one_indices, sampled_token_ids.size(1)) - 1
王敏's avatar
王敏 committed
3902

王敏's avatar
王敏 committed
3903
            spec_token_ids = self.zero_propose_draft_token_ids(
3904
3905
3906
3907
3908
3909
3910
3911
3912
3913
                scheduler_output,
                num_accepted_tokens_tensor,
                sampled_token_ids,
                sampling_metadata,
                hidden_states,
                sample_hidden_states,
                aux_hidden_states,
                spec_decode_metadata,
                attn_metadata,
            )
3914
                
3915
3916
        if max_gen_len == 1:
            # No spec decode tokens.
3917
            valid_sampled_token_ids = sampled_token_ids.tolist()
3918
3919
3920
3921
3922
3923
3924
3925
3926
3927
        else:
            # Includes spec decode tokens.
            self.spec_sampler_event.synchronize()
            valid_sampled_token_ids = self.rejection_sampler.parse_output(
                sampled_token_ids_cpu,
                self.input_batch.vocab_size,
            )
        # Mask out the sampled tokens that should not be sampled.
        for i in discard_sampled_tokens_req_indices:
            valid_sampled_token_ids[i].clear()
3928
3929
3930
        if spec_token_ids is not None:
            for i in discard_sampled_tokens_req_indices:
                spec_token_ids[i].clear()
3931
3932
3933
3934
3935
3936
3937
3938
3939
3940
3941
3942
3943
3944
3945
3946
3947
3948
3949
3950
3951
3952
3953
3954
3955
3956
3957
3958
3959
3960
3961
3962
3963
3964
3965
3966
3967
3968
3969
3970
3971
3972
3973
3974
3975
3976
3977
3978
3979
3980
3981
3982
3983
3984
3985
3986
3987
3988
3989
3990
3991
3992
3993
3994
3995
3996
3997
3998
3999
4000
4001
4002
4003
4004
4005
4006
4007
4008
4009
4010
4011
4012
4013
4014
4015
4016
4017
4018
4019
4020
4021
4022
4023
4024
4025
4026
4027
4028
4029
4030
4031
4032
4033
4034
4035
4036
4037
4038
4039
4040
4041
4042
4043
4044
4045
4046
4047
4048
4049
4050
4051
4052
4053
4054
4055
4056
4057
4058
4059
4060

        # Cache the sampled tokens in the model runner, so that the scheduler
        # doesn't need to send them back.
        # NOTE(woosuk): As an exception, when using PP, the scheduler sends
        # the sampled tokens back, because there's no direct communication
        # between the first-stage worker and the last-stage worker.
        for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
            if not sampled_ids:
                continue

            start_idx = self.input_batch.num_tokens_no_spec[req_idx]
            end_idx = start_idx + len(sampled_ids)
            assert end_idx <= self.max_model_len, (
                "Sampled token IDs exceed the max model length. "
                f"Total number of tokens: {end_idx} > max_model_len: "
                f"{self.max_model_len}")

            self.input_batch.token_ids_cpu[req_idx,
            start_idx:end_idx] = sampled_ids
            self.input_batch.num_tokens_no_spec[req_idx] = end_idx
            self.input_batch.num_tokens[req_idx] = end_idx
            req_id = self.input_batch.req_ids[req_idx]
            req_state = self.requests[req_id]
            req_state.output_token_ids.extend(sampled_ids)


        # Clear KVConnector state after all KVs are generated.
        if has_kv_transfer_group():
            get_kv_transfer_group().clear_connector_metadata()

        self.eplb_step()

        return ModelRunnerOutput(
            req_ids=self.input_batch.req_ids,
            req_id_to_index=self.input_batch.req_id_to_index,
            sampled_token_ids=valid_sampled_token_ids,
            spec_token_ids=spec_token_ids,
            logprobs=logprobs_lists,
            prompt_logprobs_dict=prompt_logprobs_dict,
            pooler_output=[],
            finished_sending=finished_sending,
            finished_recving=finished_recving,
            num_nans_in_logits=num_nans_in_logits,
        )

    def zero_propose_draft_token_ids(
            self,
            scheduler_output: "SchedulerOutput",
            num_accepted_tokens_tensor: torch.Tensor,
            sampled_token_ids: torch.Tensor,
            sampling_metadata: SamplingMetadata,
            hidden_states: torch.Tensor,
            sample_hidden_states: torch.Tensor,
            aux_hidden_states: Optional[torch.Tensor],
            spec_decode_metadata: Optional[SpecDecodeMetadata],
            attn_metadata: dict[str, Any],
    ) -> list[list[int]]:
        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
        if self.speculative_config.method == "ngram":
            assert isinstance(self.drafter, NgramProposer)
            spec_token_ids = self.propose_ngram_draft_token_ids(
                sampled_token_ids)
        elif self.speculative_config.method == "medusa":
            assert isinstance(self.drafter, MedusaProposer)
            if sample_hidden_states.shape[0] == len(sampled_token_ids):
                # The input to the target model does not include draft tokens.
                hidden_states = sample_hidden_states
            else:
                indices = []
                offset = 0
                for num_draft, tokens in zip(
                        spec_decode_metadata.num_draft_tokens,
                        sampled_token_ids):
                    indices.append(offset + len(tokens) - 1)
                    offset += num_draft + 1
                indices = torch.tensor(indices, device=self.device)
                hidden_states = sample_hidden_states[indices]

            spec_token_ids = self.drafter.propose(
                target_hidden_states=hidden_states,
                sampling_metadata=sampling_metadata,
            )
        elif self.speculative_config.use_eagle():
            assert isinstance(self.drafter, EagleProposer)
            # TODO(woosuk): Refactor the loop.
            row_indices = torch.arange(sampled_token_ids.size(0), device=sampled_token_ids.device)
            next_token_ids = sampled_token_ids[row_indices, num_accepted_tokens_tensor].flatten()
            # At this moment, we assume all eagle layers belong to the same KV
            # cache group, thus using the same attention metadata.
            eagle_attn_metadata = attn_metadata[
                self.drafter.attn_layer_names[0]]

            # NOTE: deepseek_mtp uses MLA which does not have `block_table`
            if hasattr(eagle_attn_metadata, "block_table"):
                block_table = eagle_attn_metadata.block_table
            else:
                block_table = None

            spec_scheduler_max_num_tokens = self.spec_scheduler_max_num_tokens
            if spec_decode_metadata is None:
                # input_ids can be None for multimodal models.
                target_token_ids = self.input_ids[:num_scheduled_tokens]
                # TODO(woosuk): Support M-RoPE.
                target_positions = self.positions[:num_scheduled_tokens]
                if self.use_aux_hidden_state_outputs:
                    target_hidden_states = torch.cat(
                        [h[:num_scheduled_tokens] for h in aux_hidden_states],
                        dim=-1)
                else:
                    target_hidden_states = hidden_states[:num_scheduled_tokens]
                target_slot_mapping = eagle_attn_metadata.slot_mapping
                cu_num_tokens = eagle_attn_metadata.query_start_loc
            else:
                # TODO(woosuk): Refactor this.
                cu_num_tokens, token_indices = self.drafter.prepare_inputs(
                    eagle_attn_metadata.query_start_loc,
                    num_accepted_tokens_tensor,
                )
                spec_scheduler_max_num_tokens = 1
                target_token_ids = self.input_ids[token_indices]
                # TODO(woosuk): Support M-RoPE.
                target_positions = self.positions[token_indices]
                if self.use_aux_hidden_state_outputs:
                    target_hidden_states = torch.cat(
                        [h[token_indices] for h in aux_hidden_states], dim=-1)
                else:
                    target_hidden_states = hidden_states[token_indices]
                target_slot_mapping = eagle_attn_metadata.slot_mapping[
                    token_indices]
            self.drafter.spec_scheduler_max_num_tokens = spec_scheduler_max_num_tokens
王敏's avatar
王敏 committed
4061
            draft_result = self.drafter.propose(
4062
4063
4064
4065
4066
4067
4068
4069
4070
4071
4072
4073
4074
4075
                target_token_ids=target_token_ids,
                target_positions=target_positions,
                target_hidden_states=target_hidden_states,
                target_slot_mapping=target_slot_mapping,
                next_token_ids=next_token_ids,
                cu_num_tokens=cu_num_tokens,
                block_table=block_table,
                sampling_metadata=sampling_metadata,
                decoding=spec_decode_metadata is not None,
            )
            # spec_token_ids = np.ones(draft_token_ids.shape, dtype=int).tolist()
            # self.last_draft_token_ids = draft_token_ids
            # self.last_draft_host_tokens = draft_token_ids.to('cpu', non_blocking=True)
            # self.last_draft_event.record()
王敏's avatar
王敏 committed
4076
4077
            if not envs.VLLM_REJECT_SAMPLE_OPT:
                draft_token_ids = draft_result
4078
            else:
王敏's avatar
王敏 committed
4079
                draft_token_ids, draft_probs = draft_result
王敏's avatar
王敏 committed
4080
4081
4082
4083
            spec_token_ids = draft_token_ids.tolist()

            if envs.VLLM_REJECT_SAMPLE_OPT:
                draft_req_ids = list(scheduler_output.num_scheduled_tokens.keys())
王敏's avatar
王敏 committed
4084
4085
4086
4087
4088
4089
                if self.draft_probs is None:
                    self.draft_probs = DraftProbs(
                        draft_probs, draft_req_ids)
                else:
                    self.draft_probs.update(draft_probs, draft_req_ids)

4090
4091
4092
4093
4094
        return spec_token_ids
#TODO:稳定后使用GPUModelRunnerMTP替换GPUModelRunner
if envs.VLLM_USE_ZERO_MTP:
    GPUModelRunner=GPUModelRunnerMTP
else:
laibao's avatar
laibao committed
4095
    GPUModelRunner=GPUModelRunnerBase