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

4
import gc
5
import itertools
6
import time
7
8
from collections import defaultdict
from collections.abc import Iterator
9
from contextlib import contextmanager
10
from copy import deepcopy
11
from functools import reduce
12
from itertools import product
13
from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast
14
15
16
17
18

import numpy as np
import torch
import torch.distributed
import torch.nn as nn
19
from tqdm import tqdm
20

21
import vllm.envs as envs
22
from vllm.attention import Attention, AttentionType
23
24
25
26
27
from vllm.attention.backends.abstract import (
    AttentionBackend,
    AttentionMetadata,
    MultipleOf,
)
28
from vllm.compilation.counter import compilation_counter
29
30
from vllm.compilation.cuda_graph import CUDAGraphWrapper
from vllm.compilation.monitor import set_cudagraph_capturing_enabled
31
from vllm.config import (
32
    CompilationMode,
33
34
35
36
37
    CUDAGraphMode,
    VllmConfig,
    get_layers_from_vllm_config,
    update_config,
)
38
from vllm.distributed.eplb.eplb_state import EplbState
39
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
40
from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
41
from vllm.distributed.parallel_state import (
42
    get_dcp_group,
43
44
45
46
47
48
    get_pp_group,
    get_tp_group,
    graph_capture,
    is_global_first_rank,
    prepare_communication_buffer_for_model,
)
49
from vllm.forward_context import BatchDescriptor, set_forward_context
50
from vllm.logger import init_logger
51
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
52
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
53
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
54
55
56
57
58
59
60
61
from vllm.model_executor.models.interfaces import (
    SupportsMultiModal,
    is_mixture_of_experts,
    supports_eagle3,
    supports_mrope,
    supports_multimodal_pruning,
    supports_transcription,
)
62
from vllm.model_executor.models.interfaces_base import (
63
64
65
66
    VllmModelForPooling,
    is_pooling_model,
    is_text_generation_model,
)
67
from vllm.multimodal import MULTIMODAL_REGISTRY
68
69
70
71
72
from vllm.multimodal.inputs import (
    BatchedTensorInputs,
    MultiModalKwargsItem,
    PlaceholderRange,
)
73
from vllm.multimodal.utils import group_mm_kwargs_by_modality
74
from vllm.pooling_params import PoolingParams
75
from vllm.sampling_params import SamplingType
76
from vllm.sequence import IntermediateTensors
77
from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
78
from vllm.utils import length_from_prompt_token_ids_or_embeds
79
from vllm.utils.jsontree import json_map_leaves
80
from vllm.utils.math_utils import cdiv, round_up
81
82
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.mem_utils import DeviceMemoryProfiler
83
from vllm.utils.platform_utils import is_pin_memory_available
84
85
86
87
88
from vllm.utils.torch_utils import (
    get_dtype_size,
    kv_cache_dtype_str_to_dtype,
    supports_dynamo,
)
89
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
90
from vllm.v1.attention.backends.utils import (
91
92
93
    AttentionCGSupport,
    AttentionMetadataBuilder,
    CommonAttentionMetadata,
94
    create_fast_prefill_custom_backend,
95
    get_dcp_local_seq_lens,
96
97
98
    reorder_batch_to_split_decodes_and_prefills,
    split_attn_metadata,
)
99
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
from vllm.v1.kv_cache_interface import (
    AttentionSpec,
    ChunkedLocalAttentionSpec,
    CrossAttentionSpec,
    EncoderOnlyAttentionSpec,
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
    KVCacheSpec,
    MambaSpec,
    SlidingWindowSpec,
    UniformTypeKVCacheSpecs,
)
from vllm.v1.outputs import (
    EMPTY_MODEL_RUNNER_OUTPUT,
    AsyncModelRunnerOutput,
    DraftTokenIds,
117
    KVConnectorOutput,
118
119
120
121
122
123
    LogprobsLists,
    LogprobsTensors,
    ModelRunnerOutput,
    PoolerOutput,
    SamplerOutput,
)
124
from vllm.v1.pool.metadata import PoolingMetadata
125
from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs
126
from vllm.v1.sample.metadata import SamplingMetadata
127
from vllm.v1.sample.rejection_sampler import RejectionSampler
128
from vllm.v1.sample.sampler import Sampler
129
from vllm.v1.spec_decode.eagle import EagleProposer
130
from vllm.v1.spec_decode.medusa import MedusaProposer
131
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
132
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
133
from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer
134
from vllm.v1.structured_output.utils import apply_grammar_bitmask
135
from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext
136
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
137
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
138
from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
139
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
140
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
141
142
143
144
145
from vllm.v1.worker.ubatch_utils import (
    UBatchSlice,
    UBatchSlices,
    check_ubatch_thresholds,
)
146
from vllm.v1.worker.utils import is_residual_scattered_for_sp
147

148
149
150
151
152
153
154
155
156
from .utils import (
    AttentionGroup,
    MultiModalBudget,
    add_kv_sharing_layers_to_kv_cache_groups,
    bind_kv_cache,
    gather_mm_placeholders,
    sanity_check_mm_encoder_outputs,
    scatter_mm_placeholders,
)
157

158
if TYPE_CHECKING:
159
    from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
160
    from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
161
162
163

logger = init_logger(__name__)

164
165
AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata]
# list when ubatching is enabled
166
PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict
167

168

169
170
171
172
173
174
# Wrapper for ModelRunnerOutput to support overlapped execution.
class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
    def __init__(
        self,
        model_runner_output: ModelRunnerOutput,
        sampled_token_ids: torch.Tensor,
175
        logprobs_tensors: torch.Tensor | None,
176
177
178
179
180
181
182
        invalid_req_indices: list[int],
        async_output_copy_stream: torch.cuda.Stream,
    ):
        self._model_runner_output = model_runner_output
        self._invalid_req_indices = invalid_req_indices

        # Event on the copy stream so we can synchronize the non-blocking copy.
183
        self.async_copy_ready_event = torch.cuda.Event()
184
185
186
187

        # Keep a reference to the device tensor to avoid it being
        # deallocated until we finish copying it to the host.
        self._sampled_token_ids = sampled_token_ids
188
        self._logprobs_tensors = logprobs_tensors
189
190
191
192
193

        # Initiate the copy on a separate stream, but do not synchronize it.
        default_stream = torch.cuda.current_stream()
        with torch.cuda.stream(async_output_copy_stream):
            async_output_copy_stream.wait_stream(default_stream)
194
            self.sampled_token_ids_cpu = self._sampled_token_ids.to(
195
196
                "cpu", non_blocking=True
            )
197
198
199
200
201
            self._logprobs_tensors_cpu = (
                self._logprobs_tensors.to_cpu_nonblocking()
                if self._logprobs_tensors
                else None
            )
202
            self.async_copy_ready_event.record()
203
204
205

    def get_output(self) -> ModelRunnerOutput:
        """Copy the device tensors to the host and return a ModelRunnerOutput.
206

207
208
        This function blocks until the copy is finished.
        """
209
        self.async_copy_ready_event.synchronize()
210

211
212
        # Release the device tensors once the copy has completed.
        del self._logprobs_tensors
213
214
        del self._sampled_token_ids

215
        valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist()
216
217
218
219
220
        for i in self._invalid_req_indices:
            valid_sampled_token_ids[i].clear()

        output = self._model_runner_output
        output.sampled_token_ids = valid_sampled_token_ids
221
222
223
224
        if self._logprobs_tensors_cpu:
            # NOTE(nick): this will need to be updated to use cu_num_accepted_tokens
            # for async sched + spec decode + logprobs compatibility.
            output.logprobs = self._logprobs_tensors_cpu.tolists()
225
226
227
        return output


228
229
230
231
232
233
234
235
236
237
238
239
240
241
class ExecuteModelState(NamedTuple):
    """Ephemeral cached state transferred between execute_model() and
    sample_tokens(), after execute_model() returns None."""

    scheduler_output: "SchedulerOutput"
    logits: torch.Tensor
    spec_decode_metadata: SpecDecodeMetadata | None
    spec_decode_common_attn_metadata: CommonAttentionMetadata | None
    hidden_states: torch.Tensor
    sample_hidden_states: torch.Tensor
    aux_hidden_states: list[torch.Tensor] | None
    kv_connector_output: KVConnectorOutput | None


242
class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
243
244
    def __init__(
        self,
245
        vllm_config: VllmConfig,
246
        device: torch.device,
247
    ):
248
249
250
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
        self.cache_config = vllm_config.cache_config
251
        self.compilation_config = vllm_config.compilation_config
252
253
254
255
256
257
        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.observability_config = vllm_config.observability_config
258

259
        from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
260
261

        set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3))
262

263
264
265
266
        model_config = self.model_config
        cache_config = self.cache_config
        scheduler_config = self.scheduler_config
        parallel_config = self.parallel_config
267
        self.device = device
268
269
        self.pin_memory = is_pin_memory_available()
        self.dtype = self.model_config.dtype
270
271
272
        self.kv_cache_dtype = kv_cache_dtype_str_to_dtype(
            cache_config.cache_dtype, self.model_config
        )
273

274
        self.is_pooling_model = model_config.runner_type == "pooling"
275
        self.enable_prompt_embeds = model_config.enable_prompt_embeds
276
        self.is_multimodal_raw_input_only_model = (
277
278
            model_config.is_multimodal_raw_input_only_model
        )
279
280
        # This will be overridden in load_model()
        self.is_multimodal_pruning_enabled = False
281
        self.max_model_len = model_config.max_model_len
282
283
284

        # Always set to false after the first forward pass
        self.calculate_kv_scales = self.cache_config.calculate_kv_scales
285
        self.dcp_world_size = self.parallel_config.decode_context_parallel_size
286
        self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group
287
        self.max_num_tokens = scheduler_config.max_num_batched_tokens
288
        self.max_num_reqs = scheduler_config.max_num_seqs
289

290
291
292
293
294
        # 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
        self.broadcast_pp_output = (
295
296
297
            self.parallel_config.distributed_executor_backend == "external_launcher"
            and len(get_pp_group().ranks) > 0
        )
298

299
        # Model-related.
300
        self.num_query_heads = model_config.get_num_attention_heads(parallel_config)
301
        self.hidden_size = model_config.get_hidden_size()
302
        self.attention_chunk_size = model_config.attention_chunk_size
303
        # Only relevant for models using ALiBi (e.g, MPT)
304
        self.use_alibi = model_config.uses_alibi
305

306
        self.cascade_attn_enabled = not self.model_config.disable_cascade_attn
307

308
        # Multi-modal data support
309
        self.mm_registry = MULTIMODAL_REGISTRY
310
        self.uses_mrope = model_config.uses_mrope
311
        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
312
313
            model_config
        )
314

315
316
317
        if self.model_config.is_encoder_decoder:
            # Maximum length of the encoder input, only for encoder-decoder
            # models.
318
            self.max_encoder_len = scheduler_config.max_num_encoder_input_tokens
319
320
321
        else:
            self.max_encoder_len = 0

322
        # Sampler
323
        self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
324

325
        self.eplb_state: EplbState | None = None
326
327
328
329
330
331
        """
        State of the expert parallelism load balancer.

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

332
        # Lazy initializations
333
        # self.model: nn.Module  # Set after load_model
334
        # Initialize in initialize_kv_cache
335
        self.kv_caches: list[torch.Tensor] = []
336
337
        # indexes: [kv_cache_group_id][attn_group]
        self.attn_groups: list[list[AttentionGroup]] = []
338
339
        # self.kv_cache_config: KVCacheConfig

340
341
        # mm_hash ->  encoder_output
        self.encoder_cache: dict[str, torch.Tensor] = {}
342

343
        self.use_aux_hidden_state_outputs = False
344
345
346
347
348
        # 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:
349
350
351
            self.drafter: (
                NgramProposer | SuffixDecodingProposer | EagleProposer | MedusaProposer
            )
352
353
            if self.speculative_config.method == "ngram":
                self.drafter = NgramProposer(self.vllm_config)
354
355
            elif self.speculative_config.method == "suffix":
                self.drafter = SuffixDecodingProposer(self.vllm_config)
356
            elif self.speculative_config.use_eagle():
357
                self.drafter = EagleProposer(self.vllm_config, self.device, self)
358
359
360
361
                if self.speculative_config.method == "eagle3":
                    self.use_aux_hidden_state_outputs = True
            elif self.speculative_config.method == "medusa":
                self.drafter = MedusaProposer(
362
                    vllm_config=self.vllm_config, device=self.device
363
                )
364
            else:
365
366
367
368
                raise ValueError(
                    "Unknown speculative decoding method: "
                    f"{self.speculative_config.method}"
                )
369
            self.rejection_sampler = RejectionSampler(self.sampler)
370

371
        # Request states.
372
        self.requests: dict[str, CachedRequestState] = {}
373
        self.comm_stream = torch.cuda.Stream()
374

375
376
377
378
379
380
381
382
383
        # 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.
384
        custom_logitsprocs = model_config.logits_processors
385
386
        self.input_batch = InputBatch(
            max_num_reqs=self.max_num_reqs,
387
388
389
            # We need to use the encoder length for encoder-decoer
            # because of KV cache for cross-attention.
            max_model_len=max(self.max_model_len, self.max_encoder_len),
390
391
392
            max_num_batched_tokens=self.max_num_tokens,
            device=self.device,
            pin_memory=self.pin_memory,
393
            vocab_size=self.model_config.get_vocab_size(),
394
            block_sizes=[self.cache_config.block_size],
395
            kernel_block_sizes=[self.cache_config.block_size],
396
            is_spec_decode=bool(self.vllm_config.speculative_config),
397
            logitsprocs=build_logitsprocs(
398
399
400
                self.vllm_config,
                self.device,
                self.pin_memory,
401
                self.is_pooling_model,
402
                custom_logitsprocs,
403
            ),
404
405
406
            # We currently don't know whether a particular custom logits processor
            # uses output token ids so we set this conservatively.
            logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
407
            is_pooling_model=self.is_pooling_model,
408
            dcp_kv_cache_interleave_size=self.parallel_config.dcp_kv_cache_interleave_size,
409
        )
410

411
        self.use_async_scheduling = self.scheduler_config.async_scheduling
412
413
414
415
416
417
418
419
420
        # Separate cuda stream for overlapping transfer of sampled token ids from
        # GPU to CPU when async scheduling is enabled.
        self.async_output_copy_stream: torch.cuda.Stream | None = None
        # cuda event to synchronize use of reused CPU tensors between steps
        # when async scheduling is enabled.
        self.prepare_inputs_event: torch.cuda.Event | None = None
        if self.use_async_scheduling:
            self.async_output_copy_stream = torch.cuda.Stream()
            self.prepare_inputs_event = torch.cuda.Event()
421

422
        # self.cudagraph_batch_sizes sorts in ascending order.
423
424
425
426
        if (
            self.compilation_config.cudagraph_capture_sizes
            and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
        ):
427
428
            self.cudagraph_batch_sizes = sorted(
                self.compilation_config.cudagraph_capture_sizes
429
            )
430

431
        # Cache the device properties.
432
        self._init_device_properties()
433

434
        # Persistent buffers for CUDA graphs.
435
436
437
438
439
        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
        self.positions = self._make_buffer(self.max_num_tokens, dtype=torch.int64)
        self.query_start_loc = self._make_buffer(
            self.max_num_reqs + 1, dtype=torch.int32
        )
440
        self.seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
441
442
443
444
        if self.dcp_world_size > 1:
            self.dcp_local_seq_lens = self._make_buffer(
                self.max_num_reqs, dtype=torch.int32
            )
445
446
447
        # Because inputs_embeds may be bfloat16 and we don't need a numpy
        # version of this tensor, avoid a RuntimeError by not creating a
        # numpy buffer.
448
449
450
451
452
453
454
        self.inputs_embeds = self._make_buffer(
            self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False
        )
        self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool)
        self.discard_request_indices = self._make_buffer(
            self.max_num_reqs, dtype=torch.int64
        )
455
456
        self.num_discarded_requests = 0

457
458
459
460
461
462
        self.num_decode_draft_tokens = self._make_buffer(
            self.max_num_reqs, dtype=torch.int32
        )
        self.num_accepted_tokens = self._make_buffer(
            self.max_num_reqs, dtype=torch.int64
        )
463

464
465
        # Only relevant for multimodal models
        if self.supports_mm_inputs:
466
            self.is_mm_embed = self._make_buffer(self.max_num_tokens, dtype=torch.bool)
467

468
        # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
469
        if self.uses_mrope:
Roger Wang's avatar
Roger Wang committed
470
471
472
473
            # 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
474
475
476
477
478
479

            # 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
480
            self.mrope_positions = self._make_buffer(
481
482
                (3, self.max_num_tokens + 1), dtype=torch.int64
            )
483

484
        # None in the first PP rank. The rest are set after load_model.
485
        self.intermediate_tensors: IntermediateTensors | None = None
486

487
        # OPTIMIZATION: Cache the tensors rather than creating them every step.
488
        # Keep in int64 to avoid overflow with long context
489
490
491
492
        self.arange_np = np.arange(
            max(self.max_num_reqs + 1, self.max_model_len, self.max_num_tokens),
            dtype=np.int64,
        )
493

494
495
496
497
498
        # 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] = {}
499
500
501
502
503
        self.kv_sharing_fast_prefill_eligible_layers: set[str] = set()

        self.kv_sharing_fast_prefill_logits_indices = None
        if self.cache_config.kv_sharing_fast_prefill:
            self.kv_sharing_fast_prefill_logits_indices = torch.zeros(
504
505
                self.max_num_tokens, dtype=torch.int32, device=self.device
            )
506

507
508
509
510
511
        self.uniform_decode_query_len = (
            1
            if not self.speculative_config
            else 1 + self.speculative_config.num_speculative_tokens
        )
512
513
514
515

        # Cudagraph dispatcher for runtime cudagraph dispatching.
        self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)

516
517
518
519
520
521
522
523
524
        self.mm_budget = (
            MultiModalBudget(
                self.model_config,
                self.scheduler_config,
                self.mm_registry,
            )
            if self.supports_mm_inputs
            else None
        )
525

526
        self.reorder_batch_threshold: int | None = None
527

528
529
530
531
532
        # Attention layers that are only in the KVCacheConfig of the runner
        # (e.g., KV sharing, encoder-only attention), but not in the
        # KVCacheConfig of the scheduler.
        self.runner_only_attn_layers: set[str] = set()

533
        # Cached outputs.
534
        self._draft_token_ids: list[list[int]] | torch.Tensor | None = None
535
536
        self.transfer_event = torch.cuda.Event()
        self.sampled_token_ids_pinned_cpu = torch.empty(
537
            (self.max_num_reqs, 1),
538
539
            dtype=torch.int64,
            device="cpu",
540
541
            pin_memory=self.pin_memory,
        )
542

543
544
545
        # Ephemeral state transferred between execute_model() and sample_tokens().
        self.execute_model_state: ExecuteModelState | None = None

546
547
548
549
    def reset_mm_cache(self) -> None:
        if self.mm_budget:
            self.mm_budget.reset_cache()

550
551
552
553
554
555
556
557
558
559
    def _get_positions(self, num_tokens: Any):
        if isinstance(num_tokens, int):
            if self.uses_mrope:
                return self.mrope_positions.gpu[:, :num_tokens]
            return self.positions.gpu[:num_tokens]
        else:
            if self.uses_mrope:
                return self.mrope_positions.gpu[:, num_tokens]
            return self.positions.gpu[num_tokens]

560
    def _make_buffer(
561
        self, *size: int | torch.SymInt, dtype: torch.dtype, numpy: bool = True
562
563
564
565
566
567
568
569
    ) -> CpuGpuBuffer:
        return CpuGpuBuffer(
            *size,
            dtype=dtype,
            device=self.device,
            pin_memory=self.pin_memory,
            with_numpy=numpy,
        )
570

571
572
573
    def _init_model_kwargs(self, num_tokens: int):
        model_kwargs = dict[str, Any]()

574
        if not self.is_pooling_model:
575
576
            return model_kwargs

577
578
        num_reqs = self.input_batch.num_reqs
        pooling_params = self.input_batch.get_pooling_params()
579
580
581

        token_type_id_requests = dict[int, Any]()
        for i, param in enumerate(pooling_params):
582
583
584
585
586
            if (
                param.extra_kwargs is not None
                and (token_types := param.extra_kwargs.get("compressed_token_type_ids"))
                is not None
            ):
587
588
589
590
591
                token_type_id_requests[i] = token_types

        if len(token_type_id_requests) == 0:
            return model_kwargs

592
        seq_lens = self.seq_lens.gpu[:num_reqs]
593
594
595
596
597
598
599
600
        token_type_ids = []

        for i in range(num_reqs):
            pos = token_type_id_requests.get(i, seq_lens[i])
            ids = (torch.arange(seq_lens[i]) >= pos).int()
            token_type_ids.append(ids)

        model_kwargs["token_type_ids"] = torch.concat(token_type_ids).to(
601
602
            device=self.device
        )
603
604
        return model_kwargs

605
    def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
606
607
        """
        Update the order of requests in the batch based on the attention
608
        backend's needs. For example, some attention backends (namely MLA) may
609
610
611
612
613
614
        want to separate requests based on if the attention computation will be
        compute-bound or memory-bound.

        Args:
            scheduler_output: The scheduler output.
        """
615
616
617
618
619
620
621
622
        # Attention free models have zero kv_cache_goups, however models
        # like Mamba are also attention free but use the kv_cache for
        # keeping its internal state. This is why we check the number
        # of kv_cache groups instead of solely checking
        # for self.model_config.is_attention_free.
        if len(self.kv_cache_config.kv_cache_groups) == 0:
            return

623
        if self.reorder_batch_threshold is not None:
624
625
626
            # NOTE(lucas): currently no backend supports the custom masking
            #  required for DCP with q_len > 1, so we assert here. Remove this
            #  assert once the custom mask is support is added to FA3.
627
628
629
630
            if (
                self.dcp_world_size > 1
                and envs.VLLM_ATTENTION_BACKEND != "FLASH_ATTN_MLA"
            ):
631
                assert self.reorder_batch_threshold == 1, (
632
                    "DCP not support reorder_batch_threshold > 1 now."
633
                )
634
635
636
            reorder_batch_to_split_decodes_and_prefills(
                self.input_batch,
                scheduler_output,
637
638
                decode_threshold=self.reorder_batch_threshold,
            )
639

640
641
    # Note: used for model runner override.
    def _init_device_properties(self) -> None:
642
        """Initialize attributes from torch.cuda.get_device_properties"""
643
644
645
646
647
648
649
        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()

650
    def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
651
652
653
654
655
656
        """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.

657
658
        The SamplingMetadata is updated and copied to the GPU if there is a
        new/resumed/paused/finished request in the batch.
659
660
        """
        # Remove finished requests from the cached states.
661
662
        for req_id in scheduler_output.finished_req_ids:
            self.requests.pop(req_id, None)
663
664
665
666
667
668
669
        # 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:
670
            self.input_batch.remove_request(req_id)
671
672

        # Free the cached encoder outputs.
673
674
        for mm_hash in scheduler_output.free_encoder_mm_hashes:
            self.encoder_cache.pop(mm_hash, None)
675

676
677
678
679
680
681
682
683
684
685
686
687
688
        # 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:
689
            self.input_batch.remove_request(req_id)
690

691
        reqs_to_add: list[CachedRequestState] = []
692
        # Add new requests to the cached states.
693
694
695
        for new_req_data in scheduler_output.scheduled_new_reqs:
            req_id = new_req_data.req_id
            sampling_params = new_req_data.sampling_params
696
            pooling_params = new_req_data.pooling_params
697

698
699
700
701
            if (
                sampling_params
                and sampling_params.sampling_type == SamplingType.RANDOM_SEED
            ):
702
703
704
705
706
                generator = torch.Generator(device=self.device)
                generator.manual_seed(sampling_params.seed)
            else:
                generator = None

707
708
            if self.is_pooling_model:
                assert pooling_params is not None
709
710
                task = pooling_params.task
                assert task is not None, "You did not set `task` in the API"
711

712
                model = cast(VllmModelForPooling, self.get_model())
713
                to_update = model.pooler.get_pooling_updates(task)
714
715
                to_update.apply(pooling_params)

716
            req_state = CachedRequestState(
717
                req_id=req_id,
718
                prompt_token_ids=new_req_data.prompt_token_ids,
719
                prompt_embeds=new_req_data.prompt_embeds,
720
                mm_features=new_req_data.mm_features,
721
                sampling_params=sampling_params,
722
                pooling_params=pooling_params,
723
                generator=generator,
724
725
                block_ids=new_req_data.block_ids,
                num_computed_tokens=new_req_data.num_computed_tokens,
726
                output_token_ids=[],
727
                lora_request=new_req_data.lora_request,
728
            )
729
730
            self.requests[req_id] = req_state

731
            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
732
            if self.uses_mrope:
733
                self._init_mrope_positions(req_state)
734

735
            reqs_to_add.append(req_state)
736

737
        # Update the states of the running/resumed requests.
738
        is_last_rank = get_pp_group().is_last_rank
739
740
        req_data = scheduler_output.scheduled_cached_reqs
        for i, req_id in enumerate(req_data.req_ids):
741
            req_state = self.requests[req_id]
742
743
            num_computed_tokens = req_data.num_computed_tokens[i]
            new_block_ids = req_data.new_block_ids[i]
744
            resumed_from_preemption = req_id in req_data.resumed_req_ids
745
            num_output_tokens = req_data.num_output_tokens[i]
746

747
            # Update the cached states.
748

749
            req_state.num_computed_tokens = num_computed_tokens
750
            req_index = self.input_batch.req_id_to_index.get(req_id)
751
752
753
754
755
756
757
758

            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.
759
760
761
                num_new_tokens = (
                    num_computed_tokens + len(new_token_ids) - req_state.num_tokens
                )
762
763
764
765
                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:
766
                    req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
767
768
769
770
771
            elif num_output_tokens < len(req_state.output_token_ids):
                # Some output tokens were discarded due to a sync-KV-load
                # failure. Align the cached state.
                del req_state.output_token_ids[num_output_tokens:]
                if req_index is not None:
772
773
774
775
                    end_idx = (
                        self.input_batch.num_prompt_tokens[req_index]
                        + num_output_tokens
                    )
776
777
                    self.input_batch.num_tokens[req_index] = end_idx
                    self.input_batch.num_tokens_no_spec[req_index] = end_idx
778

779
            # Update the block IDs.
780
            if not resumed_from_preemption:
781
782
                if new_block_ids is not None:
                    # Append the new blocks to the existing block IDs.
783
                    for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
784
                        block_ids.extend(new_ids)
785
            else:
786
                assert req_index is None
787
                assert new_block_ids is not None
788
789
                # The request is resumed from preemption.
                # Replace the existing block IDs with the new ones.
790
                req_state.block_ids = new_block_ids
791
792
793
794
795

            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.
796
797
798
799
800
801
802

                if self.use_async_scheduling and num_output_tokens > 0:
                    # We must recover the output token ids for resumed requests in the
                    # async scheduling case, so that correct input_ids are obtained.
                    resumed_token_ids = req_data.all_token_ids[req_id]
                    req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]

803
                reqs_to_add.append(req_state)
804
805
806
                continue

            # Update the persistent batch.
807
            self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
808
            if new_block_ids is not None:
809
                self.input_batch.block_table.append_row(new_block_ids, req_index)
810
811
812
813
814
815
816

            # 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
                end_token_index = num_computed_tokens + len(new_token_ids)
817
                self.input_batch.token_ids_cpu[
818
819
820
                    req_index, start_token_index:end_token_index
                ] = new_token_ids
                self.input_batch.num_tokens_no_spec[req_index] = end_token_index
821
                self.input_batch.num_tokens[req_index] = end_token_index
822

823
            # Add spec_token_ids to token_ids_cpu.
824
            spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
825
                req_id, []
826
            )
827
828
829
830
831
            if spec_token_ids:
                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
                self.input_batch.token_ids_cpu[
832
833
                    req_index, start_index:end_token_index
                ] = spec_token_ids
834
835
                # NOTE(woosuk): `num_tokens` here may include spec tokens.
                self.input_batch.num_tokens[req_index] += num_spec_tokens
836
837
838
839
840
841
842

            # When speculative decoding is used with structured output,
            # the scheduler can drop draft tokens that do not
            # conform to the schema. This can result in
            # scheduler_output.scheduled_spec_decode_tokens being empty,
            # even when speculative decoding is enabled.
            self.input_batch.spec_token_ids[req_index] = spec_token_ids
843

844
845
        # Add the new or resumed requests to the persistent batch.
        # The smaller empty indices are filled first.
846
847
        for request in reqs_to_add:
            self.input_batch.add_request(request)
848

849
850
851
852
853
854
        # 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)
        # Refresh batch metadata with any pending updates.
        self.input_batch.refresh_metadata()
855

856
    def _update_states_after_model_execute(
857
858
        self, output_token_ids: torch.Tensor
    ) -> None:
859
860
861
862
863
864
865
866
867
868
869
870
        """Update the cached states after model execution.

        This is used for MTP/EAGLE for hybrid models, as in linear attention,
        only the last token's state is kept. In MTP/EAGLE, for draft tokens
        the state are kept util we decide how many tokens are accepted for
        each sequence, and a shifting is done during the next iteration
        based on the number of accepted tokens.
        """
        if not self.model_config.is_hybrid or not self.speculative_config:
            return

        # Find the number of accepted tokens for each sequence.
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
        num_accepted_tokens = (
            (
                torch.cat(
                    [
                        output_token_ids,
                        torch.full(
                            (output_token_ids.size(0), 1),
                            -1,
                            device=output_token_ids.device,
                        ),
                    ],
                    dim=1,
                )
                == -1
            )
            .int()
            .argmax(-1)
            .cpu()
            .numpy()
        )
891
892
893
        for i, num_tokens in enumerate(num_accepted_tokens):
            self.input_batch.num_accepted_tokens_cpu[i] = num_tokens

894
    def _init_mrope_positions(self, req_state: CachedRequestState):
895
896
        model = self.get_model()
        assert supports_mrope(model), "M-RoPE support is not implemented."
897
898

        req_state.mrope_positions, req_state.mrope_position_delta = (
899
            model.get_mrope_input_positions(
900
                req_state.prompt_token_ids,
901
                req_state.mm_features,
902
            )
903
        )
904

905
    def _extract_mm_kwargs(
906
        self,
907
908
        scheduler_output: "SchedulerOutput",
    ) -> BatchedTensorInputs:
909
        if not scheduler_output or not self.is_multimodal_raw_input_only_model:
910
            return {}
911

912
913
        mm_kwargs = list[MultiModalKwargsItem]()
        for req in scheduler_output.scheduled_new_reqs:
914
915
916
            for feature in req.mm_features:
                if feature.data is not None:
                    mm_kwargs.append(feature.data)
917

918
        # Input all modalities at once
919
        model = cast(SupportsMultiModal, self.model)
920
921
        mm_kwargs_combined: BatchedTensorInputs = {}
        for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
922
923
924
925
            mm_kwargs,
            device=self.device,
            pin_memory=self.pin_memory,
            merge_by_field_config=model.merge_by_field_config,
926
            multimodal_cpu_fields=model.multimodal_cpu_fields,
927
928
        ):
            mm_kwargs_combined.update(mm_kwargs_group)
929

930
        return mm_kwargs_combined
931

932
    def _dummy_mm_kwargs(self, num_seqs: int) -> BatchedTensorInputs:
933
        if not self.is_multimodal_raw_input_only_model:
934
            return {}
935

936
937
938
939
940
        mm_budget = self.mm_budget
        assert mm_budget is not None

        dummy_modality = mm_budget.get_modality_with_max_tokens()
        return self._get_mm_dummy_batch(dummy_modality, num_seqs)
941

942
943
944
    def _get_cumsum_and_arange(
        self,
        num_tokens: np.ndarray,
945
        cumsum_dtype: np.dtype | None = None,
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
    ) -> 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

962
963
964
    def _prepare_input_ids(
        self, total_num_scheduled_tokens: int, cu_num_tokens: np.ndarray
    ) -> None:
965
        """Prepare the input IDs for the current batch.
966

967
968
969
970
971
972
973
        Carefully handles the `prev_sampled_token_ids` which can be cached
        from the previous engine iteration, in which case those tokens on the
        GPU need to be copied into the corresponding slots into input_ids."""

        if self.input_batch.prev_sampled_token_ids is None:
            # Normal scheduling case
            self.input_ids.copy_to_gpu(total_num_scheduled_tokens)
974
975
976
            if self.enable_prompt_embeds:
                self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)
                self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
            return

        # Async scheduling case, where some decode requests from the previous
        # iteration won't have entries in input_ids_cpu and need to be copied
        # on the GPU from prev_sampled_token_ids.
        prev_req_id_to_index = self.input_batch.prev_req_id_to_index
        assert prev_req_id_to_index is not None
        flattened_indices = []
        prev_common_req_indices = []
        indices_match = True
        max_flattened_index = -1
        for req_id, cur_index in self.input_batch.req_id_to_index.items():
            if (prev_index := prev_req_id_to_index.get(req_id)) is not None:
                prev_common_req_indices.append(prev_index)
                # We need to compute the flattened input_ids index of the
                # last token in each common request.
                flattened_index = cu_num_tokens[cur_index].item() - 1
                flattened_indices.append(flattened_index)
995
                indices_match &= prev_index == flattened_index
996
997
998
999
1000
1001
                max_flattened_index = max(max_flattened_index, flattened_index)
        num_commmon_tokens = len(flattened_indices)
        if num_commmon_tokens < total_num_scheduled_tokens:
            # If not all requests are decodes from the last iteration,
            # We need to copy the input_ids_cpu to the GPU first.
            self.input_ids.copy_to_gpu(total_num_scheduled_tokens)
1002
1003
1004
            if self.enable_prompt_embeds:
                self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)
                self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
1005
1006
        if num_commmon_tokens == 0:
            # No requests in common with the previous iteration
1007
            # So input_ids.cpu will have all the input ids.
1008
1009
1010
1011
1012
1013
1014
            return
        if indices_match and max_flattened_index == (num_commmon_tokens - 1):
            # Common-case optimization: the batch is unchanged
            # and no reordering happened.
            # The indices are both the same permutation of 0..N-1 so
            # we can copy directly using a single slice.
            self.input_ids.gpu[:num_commmon_tokens].copy_(
1015
1016
1017
                self.input_batch.prev_sampled_token_ids[:num_commmon_tokens, 0],
                non_blocking=True,
            )
1018
1019
            if self.enable_prompt_embeds:
                self.is_token_ids.gpu[:num_commmon_tokens] = True
1020
            return
1021
        # Upload the index tensors asynchronously so the scatter can be non-blocking.
1022
1023
1024
        input_ids_index_tensor = torch.tensor(
            flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory
        ).to(self.device, non_blocking=True)
1025
        prev_common_req_indices_tensor = torch.tensor(
1026
1027
            prev_common_req_indices, dtype=torch.int64, pin_memory=self.pin_memory
        ).to(self.device, non_blocking=True)
1028
1029
1030
1031
        self.input_ids.gpu.scatter_(
            dim=0,
            index=input_ids_index_tensor,
            src=self.input_batch.prev_sampled_token_ids[
1032
1033
1034
                prev_common_req_indices_tensor, 0
            ],
        )
1035

1036
1037
    def _get_encoder_seq_lens(
        self,
1038
        scheduled_encoder_inputs: dict[str, list[int]],
1039
1040
        kv_cache_spec: KVCacheSpec,
        num_reqs: int,
1041
    ) -> np.ndarray | None:
1042
1043
1044
1045
1046
1047
        if not isinstance(kv_cache_spec, CrossAttentionSpec):
            return None

        # Build encoder_seq_lens array mapping request indices to
        # encoder lengths for inputs scheduled in this batch
        encoder_seq_lens = np.zeros(num_reqs, dtype=np.int32)
1048
        for req_id in scheduled_encoder_inputs:
1049
1050
1051
1052
1053
            req_index = self.input_batch.req_id_to_index[req_id]
            encoder_seq_lens[req_index] = self.max_encoder_len

        return encoder_seq_lens

1054
    def _prepare_inputs(
1055
1056
1057
1058
        self,
        scheduler_output: "SchedulerOutput",
        num_scheduled_tokens: np.ndarray,
        max_num_scheduled_tokens: int,
1059
1060
    ) -> tuple[
        torch.Tensor,
1061
1062
1063
        SpecDecodeMetadata | None,
        UBatchSlices | None,
        torch.Tensor | None,
1064
    ]:
1065
1066
        """
        :return: tuple[
1067
            logits_indices, spec_decode_metadata,
1068
            ubatch_slices, num_tokens_across_dp,
1069
1070
        ]
        """
1071
1072
1073
1074
1075
1076
1077
        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.
1078
        self.input_batch.block_table.commit_block_table(num_reqs)
1079
1080
1081

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

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

        # Get positions.
1089
        positions_np = self.positions.np[:total_num_scheduled_tokens]
1090
1091
1092
1093
1094
        np.add(
            self.input_batch.num_computed_tokens_cpu[req_indices],
            arange,
            out=positions_np,
        )
1095

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

1101
1102
1103
1104
        # 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.
1105
1106
1107
        token_indices = (
            positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]
        )
1108
        token_indices_tensor = torch.from_numpy(token_indices)
1109

1110
1111
1112
        # 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.
1113
1114
1115
1116
1117
1118
        torch.index_select(
            self.input_batch.token_ids_cpu_tensor.flatten(),
            0,
            token_indices_tensor,
            out=self.input_ids.cpu[:total_num_scheduled_tokens],
        )
1119
        if self.enable_prompt_embeds:
1120
            is_token_ids = self.input_batch.is_token_ids_tensor.flatten()
1121
1122
1123
1124
            torch.index_select(
                is_token_ids,
                0,
                token_indices_tensor,
1125
1126
                out=self.is_token_ids.cpu[:total_num_scheduled_tokens],
            )
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159

        # Because we did not pre-allocate a massive prompt_embeds CPU tensor on
        # the InputBatch, we need to fill in the prompt embeds into the expected
        # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor.
        if self.input_batch.req_prompt_embeds:
            output_idx = 0
            for req_idx in range(num_reqs):
                num_sched = num_scheduled_tokens[req_idx]

                # Skip if this request doesn't have embeddings
                if req_idx not in self.input_batch.req_prompt_embeds:
                    output_idx += num_sched
                    continue

                # Skip if no tokens scheduled
                if num_sched <= 0:
                    output_idx += num_sched
                    continue

                req_embeds = self.input_batch.req_prompt_embeds[req_idx]
                start_pos = self.input_batch.num_computed_tokens_cpu[req_idx]

                # Skip if trying to read beyond available embeddings
                if start_pos >= req_embeds.shape[0]:
                    output_idx += num_sched
                    continue

                # Copy available embeddings
                end_pos = start_pos + num_sched
                actual_end = min(end_pos, req_embeds.shape[0])
                actual_num_sched = actual_end - start_pos

                if actual_num_sched > 0:
1160
1161
1162
                    self.inputs_embeds.cpu[
                        output_idx : output_idx + actual_num_sched
                    ].copy_(req_embeds[start_pos:actual_end])
1163
1164

                output_idx += num_sched
1165

1166
1167
        self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)
        self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)
1168
1169

        # Prepare the attention metadata.
1170
        self.query_start_loc.np[0] = 0
1171
        self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens
1172
1173
        # Note: pad query_start_loc to be non-decreasing, as kernels
        # like FlashAttention requires that
1174
        self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1])
1175
        self.query_start_loc.copy_to_gpu()
1176
        query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]
1177

1178
        num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
1179
        num_tokens_padded = self._get_num_input_tokens(num_tokens_unpadded)
1180
1181
1182
        uniform_decode = (
            max_num_scheduled_tokens == self.uniform_decode_query_len
        ) and (total_num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)
1183
1184
1185
1186
1187
1188
1189

        # Disable DP padding when running eager to avoid excessive padding when
        # running prefills. 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.
        allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

1190
        ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(
1191
1192
1193
1194
1195
1196
1197
            num_tokens_unpadded=num_tokens_unpadded,
            parallel_config=self.parallel_config,
            allow_microbatching=True,
            allow_dp_padding=allow_dp_padding,
            num_tokens_padded=num_tokens_padded,
            uniform_decode=uniform_decode,
            num_scheduled_tokens_per_request=num_scheduled_tokens,
1198
        )
1199

1200
        self.seq_lens.np[:num_reqs] = (
1201
1202
            self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens
        )
1203
        # Fill unused with 0 for full cuda graph mode.
1204
1205
        self.seq_lens.np[num_reqs:].fill(0)
        self.seq_lens.copy_to_gpu()
1206

1207
        num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids]
1208
1209
1210
1211
1212
1213
1214
        num_tokens_np = np.array(num_tokens, dtype=np.int32)

        # Record the index of requests that should not be sampled,
        # so that we could clear the sampled tokens before returning
        discard_requests_mask = self.seq_lens.np[:num_reqs] < num_tokens_np
        discard_request_indices = np.nonzero(discard_requests_mask)[0]
        self.num_discarded_requests = len(discard_request_indices)
1215
1216
1217
        self.discard_request_indices.np[: self.num_discarded_requests] = (
            discard_request_indices
        )
1218
1219
1220

        self.discard_request_indices.copy_to_gpu(self.num_discarded_requests)

1221
        # Copy the tensors to the GPU.
1222
1223
        self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens)

1224
        if self.uses_mrope:
1225
            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
1226
1227
            self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(
                self.mrope_positions.cpu[:, :total_num_scheduled_tokens],
1228
1229
                non_blocking=True,
            )
1230
1231
        else:
            # Common case (1D positions)
1232
            self.positions.copy_to_gpu(total_num_scheduled_tokens)
1233

1234
        use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
1235
1236
1237
1238
1239
1240
1241
        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
1242
            num_draft_tokens = None
1243
            spec_decode_metadata = None
1244
            num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
1245
1246
1247
1248
1249
        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)
1250
1251
1252
            # For chunked prefills, use -1 as mask rather than 0, as guided
            # decoding may rollback speculative tokens.
            num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32)
1253
1254
1255
1256
            for (
                req_id,
                draft_token_ids,
            ) in scheduler_output.scheduled_spec_decode_tokens.items():
1257
1258
                req_idx = self.input_batch.req_id_to_index[req_id]
                num_draft_tokens[req_idx] = len(draft_token_ids)
1259
1260
1261
1262
1263
1264
1265
1266
                num_decode_draft_tokens[req_idx] = (
                    len(draft_token_ids)
                    if (
                        self.input_batch.num_computed_tokens_cpu[req_idx]
                        >= self.input_batch.num_prompt_tokens[req_idx]
                    )
                    else -1
                )
1267
            spec_decode_metadata = self._calc_spec_decode_metadata(
1268
1269
                num_draft_tokens, cu_num_tokens
            )
1270
            logits_indices = spec_decode_metadata.logits_indices
1271
            num_sampled_tokens = num_draft_tokens + 1
1272
            # For DECODE only cuda graph of some attention backends (e.g., GDN).
1273
            self.num_decode_draft_tokens.np[:num_reqs] = num_decode_draft_tokens
1274
1275
            self.num_decode_draft_tokens.np[num_reqs:].fill(-1)
            self.num_decode_draft_tokens.copy_to_gpu()
1276

1277
1278
1279
1280
1281
        # Hot-Swap lora model
        if self.lora_config:
            assert (
                np.sum(num_sampled_tokens)
                <= self.vllm_config.scheduler_config.max_num_batched_tokens
1282
            )
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
            self.set_active_loras(
                self.input_batch, num_scheduled_tokens, num_sampled_tokens
            )

        return (
            logits_indices,
            spec_decode_metadata,
            ubatch_slices,
            num_tokens_across_dp,
        )

    def _build_attention_metadata(
        self,
        total_num_scheduled_tokens: int,
        max_num_scheduled_tokens: int,
        num_reqs: int,
        ubatch_slices: UBatchSlices | None = None,
        logits_indices: torch.Tensor | None = None,
        use_spec_decode: bool = False,
        for_cudagraph_capture: bool = False,
        scheduled_encoder_inputs: dict[str, list[int]] | None = None,
        cascade_attn_prefix_lens: list[list[int]] | None = None,
    ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]:
        """
        :return: tuple[attn_metadata, spec_decode_common_attn_metadata]
        """
        logits_indices_padded = None
        num_logits_indices = 0
        if logits_indices is not None:
            num_logits_indices = logits_indices.size(0)
            if self.cache_config.kv_sharing_fast_prefill:
                logits_indices_padded = self._prepare_kv_sharing_fast_prefill(
                    logits_indices
                )
1317

1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
        # update seq_lens of decode reqs under DCP.
        if self.dcp_world_size > 1:
            self.dcp_local_seq_lens.cpu[:num_reqs] = get_dcp_local_seq_lens(
                self.seq_lens.cpu[:num_reqs],
                self.dcp_world_size,
                self.dcp_rank,
                self.parallel_config.dcp_kv_cache_interleave_size,
            )
            self.dcp_local_seq_lens.copy_to_gpu(num_reqs)

1328
1329
1330
        attn_metadata: PerLayerAttnMetadata = {}
        if ubatch_slices is not None:
            attn_metadata = [dict() for _ in range(len(ubatch_slices))]
1331

1332
1333
        # Used in the below loop
        query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]
1334
        query_start_loc_cpu = self.query_start_loc.cpu[: num_reqs + 1]
1335
        seq_lens = self.seq_lens.gpu[:num_reqs]
1336
        seq_lens_cpu = self.seq_lens.cpu[:num_reqs]
1337
1338
1339
        num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[
            :num_reqs
        ]
1340
1341
1342
        dcp_local_seq_lens = (
            self.dcp_local_seq_lens.gpu[:num_reqs] if self.dcp_world_size > 1 else None
        )
1343
        spec_decode_common_attn_metadata = None
1344
1345
1346
1347
1348
1349
1350
1351
1352

        if for_cudagraph_capture:
            # For some attention backends (e.g. FA) with sliding window models we need
            # to make sure the backend see a max_seq_len that is larger to the sliding
            # window size when capturing to make sure the correct kernel is selected.
            max_seq_len = self.max_model_len
        else:
            max_seq_len = self.seq_lens.np[:num_reqs].max().item()

1353
1354
        if use_spec_decode:
            self.num_accepted_tokens.np[:num_reqs] = (
1355
1356
                self.input_batch.num_accepted_tokens_cpu[:num_reqs]
            )
1357
1358
            self.num_accepted_tokens.np[num_reqs:].fill(1)
            self.num_accepted_tokens.copy_to_gpu()
1359

1360
1361
        # Prepare the attention metadata for each KV cache group and make layers
        # in the same group share the same metadata.
1362
        for kv_cache_gid, kv_cache_group in enumerate(
1363
1364
            self.kv_cache_config.kv_cache_groups
        ):
1365
            encoder_seq_lens = self._get_encoder_seq_lens(
1366
1367
1368
                scheduled_encoder_inputs or {},
                kv_cache_group.kv_cache_spec,
                num_reqs,
1369
            )
1370

1371
            if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):
1372
1373
1374
1375
1376
                # Encoder-only layers do not have KV cache, so we need to
                # create a dummy block table and slot mapping for them.
                blk_table_tensor = torch.zeros(
                    (num_reqs, 1),
                    dtype=torch.int32,
1377
1378
1379
                    device=self.device,
                )
                slot_mapping = torch.zeros(
1380
                    (total_num_scheduled_tokens,),
1381
1382
1383
                    dtype=torch.int64,
                    device=self.device,
                )
1384
            else:
1385
                blk_table = self.input_batch.block_table[kv_cache_gid]
1386
                blk_table_tensor = blk_table.get_device_tensor(num_reqs)
1387
                slot_mapping = blk_table.slot_mapping.gpu[:total_num_scheduled_tokens]
1388
1389
1390

                # Fill unused with -1. Needed for reshape_and_cache in full cuda
                # graph mode.
1391
                blk_table.slot_mapping.gpu[total_num_scheduled_tokens:].fill_(-1)
1392

1393
            common_attn_metadata = CommonAttentionMetadata(
1394
1395
1396
1397
1398
                query_start_loc=query_start_loc,
                query_start_loc_cpu=query_start_loc_cpu,
                seq_lens=seq_lens,
                seq_lens_cpu=seq_lens_cpu,
                num_computed_tokens_cpu=num_computed_tokens_cpu,
1399
1400
1401
                num_reqs=num_reqs,
                num_actual_tokens=total_num_scheduled_tokens,
                max_query_len=max_num_scheduled_tokens,
1402
                max_seq_len=max_seq_len,
1403
1404
                block_table_tensor=blk_table_tensor,
                slot_mapping=slot_mapping,
1405
                logits_indices_padded=logits_indices_padded,
1406
                num_logits_indices=num_logits_indices,
1407
                causal=True,
1408
                encoder_seq_lens=encoder_seq_lens,
1409
                dcp_local_seq_lens=dcp_local_seq_lens,
1410
1411
            )

1412
            if self.speculative_config and spec_decode_common_attn_metadata is None:
1413
                if isinstance(self.drafter, EagleProposer):
1414
                    if self.drafter.attn_layer_names[0] in kv_cache_group.layer_names:
1415
1416
1417
                        spec_decode_common_attn_metadata = common_attn_metadata
                else:
                    spec_decode_common_attn_metadata = common_attn_metadata
1418

1419
1420
1421
1422
1423
1424
            for attn_gid, attn_group in enumerate(self.attn_groups[kv_cache_gid]):
                cascade_attn_prefix_len = (
                    cascade_attn_prefix_lens[kv_cache_gid][attn_gid]
                    if cascade_attn_prefix_lens
                    else 0
                )
1425
                builder = attn_group.get_metadata_builder()
1426

1427
                extra_attn_metadata_args = {}
1428
                if use_spec_decode and isinstance(builder, GDNAttentionMetadataBuilder):
1429
                    extra_attn_metadata_args = dict(
1430
1431
1432
1433
                        num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs],
                        num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[
                            :num_reqs
                        ],
1434
1435
                    )

1436
1437
                if ubatch_slices is not None:
                    common_attn_metadata_list = split_attn_metadata(
1438
1439
                        ubatch_slices, common_attn_metadata
                    )
1440
                    for ubid, common_attn_metadata in enumerate(
1441
1442
                        common_attn_metadata_list
                    ):
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
                        builder = attn_group.get_metadata_builder(ubatch_id=ubid)
                        if for_cudagraph_capture:
                            attn_metadata_i = builder.build_for_cudagraph_capture(
                                common_attn_metadata
                            )
                        else:
                            attn_metadata_i = builder.build(
                                common_prefix_len=cascade_attn_prefix_len,
                                common_attn_metadata=common_attn_metadata,
                            )
                        for layer_name in kv_cache_group.layer_names:
1454
1455
1456
1457
                            assert type(attn_metadata) is list
                            attn_metadata[ubid][layer_name] = attn_metadata_i
                else:
                    assert isinstance(attn_metadata, dict)
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
                    if for_cudagraph_capture:
                        attn_metadata_i = builder.build_for_cudagraph_capture(
                            common_attn_metadata
                        )
                    else:
                        attn_metadata_i = builder.build(
                            common_prefix_len=cascade_attn_prefix_len,
                            common_attn_metadata=common_attn_metadata,
                            **extra_attn_metadata_args,
                        )
1468
1469
                    for layer_name in attn_group.layer_names:
                        attn_metadata[layer_name] = attn_metadata_i
1470

1471
        return attn_metadata, spec_decode_common_attn_metadata
1472

1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
    def _compute_cascade_attn_prefix_lens(
        self,
        num_scheduled_tokens: np.ndarray,
        num_common_prefix_blocks: list[int],
    ) -> list[list[int]] | None:
        """
        :return: Optional[cascade_attn_prefix_lens]
            cascade_attn_prefix_lens is 2D: ``[kv_cache_group_id][attn_group_idx]``,
            None if we should not use cascade attention
        """
1483

1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
        use_cascade_attn = False
        num_kv_cache_groups = len(self.kv_cache_config.kv_cache_groups)
        cascade_attn_prefix_lens: list[list[int]] = [
            [] for _ in range(num_kv_cache_groups)
        ]

        for kv_cache_gid in range(num_kv_cache_groups):
            for attn_group in self.attn_groups[kv_cache_gid]:
                if isinstance(attn_group.kv_cache_spec, EncoderOnlyAttentionSpec):
                    cascade_attn_prefix_len = 0
                else:
                    # 0 if cascade attention should not be used
                    cascade_attn_prefix_len = self._compute_cascade_attn_prefix_len(
                        num_scheduled_tokens,
                        num_common_prefix_blocks[kv_cache_gid],
                        attn_group.kv_cache_spec,
                        attn_group.get_metadata_builder(),
                    )
                cascade_attn_prefix_lens[kv_cache_gid].append(cascade_attn_prefix_len)
                use_cascade_attn |= cascade_attn_prefix_len > 0

        return cascade_attn_prefix_lens if use_cascade_attn else None
1506

1507
1508
1509
1510
    def _compute_cascade_attn_prefix_len(
        self,
        num_scheduled_tokens: np.ndarray,
        num_common_prefix_blocks: int,
1511
1512
        kv_cache_spec: KVCacheSpec,
        attn_metadata_builder: AttentionMetadataBuilder,
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
    ) -> 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.
        """
1531

1532
        common_prefix_len = num_common_prefix_blocks * kv_cache_spec.block_size
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
        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]
1570
        # Request 3's num_computed_tokens: 3 (i.e., [A, B, C])
1571
1572
1573
1574
1575
1576
1577
        # 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(
1578
1579
            common_prefix_len, self.input_batch.num_computed_tokens_cpu[:num_reqs].min()
        )
1580
        # common_prefix_len should be a multiple of the block size.
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
        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
        )
        use_local_attention = isinstance(kv_cache_spec, ChunkedLocalAttentionSpec) or (
            isinstance(kv_cache_spec, FullAttentionSpec)
            and kv_cache_spec.attention_chunk_size is not None
        )
1592
1593
        assert isinstance(kv_cache_spec, AttentionSpec)
        use_cascade = attn_metadata_builder.use_cascade_attention(
1594
1595
1596
            common_prefix_len=common_prefix_len,
            query_lens=num_scheduled_tokens,
            num_query_heads=self.num_query_heads,
1597
            num_kv_heads=kv_cache_spec.num_kv_heads,
1598
            use_alibi=self.use_alibi,
1599
            use_sliding_window=use_sliding_window,
1600
            use_local_attention=use_local_attention,
1601
            num_sms=self.num_sms,
1602
            dcp_world_size=self.dcp_world_size,
1603
1604
1605
        )
        return common_prefix_len if use_cascade else 0

1606
1607
    def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"):
        mrope_pos_ptr = 0
1608
        for index, req_id in enumerate(self.input_batch.req_ids):
1609
1610
1611
            req = self.requests[req_id]
            assert req.mrope_positions is not None

1612
1613
            num_computed_tokens = self.input_batch.num_computed_tokens_cpu[index]
            num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
1614
            num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
1615
1616
                req.prompt_token_ids, req.prompt_embeds
            )
1617
1618

            if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens:
1619
1620
                prompt_part_len = max(0, num_prompt_tokens - num_computed_tokens)
                completion_part_len = max(0, num_scheduled_tokens - prompt_part_len)
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
            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

1634
1635
1636
                self.mrope_positions.cpu[:, dst_start:dst_end] = req.mrope_positions[
                    :, src_start:src_end
                ]
1637
1638
1639
1640
1641
1642
1643
                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

1644
                MRotaryEmbedding.get_next_input_positions_tensor(
1645
                    out=self.mrope_positions.np,
1646
1647
1648
1649
1650
                    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,
                )
1651
1652
1653

                mrope_pos_ptr += completion_part_len

1654
1655
    def _calc_spec_decode_metadata(
        self,
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
        num_draft_tokens: np.ndarray,
        cu_num_scheduled_tokens: np.ndarray,
    ) -> 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
1672
1673
1674
1675

        # 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(
1676
1677
            num_sampled_tokens, cumsum_dtype=np.int32
        )
1678
        # Step 2. [0, 0, 0, 0, 103, 104, 104, 104, 206, 207, 207]
1679
        logits_indices = np.repeat(
1680
1681
            cu_num_scheduled_tokens - num_sampled_tokens, num_sampled_tokens
        )
1682
        # Step 3. [0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
1683
1684
1685
1686
1687
1688
        logits_indices += arange

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

        # Compute the draft logits indices.
1689
1690
1691
        # 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(
1692
1693
            num_draft_tokens, cumsum_dtype=np.int32
        )
1694
1695
        # [0, 0, 0, 5, 5, 9]
        target_logits_indices = np.repeat(
1696
1697
            cu_num_sampled_tokens - num_sampled_tokens, num_draft_tokens
        )
1698
1699
1700
1701
1702
        # [0, 1, 2, 5, 6, 9]
        target_logits_indices += arange

        # TODO: Optimize the CPU -> GPU copy.
        cu_num_draft_tokens = torch.from_numpy(cu_num_draft_tokens).to(
1703
1704
            self.device, non_blocking=True
        )
1705
1706
1707
        cu_num_sampled_tokens = torch.from_numpy(cu_num_sampled_tokens).to(
            self.device, non_blocking=True
        )
1708
1709
1710
        logits_indices = torch.from_numpy(logits_indices).to(
            self.device, non_blocking=True
        )
1711
        target_logits_indices = torch.from_numpy(target_logits_indices).to(
1712
1713
            self.device, non_blocking=True
        )
1714
        bonus_logits_indices = torch.from_numpy(bonus_logits_indices).to(
1715
1716
            self.device, non_blocking=True
        )
1717

1718
1719
        # Compute the draft token ids.
        # draft_token_indices:      [  1,   2,   3, 105, 106, 208]
1720
        draft_token_ids = self.input_ids.gpu[logits_indices]
1721
1722
        draft_token_ids = draft_token_ids[target_logits_indices + 1]

1723
        return SpecDecodeMetadata(
1724
1725
1726
            draft_token_ids=draft_token_ids,
            num_draft_tokens=num_draft_tokens.tolist(),
            cu_num_draft_tokens=cu_num_draft_tokens,
1727
            cu_num_sampled_tokens=cu_num_sampled_tokens,
1728
1729
1730
1731
1732
            target_logits_indices=target_logits_indices,
            bonus_logits_indices=bonus_logits_indices,
            logits_indices=logits_indices,
        )

1733
1734
1735
1736
1737
1738
1739
    def _prepare_kv_sharing_fast_prefill(
        self,
        logits_indices: torch.Tensor,
    ) -> torch.Tensor:
        assert self.kv_sharing_fast_prefill_logits_indices is not None
        num_logits = logits_indices.shape[0]
        assert num_logits > 0
1740
        self.kv_sharing_fast_prefill_logits_indices[:num_logits].copy_(logits_indices)
1741
1742
1743
1744
1745
        # There might have leftover indices in logits_indices[num_logits:]
        # from previous iterations, whose values may be greater than the
        # batch size in the current iteration. To ensure indices are always
        # valid, we fill the padded indices with the last index.
        self.kv_sharing_fast_prefill_logits_indices[num_logits:].fill_(
1746
1747
1748
1749
1750
1751
            logits_indices[-1].item()
        )
        if (
            self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
            and num_logits <= self.cudagraph_batch_sizes[-1]
        ):
1752
1753
1754
1755
1756
            # Use piecewise CUDA graphs.
            # Add padding to the batch size.
            num_logits_padded = self.vllm_config.pad_for_cudagraph(num_logits)
        else:
            num_logits_padded = num_logits
1757
1758
1759
        logits_indices_padded = self.kv_sharing_fast_prefill_logits_indices[
            :num_logits_padded
        ]
1760
1761
        return logits_indices_padded

1762
1763
1764
1765
1766
1767
1768
1769
    def _batch_mm_kwargs_from_scheduler(
        self,
        scheduler_output: "SchedulerOutput",
    ) -> tuple[list[MultiModalKwargsItem], list[tuple[str, PlaceholderRange]]]:
        """Batch multimodal kwargs from scheduled encoder inputs.

        Args:
            scheduler_output: The scheduler output containing scheduled encoder
1770
                inputs.
1771
1772
1773
1774
1775
1776

        Returns:
            A tuple of (mm_kwargs, req_ids_pos) where:
            - mm_kwargs: List of multimodal kwargs items to be batched
            - mm_hashes_pos: List of (mm_hash, position_info) tuples
        """
1777
1778
        scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
        if not scheduled_encoder_inputs:
1779
            return [], []
1780
        # Batch the multi-modal inputs.
1781
        mm_kwargs = list[MultiModalKwargsItem]()
1782
1783
        # list of tuple (mm_hash, position_info)
        mm_hashes_pos = list[tuple[str, PlaceholderRange]]()
1784
1785
        for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
            req_state = self.requests[req_id]
1786
1787

            for mm_input_id in encoder_input_ids:
1788
1789
1790
1791
                mm_feature = req_state.mm_features[mm_input_id]
                mm_hash = mm_feature.identifier
                mm_kwargs.append(mm_feature.data)
                mm_hashes_pos.append((mm_hash, mm_feature.mm_position))
1792

1793
1794
1795
1796
1797
        return mm_kwargs, mm_hashes_pos

    def _execute_mm_encoder(self, scheduler_output: "SchedulerOutput"):
        # Batch the multi-modal inputs using the helper method.
        mm_kwargs, mm_hashes_pos = self._batch_mm_kwargs_from_scheduler(
1798
1799
            scheduler_output
        )
1800
1801
1802
1803

        if not mm_kwargs:
            return

1804
1805
1806
1807
1808
1809
1810
        # 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.
1811
        model = cast(SupportsMultiModal, self.model)
1812
        encoder_outputs = []
1813
        for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
1814
1815
1816
1817
            mm_kwargs,
            device=self.device,
            pin_memory=self.pin_memory,
            merge_by_field_config=model.merge_by_field_config,
1818
            multimodal_cpu_fields=model.multimodal_cpu_fields,
1819
        ):
1820
1821
1822
            curr_group_outputs = []

            # EVS-related change.
1823
            # (ekhvedchenia): Temporary hack to limit peak memory usage when
1824
            # processing multimodal data. This solves the issue with scheduler
1825
1826
1827
1828
            # putting too many video samples into a single batch. Scheduler
            # uses pruned vision tokens count to compare it versus compute
            # budget which is incorrect (Either input media size or non-pruned
            # output vision tokens count should be considered)
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
            # TODO(ywang96): Fix memory profiling to take EVS into account and
            # remove this hack.
            if (
                self.is_multimodal_pruning_enabled
                and modality == "video"
                and num_items > 1
            ):
                for video_mm_kwargs_item in filter(
                    lambda item: item.modality == "video", mm_kwargs
                ):
                    _, _, micro_batch_mm_inputs = next(
                        group_mm_kwargs_by_modality(
                            [video_mm_kwargs_item],
                            device=self.device,
                            pin_memory=self.pin_memory,
                            merge_by_field_config=model.merge_by_field_config,
1845
                            multimodal_cpu_fields=model.multimodal_cpu_fields,
1846
                        )
1847
                    )
1848
1849

                    micro_batch_outputs = model.get_multimodal_embeddings(
1850
1851
                        **micro_batch_mm_inputs
                    )
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861

                    curr_group_outputs.extend(micro_batch_outputs)
            else:
                # 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.
1862
                curr_group_outputs = model.get_multimodal_embeddings(**mm_kwargs_group)
1863

1864
1865
            sanity_check_mm_encoder_outputs(
                curr_group_outputs,
1866
                expected_num_items=num_items,
1867
            )
1868
            encoder_outputs.extend(curr_group_outputs)
1869

1870
1871
1872
        # Cache the encoder outputs by mm_hash
        for (mm_hash, pos_info), output in zip(mm_hashes_pos, encoder_outputs):
            self.encoder_cache[mm_hash] = scatter_mm_placeholders(
1873
1874
1875
1876
1877
                output,
                is_embed=pos_info.is_embed,
            )

    def _gather_mm_embeddings(
1878
1879
        self,
        scheduler_output: "SchedulerOutput",
1880
        shift_computed_tokens: int = 0,
1881
1882
1883
1884
1885
1886
1887
1888
    ) -> tuple[list[torch.Tensor], torch.Tensor]:
        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

        mm_embeds = list[torch.Tensor]()
        is_mm_embed = self.is_mm_embed.cpu
        is_mm_embed[:total_num_scheduled_tokens] = False

        req_start_idx = 0
1889
        should_sync_mrope_positions = False
1890

1891
        for req_id in self.input_batch.req_ids:
1892
1893
            mm_embeds_req: list[torch.Tensor] = []

1894
            num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
1895
            req_state = self.requests[req_id]
1896
            num_computed_tokens = req_state.num_computed_tokens + shift_computed_tokens
1897

1898
1899
            for mm_feature in req_state.mm_features:
                pos_info = mm_feature.mm_position
1900
1901
                start_pos = pos_info.offset
                num_encoder_tokens = pos_info.length
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917

                # 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,
1918
1919
                    num_encoder_tokens,
                )
1920
                assert start_idx < end_idx
1921

1922
                mm_hash = mm_feature.identifier
1923
                encoder_output = self.encoder_cache.get(mm_hash, None)
1924
                assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."
1925
1926
1927
1928

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

1929
                req_start_pos = req_start_idx + start_pos - num_computed_tokens
1930
1931
1932
                is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] = (
                    True if is_embed is None else is_embed
                )
1933

1934
1935
1936
1937
                mm_embeds_item = gather_mm_placeholders(
                    encoder_output[start_idx:end_idx],
                    is_embed=is_embed,
                )
1938
1939
1940
                mm_embeds_req.append(mm_embeds_item)

            if self.is_multimodal_pruning_enabled and self.uses_mrope:
1941
                assert req_state.mrope_positions is not None
1942
1943
1944
1945
1946
1947
1948
                should_sync_mrope_positions = True
                mm_embeds_req, new_mrope_positions, new_delta = (
                    self.model.recompute_mrope_positions(
                        input_ids=req_state.prompt_token_ids,
                        multimodal_embeddings=mm_embeds_req,
                        mrope_positions=req_state.mrope_positions,
                        num_computed_tokens=req_state.num_computed_tokens,
1949
1950
                    )
                )
1951
1952
1953
1954
                req_state.mrope_positions.copy_(new_mrope_positions)
                req_state.mrope_position_delta = new_delta

            mm_embeds.extend(mm_embeds_req)
1955
1956
1957
            req_start_idx += num_scheduled_tokens

        is_mm_embed = self.is_mm_embed.copy_to_gpu(total_num_scheduled_tokens)
1958
1959
1960

        if should_sync_mrope_positions:
            self._calc_mrope_positions(scheduler_output)
1961
            self.mrope_positions.copy_to_gpu(total_num_scheduled_tokens)
1962

1963
        return mm_embeds, is_mm_embed
1964

1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
    def _extract_encoder_inputs(
        self,
        scheduler_output: "SchedulerOutput",
    ) -> dict[str, torch.Tensor]:
        """Extract encoder inputs for encoder-decoder models.

        This method extracts multimodal input features from scheduled encoder
        inputs and formats them for the encoder-decoder model forward pass.
        """
        # Batch the multi-modal inputs using the helper method.
        mm_kwargs, _ = self._batch_mm_kwargs_from_scheduler(scheduler_output)

        if not mm_kwargs:
            return {}

        # Group MM kwargs by modality and extract features
1981
        model = cast(SupportsMultiModal, self.model)
1982
1983
        encoder_features = {}
        for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
1984
1985
1986
1987
            mm_kwargs,
            device=self.device,
            pin_memory=self.pin_memory,
            merge_by_field_config=model.merge_by_field_config,
1988
            multimodal_cpu_fields=model.multimodal_cpu_fields,
1989
1990
1991
1992
1993
1994
1995
1996
        ):
            # Add the grouped features to encoder_features dict
            # This allows the model to receive them as kwargs (e.g.,
            # input_features=...)
            encoder_features.update(mm_kwargs_group)

        return encoder_features

1997
    def get_model(self) -> nn.Module:
1998
        # get raw model out of the cudagraph wrapper.
1999
        if isinstance(self.model, (CUDAGraphWrapper, UBatchWrapper)):
2000
            return self.model.unwrap()
2001
2002
        return self.model

2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
    def get_supported_generation_tasks(self) -> list[GenerationTask]:
        model = self.get_model()
        supported_tasks = list[GenerationTask]()

        if is_text_generation_model(model):
            supported_tasks.append("generate")

        if supports_transcription(model):
            if model.supports_transcription_only:
                return ["transcription"]

            supported_tasks.append("transcription")

        return supported_tasks

2018
2019
2020
2021
2022
    def get_supported_pooling_tasks(self) -> list[PoolingTask]:
        model = self.get_model()
        if not is_pooling_model(model):
            return []

2023
2024
        supported_tasks = list(model.pooler.get_supported_tasks())

2025
2026
2027
2028
2029
        if self.scheduler_config.chunked_prefill_enabled:
            if "token_embed" in supported_tasks:
                supported_tasks.remove("token_embed")
            if "token_classify" in supported_tasks:
                supported_tasks.remove("token_classify")
2030

2031
2032
            logger.debug_once(
                "Chunked prefill is not supported with "
2033
2034
                "token_embed and token_classify tasks "
                "which using ALL pooling. "
2035
2036
2037
                "Please turn off chunked prefill by "
                "`--no-enable-chunked-prefill` before using it."
            )
2038
2039
2040
2041
2042

        if "score" in supported_tasks:
            num_labels = getattr(self.model_config.hf_config, "num_labels", 0)
            if num_labels != 1:
                supported_tasks.remove("score")
2043
                logger.debug_once("Score API is only enabled for num_labels == 1.")
2044
2045

        return supported_tasks
2046

2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
        tasks = list[SupportedTask]()

        if self.model_config.runner_type == "generate":
            tasks.extend(self.get_supported_generation_tasks())
        if self.model_config.runner_type == "pooling":
            tasks.extend(self.get_supported_pooling_tasks())

        return tuple(tasks)

2057
    def sync_and_slice_intermediate_tensors(
2058
2059
2060
2061
2062
        self,
        num_tokens: int,
        intermediate_tensors: IntermediateTensors,
        sync_self: bool,
    ) -> IntermediateTensors:
2063
2064
2065
        assert self.intermediate_tensors is not None

        tp = self.vllm_config.parallel_config.tensor_parallel_size
2066
        is_rs = is_residual_scattered_for_sp(self.vllm_config, num_tokens)
2067
2068
2069
2070
2071
2072

        # 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():
2073
                is_scattered = k == "residual" and is_rs
2074
                copy_len = num_tokens // tp if is_scattered else num_tokens
2075
                self.intermediate_tensors[k][:copy_len].copy_(
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
                    v[:copy_len], non_blocking=True
                )

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

    def eplb_step(self, is_dummy: bool = False, is_profile: bool = False) -> None:
2089
2090
2091
2092
2093
2094
2095
        """
        Step for the EPLB (Expert Parallelism Load Balancing) state.
        """
        if not self.parallel_config.enable_eplb:
            return

        assert self.eplb_state is not None
2096
2097
        model = self.get_model()
        assert is_mixture_of_experts(model)
2098
2099
2100
        self.eplb_state.step(
            is_dummy,
            is_profile,
2101
            log_stats=self.parallel_config.eplb_config.log_balancedness,
2102
2103
        )

2104
2105
2106
2107
    # This is where the second ubatch is adjusted to account for the padding.
    # Should be called after attention metadata creation. This just pads
    # the second ubatch slice out to the total number of tokens
    # (num_tokens + padding)
2108
2109
    @staticmethod
    def pad_out_ubatch_slice(ubatch_slices: UBatchSlices, num_total_tokens: int):
2110
2111
2112
2113
2114
2115
        padded_second_ubatch_slice = slice(
            ubatch_slices[1].token_slice.start, num_total_tokens
        )
        ubatch_slices[1] = UBatchSlice(
            padded_second_ubatch_slice, padded_second_ubatch_slice
        )
2116

2117
2118
2119
2120
2121
2122
    def _pool(
        self,
        hidden_states: torch.Tensor,
        num_scheduled_tokens: int,
        num_scheduled_tokens_np: np.ndarray,
    ) -> ModelRunnerOutput:
2123
2124
2125
        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"
        )
2126

2127
        hidden_states = hidden_states[:num_scheduled_tokens]
2128
        pooling_metadata = self.input_batch.get_pooling_metadata()
2129
2130
2131
2132
        pooling_metadata.build_pooling_cursor(
            num_scheduled_tokens_np.tolist(), device=hidden_states.device
        )
        seq_lens_cpu = self.seq_lens.cpu[: self.input_batch.num_reqs]
2133

2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
        model = cast(VllmModelForPooling, self.model)
        raw_pooler_output: PoolerOutput = model.pooler(
            hidden_states=hidden_states,
            pooling_metadata=pooling_metadata,
        )
        raw_pooler_output = json_map_leaves(
            lambda x: x.to("cpu", non_blocking=True),
            raw_pooler_output,
        )
        self._sync_device()
2144

2145
        pooler_output: list[torch.Tensor | None] = []
2146
        for raw_output, seq_len, prompt_len in zip(
2147
2148
            raw_pooler_output, seq_lens_cpu, pooling_metadata.prompt_lens
        ):
2149
            output = raw_output if seq_len == prompt_len else None
2150
            pooler_output.append(output)
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160

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

2161
    def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:
2162
2163
2164
2165
2166
2167
        if (
            self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
            and hasattr(self, "cudagraph_batch_sizes")
            and self.cudagraph_batch_sizes
            and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]
        ):
2168
2169
2170
2171
2172
2173
2174
2175
            # Use CUDA graphs.
            # Add padding to the batch size.
            return self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)

        # 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
2176
2177
2178
2179
        if (
            self.compilation_config.pass_config.enable_sequence_parallelism
            and tp_size > 1
        ):
2180
2181
2182
            return round_up(num_scheduled_tokens, tp_size)
        return num_scheduled_tokens

2183
    def _preprocess(
2184
2185
        self,
        scheduler_output: "SchedulerOutput",
2186
        num_input_tokens: int,  # Padded
2187
        intermediate_tensors: IntermediateTensors | None = None,
2188
    ) -> tuple[
2189
2190
        torch.Tensor | None,
        torch.Tensor | None,
2191
        torch.Tensor,
2192
        IntermediateTensors | None,
2193
2194
        dict[str, Any],
    ]:
2195
        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2196
        is_first_rank = get_pp_group().is_first_rank
2197

2198
2199
        # _prepare_inputs may reorder the batch, so we must gather multi
        # modal outputs after that to ensure the correct order
2200
2201
        if (
            self.supports_mm_inputs
2202
            and is_first_rank
2203
2204
            and not self.model_config.is_encoder_decoder
        ):
2205
2206
            # Run the multimodal encoder if any.
            self._execute_mm_encoder(scheduler_output)
2207
            mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output)
2208

2209
2210
2211
            # 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.
2212
            inputs_embeds_scheduled = self.model.get_input_embeddings(
2213
2214
2215
                self.input_ids.gpu[:num_scheduled_tokens],
                multimodal_embeddings=mm_embeds,
                is_multimodal=is_mm_embed,
2216
            )
2217

2218
            # TODO(woosuk): Avoid the copy. Optimize.
2219
            self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
2220

2221
            input_ids = None
2222
            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
2223
2224
2225
2226
            model_kwargs = {
                **self._init_model_kwargs(num_scheduled_tokens),
                **self._extract_mm_kwargs(scheduler_output),
            }
2227
        elif self.enable_prompt_embeds and is_first_rank:
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
            # Get the input embeddings for the tokens that are not input embeds,
            # then put them into the appropriate positions.
            # TODO(qthequartermasterman): Since even when prompt embeds are
            # enabled, (a) not all requests will use prompt embeds, and (b)
            # after the initial prompt is processed, the rest of the generated
            # tokens will be token ids, it is not desirable to have the
            # embedding layer outside of the CUDA graph all the time. The v0
            # engine avoids this by "double compiling" the CUDA graph, once
            # with input_ids and again with inputs_embeds, for all num_tokens.
            # If a batch only has token ids, then including the embedding layer
            # in the CUDA graph will be more performant (like in the else case
            # below).
2240
2241
2242
            token_ids_idx = (
                self.is_token_ids.gpu[:num_scheduled_tokens]
                .nonzero(as_tuple=False)
2243
                .squeeze(1)
2244
            )
2245
2246
2247
            # Some tokens ids may need to become embeds
            if token_ids_idx.numel() > 0:
                token_ids = self.input_ids.gpu[token_ids_idx]
2248
                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
2249
2250
2251
2252
2253
                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds

            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
            model_kwargs = self._init_model_kwargs(num_input_tokens)
            input_ids = None
2254
        else:
2255
2256
2257
2258
            # 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.
2259
            input_ids = self.input_ids.gpu[:num_input_tokens]
2260
            inputs_embeds = None
2261
            model_kwargs = self._init_model_kwargs(num_input_tokens)
2262
        if self.uses_mrope:
2263
            positions = self.mrope_positions.gpu[:, :num_input_tokens]
2264
        else:
2265
            positions = self.positions.gpu[:num_input_tokens]
2266

2267
        if is_first_rank:
2268
2269
            intermediate_tensors = None
        else:
2270
            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
2271
2272
                num_input_tokens, intermediate_tensors, True
            )
2273

2274
2275
2276
2277
        if (
            self.model_config.is_encoder_decoder
            and scheduler_output.scheduled_encoder_inputs
        ):
2278
2279
2280
            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
            model_kwargs.update(encoder_inputs)

2281
2282
2283
2284
2285
2286
2287
        return (
            input_ids,
            inputs_embeds,
            positions,
            intermediate_tensors,
            model_kwargs,
        )
2288

2289
    def _sample(
2290
        self,
2291
2292
        logits: torch.Tensor | None,
        spec_decode_metadata: SpecDecodeMetadata | None,
2293
    ) -> SamplerOutput:
2294
        # Sample the next token and get logprobs if needed.
2295
        sampling_metadata = self.input_batch.sampling_metadata
2296
        if spec_decode_metadata is None:
2297
2298
2299
            # Update output token ids with tokens sampled in last step
            # if async scheduling and required by current sampling params.
            self.input_batch.update_async_output_token_ids()
2300
            return self.sampler(
2301
2302
2303
                logits=logits,
                sampling_metadata=sampling_metadata,
            )
2304

2305
        sampler_output = self.rejection_sampler(
2306
2307
            spec_decode_metadata,
            None,  # draft_probs
2308
            logits,
2309
2310
            sampling_metadata,
        )
2311
        self._update_states_after_model_execute(sampler_output.sampled_token_ids)
2312
2313
2314
        return sampler_output

    def _bookkeeping_sync(
2315
2316
2317
        self,
        scheduler_output: "SchedulerOutput",
        sampler_output: SamplerOutput,
2318
        logits: torch.Tensor | None,
2319
2320
        hidden_states: torch.Tensor,
        num_scheduled_tokens: int,
2321
        spec_decode_metadata: SpecDecodeMetadata | None,
2322
    ) -> tuple[
2323
        dict[str, int],
2324
        LogprobsLists | None,
2325
        list[list[int]],
2326
        dict[str, LogprobsTensors | None],
2327
2328
2329
        list[str],
        dict[str, int],
        list[int],
2330
    ]:
2331
2332
2333
2334
        num_nans_in_logits = {}
        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
            num_nans_in_logits = self._get_nans_in_logits(logits)

2335
2336
2337
        discard_sampled_tokens_req_indices = self.discard_request_indices.np[
            : self.num_discarded_requests
        ]
2338
2339
2340
2341
        for i in discard_sampled_tokens_req_indices:
            gen = self.input_batch.generators.get(int(i))
            if gen is not None:
                gen.set_offset(gen.get_offset() - 4)
2342

2343
2344
2345
        # Copy some objects so they don't get modified after returning.
        # This is important when using async scheduling.
        req_ids_output_copy = self.input_batch.req_ids.copy()
2346
        req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
2347
2348

        num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]
2349
        sampled_token_ids = sampler_output.sampled_token_ids
2350
        invalid_req_indices = []
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
        if not self.use_async_scheduling:
            # Get the valid generated tokens.
            max_gen_len = sampled_token_ids.shape[-1]
            if max_gen_len == 1:
                # No spec decode tokens.
                valid_sampled_token_ids = self._to_list(sampled_token_ids)
            else:
                # Includes spec decode tokens.
                valid_sampled_token_ids = self.rejection_sampler.parse_output(
                    sampled_token_ids,
                    self.input_batch.vocab_size,
                )
            # Mask out the sampled tokens that should not be sampled.
            for i in discard_sampled_tokens_req_indices:
2365
                valid_sampled_token_ids[int(i)].clear()
2366
        else:
2367
            valid_sampled_token_ids = []
2368
            invalid_req_indices = discard_sampled_tokens_req_indices.tolist()
2369
2370
2371
2372
2373
2374
            invalid_req_indices_set = set(invalid_req_indices)
            assert sampled_token_ids.shape[-1] == 1

            # Cache the sampled tokens on the GPU and avoid CPU sync.
            # These will be copied into input_ids in the next step
            # when preparing inputs.
2375
            self.input_batch.prev_sampled_token_ids = sampled_token_ids
2376
2377
2378
2379
2380
            self.input_batch.prev_req_id_to_index = {
                req_id: i
                for i, req_id in enumerate(self.input_batch.req_ids)
                if i not in invalid_req_indices_set
            }
2381

2382
2383
2384
2385
2386
        # 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.
2387
        req_ids = self.input_batch.req_ids
2388
2389
2390
2391
        logprobs_tensors = sampler_output.logprobs_tensors
        cu_num_accepted_tokens = (
            [0] if spec_decode_metadata and logprobs_tensors else None
        )
2392
2393
        for req_idx in range(num_sampled_tokens):
            if self.use_async_scheduling:
2394
                sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None
2395
2396
            else:
                sampled_ids = valid_sampled_token_ids[req_idx]
2397
2398
2399
2400
2401
2402
2403
2404

            num_sampled_ids: int = len(sampled_ids) if sampled_ids else 0

            if cu_num_accepted_tokens is not None:
                cu_num_accepted_tokens.append(
                    cu_num_accepted_tokens[-1] + num_sampled_ids
                )

2405
2406
2407
2408
            if not sampled_ids:
                continue

            start_idx = self.input_batch.num_tokens_no_spec[req_idx]
2409
            end_idx = start_idx + num_sampled_ids
2410
2411
2412
2413
            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}"
2414
            )
2415

2416
2417
            self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids
            self.input_batch.is_token_ids[req_idx, start_idx:end_idx] = True
2418
2419
            self.input_batch.num_tokens_no_spec[req_idx] = end_idx
            self.input_batch.num_tokens[req_idx] = end_idx
2420

2421
            req_id = req_ids[req_idx]
2422
2423
2424
            req_state = self.requests[req_id]
            req_state.output_token_ids.extend(sampled_ids)

2425
2426
        logprobs_lists = (
            logprobs_tensors.tolists(cu_num_accepted_tokens)
2427
            if not self.use_async_scheduling and logprobs_tensors is not None
2428
2429
2430
2431
2432
2433
2434
2435
2436
            else None
        )

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

2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
        return (
            num_nans_in_logits,
            logprobs_lists,
            valid_sampled_token_ids,
            prompt_logprobs_dict,
            req_ids_output_copy,
            req_id_to_index_output_copy,
            invalid_req_indices,
        )

2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
    @contextmanager
    def synchronize_input_prep(self):
        if self.prepare_inputs_event is None:
            yield
            return

        # Ensure prior step has finished with reused CPU tensors.
        # This is required in the async scheduling case because
        # the CPU->GPU transfer happens async.
        self.prepare_inputs_event.synchronize()
        try:
            yield
        finally:
            self.prepare_inputs_event.record()

2462
2463
    def _model_forward(
        self,
2464
2465
2466
2467
        input_ids: torch.Tensor | None = None,
        positions: torch.Tensor | None = None,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
2468
2469
2470
2471
2472
        **model_kwargs: dict[str, Any],
    ) -> Any:
        """Helper method to call the model forward pass.

        This method can be overridden by subclasses for model execution.
2473
        Motivation: We can inspect only this method versus
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
        the whole execute_model, which has additional logic.

        Args:
            input_ids: Input token IDs
            positions: Token positions
            intermediate_tensors: Tensors from previous pipeline stages
            inputs_embeds: Input embeddings (alternative to input_ids)
            **model_kwargs: Additional model arguments

        Returns:
            Model output tensor
        """
        return self.model(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
            **model_kwargs,
        )

2494
2495
2496
2497
    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: "SchedulerOutput",
2498
        intermediate_tensors: IntermediateTensors | None = None,
2499
2500
2501
2502
2503
2504
2505
    ) -> ModelRunnerOutput | IntermediateTensors | None:
        if self.execute_model_state is not None:
            raise RuntimeError(
                "State error: sample_tokens() must be called "
                "after execute_model() returns None."
            )
        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2506
        with record_function_or_nullcontext("gpu_model_runner: preprocess"):
2507
2508
2509
2510
            with self.synchronize_input_prep():
                # Update persistent batch states.
                self._update_states(scheduler_output)

2511
                if not num_scheduled_tokens:
2512
2513
2514
2515
                    if not has_kv_transfer_group():
                        # Return empty ModelRunnerOutput if no work to do.
                        return EMPTY_MODEL_RUNNER_OUTPUT
                    return self.kv_connector_no_forward(
2516
2517
                        scheduler_output, self.vllm_config
                    )
2518
2519
2520
2521
                if self.cache_config.kv_sharing_fast_prefill:
                    assert not self.input_batch.num_prompt_logprobs, (
                        "--kv-sharing-fast-prefill produces incorrect "
                        "logprobs for prompt tokens, tokens, please disable "
2522
2523
                        "it when the requests need prompt logprobs"
                    )
2524

2525
2526
2527
2528
2529
2530
                num_reqs = self.input_batch.num_reqs
                req_ids = self.input_batch.req_ids
                tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
                num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
                max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())

2531
2532
2533
2534
                (
                    logits_indices,
                    spec_decode_metadata,
                    ubatch_slices,
2535
                    num_tokens_across_dp,
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
                ) = self._prepare_inputs(
                    scheduler_output, num_scheduled_tokens_np, max_num_scheduled_tokens
                )

                cascade_attn_prefix_lens = None
                # Disable cascade attention when using microbatching (DBO)
                if self.cascade_attn_enabled and ubatch_slices is None:
                    # Pre-compute cascade attention prefix lengths
                    # NOTE: Must be AFTER _prepare_inputs uses self.input_batch state
                    cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
                        num_scheduled_tokens_np,
                        scheduler_output.num_common_prefix_blocks,
                    )

                # TODO(lucas): move cudagraph dispatching here:
                #   https://github.com/vllm-project/vllm/issues/23789

                total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
                use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
                attn_metadata, spec_decode_common_attn_metadata = (
                    self._build_attention_metadata(
                        total_num_scheduled_tokens=total_num_scheduled_tokens,
                        max_num_scheduled_tokens=max_num_scheduled_tokens,
                        num_reqs=num_reqs,
                        ubatch_slices=ubatch_slices,
                        logits_indices=logits_indices,
                        use_spec_decode=use_spec_decode,
                        scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,
                        cascade_attn_prefix_lens=cascade_attn_prefix_lens,
                    )
                )
2567

2568
            dp_rank = self.parallel_config.data_parallel_rank
2569
2570
            if ubatch_slices:
                assert num_tokens_across_dp is not None
2571
                num_input_tokens = int(num_tokens_across_dp[dp_rank].item())
2572
2573
                self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
            elif num_tokens_across_dp is not None:
2574
                num_input_tokens = int(num_tokens_across_dp[dp_rank].item())
2575
2576
2577
2578
2579
            else:
                num_input_tokens = self._get_num_input_tokens(
                    scheduler_output.total_num_scheduled_tokens
                )

2580
2581
2582
2583
2584
2585
            (
                input_ids,
                inputs_embeds,
                positions,
                intermediate_tensors,
                model_kwargs,
2586
            ) = self._preprocess(
2587
                scheduler_output, num_input_tokens, intermediate_tensors
2588
2589
            )

2590
2591
2592
            uniform_decode = (
                max_num_scheduled_tokens == self.uniform_decode_query_len
            ) and (num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)
2593
            batch_descriptor = BatchDescriptor(
2594
2595
2596
                num_tokens=num_input_tokens,
                uniform_decode=uniform_decode,
                has_lora=len(self.input_batch.lora_id_to_lora_request) > 0,
2597
2598
            )
            cudagraph_runtime_mode, batch_descriptor = (
2599
2600
2601
2602
                self.cudagraph_dispatcher.dispatch(
                    batch_descriptor,
                    use_cascade_attn=cascade_attn_prefix_lens is not None,
                )
2603
            )
2604

2605
        # Set cudagraph mode to none if calc_kv_scales is true.
2606
2607
2608
2609
2610
2611
        # KV scales calculation involves dynamic operations that are incompatible
        # with CUDA graph capture.
        if self.calculate_kv_scales:
            cudagraph_runtime_mode = CUDAGraphMode.NONE
            # Mark KV scales as calculated after the first forward pass
            self.calculate_kv_scales = False
2612

2613
2614
        # Run the model.
        # Use persistent buffers for CUDA graphs.
2615
2616
        with (
            set_forward_context(
2617
2618
2619
2620
2621
2622
                attn_metadata,
                self.vllm_config,
                num_tokens=num_input_tokens,
                num_tokens_across_dp=num_tokens_across_dp,
                cudagraph_runtime_mode=cudagraph_runtime_mode,
                batch_descriptor=batch_descriptor,
2623
                ubatch_slices=ubatch_slices,
2624
            ),
2625
            record_function_or_nullcontext("gpu_model_runner: forward"),
2626
2627
            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
        ):
2628
            model_output = self._model_forward(
2629
2630
2631
2632
2633
2634
2635
                input_ids=input_ids,
                positions=positions,
                intermediate_tensors=intermediate_tensors,
                inputs_embeds=inputs_embeds,
                **model_kwargs,
            )

2636
        with record_function_or_nullcontext("gpu_model_runner: postprocess"):
2637
            if self.use_aux_hidden_state_outputs:
2638
                # True when EAGLE 3 is used.
2639
2640
                hidden_states, aux_hidden_states = model_output
            else:
2641
                # Common case.
2642
2643
2644
                hidden_states = model_output
                aux_hidden_states = None

2645
2646
2647
2648
2649
            if not self.broadcast_pp_output:
                # Common case.
                if not get_pp_group().is_last_rank:
                    # Return the intermediate tensors.
                    assert isinstance(hidden_states, IntermediateTensors)
2650
2651
                    hidden_states.kv_connector_output = kv_connector_output
                    return hidden_states
2652

2653
                if self.is_pooling_model:
2654
                    # Return the pooling output.
2655
2656
2657
                    output = self._pool(
                        hidden_states, num_scheduled_tokens, num_scheduled_tokens_np
                    )
2658
2659
                    output.kv_connector_output = kv_connector_output
                    return output
2660
2661

                sample_hidden_states = hidden_states[logits_indices]
2662
                logits = self.model.compute_logits(sample_hidden_states)
2663
2664
2665
2666
            else:
                # Rare case.
                assert not self.is_pooling_model

2667
                sample_hidden_states = hidden_states[logits_indices]
2668
                if not get_pp_group().is_last_rank:
2669
                    all_gather_tensors = {
2670
2671
2672
                        "residual": not is_residual_scattered_for_sp(
                            self.vllm_config, num_input_tokens
                        )
2673
                    }
2674
                    get_pp_group().send_tensor_dict(
2675
2676
                        hidden_states.tensors,
                        all_gather_group=get_tp_group(),
2677
2678
                        all_gather_tensors=all_gather_tensors,
                    )
2679
2680
                    logits = None
                else:
2681
                    logits = self.model.compute_logits(sample_hidden_states)
2682
2683
2684
2685
2686

                model_output_broadcast_data = {}
                if logits is not None:
                    model_output_broadcast_data["logits"] = logits.contiguous()

2687
2688
2689
                model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
                    model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
                )
2690
2691
2692
                assert model_output_broadcast_data is not None
                logits = model_output_broadcast_data["logits"]

2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
        self.execute_model_state = ExecuteModelState(
            scheduler_output,
            logits,
            spec_decode_metadata,
            spec_decode_common_attn_metadata,
            hidden_states,
            sample_hidden_states,
            aux_hidden_states,
            kv_connector_output,
        )
        return None

    @torch.inference_mode
    def sample_tokens(
        self, grammar_output: "GrammarOutput | None"
    ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
        if self.execute_model_state is None:
            # Nothing to do (PP non-final rank case), output isn't used.
            return None  # noqa

        # Unpack ephemeral state.
        (
            scheduler_output,
            logits,
            spec_decode_metadata,
            spec_decode_common_attn_metadata,
            hidden_states,
            sample_hidden_states,
            aux_hidden_states,
            kv_connector_output,
        ) = self.execute_model_state
        # Clear ephemeral state.
        self.execute_model_state = None

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

2733
        with record_function_or_nullcontext("gpu_model_runner: sample"):
2734
2735
            sampler_output = self._sample(logits, spec_decode_metadata)

2736
2737
        def propose_draft_token_ids(sampled_token_ids):
            assert spec_decode_common_attn_metadata is not None
2738
            with record_function_or_nullcontext("gpu_model_runner: draft"):
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
                self._draft_token_ids = self.propose_draft_token_ids(
                    scheduler_output,
                    sampled_token_ids,
                    self.input_batch.sampling_metadata,
                    hidden_states,
                    sample_hidden_states,
                    aux_hidden_states,
                    spec_decode_metadata,
                    spec_decode_common_attn_metadata,
                )

2750
2751
2752
2753
2754
        use_padded_batch_for_eagle = (
            self.speculative_config
            and self.speculative_config.use_eagle()
            and not self.speculative_config.disable_padded_drafter_batch
        )
2755
2756
2757
        effective_drafter_max_model_len = self.max_model_len
        if effective_drafter_max_model_len is None:
            effective_drafter_max_model_len = self.model_config.max_model_len
2758
2759
2760
2761
2762
        if (
            self.speculative_config
            and self.speculative_config.draft_model_config is not None
            and self.speculative_config.draft_model_config.max_model_len is not None
        ):
2763
            effective_drafter_max_model_len = (
2764
2765
                self.speculative_config.draft_model_config.max_model_len
            )
2766
        input_fits_in_drafter = spec_decode_common_attn_metadata and (
2767
2768
2769
2770
            spec_decode_common_attn_metadata.max_seq_len
            + self.speculative_config.num_speculative_tokens
            <= effective_drafter_max_model_len
        )
2771
        if use_padded_batch_for_eagle and input_fits_in_drafter:
2772
2773
2774
2775
            # EAGLE speculative decoding can use the GPU sampled tokens
            # as inputs, and does not need to wait for bookkeeping to finish.
            propose_draft_token_ids(sampler_output.sampled_token_ids)

2776
        with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
2777
2778
2779
2780
2781
2782
2783
2784
            (
                num_nans_in_logits,
                logprobs_lists,
                valid_sampled_token_ids,
                prompt_logprobs_dict,
                req_ids_output_copy,
                req_id_to_index_output_copy,
                invalid_req_indices,
2785
2786
2787
2788
2789
            ) = self._bookkeeping_sync(
                scheduler_output,
                sampler_output,
                logits,
                hidden_states,
2790
                scheduler_output.total_num_scheduled_tokens,
2791
                spec_decode_metadata,
2792
            )
2793

2794
2795
2796
2797
2798
        if (
            self.speculative_config
            and not use_padded_batch_for_eagle
            and input_fits_in_drafter
        ):
2799
2800
2801
            # ngram and other speculative decoding methods use the sampled
            # tokens on the CPU, so they are run after bookkeeping.
            propose_draft_token_ids(valid_sampled_token_ids)
2802

2803
        with record_function_or_nullcontext("gpu_model_runner: eplb"):
2804
            self.eplb_step()
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
        with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
            output = ModelRunnerOutput(
                req_ids=req_ids_output_copy,
                req_id_to_index=req_id_to_index_output_copy,
                sampled_token_ids=valid_sampled_token_ids,
                logprobs=logprobs_lists,
                prompt_logprobs_dict=prompt_logprobs_dict,
                pooler_output=[],
                kv_connector_output=kv_connector_output,
                num_nans_in_logits=num_nans_in_logits,
            )
2816

2817
2818
        if not self.use_async_scheduling:
            return output
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
        with record_function_or_nullcontext(
            "gpu_model_runner: AsyncGPUModelRunnerOutput"
        ):
            async_output = AsyncGPUModelRunnerOutput(
                model_runner_output=output,
                sampled_token_ids=sampler_output.sampled_token_ids,
                logprobs_tensors=sampler_output.logprobs_tensors,
                invalid_req_indices=invalid_req_indices,
                async_output_copy_stream=self.async_output_copy_stream,
            )
        with record_function_or_nullcontext(
            "gpu_model_runner: set_async_sampled_token_ids"
        ):
            # Save ref of sampled_token_ids CPU tensor if the batch contains
            # any requests with sampling params that that require output ids.
            self.input_batch.set_async_sampled_token_ids(
                async_output.sampled_token_ids_cpu,
                async_output.async_copy_ready_event,
            )
2838
2839
2840

        return async_output

2841
    def take_draft_token_ids(self) -> DraftTokenIds | None:
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
        if self._draft_token_ids is None:
            return None
        req_ids = self.input_batch.req_ids
        if isinstance(self._draft_token_ids, torch.Tensor):
            draft_token_ids = self._draft_token_ids.tolist()
        else:
            draft_token_ids = self._draft_token_ids
        self._draft_token_ids = None
        return DraftTokenIds(req_ids, draft_token_ids)

2852
2853
2854
    def propose_draft_token_ids(
        self,
        scheduler_output: "SchedulerOutput",
2855
        sampled_token_ids: torch.Tensor | list[list[int]],
2856
2857
2858
        sampling_metadata: SamplingMetadata,
        hidden_states: torch.Tensor,
        sample_hidden_states: torch.Tensor,
2859
2860
        aux_hidden_states: list[torch.Tensor] | None,
        spec_decode_metadata: SpecDecodeMetadata | None,
2861
        common_attn_metadata: CommonAttentionMetadata,
2862
    ) -> list[list[int]] | torch.Tensor:
2863
2864
        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
        if self.speculative_config.method == "ngram":
2865
            assert isinstance(sampled_token_ids, list)
2866
            assert isinstance(self.drafter, NgramProposer)
2867
            draft_token_ids = self.drafter.propose(
2868
2869
                sampled_token_ids,
                self.input_batch.req_ids,
2870
2871
                self.input_batch.num_tokens_no_spec,
                self.input_batch.token_ids_cpu,
2872
2873
                self.input_batch.spec_decode_unsupported_reqs,
            )
2874
2875
2876
2877
        elif self.speculative_config.method == "suffix":
            assert isinstance(sampled_token_ids, list)
            assert isinstance(self.drafter, SuffixDecodingProposer)
            draft_token_ids = self.drafter.propose(self.input_batch, sampled_token_ids)
2878
        elif self.speculative_config.method == "medusa":
2879
            assert isinstance(sampled_token_ids, list)
2880
            assert isinstance(self.drafter, MedusaProposer)
2881

2882
2883
            if sample_hidden_states.shape[0] == len(sampled_token_ids):
                # The input to the target model does not include draft tokens.
2884
2885
2886
2887
                hidden_states = sample_hidden_states
            else:
                indices = []
                offset = 0
2888
2889
2890
                assert spec_decode_metadata is not None, (
                    "No spec decode metadata for medusa"
                )
2891
                for num_draft, tokens in zip(
2892
2893
                    spec_decode_metadata.num_draft_tokens, sampled_token_ids
                ):
2894
2895
                    indices.append(offset + len(tokens) - 1)
                    offset += num_draft + 1
2896
                indices = torch.tensor(indices, device=self.device)
2897
2898
                hidden_states = sample_hidden_states[indices]

2899
            draft_token_ids = self.drafter.propose(
2900
2901
2902
                target_hidden_states=hidden_states,
                sampling_metadata=sampling_metadata,
            )
2903
        elif self.speculative_config.use_eagle():
2904
            assert isinstance(self.drafter, EagleProposer)
2905
2906
2907
2908
2909

            if self.speculative_config.disable_padded_drafter_batch:
                # When padded-batch is disabled, the sampled_token_ids should be
                # the cpu-side list[list[int]] of valid sampled tokens for each
                # request, with invalid requests having empty lists.
2910
2911
                assert isinstance(sampled_token_ids, list), (
                    "sampled_token_ids should be a python list when"
2912
                    "padded-batch is disabled."
2913
                )
2914
                next_token_ids = self.drafter.prepare_next_token_ids_cpu(
2915
2916
2917
2918
2919
                    sampled_token_ids,
                    self.requests,
                    self.input_batch,
                    scheduler_output.num_scheduled_tokens,
                )
2920
2921
2922
2923
2924
            else:
                # When using padded-batch, the sampled_token_ids should be
                # the gpu tensor of sampled tokens for each request, of shape
                # (num_reqs, num_spec_tokens + 1) with rejected tokens having
                # value -1.
2925
2926
                assert isinstance(sampled_token_ids, torch.Tensor), (
                    "sampled_token_ids should be a torch.Tensor when"
2927
                    "padded-batch is enabled."
2928
2929
                )
                next_token_ids, valid_sampled_tokens_count = (
2930
2931
2932
2933
2934
2935
                    self.drafter.prepare_next_token_ids_padded(
                        common_attn_metadata,
                        sampled_token_ids,
                        self.requests,
                        self.input_batch,
                        self.discard_request_indices.gpu,
2936
                        self.num_discarded_requests,
2937
                    )
2938
                )
Jiayi Yao's avatar
Jiayi Yao committed
2939

2940
            if spec_decode_metadata is None:
2941
                token_indices_to_sample = None
2942
                # input_ids can be None for multimodal models.
2943
                target_token_ids = self.input_ids.gpu[:num_scheduled_tokens]
2944
                target_positions = self._get_positions(num_scheduled_tokens)
2945
                if self.use_aux_hidden_state_outputs:
Wentao Ye's avatar
Wentao Ye committed
2946
                    assert aux_hidden_states is not None
2947
                    target_hidden_states = torch.cat(
2948
2949
                        [h[:num_scheduled_tokens] for h in aux_hidden_states], dim=-1
                    )
2950
2951
                else:
                    target_hidden_states = hidden_states[:num_scheduled_tokens]
2952
            else:
2953
2954
                if self.speculative_config.disable_padded_drafter_batch:
                    token_indices_to_sample = None
2955
2956
2957
2958
2959
                    common_attn_metadata, token_indices = self.drafter.prepare_inputs(
                        common_attn_metadata,
                        sampled_token_ids,
                        spec_decode_metadata.num_draft_tokens,
                    )
2960
                else:
2961
                    common_attn_metadata, token_indices, token_indices_to_sample = (
2962
2963
2964
                        self.drafter.prepare_inputs_padded(
                            common_attn_metadata,
                            spec_decode_metadata,
2965
2966
2967
                            valid_sampled_tokens_count,
                        )
                    )
2968

2969
                target_token_ids = self.input_ids.gpu[token_indices]
2970
                target_positions = self._get_positions(token_indices)
2971
                if self.use_aux_hidden_state_outputs:
Wentao Ye's avatar
Wentao Ye committed
2972
                    assert aux_hidden_states is not None
2973
                    target_hidden_states = torch.cat(
2974
2975
                        [h[token_indices] for h in aux_hidden_states], dim=-1
                    )
2976
2977
                else:
                    target_hidden_states = hidden_states[token_indices]
2978

2979
            if self.supports_mm_inputs:
2980
2981
2982
2983
2984
2985
                mm_embed_inputs = self._gather_mm_embeddings(
                    scheduler_output,
                    shift_computed_tokens=1,
                )
            else:
                mm_embed_inputs = None
2986

2987
            draft_token_ids = self.drafter.propose(
2988
2989
2990
2991
                target_token_ids=target_token_ids,
                target_positions=target_positions,
                target_hidden_states=target_hidden_states,
                next_token_ids=next_token_ids,
2992
                last_token_indices=token_indices_to_sample,
2993
                sampling_metadata=sampling_metadata,
2994
                common_attn_metadata=common_attn_metadata,
2995
                mm_embed_inputs=mm_embed_inputs,
2996
            )
2997

2998
        return draft_token_ids
2999

3000
3001
3002
    def update_config(self, overrides: dict[str, Any]) -> None:
        allowed_config_names = {"load_config", "model_config"}
        for config_name, config_overrides in overrides.items():
3003
3004
            assert config_name in allowed_config_names, (
                f"Config `{config_name}` not supported. "
3005
                f"Allowed configs: {allowed_config_names}"
3006
            )
3007
3008
3009
3010
            config = getattr(self, config_name)
            new_config = update_config(config, config_overrides)
            setattr(self, config_name, new_config)

3011
3012
3013
3014
3015
    def load_model(self, eep_scale_up: bool = False) -> None:
        """
        Args:
            eep_scale_up: the model loading is for elastic EP scale up.
        """
3016
3017
3018
3019
3020
        logger.info_once(
            "Starting to load model %s...",
            self.model_config.model,
            scope="global",
        )
3021
3022
3023
3024
3025
        global_expert_loads, old_global_expert_indices_per_model, rank_mapping = (
            EplbState.get_eep_state(self.parallel_config)
            if eep_scale_up
            else (None, None, None)
        )
3026

3027
3028
3029
        if self.parallel_config.enable_eplb:
            self.eplb_state = EplbState(self.parallel_config, self.device)
            eplb_models = 0
3030
        with DeviceMemoryProfiler() as m:
3031
            time_before_load = time.perf_counter()
3032
            model_loader = get_model_loader(self.load_config)
3033
            self.model = model_loader.load_model(
3034
3035
                vllm_config=self.vllm_config, model_config=self.model_config
            )
3036
            if self.lora_config:
3037
3038
3039
                self.model = self.load_lora_model(
                    self.model, self.vllm_config, self.device
                )
3040
            if hasattr(self, "drafter"):
3041
                logger.info_once("Loading drafter model...")
3042
                self.drafter.load_model(self.model)
3043
3044
3045
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
3073
                if (
                    hasattr(self.drafter, "model")
                    and is_mixture_of_experts(self.drafter.model)
                    and self.parallel_config.enable_eplb
                ):
                    logger.info_once(
                        "EPLB is enabled for drafter model %s.",
                        self.vllm_config.speculative_config.draft_model_config.model,
                    )

                    global_expert_load = (
                        global_expert_loads[eplb_models]
                        if global_expert_loads
                        else None
                    )
                    old_global_expert_indices = (
                        old_global_expert_indices_per_model[eplb_models]
                        if old_global_expert_indices_per_model
                        else None
                    )
                    if self.eplb_state is None:
                        self.eplb_state = EplbState(self.parallel_config, self.device)
                    self.eplb_state.add_model(
                        self.drafter.model,
                        self.vllm_config.speculative_config.draft_model_config,
                        global_expert_load,
                        old_global_expert_indices,
                        rank_mapping,
                    )
                    eplb_models += 1

3074
            if self.use_aux_hidden_state_outputs:
3075
                if not supports_eagle3(self.get_model()):
3076
3077
                    raise RuntimeError(
                        "Model does not support EAGLE3 interface but "
3078
3079
                        "aux_hidden_state_outputs was requested"
                    )
3080
3081
3082
3083
3084
3085
3086
3087
3088
3089
3090
3091
3092

                # Try to get auxiliary layers from speculative config,
                # otherwise use model's default layers
                aux_layers = self._get_eagle3_aux_layers_from_config()
                if aux_layers:
                    logger.info(
                        "Using auxiliary layers from speculative config: %s",
                        aux_layers,
                    )
                else:
                    aux_layers = self.model.get_eagle3_aux_hidden_state_layers()

                self.model.set_aux_hidden_state_layers(aux_layers)
3093
            time_after_load = time.perf_counter()
3094
        self.model_memory_usage = m.consumed_memory
3095
        logger.info_once(
3096
            "Model loading took %.4f GiB memory and %.6f seconds",
3097
3098
            self.model_memory_usage / GiB_bytes,
            time_after_load - time_before_load,
3099
            scope="local",
3100
        )
3101
        prepare_communication_buffer_for_model(self.model)
3102
        self.is_multimodal_pruning_enabled = (
3103
            supports_multimodal_pruning(self.get_model())
3104
3105
            and self.model_config.multimodal_config.is_multimodal_pruning_enabled()
        )
3106

3107
        if is_mixture_of_experts(self.model) and self.parallel_config.enable_eplb:
3108
3109
3110
3111
3112
3113
3114
3115
3116
3117
3118
            logger.info_once("EPLB is enabled for model %s.", self.model_config.model)
            global_expert_load = (
                global_expert_loads[eplb_models] if global_expert_loads else None
            )
            old_global_expert_indices = (
                old_global_expert_indices_per_model[eplb_models]
                if old_global_expert_indices_per_model
                else None
            )
            assert self.eplb_state is not None
            self.eplb_state.add_model(
3119
                self.model,
3120
                self.model_config,
3121
3122
3123
                global_expert_load,
                old_global_expert_indices,
                rank_mapping,
3124
3125
            )

3126
        if (
3127
3128
            self.vllm_config.compilation_config.mode
            == CompilationMode.STOCK_TORCH_COMPILE
3129
            and supports_dynamo()
3130
        ):
3131
            backend = self.vllm_config.compilation_config.init_backend(self.vllm_config)
3132
            compilation_counter.stock_torch_compile_count += 1
3133
            self.model.compile(fullgraph=True, backend=backend)
3134
            return
3135
        # for other compilation modes, cudagraph behavior is controlled by
3136
3137
3138
        # CudagraphWraper and CudagraphDispatcher of vllm.

        # wrap the model with full cudagraph wrapper if needed.
3139
3140
3141
3142
3143
3144
3145
        if (
            self.compilation_config.cudagraph_mode.has_full_cudagraphs()
            and not self.parallel_config.enable_dbo
        ):
            self.model = CUDAGraphWrapper(
                self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
            )
3146
3147
        elif self.parallel_config.enable_dbo:
            if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
3148
3149
3150
                self.model = UBatchWrapper(
                    self.model, self.vllm_config, CUDAGraphMode.FULL, self.device
                )
3151
            else:
3152
3153
3154
                self.model = UBatchWrapper(
                    self.model, self.vllm_config, CUDAGraphMode.NONE, self.device
                )
3155

3156
    def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None:
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
3175
3176
3177
3178
3179
        """Extract Eagle3 auxiliary layer indices from speculative config.

        These indices specify which hidden states from the base model should
        be used as auxiliary inputs for the Eagle3 drafter model during
        speculative decoding.

        Returns:
            Tuple of layer indices if found in draft model config,
            None otherwise.
        """
        if not (self.speculative_config and self.speculative_config.draft_model_config):
            return None

        hf_config = self.speculative_config.draft_model_config.hf_config
        if not hasattr(hf_config, "eagle_aux_hidden_state_layer_ids"):
            return None

        layer_ids = hf_config.eagle_aux_hidden_state_layer_ids
        if layer_ids and isinstance(layer_ids, (list, tuple)):
            return tuple(layer_ids)

        return None

3180
    def reload_weights(self) -> None:
3181
        assert getattr(self, "model", None) is not None, (
3182
            "Cannot reload weights before model is loaded."
3183
        )
3184
3185
        model_loader = get_model_loader(self.load_config)
        logger.info("Reloading weights inplace...")
3186
        model_loader.load_weights(self.get_model(), model_config=self.model_config)
3187

3188
3189
3190
3191
3192
    def save_tensorized_model(
        self,
        tensorizer_config: "TensorizerConfig",
    ) -> None:
        TensorizerLoader.save_model(
3193
            self.get_model(),
3194
            tensorizer_config=tensorizer_config,
3195
            model_config=self.model_config,
3196
3197
        )

3198
3199
3200
    def _get_prompt_logprobs_dict(
        self,
        hidden_states: torch.Tensor,
3201
        num_scheduled_tokens: dict[str, int],
3202
    ) -> dict[str, LogprobsTensors | None]:
3203
3204
3205
3206
        num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs
        if not num_prompt_logprobs_dict:
            return {}

3207
        in_progress_dict = self.input_batch.in_progress_prompt_logprobs_cpu
3208
        prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
3209
3210
3211
3212
3213

        # 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():
3214
            num_tokens = num_scheduled_tokens[req_id]
3215
3216
3217

            # Get metadata for this request.
            request = self.requests[req_id]
3218
3219
3220
3221
            if request.prompt_token_ids is None:
                # Prompt logprobs is incompatible with prompt embeddings
                continue

3222
3223
            num_prompt_tokens = len(request.prompt_token_ids)
            prompt_token_ids = torch.tensor(request.prompt_token_ids).to(
3224
3225
                self.device, non_blocking=True
            )
3226

3227
3228
3229
3230
3231
3232
            # 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(
3233
3234
                    num_prompt_tokens - 1, num_prompt_logprobs + 1
                )
3235
3236
                in_progress_dict[req_id] = logprobs_tensors

3237
            # Determine number of logits to retrieve.
3238
3239
            start_idx = request.num_computed_tokens
            start_tok = start_idx + 1
3240
            num_remaining_tokens = num_prompt_tokens - start_tok
3241
            if num_tokens <= num_remaining_tokens:
3242
                # This is a chunk, more tokens remain.
3243
3244
3245
                # 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.
3246
3247
3248
3249
3250
                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)
3251
3252
3253
3254
3255
3256
3257
                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
3258
3259
3260
3261
3262

            # 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]
3263
            offset = self.query_start_loc.np[req_idx].item()
3264
            prompt_hidden_states = hidden_states[offset : offset + num_logits]
3265
            logits = self.model.compute_logits(prompt_hidden_states)
3266
3267
3268
3269

            # 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.
3270
            tgt_token_ids = prompt_token_ids[start_tok : start_tok + num_logits]
3271
3272

            # Compute prompt logprobs.
3273
3274
            logprobs = self.sampler.compute_logprobs(logits)
            token_ids, logprobs, ranks = self.sampler.gather_logprobs(
3275
3276
                logprobs, num_prompt_logprobs, tgt_token_ids
            )
3277
3278

            # Transfer GPU->CPU async.
3279
3280
            chunk_slice = slice(start_idx, start_idx + num_logits)
            logprobs_tensors.logprob_token_ids[chunk_slice].copy_(
3281
3282
3283
                token_ids, non_blocking=True
            )
            logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, non_blocking=True)
3284
            logprobs_tensors.selected_token_ranks[chunk_slice].copy_(
3285
3286
                ranks, non_blocking=True
            )
3287
3288
3289
3290
3291

        # 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]
3292
            del in_progress_dict[req_id]
3293
3294

        # Must synchronize the non-blocking GPU->CPU transfers.
3295
        if prompt_logprobs_dict:
3296
            self._sync_device()
3297
3298
3299

        return prompt_logprobs_dict

3300
3301
    def _get_nans_in_logits(
        self,
3302
        logits: torch.Tensor | None,
3303
3304
3305
3306
3307
3308
3309
3310
3311
3312
3313
    ) -> 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])
3314
3315
3316
                    if num_nans_for_index is not None and req_index < logits.shape[0]
                    else 0
                )
3317
3318
3319
3320
            return num_nans_in_logits
        except IndexError:
            return {}

3321
3322
3323
3324
3325
3326
    @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
3327
         - during DP rank dummy run
3328
3329
3330
3331
3332
3333
3334
3335
3336
3337
3338
        """
        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(
3339
                    self.input_ids.gpu,
3340
3341
                    low=0,
                    high=self.model_config.get_vocab_size(),
3342
3343
                    dtype=input_ids.dtype,
                )
3344

3345
            logger.debug_once("Randomizing dummy data for DP Rank")
3346
            input_ids.copy_(rand_input_ids()[: input_ids.size(0)], non_blocking=True)
3347
3348
3349
            yield
            input_ids.fill_(0)

3350
3351
3352
3353
3354
3355
    def _get_mm_dummy_batch(
        self,
        modality: str,
        max_items_per_batch: int,
    ) -> BatchedTensorInputs:
        """Dummy data for profiling and precompiling multimodal models."""
3356
3357
        assert self.mm_budget is not None

3358
3359
        dummy_decoder_data = self.mm_registry.get_decoder_dummy_data(
            model_config=self.model_config,
3360
            seq_len=self.max_model_len,
3361
            mm_counts={modality: 1},
3362
            cache=self.mm_budget.cache,
3363
3364
3365
3366
        )
        dummy_mm_data = dummy_decoder_data.multi_modal_data

        # Result in the maximum GPU consumption of the model
3367
3368
        dummy_mm_item = dummy_mm_data[modality][0]
        dummy_mm_items = [dummy_mm_item] * max_items_per_batch
3369

3370
        model = cast(SupportsMultiModal, self.model)
3371
3372
3373
3374
3375
3376
3377
        return next(
            mm_kwargs_group
            for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(
                dummy_mm_items,
                device=self.device,
                pin_memory=self.pin_memory,
                merge_by_field_config=model.merge_by_field_config,
3378
                multimodal_cpu_fields=model.multimodal_cpu_fields,
3379
3380
            )
        )
3381

3382
3383
3384
3385
    @torch.inference_mode()
    def _dummy_run(
        self,
        num_tokens: int,
3386
        cudagraph_runtime_mode: CUDAGraphMode | None = None,
3387
3388
        force_attention: bool = False,
        uniform_decode: bool = False,
3389
        allow_microbatching: bool = True,
3390
3391
        skip_eplb: bool = False,
        is_profile: bool = False,
3392
        create_mixed_batch: bool = False,
3393
        remove_lora: bool = True,
3394
        activate_lora: bool = False,
3395
    ) -> tuple[torch.Tensor, torch.Tensor]:
3396
3397
3398
3399
3400
3401
3402
        """
        Run a dummy forward pass to warm up/profile run or capture the
        CUDA graph for the model.

        Args:
            num_tokens: Number of tokens to run the dummy forward pass.
            cudagraph_runtime_mode: used to control the behavior.
3403
                - if not set will determine the cudagraph mode based on using
3404
                    the self.cudagraph_dispatcher.
3405
3406
3407
3408
                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
                    needed.
3409
            force_attention: If True, always create attention metadata. Used to
3410
3411
3412
3413
                warm up attention backend when mode is NONE.
            uniform_decode: If True, the batch is a uniform decode batch.
            skip_eplb: If True, skip EPLB state update.
            is_profile: If True, this is a profile run.
3414
3415
            create_mixed_batch: If True, create a mixed batch with both decode
                (1 token) and prefill (multiple tokens) requests.
3416
            remove_lora: If False, dummy LoRAs are not destroyed after the run
3417
            activate_lora: If False, dummy_run is performed without LoRAs.
3418
        """
3419
3420
3421
3422
        assert (
            cudagraph_runtime_mode is None
            or cudagraph_runtime_mode.valid_runtime_modes()
        )
3423

3424
        # If cudagraph_mode.decode_mode() == FULL and
3425
        # cudagraph_mode.separate_routine(). This means that we are using
3426
3427
3428
3429
3430
3431
3432
3433
3434
3435
3436
        # different graphs and/or modes for mixed prefill-decode batches vs.
        # uniform decode batches. A uniform decode batch means that all
        # requests have identical query length, except a potential virtual
        # request (shorter) in the batch account for padding.
        # Uniform decode batch could either be common pure decode, where
        # max_query_len == 1, or speculative decode, where
        # max_query_len == 1 + num_spec_decode_tokens.

        # When setting max_query_len = 1, we switch to and capture the optimized
        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
        # for GQA/MQA.
3437
        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
3438

3439
3440
3441
3442
3443
        # 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.
        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
        max_num_reqs = self.scheduler_config.max_num_seqs
3444
3445
3446
3447
        if create_mixed_batch:
            assert not uniform_decode
            # Create mixed batch:
            # first half decode tokens, second half one prefill
3448
            num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2)
3449
3450
3451
3452
            num_prefill_tokens = num_tokens - num_decode_tokens
            num_reqs = num_decode_tokens + 1

            # Create decode requests (1 token each) followed by prefill request
3453
            num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]
3454
3455
3456
            # Note: Overriding max_query_len to be the prefill tokens
            max_query_len = num_prefill_tokens
        elif uniform_decode:
3457
            assert not create_mixed_batch
3458
            num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len))
3459
3460
            num_scheduled_tokens_list = [max_query_len] * num_reqs
            if num_tokens % max_query_len != 0:
3461
                num_scheduled_tokens_list[-1] = num_tokens % max_query_len
3462
3463
3464
3465
3466
3467
        else:
            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

3468
3469
        assert sum(num_scheduled_tokens_list) == num_tokens
        assert len(num_scheduled_tokens_list) == num_reqs
3470
        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
3471
        total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
3472
        num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
3473

3474
3475
3476
        # Disable DP padding when running eager
        allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

3477
3478
        # We currently only microbatch if the number of tokens is
        # over a certain threshold.
3479
        ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(
3480
3481
3482
3483
3484
3485
3486
            num_tokens_unpadded=total_num_scheduled_tokens,
            parallel_config=self.vllm_config.parallel_config,
            allow_microbatching=allow_microbatching,
            allow_dp_padding=allow_dp_padding,
            num_tokens_padded=total_num_scheduled_tokens,
            uniform_decode=uniform_decode,
            num_scheduled_tokens_per_request=num_scheduled_tokens,
3487
3488
3489
        )
        num_tokens_after_padding = num_tokens
        if num_tokens_across_dp is not None:
3490
3491
            dp_rank = self.parallel_config.data_parallel_rank
            num_tokens_after_padding = int(num_tokens_across_dp[dp_rank])
3492

3493
        attn_metadata: PerLayerAttnMetadata | None = None
3494
3495
3496

        # If force_attention is True, we always capture attention. Otherwise,
        # it only happens for cudagraph_runtime_mode=FULL.
3497
        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
3498
3499
3500
3501
3502
3503
            if create_mixed_batch:
                # In the mixed batch mode (used for FI warmup), we use
                # shorter sequence lengths to run faster.
                # TODO(luka) better system for describing dummy batches
                seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
            else:
3504
                seq_lens = max_query_len
3505
            self.seq_lens.np[:num_reqs] = seq_lens
3506
3507
            self.seq_lens.np[num_reqs:] = 0
            self.seq_lens.copy_to_gpu()
3508

3509
3510
            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
3511
3512
            self.query_start_loc.copy_to_gpu()

3513
3514
3515
3516
3517
3518
3519
            attn_metadata, _ = self._build_attention_metadata(
                total_num_scheduled_tokens=num_tokens,
                max_num_scheduled_tokens=max_query_len,
                num_reqs=num_reqs,
                ubatch_slices=ubatch_slices,
                for_cudagraph_capture=True,
            )
3520

3521
        with self.maybe_dummy_run_with_lora(
3522
3523
3524
3525
3526
            self.lora_config,
            num_scheduled_tokens,
            num_sampled_tokens,
            activate_lora,
            remove_lora,
3527
        ):
3528
3529
3530
            # Make sure padding doesn't exceed max_num_tokens
            assert num_tokens_after_padding <= self.max_num_tokens
            model_kwargs = self._init_model_kwargs(num_tokens_after_padding)
3531
            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
3532
                input_ids = None
3533
                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_after_padding]
3534
                model_kwargs = {
3535
                    **model_kwargs,
3536
3537
                    **self._dummy_mm_kwargs(num_reqs),
                }
3538
3539
            elif self.enable_prompt_embeds:
                input_ids = None
3540
3541
                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_after_padding]
                model_kwargs = self._init_model_kwargs(num_tokens_after_padding)
3542
            else:
3543
                input_ids = self.input_ids.gpu[:num_tokens_after_padding]
3544
                inputs_embeds = None
3545

3546
            if self.uses_mrope:
3547
                positions = self.mrope_positions.gpu[:, :num_tokens_after_padding]
3548
            else:
3549
                positions = self.positions.gpu[:num_tokens_after_padding]
3550
3551
3552
3553
3554
3555
3556
3557
3558

            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,
3559
3560
3561
                            device=self.device,
                        )
                    )
3562
3563

                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
3564
                    num_tokens_after_padding, None, False
3565
                )
3566
3567

            # filter out the valid batch descriptor
3568
3569
3570
3571
3572
            _cg_mode, batch_descriptor = (
                self.cudagraph_dispatcher.dispatch(
                    BatchDescriptor(
                        num_tokens=num_tokens_after_padding,
                        uniform_decode=uniform_decode,
3573
                        has_lora=activate_lora and self.lora_config is not None,
3574
3575
3576
3577
3578
                    )
                )
                if not is_profile
                else (CUDAGraphMode.NONE, None)
            )
3579
3580
3581
            if cudagraph_runtime_mode is not None:
                # we allow forcing NONE when the dispatcher disagrees to support
                # warm ups for cudagraph capture
3582
3583
3584
3585
                assert (
                    cudagraph_runtime_mode == CUDAGraphMode.NONE
                    or cudagraph_runtime_mode == _cg_mode
                ), (
3586
                    f"Cudagraph runtime mode mismatch at dummy_run. "
3587
3588
                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
                )
3589
3590
            else:
                cudagraph_runtime_mode = _cg_mode
3591

3592
            if ubatch_slices is not None:
3593
3594
3595
3596
3597
3598
3599
                # Adjust values to reflect a single ubatch.
                # TODO(sage,lucas): this is cruft that should be addressed in
                #  the padding refactor.
                num_tokens_after_padding = ubatch_slices[0].num_tokens
                if num_tokens_across_dp is not None:
                    num_tokens_across_dp[:] = num_tokens_after_padding

3600
3601
3602
            with (
                self.maybe_randomize_inputs(input_ids),
                set_forward_context(
3603
3604
                    attn_metadata,
                    self.vllm_config,
3605
                    num_tokens=num_tokens_after_padding,
3606
3607
                    num_tokens_across_dp=num_tokens_across_dp,
                    cudagraph_runtime_mode=cudagraph_runtime_mode,
3608
                    batch_descriptor=batch_descriptor,
3609
3610
3611
                    ubatch_slices=ubatch_slices,
                ),
            ):
3612
                outputs = self.model(
3613
3614
3615
3616
                    input_ids=input_ids,
                    positions=positions,
                    intermediate_tensors=intermediate_tensors,
                    inputs_embeds=inputs_embeds,
3617
                    **model_kwargs,
3618
                )
3619

3620
3621
3622
3623
            if self.use_aux_hidden_state_outputs:
                hidden_states, _ = outputs
            else:
                hidden_states = outputs
3624

3625
            if self.speculative_config and self.speculative_config.use_eagle():
3626
                assert isinstance(self.drafter, EagleProposer)
3627
3628
3629
3630
                use_cudagraphs = (
                    cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE
                    and not self.speculative_config.enforce_eager
                )
3631
3632
3633
3634
3635
3636
3637
3638
3639
3640
3641
3642

                # Note(gnovack) - We need to disable cudagraphs for one of the two
                # lora cases when cudagraph_specialize_lora is enabled. This is a
                # short term mitigation for issue mentioned in
                # https://github.com/vllm-project/vllm/issues/28334
                if self.compilation_config.cudagraph_specialize_lora and activate_lora:
                    use_cudagraphs = False

                self.drafter.dummy_run(
                    num_tokens,
                    use_cudagraphs=use_cudagraphs,
                )
3643

3644
3645
3646
3647
3648
3649
3650
3651
3652
3653
        # 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)

3654
        logit_indices = np.cumsum(num_scheduled_tokens) - 1
3655
3656
3657
3658
        logit_indices_device = torch.from_numpy(logit_indices).to(
            self.device, non_blocking=True
        )
        return hidden_states, hidden_states[logit_indices_device]
3659
3660
3661
3662
3663
3664

    @torch.inference_mode()
    def _dummy_sampler_run(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
3665
3666
3667
3668
        # 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)
3669

3670
        logits = self.model.compute_logits(hidden_states)
3671
3672
        num_reqs = logits.size(0)

3673
        dummy_tensors = lambda v: torch.full((num_reqs,), v, device=self.device)
3674
3675
3676
3677
3678
3679
3680
3681
3682
3683
3684
3685
3686
3687
3688

        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)],
3689
            spec_token_ids=[[] for _ in range(num_reqs)],
3690
3691
            allowed_token_ids_mask=None,
            bad_words_token_ids={},
3692
            logitsprocs=LogitsProcessors(),
3693
        )
3694
        try:
3695
3696
3697
            sampler_output = self.sampler(
                logits=logits, sampling_metadata=dummy_metadata
            )
3698
        except RuntimeError as e:
3699
            if "out of memory" in str(e):
3700
3701
3702
3703
                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 "
3704
3705
                    "initializing the engine."
                ) from e
3706
3707
            else:
                raise e
3708
        if self.speculative_config:
3709
3710
            draft_token_ids = [[0] for _ in range(num_reqs)]
            dummy_spec_decode_metadata = SpecDecodeMetadata.make_dummy(
3711
3712
                draft_token_ids, self.device
            )
3713
3714
3715
3716
3717
3718

            num_tokens = sum(len(ids) for ids in draft_token_ids)
            # draft_probs = torch.randn(
            #     num_tokens, logits.shape[-1], device=self.device,
            #     dtype=logits.dtype)
            draft_probs = None
3719
3720
3721
3722
3723
            logits = torch.randn(
                num_tokens + num_reqs,
                logits.shape[-1],
                device=self.device,
                dtype=logits.dtype,
3724
            )
3725
3726
3727
            self.rejection_sampler(
                dummy_spec_decode_metadata,
                draft_probs,
3728
                logits,
3729
3730
                dummy_metadata,
            )
3731
        return sampler_output
3732

3733
    def _dummy_pooler_run_task(
3734
3735
        self,
        hidden_states: torch.Tensor,
3736
3737
        task: PoolingTask,
    ) -> PoolerOutput:
3738
3739
3740
3741
3742
3743
3744
3745
3746
3747
3748
        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

        req_num_tokens = num_tokens // num_reqs

3749
        dummy_prompt_lens = torch.tensor(
3750
3751
            num_scheduled_tokens_list,
            device="cpu",
3752
        )
3753
3754
3755
        dummy_token_ids = torch.zeros(
            (num_reqs, req_num_tokens), dtype=torch.int32, device=self.device
        )
3756

3757
        model = cast(VllmModelForPooling, self.get_model())
3758
        dummy_pooling_params = PoolingParams(task=task)
3759
        dummy_pooling_params.verify(task=task, model_config=self.model_config)
3760
        to_update = model.pooler.get_pooling_updates(task)
3761
3762
        to_update.apply(dummy_pooling_params)

3763
        dummy_metadata = PoolingMetadata(
3764
3765
3766
3767
            prompt_lens=dummy_prompt_lens,
            prompt_token_ids=dummy_token_ids,
            pooling_params=[dummy_pooling_params] * num_reqs,
        )
3768

3769
3770
3771
        dummy_metadata.build_pooling_cursor(
            num_scheduled_tokens_list, device=hidden_states.device
        )
3772

3773
        try:
3774
3775
3776
            return model.pooler(
                hidden_states=hidden_states, pooling_metadata=dummy_metadata
            )
3777
        except RuntimeError as e:
3778
            if "out of memory" in str(e):
3779
                raise RuntimeError(
3780
3781
3782
                    "CUDA out of memory occurred when warming up pooler "
                    f"({task=}) with {num_reqs} dummy requests. Please try "
                    "lowering `max_num_seqs` or `gpu_memory_utilization` when "
3783
3784
                    "initializing the engine."
                ) from e
3785
3786
            else:
                raise e
3787
3788
3789
3790
3791
3792
3793

    @torch.inference_mode()
    def _dummy_pooler_run(
        self,
        hidden_states: torch.Tensor,
    ) -> PoolerOutput:
        # Find the task that has the largest output for subsequent steps
3794
3795
3796
3797
3798
3799
3800
3801
3802
3803
3804
3805
3806
3807
3808
3809
3810
3811
3812
3813
        supported_pooling_tasks = self.get_supported_pooling_tasks()

        if not supported_pooling_tasks:
            if self.scheduler_config.chunked_prefill_enabled:
                raise RuntimeError(
                    f"Model {self.model_config.model} does not support "
                    "any pooling tasks with chunked prefill enabled. "
                    "Please add --no-enable-chunked-prefill to your "
                    "config or CLI args. See "
                    "https://docs.vllm.ai/en/latest/models/pooling_models.html "
                    "to learn more."
                )
            else:
                raise RuntimeError(
                    f"Model {self.model_config.model} does not support "
                    "any pooling tasks. See "
                    "https://docs.vllm.ai/en/latest/models/pooling_models.html "
                    "to learn more."
                )

3814
        output_size = dict[PoolingTask, float]()
3815
        for task in supported_pooling_tasks:
3816
3817
            # Run a full batch with each task to ensure none of them OOMs
            output = self._dummy_pooler_run_task(hidden_states, task)
3818
            output_size[task] = sum(o.nbytes for o in output)
3819
3820
3821
3822
            del output  # Allow GC

        max_task = max(output_size.items(), key=lambda x: x[1])[0]
        return self._dummy_pooler_run_task(hidden_states, max_task)
3823

3824
    def profile_run(self) -> None:
3825
        # Profile with multimodal encoder & encoder cache.
3826
        if self.supports_mm_inputs:
3827
            if self.model_config.multimodal_config.skip_mm_profiling:
3828
                logger.info(
3829
                    "Skipping memory profiling for multimodal encoder and "
3830
3831
                    "encoder cache."
                )
3832
3833
3834
3835
3836
3837
3838
3839
            else:
                mm_budget = self.mm_budget
                assert mm_budget is not None

                if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
                    # NOTE: Currently model is profiled with a single non-text
                    # modality with the max possible input tokens even when
                    # it supports multiple.
3840
                    dummy_modality = mm_budget.get_modality_with_max_tokens()
3841
3842
3843
                    max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[
                        dummy_modality
                    ]
3844
3845
3846
3847
3848
3849
3850
3851
3852

                    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_mm_items_per_batch,
                        dummy_modality,
                    )
3853

3854
3855
3856
3857
3858
                    # Create dummy batch of multimodal inputs.
                    batched_dummy_mm_inputs = self._get_mm_dummy_batch(
                        dummy_modality,
                        max_mm_items_per_batch,
                    )
3859

3860
                    # Run multimodal encoder.
3861
3862
3863
                    dummy_encoder_outputs = self.model.get_multimodal_embeddings(
                        **batched_dummy_mm_inputs
                    )
3864

3865
3866
3867
3868
                    sanity_check_mm_encoder_outputs(
                        dummy_encoder_outputs,
                        expected_num_items=max_mm_items_per_batch,
                    )
3869

3870
3871
3872
3873
3874
3875
3876
3877
3878
3879
                    # NOTE: This happens when encoder cache needs to store
                    # the embeddings that encoder outputs are scattered onto.
                    # In this case we create dummy embeddings of size
                    # (encode_budget, hidden_size) and scatter encoder
                    # output into it.
                    encoder_output_shape = dummy_encoder_outputs[0].shape
                    if encoder_output_shape[0] < encoder_budget:
                        expanded_outputs = []
                        for output in dummy_encoder_outputs:
                            expanded = output.new_zeros(
3880
3881
                                (encoder_budget, encoder_output_shape[-1])
                            )
3882
3883
3884
3885
3886
3887
                            num_tokens = output.shape[0]
                            expanded[:num_tokens].copy_(output)
                            expanded_outputs.append(expanded)

                        dummy_encoder_outputs = expanded_outputs

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

3891
        # Add `is_profile` here to pre-allocate communication buffers
3892
3893
3894
        hidden_states, last_hidden_states = self._dummy_run(
            self.max_num_tokens, is_profile=True
        )
3895
        if get_pp_group().is_last_rank:
3896
3897
3898
3899
            if self.is_pooling_model:
                output = self._dummy_pooler_run(hidden_states)
            else:
                output = self._dummy_sampler_run(last_hidden_states)
3900
        else:
3901
            output = None
3902
        self._sync_device()
3903
        del hidden_states, output
3904
        self.encoder_cache.clear()
3905
        gc.collect()
3906

3907
    def capture_model(self) -> int:
3908
        if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE:
3909
            logger.warning(
3910
                "Skipping CUDA graph capture. To turn on CUDA graph capture, "
3911
3912
                "ensure `cudagraph_mode` was not manually set to `NONE`"
            )
3913
            return 0
3914

3915
3916
        compilation_counter.num_gpu_runner_capture_triggers += 1

3917
3918
        start_time = time.perf_counter()

3919
3920
3921
3922
3923
3924
3925
3926
3927
3928
3929
3930
3931
3932
        @contextmanager
        def freeze_gc():
            # Optimize garbage collection during CUDA graph capture.
            # Clean up, then freeze all remaining objects from being included
            # in future collections.
            gc.collect()
            should_freeze = not envs.VLLM_ENABLE_CUDAGRAPH_GC
            if should_freeze:
                gc.freeze()
            try:
                yield
            finally:
                if should_freeze:
                    gc.unfreeze()
3933
                    gc.collect()
3934

3935
3936
3937
        # 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.
3938
        set_cudagraph_capturing_enabled(True)
3939
        with freeze_gc(), graph_capture(device=self.device):
3940
            start_free_gpu_memory = torch.cuda.mem_get_info()[0]
3941
            cudagraph_mode = self.compilation_config.cudagraph_mode
3942
            assert cudagraph_mode is not None
3943
3944
3945
3946
3947
3948
3949
3950
3951

            if self.lora_config:
                if self.compilation_config.cudagraph_specialize_lora:
                    lora_cases = [True, False]
                else:
                    lora_cases = [True]
            else:
                lora_cases = [False]

3952
3953
            if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:
                cudagraph_runtime_mode = cudagraph_mode.mixed_mode()
3954
                # make sure we capture the largest batch size first
3955
3956
3957
                compilation_cases = list(
                    product(reversed(self.cudagraph_batch_sizes), lora_cases)
                )
3958
3959
3960
                self._capture_cudagraphs(
                    compilation_cases,
                    cudagraph_runtime_mode=cudagraph_runtime_mode,
3961
3962
                    uniform_decode=False,
                )
3963

3964
3965
            # Capture full cudagraph for uniform decode batches if we
            # don't already have full mixed prefill-decode cudagraphs.
3966
3967
3968
3969
3970
3971
3972
            if (
                cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
                and cudagraph_mode.separate_routine()
            ):
                max_num_tokens = (
                    self.scheduler_config.max_num_seqs * self.uniform_decode_query_len
                )
3973
                decode_cudagraph_batch_sizes = [
3974
3975
                    x
                    for x in self.cudagraph_batch_sizes
3976
                    if max_num_tokens >= x >= self.uniform_decode_query_len
3977
                ]
3978
3979
3980
                compilation_cases_decode = list(
                    product(reversed(decode_cudagraph_batch_sizes), lora_cases)
                )
3981
3982
3983
                self._capture_cudagraphs(
                    compilation_cases=compilation_cases_decode,
                    cudagraph_runtime_mode=CUDAGraphMode.FULL,
3984
3985
                    uniform_decode=True,
                )
3986

3987
3988
3989
            torch.cuda.synchronize()
            end_free_gpu_memory = torch.cuda.mem_get_info()[0]

3990
3991
3992
        # Disable cudagraph capturing globally, so any unexpected cudagraph
        # capturing will be detected and raise an error after here.
        # Note: We don't put it into graph_capture context manager because
3993
        # we may do lazy capturing in future that still allows capturing
3994
3995
        # after here.
        set_cudagraph_capturing_enabled(False)
3996
3997
3998
3999
4000

        end_time = time.perf_counter()
        elapsed_time = end_time - start_time
        cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
        # This usually takes 5~20 seconds.
4001
        logger.info_once(
4002
4003
4004
            "Graph capturing finished in %.0f secs, took %.2f GiB",
            elapsed_time,
            cuda_graph_size / (1 << 30),
4005
            scope="local",
4006
        )
4007
        return cuda_graph_size
4008

4009
4010
    def _capture_cudagraphs(
        self,
4011
        compilation_cases: list[tuple[int, bool]],
4012
4013
4014
4015
4016
4017
4018
        cudagraph_runtime_mode: CUDAGraphMode,
        uniform_decode: bool,
    ):
        assert (
            cudagraph_runtime_mode != CUDAGraphMode.NONE
            and cudagraph_runtime_mode.valid_runtime_modes()
        ), f"Invalid cudagraph runtime mode: {cudagraph_runtime_mode}"
4019
4020
4021
4022
4023
4024
4025
4026

        # Only rank 0 should print progress bar during capture
        if is_global_first_rank():
            compilation_cases = tqdm(
                compilation_cases,
                disable=not self.load_config.use_tqdm_on_load,
                desc="Capturing CUDA graphs ({}, {})".format(
                    "decode" if uniform_decode else "mixed prefill-decode",
4027
4028
4029
                    cudagraph_runtime_mode.name,
                ),
            )
4030

4031
        # We skip EPLB here since we don't want to record dummy metrics
4032
        for num_tokens, activate_lora in compilation_cases:
4033
            # We currently only capture ubatched graphs when its a FULL
4034
4035
4036
            # cudagraph, a uniform decode batch, and the number of tokens
            # is above the threshold. Otherwise we just capture a non-ubatched
            # version of the graph
4037
4038
4039
4040
            allow_microbatching = (
                self.parallel_config.enable_dbo
                and cudagraph_runtime_mode == CUDAGraphMode.FULL
                and uniform_decode
4041
4042
4043
4044
4045
                and check_ubatch_thresholds(
                    config=self.vllm_config.parallel_config,
                    num_tokens=num_tokens,
                    uniform_decode=uniform_decode,
                )
4046
            )
4047

4048
4049
4050
4051
4052
4053
            for _ in range(self.compilation_config.cudagraph_num_of_warmups):
                # Use CUDAGraphRuntimeStyle.NONE (default) for warmup.
                # But be careful, warm up with `NONE`is orthogonal to
                # if we want to warm up attention or not. This is
                # different from the case where `FULL` implies capture
                # attention while `PIECEWISE` implies no attention.
4054
4055
4056
4057
4058
4059
4060
4061
4062
                force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL
                self._dummy_run(
                    num_tokens,
                    cudagraph_runtime_mode=CUDAGraphMode.NONE,
                    force_attention=force_attention,
                    uniform_decode=uniform_decode,
                    allow_microbatching=allow_microbatching,
                    skip_eplb=True,
                    remove_lora=False,
4063
                    activate_lora=activate_lora,
4064
4065
4066
4067
4068
4069
4070
4071
                )
            self._dummy_run(
                num_tokens,
                cudagraph_runtime_mode=cudagraph_runtime_mode,
                uniform_decode=uniform_decode,
                allow_microbatching=allow_microbatching,
                skip_eplb=True,
                remove_lora=False,
4072
                activate_lora=activate_lora,
4073
            )
4074
        self.maybe_remove_all_loras(self.lora_config)
4075

4076
4077
4078
4079
    def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
        """
        Initialize the attention backends and attention metadata builders.
        """
4080
        assert len(self.attn_groups) == 0, "Attention backends are already initialized"
4081

4082
4083
4084
4085
4086
4087
        class AttentionGroupKey(NamedTuple):
            attn_backend: type[AttentionBackend]
            kv_cache_spec: KVCacheSpec

        def get_attn_backends_for_group(
            kv_cache_group_spec: KVCacheGroupSpec,
4088
        ) -> tuple[dict[AttentionGroupKey, list[str]], set[type[AttentionBackend]]]:
4089
            layers = get_layers_from_vllm_config(
4090
4091
                self.vllm_config, AttentionLayerBase, kv_cache_group_spec.layer_names
            )
4092
4093
            attn_backends = {}
            attn_backend_layers = defaultdict(list)
4094
            # Dedupe based on full class name; this is a bit safer than
4095
4096
4097
4098
            # using the class itself as the key because when we create dynamic
            # attention backend subclasses (e.g. ChunkedLocalAttention) unless
            # they are cached correctly, there will be different objects per
            # layer.
4099
            for layer_name in kv_cache_group_spec.layer_names:
4100
                attn_backend = layers[layer_name].get_attn_backend()
4101
4102
4103
4104
4105
4106
4107

                if layer_name in self.kv_sharing_fast_prefill_eligible_layers:
                    attn_backend = create_fast_prefill_custom_backend(
                        "FastPrefill",
                        attn_backend,
                    )

4108
4109
4110
                full_cls_name = attn_backend.full_cls_name()
                layer_kv_cache_spec = kv_cache_group_spec.kv_cache_spec
                if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
4111
                    layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name]
4112
                key = (full_cls_name, layer_kv_cache_spec)
4113
4114
4115
                attn_backends[key] = AttentionGroupKey(
                    attn_backend, layer_kv_cache_spec
                )
4116
                attn_backend_layers[key].append(layer_name)
4117
4118
4119
4120
            return (
                {attn_backends[k]: v for k, v in attn_backend_layers.items()},
                set(group_key.attn_backend for group_key in attn_backends.values()),
            )
4121
4122

        def create_attn_groups(
4123
            attn_backends_map: dict[AttentionGroupKey, list[str]],
4124
            kv_cache_group_id: int,
4125
4126
        ) -> list[AttentionGroup]:
            attn_groups: list[AttentionGroup] = []
4127
            for (attn_backend, kv_cache_spec), layer_names in attn_backends_map.items():
4128
                attn_group = AttentionGroup(
4129
                    attn_backend,
4130
                    layer_names,
4131
                    kv_cache_spec,
4132
                    kv_cache_group_id,
4133
4134
                )

4135
4136
4137
                attn_groups.append(attn_group)
            return attn_groups

4138
4139
        attention_backend_maps = []
        attention_backend_set: set[type[AttentionBackend]] = set()
4140
        for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
4141
            attn_backends = get_attn_backends_for_group(kv_cache_group_spec)
4142
4143
4144
4145
4146
4147
            attention_backend_maps.append(attn_backends[0])
            attention_backend_set.update(attn_backends[1])

        # Resolve cudagraph_mode before actually initialize metadata_builders
        self._check_and_update_cudagraph_mode(attention_backend_set)

4148
4149
        for i, attn_backend_map in enumerate(attention_backend_maps):
            self.attn_groups.append(create_attn_groups(attn_backend_map, i))
4150

4151
4152
4153
4154
4155
4156
4157
4158
4159
4160
4161
4162
4163
4164
4165
4166
4167
4168
    def initialize_metadata_builders(
        self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
    ) -> None:
        """
        Create the metadata builders for all KV cache groups and attn groups.
        """
        for kv_cache_group_id in range(len(kv_cache_config.kv_cache_groups)):
            for attn_group in self.attn_groups[kv_cache_group_id]:
                attn_group.create_metadata_builders(
                    self.vllm_config,
                    self.device,
                    kernel_block_sizes[kv_cache_group_id]
                    if kv_cache_group_id < len(kernel_block_sizes)
                    else None,
                    num_metadata_builders=1
                    if not self.parallel_config.enable_dbo
                    else 2,
                )
co63oc's avatar
co63oc committed
4169
        # Calculate reorder batch threshold (if needed)
4170
4171
        # Note (tdoublep): do this *after* constructing builders,
        # because some of them change the threshold at init time.
4172
4173
        self.calculate_reorder_batch_threshold()

4174
4175
4176
    def _check_and_update_cudagraph_mode(
        self, attention_backends: set[type[AttentionBackend]]
    ) -> None:
4177
        """
4178
        Resolve the cudagraph_mode when there are multiple attention
4179
4180
4181
4182
        backends with potential conflicting CUDA graph support.
        Then initialize the cudagraph_dispatcher based on the resolved
        cudagraph_mode.
        """
4183
        min_cg_support = AttentionCGSupport.ALWAYS
4184
        min_cg_backend_name = None
4185

4186
4187
4188
4189
4190
        for attn_backend in attention_backends:
            builder_cls = attn_backend.get_builder_cls()
            if builder_cls.cudagraph_support.value < min_cg_support.value:
                min_cg_support = builder_cls.cudagraph_support
                min_cg_backend_name = attn_backend.__name__
4191
4192
4193
        # Flexible resolve the cudagraph mode
        cudagraph_mode = self.compilation_config.cudagraph_mode
        # check cudagraph for mixed batch is supported
4194
4195
4196
4197
4198
4199
        if (
            cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL
            and min_cg_support != AttentionCGSupport.ALWAYS
        ):
            msg = (
                f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
4200
                f"with {min_cg_backend_name} backend (support: "
4201
4202
                f"{min_cg_support})"
            )
4203
4204
            if min_cg_support == AttentionCGSupport.NEVER:
                # if not supported any full cudagraphs, just raise it.
4205
4206
                msg += (
                    "; please try cudagraph_mode=PIECEWISE, and "
4207
                    "make sure compilation mode is VLLM_COMPILE"
4208
                )
4209
4210
4211
4212
4213
                raise ValueError(msg)

            # attempt to resolve the full cudagraph related mode
            if self.compilation_config.splitting_ops_contain_attention():
                msg += "; setting cudagraph_mode=FULL_AND_PIECEWISE"
4214
                cudagraph_mode = self.compilation_config.cudagraph_mode = (
4215
                    CUDAGraphMode.FULL_AND_PIECEWISE
4216
                )
4217
4218
            else:
                msg += "; setting cudagraph_mode=FULL_DECODE_ONLY"
4219
                cudagraph_mode = self.compilation_config.cudagraph_mode = (
4220
                    CUDAGraphMode.FULL_DECODE_ONLY
4221
                )
4222
4223
            logger.warning(msg)

4224
        # check that if we are doing decode full-cudagraphs it is supported
4225
4226
4227
4228
4229
4230
        if (
            cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
            and min_cg_support == AttentionCGSupport.NEVER
        ):
            msg = (
                f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
4231
                f"with {min_cg_backend_name} backend (support: "
4232
4233
                f"{min_cg_support})"
            )
4234
            if self.compilation_config.mode == CompilationMode.VLLM_COMPILE and (
4235
4236
4237
4238
4239
                self.compilation_config.splitting_ops_contain_attention()
                or self.compilation_config.use_inductor_graph_partition
            ):
                msg += (
                    "; setting cudagraph_mode=PIECEWISE because "
4240
                    "attention is compiled piecewise"
4241
4242
                )
                cudagraph_mode = self.compilation_config.cudagraph_mode = (
4243
                    CUDAGraphMode.PIECEWISE
4244
                )
4245
            else:
4246
4247
                msg += (
                    "; setting cudagraph_mode=NONE because "
4248
                    "attention is not compiled piecewise"
4249
4250
                )
                cudagraph_mode = self.compilation_config.cudagraph_mode = (
4251
                    CUDAGraphMode.NONE
4252
                )
4253
4254
            logger.warning(msg)

4255
4256
        # check that if we are doing spec-decode + decode full-cudagraphs it is
        # supported
4257
4258
4259
4260
4261
4262
4263
4264
        if (
            cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
            and self.uniform_decode_query_len > 1
            and min_cg_support.value < AttentionCGSupport.UNIFORM_BATCH.value
        ):
            msg = (
                f"CUDAGraphMode.{cudagraph_mode.name} is not supported"
                f" with spec-decode for attention backend "
4265
                f"{min_cg_backend_name} (support: {min_cg_support})"
4266
            )
4267
4268
            if self.compilation_config.splitting_ops_contain_attention():
                msg += "; setting cudagraph_mode=PIECEWISE"
4269
                cudagraph_mode = self.compilation_config.cudagraph_mode = (
4270
                    CUDAGraphMode.PIECEWISE
4271
                )
4272
4273
            else:
                msg += "; setting cudagraph_mode=NONE"
4274
                cudagraph_mode = self.compilation_config.cudagraph_mode = (
4275
                    CUDAGraphMode.NONE
4276
                )
4277
4278
4279
4280
            logger.warning(msg)

        # double check that we can support full cudagraph if they are requested
        # even after automatic downgrades
4281
4282
4283
4284
4285
4286
        if (
            cudagraph_mode.has_full_cudagraphs()
            and min_cg_support == AttentionCGSupport.NEVER
        ):
            raise ValueError(
                f"CUDAGraphMode.{cudagraph_mode.name} is not "
4287
                f"supported with {min_cg_backend_name} backend ("
4288
4289
                f"support:{min_cg_support}) "
                "; please try cudagraph_mode=PIECEWISE, "
4290
                "and make sure compilation mode is VLLM_COMPILE"
4291
            )
4292

4293
4294
        # Trigger cudagraph dispatching keys initialization after
        # resolved cudagraph mode.
4295
        self.cudagraph_dispatcher.initialize_cudagraph_keys(
4296
4297
            self.compilation_config.cudagraph_mode, self.uniform_decode_query_len
        )
4298

4299
4300
    def calculate_reorder_batch_threshold(self) -> None:
        """
4301
4302
4303
4304
        Choose the minimum reorder batch threshold from all attention groups.
        Backends should be able to support lower threshold then what they request
        just may have a performance penalty due to that backend treating decodes
        as prefills.
4305
        """
4306
4307
4308
4309
4310
4311
        min_none_high = lambda a, b: a if b is None else b if a is None else min(a, b)

        reorder_batch_thresholds = [
            group.get_metadata_builder().reorder_batch_threshold
            for group in self._attn_group_iterator()
        ]
4312
4313
4314
4315
4316
        # If there are no attention groups (attention-free model) or no backend
        # reports a threshold, leave reordering disabled.
        if len(reorder_batch_thresholds) == 0:
            self.reorder_batch_threshold = None
            return
4317
        self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds)
4318

4319
4320
4321
    @staticmethod
    def select_common_block_size(
        kv_manager_block_size: int, attn_groups: list[AttentionGroup]
4322
4323
    ) -> int:
        """
4324
4325
4326
4327
4328
        Select a block size that is supported by all backends and is a factor of
        kv_manager_block_size.

        If kv_manager_block_size is supported by all backends, return it directly.
        Otherwise, return the max supported size.
4329
4330
4331
4332
4333
4334

        Args:
            kv_manager_block_size: Block size of KV cache
            attn_groups: List of attention groups

        Returns:
4335
            The selected block size
4336
4337

        Raises:
4338
            ValueError: If no valid block size found
4339
4340
        """

4341
4342
4343
4344
4345
4346
4347
4348
        def block_size_is_supported(
            backends: list[type[AttentionBackend]], block_size: int
        ) -> bool:
            """
            Check if the block size is supported by all backends.
            """
            for backend in backends:
                is_supported = False
4349
                for supported_size in backend.supported_kernel_block_sizes:
4350
4351
4352
4353
4354
4355
4356
4357
4358
4359
4360
4361
4362
4363
4364
4365
4366
4367
4368
4369
4370
4371
4372
4373
4374
4375
4376
4377
4378
4379
                    if isinstance(supported_size, int):
                        if block_size == supported_size:
                            is_supported = True
                    elif isinstance(supported_size, MultipleOf):
                        if block_size % supported_size.base == 0:
                            is_supported = True
                    else:
                        raise ValueError(f"Unknown supported size: {supported_size}")
                if not is_supported:
                    return False
            return True

        backends = [group.backend for group in attn_groups]

        # Case 1: if the block_size of kv cache manager is supported by all backends,
        # return it directly
        if block_size_is_supported(backends, kv_manager_block_size):
            return kv_manager_block_size

        # Case 2: otherwise, the block_size must be an `int`-format supported size of
        # at least one backend. Iterate over all `int`-format supported sizes in
        # descending order and return the first one that is supported by all backends.
        # Simple proof:
        # If the supported size b is in MultipleOf(x_i) format for all attention
        # backends i, and b a factor of kv_manager_block_size, then
        # kv_manager_block_size also satisfies MultipleOf(x_i) for all i. We will
        # return kv_manager_block_size in case 1.
        all_int_supported_sizes = set(
            supported_size
            for backend in backends
4380
            for supported_size in backend.supported_kernel_block_sizes
4381
4382
            if isinstance(supported_size, int)
        )
4383

4384
4385
4386
4387
4388
4389
        for supported_size in sorted(all_int_supported_sizes, reverse=True):
            if kv_manager_block_size % supported_size != 0:
                continue
            if block_size_is_supported(backends, supported_size):
                return supported_size
        raise ValueError(f"No common block size for {kv_manager_block_size}. ")
4390

4391
4392
4393
    def may_reinitialize_input_batch(
        self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
    ) -> None:
4394
4395
4396
4397
4398
4399
4400
        """
        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.
4401
            kernel_block_sizes: The kernel block sizes for each KV cache group.
4402
4403
4404
4405
        """
        block_sizes = [
            kv_cache_group.kv_cache_spec.block_size
            for kv_cache_group in kv_cache_config.kv_cache_groups
4406
            if not isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec)
4407
        ]
4408
4409
4410
4411

        if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [
            self.cache_config.block_size
        ]:
4412
4413
4414
            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
4415
4416
                "for more details."
            )
4417
4418
            self.input_batch = InputBatch(
                max_num_reqs=self.max_num_reqs,
4419
                max_model_len=max(self.max_model_len, self.max_encoder_len),
4420
4421
4422
4423
4424
                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,
4425
                kernel_block_sizes=kernel_block_sizes,
4426
                is_spec_decode=bool(self.vllm_config.speculative_config),
4427
                logitsprocs=self.input_batch.logitsprocs,
4428
                logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids,
4429
                is_pooling_model=self.is_pooling_model,
4430
4431
                num_speculative_tokens=(
                    self.vllm_config.speculative_config.num_speculative_tokens
4432
4433
4434
                    if self.vllm_config.speculative_config
                    else 0
                ),
4435
4436
            )

4437
    def _allocate_kv_cache_tensors(
4438
4439
        self, kv_cache_config: KVCacheConfig
    ) -> dict[str, torch.Tensor]:
4440
        """
4441
4442
4443
        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.

4444
        Args:
4445
            kv_cache_config: The KV cache config
4446
        Returns:
4447
            dict[str, torch.Tensor]: A map between layer names to their
4448
            corresponding memory buffer for KV cache.
4449
        """
4450
4451
        kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
        for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
4452
4453
4454
            tensor = torch.zeros(
                kv_cache_tensor.size, dtype=torch.int8, device=self.device
            )
4455
4456
4457
4458
4459
            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:
4460
4461
4462
4463
            for layer_name in group.layer_names:
                if layer_name in self.runner_only_attn_layers:
                    continue
                layer_names.add(layer_name)
4464
4465
4466
        assert layer_names == set(kv_cache_raw_tensors.keys()), (
            "Some layers are not correctly initialized"
        )
4467
4468
        return kv_cache_raw_tensors

4469
4470
4471
    def _attn_group_iterator(self) -> Iterator[AttentionGroup]:
        return itertools.chain.from_iterable(self.attn_groups)

4472
    def _kv_cache_spec_attn_group_iterator(self) -> Iterator[AttentionGroup]:
4473
4474
        if not self.kv_cache_config.kv_cache_groups:
            return
4475
4476
        for attn_groups in self.attn_groups:
            yield from attn_groups
4477

4478
4479
4480
4481
4482
4483
4484
4485
4486
4487
4488
4489
4490
4491
4492
    def _prepare_kernel_block_sizes(self, kv_cache_config: KVCacheConfig) -> list[int]:
        """
        Generate kernel_block_sizes that matches each block_size.

        For attention backends that support virtual block splitting,
        use the supported block sizes from the backend.
        For other backends (like Mamba), use the same block size (no splitting).

        Args:
            kv_cache_config: The KV cache configuration.

        Returns:
            list[int]: List of kernel block sizes for each cache group.
        """
        kernel_block_sizes = []
4493
        for kv_cache_gid, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
4494
4495
4496
4497
4498
4499
            kv_cache_spec = kv_cache_group.kv_cache_spec
            if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
                # All layers in the UniformTypeKVCacheSpecs have the same type,
                # Pick an arbitrary one to dispatch.
                kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values()))
            if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
4500
                continue
4501
            elif isinstance(kv_cache_spec, AttentionSpec):
4502
4503
4504
                # This is an attention backend that supports virtual
                # block splitting. Get the supported block sizes from
                # all backends in the group.
4505
                attn_groups = self.attn_groups[kv_cache_gid]
4506
                kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size
4507
                selected_kernel_size = self.select_common_block_size(
4508
4509
4510
                    kv_manager_block_size, attn_groups
                )
                kernel_block_sizes.append(selected_kernel_size)
4511
            elif isinstance(kv_cache_spec, MambaSpec):
4512
4513
                # This is likely Mamba or other non-attention cache,
                # no splitting.
4514
                kernel_block_sizes.append(kv_cache_spec.block_size)
4515
4516
4517
4518
4519
4520
            else:
                raise NotImplementedError(
                    f"unknown kv cache spec {kv_cache_group.kv_cache_spec}"
                )
        return kernel_block_sizes

4521
4522
4523
4524
    def _reshape_kv_cache_tensors(
        self,
        kv_cache_config: KVCacheConfig,
        kv_cache_raw_tensors: dict[str, torch.Tensor],
4525
        kernel_block_sizes: list[int],
4526
    ) -> dict[str, torch.Tensor]:
4527
        """
4528
        Reshape the KV cache tensors to the desired shape and dtype.
4529

4530
        Args:
4531
4532
            kv_cache_config: The KV cache config
            kv_cache_raw_tensors: The KV cache buffer of each layer, with
4533
                correct size but uninitialized shape.
4534
            kernel_block_sizes: The kernel block sizes for each KV cache group.
4535
        Returns:
4536
            Dict[str, torch.Tensor]: A map between layer names to their
4537
4538
            corresponding memory buffer for KV cache.
        """
4539
        kv_caches: dict[str, torch.Tensor] = {}
4540
        has_attn, has_mamba = False, False
4541
4542
        for group in self._kv_cache_spec_attn_group_iterator():
            kv_cache_spec = group.kv_cache_spec
4543
            attn_backend = group.backend
4544
4545
4546
4547
            if group.kv_cache_group_id == len(kernel_block_sizes):
                # There may be a last group for layers without kv cache.
                continue
            kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
4548
            for layer_name in group.layer_names:
4549
4550
                if layer_name in self.runner_only_attn_layers:
                    continue
4551
4552
                raw_tensor = kv_cache_raw_tensors[layer_name]
                assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
4553
                num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
4554
                if isinstance(kv_cache_spec, AttentionSpec):
4555
                    has_attn = True
4556
4557
                    num_blocks_per_kv_block = (
                        kv_cache_spec.block_size // kernel_block_size
4558
4559
4560
                    )
                    kernel_num_blocks = num_blocks * num_blocks_per_kv_block

4561
                    kv_cache_shape = attn_backend.get_kv_cache_shape(
4562
                        kernel_num_blocks,
4563
                        kernel_block_size,
4564
4565
                        kv_cache_spec.num_kv_heads,
                        kv_cache_spec.head_size,
4566
4567
                        cache_dtype_str=self.cache_config.cache_dtype,
                    )
4568
                    dtype = kv_cache_spec.dtype
4569
                    try:
4570
                        kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
4571
                        assert len(kv_cache_stride_order) == len(kv_cache_shape)
4572
                    except (AttributeError, NotImplementedError):
4573
                        kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
4574
4575
4576
4577
4578
                    # 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.
4579
4580
4581
                    kv_cache_shape = tuple(
                        kv_cache_shape[i] for i in kv_cache_stride_order
                    )
4582
4583
4584
4585
4586
                    # Maintain original KV shape view.
                    inv_order = [
                        kv_cache_stride_order.index(i)
                        for i in range(len(kv_cache_stride_order))
                    ]
4587
4588
4589
4590
4591
4592
                    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
4593
                elif isinstance(kv_cache_spec, MambaSpec):
4594
                    has_mamba = True
Chen Zhang's avatar
Chen Zhang committed
4595
4596
                    raw_tensor = kv_cache_raw_tensors[layer_name]
                    state_tensors = []
4597
                    storage_offset_bytes = 0
4598
                    for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):
4599
4600
                        dtype_size = get_dtype_size(dtype)
                        num_element_per_page = (
4601
4602
                            kv_cache_spec.page_size_bytes // dtype_size
                        )
Chen Zhang's avatar
Chen Zhang committed
4603
                        target_shape = (num_blocks, *shape)
4604
4605
                        stride = torch.empty(target_shape).stride()
                        target_stride = (num_element_per_page, *stride[1:])
4606
                        assert storage_offset_bytes % dtype_size == 0
4607
4608
4609
4610
                        tensor = torch.as_strided(
                            raw_tensor.view(dtype),
                            size=target_shape,
                            stride=target_stride,
4611
                            storage_offset=storage_offset_bytes // dtype_size,
4612
                        )
Chen Zhang's avatar
Chen Zhang committed
4613
                        state_tensors.append(tensor)
4614
                        storage_offset_bytes += stride[0] * dtype_size
4615
4616

                    kv_caches[layer_name] = state_tensors
4617
                else:
4618
                    raise NotImplementedError
4619
4620

        if has_attn and has_mamba:
4621
            self._update_hybrid_attention_mamba_layout(kv_caches)
4622

4623
4624
        return kv_caches

4625
    def _update_hybrid_attention_mamba_layout(
4626
4627
        self, kv_caches: dict[str, torch.Tensor]
    ) -> None:
4628
        """
4629
4630
        Update the layout of attention layers from (2, num_blocks, ...) to
        (num_blocks, 2, ...).
4631
4632

        Args:
4633
            kv_caches: The KV cache buffer of each layer.
4634
4635
        """

4636
4637
        for group in self._kv_cache_spec_attn_group_iterator():
            kv_cache_spec = group.kv_cache_spec
4638
            for layer_name in group.layer_names:
4639
                kv_cache = kv_caches[layer_name]
4640
4641
4642
4643
                if isinstance(kv_cache_spec, AttentionSpec) and kv_cache.shape[0] == 2:
                    assert kv_cache.shape[1] != 2, (
                        "Fail to determine whether the layout is "
                        "(2, num_blocks, ...) or (num_blocks, 2, ...) for "
4644
                        f"a tensor of shape {kv_cache.shape}"
4645
                    )
4646
                    hidden_size = kv_cache.shape[2:].numel()
4647
4648
4649
4650
                    kv_cache.as_strided_(
                        size=kv_cache.shape,
                        stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]),
                    )
4651

4652
    def initialize_kv_cache_tensors(
4653
        self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
4654
    ) -> dict[str, torch.Tensor]:
4655
4656
4657
4658
4659
        """
        Initialize the memory buffer for KV cache.

        Args:
            kv_cache_config: The KV cache config
4660
4661
            kernel_block_sizes: The kernel block sizes for each KV cache group.

4662
        Returns:
4663
            Dict[str, torch.Tensor]: A map between layer names to their
4664
4665
4666
4667
4668
            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
4669
        kv_caches = self._reshape_kv_cache_tensors(
4670
            kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes
4671
        )
4672

4673
        # Set up cross-layer KV cache sharing
4674
4675
        for layer_name, target_layer_name in self.shared_kv_cache_layers.items():
            logger.debug("%s reuses KV cache of %s", layer_name, target_layer_name)
4676
4677
            kv_caches[layer_name] = kv_caches[target_layer_name]

4678
4679
4680
4681
4682
4683
4684
4685
4686
        num_attn_module = (
            2 if self.model_config.hf_config.model_type == "longcat_flash" else 1
        )
        bind_kv_cache(
            kv_caches,
            self.compilation_config.static_forward_context,
            self.kv_caches,
            num_attn_module,
        )
4687
4688
4689
        return kv_caches

    def maybe_add_kv_sharing_layers_to_kv_cache_groups(
4690
4691
        self, kv_cache_config: KVCacheConfig
    ) -> None:
4692
4693
4694
4695
4696
4697
4698
4699
4700
4701
4702
4703
4704
4705
4706
4707
4708
4709
        """
        Add layers that re-use KV cache to KV cache group of its target layer.
        Mapping of KV cache tensors happens in `initialize_kv_cache_tensors()`
        """
        if not self.shared_kv_cache_layers:
            # No cross-layer KV sharing, return
            return

        add_kv_sharing_layers_to_kv_cache_groups(
            self.shared_kv_cache_layers,
            kv_cache_config.kv_cache_groups,
            self.runner_only_attn_layers,
        )

        if self.cache_config.kv_sharing_fast_prefill:
            # In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other
            # similar KV sharing setups, only the layers that generate KV caches
            # are involved in the prefill phase, enabling prefill to early exit.
4710
            attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
4711
4712
            for layer_name in reversed(attn_layers):
                if layer_name in self.shared_kv_cache_layers:
4713
                    self.kv_sharing_fast_prefill_eligible_layers.add(layer_name)
4714
4715
                else:
                    break
4716

4717
4718
4719
4720
4721
4722
4723
    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
        """
4724
        kv_cache_config = deepcopy(kv_cache_config)
4725
        self.kv_cache_config = kv_cache_config
4726
        self.may_add_encoder_only_layers_to_kv_cache_config()
4727
        self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
4728
        self.initialize_attn_backend(kv_cache_config)
4729
4730
4731
4732
4733
4734
        # The kernel block size for all KV cache groups. For example, if
        # kv_cache_manager uses block_size 256 for a given group, but the attention
        # backends for that group only supports block_size 64, we will return
        # kernel_block_size 64 and split the 256-token-block to 4 blocks with 64
        # tokens each.
        kernel_block_sizes = self._prepare_kernel_block_sizes(kv_cache_config)
4735
4736
4737
4738

        # create metadata builders
        self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes)

4739
        # Reinitialize need to after initialize_attn_backend
4740
4741
4742
4743
        self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
        kv_caches = self.initialize_kv_cache_tensors(
            kv_cache_config, kernel_block_sizes
        )
4744

4745
4746
4747
4748
4749
4750
        if self.speculative_config and self.speculative_config.use_eagle():
            assert isinstance(self.drafter, EagleProposer)
            # validate all draft model layers belong to the same kv cache
            # group
            self.drafter.validate_same_kv_cache_group(kv_cache_config)

Robert Shaw's avatar
Robert Shaw committed
4751
        if has_kv_transfer_group():
4752
4753
4754
            kv_transfer_group = get_kv_transfer_group()
            kv_transfer_group.register_kv_caches(kv_caches)
            kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
Robert Shaw's avatar
Robert Shaw committed
4755

4756
        if self.dcp_world_size > 1:
4757
            layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)
4758
4759
4760
4761
4762
            for layer in layers.values():
                assert layer.impl.need_to_return_lse_for_decode, (
                    "DCP requires attention impls to return"
                    " the softmax lse for decode, but the impl "
                    f"{layer.impl.__class__.__name__} "
4763
4764
                    "does not return the softmax lse for decode."
                )
4765

4766
4767
4768
4769
4770
    def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
        """
        Add encoder-only layers to the KV cache config.
        """
        block_size = self.vllm_config.cache_config.block_size
4771
        encoder_only_attn_specs: dict[AttentionSpec, list[str]] = defaultdict(list)
4772
4773
4774
        attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)
        for layer_name, attn_module in attn_layers.items():
            if attn_module.attn_type == AttentionType.ENCODER_ONLY:
4775
                attn_spec: AttentionSpec = EncoderOnlyAttentionSpec(
4776
4777
4778
                    block_size=block_size,
                    num_kv_heads=attn_module.num_kv_heads,
                    head_size=attn_module.head_size,
4779
4780
                    dtype=self.kv_cache_dtype,
                )
4781
4782
4783
                encoder_only_attn_specs[attn_spec].append(layer_name)
                self.runner_only_attn_layers.add(layer_name)
        if len(encoder_only_attn_specs) > 0:
4784
4785
4786
            assert len(encoder_only_attn_specs) == 1, (
                "Only support one encoder-only attention spec now"
            )
4787
4788
            spec, layer_names = encoder_only_attn_specs.popitem()
            self.kv_cache_config.kv_cache_groups.append(
4789
4790
                KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
            )
4791

4792
    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
4793
        """
4794
        Generates the KVCacheSpec by parsing the kv cache format from each
4795
4796
        Attention module in the static forward context.
        Returns:
4797
            KVCacheSpec: A dictionary mapping layer names to their KV cache
4798
4799
4800
            format. Layers that do not need KV cache are not included.
        """

4801
        kv_cache_spec: dict[str, KVCacheSpec] = {}
4802
        attn_layers = get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase)
Chen Zhang's avatar
Chen Zhang committed
4803
        for layer_name, attn_module in attn_layers.items():
4804
4805
4806
4807
4808
4809
4810
4811
4812
4813
4814
4815
4816
4817
4818
            if isinstance(attn_module, Attention) and (
                kv_tgt_layer := attn_module.kv_sharing_target_layer_name
            ):
                # 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
            # Skip modules that don't need KV cache (eg encoder-only attention)
            if spec := attn_module.get_kv_cache_spec(self.vllm_config):
                kv_cache_spec[layer_name] = spec
4819

4820
        return kv_cache_spec
4821
4822
4823
4824
4825
4826
4827
4828
4829
4830

    def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
        # This is a short term mitigation for issue mentioned in
        # https://github.com/vllm-project/vllm/issues/22754.
        # `tolist` would trigger a cuda wise stream sync, which
        # would block other copy ops from other cuda streams.
        # A cuda event sync would avoid such a situation. Since
        # this is in the critical path of every single model
        # forward loop, this has caused perf issue for a disagg
        # setup.
4831
        pinned = self.sampled_token_ids_pinned_cpu[: sampled_token_ids.shape[0]]
4832
4833
4834
4835
        pinned.copy_(sampled_token_ids, non_blocking=True)
        self.transfer_event.record()
        self.transfer_event.synchronize()
        return pinned.tolist()