model_runner.py 56 KB
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
2
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
"""
NOTE: Coding style guide for this file:
This model runner is shared by all models: text and multimodal, generative
and embedding, public and private. As a result, this file must only contain
code that is common to every model. Model-specific behavior belongs in the
appropriate model-specific files.

In other words:
* Be paranoid about changing this file. It should remain stable.
* Be even more paranoid about adding new lines. It should remain minimal.

Even for shared features (for example, different parallelism modes), keep the
complexity out of this path. The less common the feature, the more it should be
hidden. Prefer utility functions defined elsewhere and call them from here,
instead of embedding feature-specific logic directly.
"""

20
import functools
Woosuk Kwon's avatar
Woosuk Kwon committed
21
22
23
import gc
import time
from copy import deepcopy
24
from typing import Any, NamedTuple
Woosuk Kwon's avatar
Woosuk Kwon committed
25
26
27
28
29
30
31

import numpy as np
import torch
import torch.nn as nn

from vllm.config import VllmConfig
from vllm.config.compilation import CUDAGraphMode
32
from vllm.distributed.parallel_state import (
33
    get_dcp_group,
34
35
36
    get_pp_group,
    prepare_communication_buffer_for_model,
)
37
from vllm.forward_context import BatchDescriptor, set_forward_context
Woosuk Kwon's avatar
Woosuk Kwon committed
38
from vllm.logger import init_logger
39
40
41
from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
    initialize_mamba_ssu_backend,
)
Woosuk Kwon's avatar
Woosuk Kwon committed
42
from vllm.model_executor.model_loader import get_model_loader
43
from vllm.multimodal import MULTIMODAL_REGISTRY
44
from vllm.sequence import IntermediateTensors
45
from vllm.tasks import SupportedTask
46
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
Woosuk Kwon's avatar
Woosuk Kwon committed
47
48
49
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.kv_cache_interface import KVCacheConfig
50
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
51
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
52
from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
Woosuk Kwon's avatar
Woosuk Kwon committed
53
from vllm.v1.worker.gpu.attn_utils import (
54
    build_slot_mappings_by_layer,
Woosuk Kwon's avatar
Woosuk Kwon committed
55
56
57
58
59
    get_kv_cache_spec,
    init_attn_backend,
    init_kv_cache,
)
from vllm.v1.worker.gpu.block_table import BlockTables
60
from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
61
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
62
63
64
65
66
from vllm.v1.worker.gpu.cudagraph_utils import (
    BatchExecutionDescriptor,
    ModelCudaGraphManager,
    get_uniform_token_count,
)
67
from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
68
from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after
Woosuk Kwon's avatar
Woosuk Kwon committed
69
70
71
from vllm.v1.worker.gpu.input_batch import (
    InputBatch,
    InputBuffers,
72
    combine_sampled_and_draft_tokens,
73
    expand_idx_mapping,
74
    get_num_sampled_and_rejected,
75
    post_update,
76
    post_update_pool,
77
78
    prepare_pos_seq_lens,
    prepare_prefill_inputs,
Woosuk Kwon's avatar
Woosuk Kwon committed
79
)
80
81
82
83
84
from vllm.v1.worker.gpu.kv_connector import (
    NO_OP_KV_CONNECTOR,
    KVConnector,
    get_kv_connector,
)
85
from vllm.v1.worker.gpu.lora_utils import LoraState
86
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
87
from vllm.v1.worker.gpu.model_states import init_model_state
88
from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
89
from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive
90
from vllm.v1.worker.gpu.sample.output import SamplerOutput
91
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
92
from vllm.v1.worker.gpu.sample.sampler import Sampler
93
from vllm.v1.worker.gpu.spec_decode import init_speculator
94
95
96
from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
    set_eagle3_aux_hidden_state_layers,
)
97
from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler
98
from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
99
from vllm.v1.worker.gpu.states import RequestState
100
from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
Woosuk Kwon's avatar
Woosuk Kwon committed
101
102
103
104
105
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

logger = init_logger(__name__)


106
class GPUModelRunner(LoRAModelRunnerMixin):
107
    def __init__(self, vllm_config: VllmConfig, device: torch.device):
Woosuk Kwon's avatar
Woosuk Kwon committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
        self.cache_config = vllm_config.cache_config
        self.compilation_config = vllm_config.compilation_config
        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

        self.device = device
        self.dtype = self.model_config.dtype
        self.kv_cache_dtype = self.dtype
        if self.cache_config.cache_dtype != "auto":
            # Quantized KV cache.
            self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
                self.cache_config.cache_dtype
            ]

        self.vocab_size = self.model_config.get_vocab_size()
        self.max_model_len = self.model_config.max_model_len
        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
        self.max_num_reqs = self.scheduler_config.max_num_seqs
132
        self.is_encoder_decoder = self.model_config.is_encoder_decoder
133

Woosuk Kwon's avatar
Woosuk Kwon committed
134
135
136
        self.use_async_scheduling = self.scheduler_config.async_scheduling
        self.output_copy_stream = torch.cuda.Stream(self.device)

137
        # Pipeline parallelism.
138
139
140
141
        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
        self.is_first_pp_rank = get_pp_group().is_first_rank
        self.is_last_pp_rank = get_pp_group().is_last_rank

142
143
        # Persistent buffer for intermediate tensors (non-first PP ranks).
        self.intermediate_tensors: IntermediateTensors | None = None
144

145
146
147
148
        # Data parallelism.
        self.dp_size = self.parallel_config.data_parallel_size
        self.dp_rank = self.parallel_config.data_parallel_rank

149
150
151
152
153
154
        # Decode context parallelism.
        self.dcp_size = self.parallel_config.decode_context_parallel_size
        self.use_dcp = self.dcp_size > 1
        self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0
        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size

155
156
157
158
159
160
161
162
163
        # Multimodal
        self.mm_registry = MULTIMODAL_REGISTRY
        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
            self.model_config
        )
        self.encoder_cache = None
        if self.supports_mm_inputs and self.is_first_pp_rank:
            self.encoder_cache = EncoderCache()

164
        # Speculative decoding.
165
        self.speculator = None
166
        self.num_speculative_steps = 0
167
        self.use_aux_hidden_state_outputs = False
168
169
        if self.speculative_config is not None:
            self.num_speculative_steps = self.speculative_config.num_speculative_tokens
170

171
172
173
174
175
176
            if self.is_last_pp_rank:
                self.speculator = init_speculator(self.vllm_config, self.device)

            if self.speculative_config.method == "eagle3":
                # EAGLE3 may require auxiliary hidden states from target model outputs.
                self.use_aux_hidden_state_outputs = True
177
                if self.use_pp:
178
179
180
181
                    raise ValueError("EAGLE3 with pipeline parallel is not supported.")

        # Draft tokens propagation - for spec-dec + struct outputs.
        self.draft_tokens_handler = DraftTokensHandler(self.device)
182
        self.uniform_decode_query_len = 1 + self.num_speculative_steps
183

184
185
186
187
        # Pooling models.
        self.is_pooling_model = self.model_config.runner_type == "pooling"
        self.pooling_runner: PoolingRunner | None = None

188
        # General request states.
Woosuk Kwon's avatar
Woosuk Kwon committed
189
190
191
192
        self.req_states = RequestState(
            max_num_reqs=self.max_num_reqs,
            max_model_len=self.max_model_len,
            max_num_batched_tokens=self.max_num_tokens,
193
            num_speculative_steps=self.num_speculative_steps,
Woosuk Kwon's avatar
Woosuk Kwon committed
194
195
196
197
198
199
200
201
            vocab_size=self.vocab_size,
            device=self.device,
        )
        self.input_buffers = InputBuffers(
            max_num_reqs=self.max_num_reqs,
            max_num_tokens=self.max_num_tokens,
            device=self.device,
        )
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218

        self.sampler: Sampler | None = None
        self.rejection_sampler: RejectionSampler | None = None
        self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
        self.structured_outputs_worker: StructuredOutputsWorker | None = None
        if self.is_last_pp_rank and not self.is_pooling_model:
            # Initialize sampling-related workers.
            # These components are only set up on the last PP rank and
            # for generative (non-pooling) models.
            self.sampler = Sampler(
                max_num_reqs=self.max_num_reqs,
                vocab_size=self.vocab_size,
                device=self.device,
                req_states=self.req_states,
                logprobs_mode=self.model_config.logprobs_mode,
                num_speculative_tokens=self.num_speculative_steps + 1,
            )
219
220
221
222
223
            if self.speculative_config is not None:
                self.rejection_sampler = RejectionSampler(
                    self.sampler,
                    self.speculative_config,
                )
224
225
226
227
228
229
            self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
            self.structured_outputs_worker = StructuredOutputsWorker(
                max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
                vocab_size=self.vocab_size,
                device=self.device,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
230

231
        # For CUDA graphs, and will init cudagraph_manager after init_attn_backend.
232
        self.decode_query_len = self.num_speculative_steps + 1
233
        self.cudagraph_manager: ModelCudaGraphManager | None = None
234
235
        # LoRA-related workers.
        self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
236
        # KV Connector if configured.
237
238
        self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR

239
        # For transferring state from execute_model to subsequent sample_tokens call.
240
        self.execute_model_state: ExecuteModelState | None = None
241

242
243
244
        # Expert parallelism load balancer.
        self.eplb = EPLBController(self.parallel_config, self.device)

245
246
247
248
    def update_max_model_len(self, max_model_len: int) -> None:
        self.max_model_len = max_model_len
        self.req_states.max_model_len = max_model_len

249
250
251
    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
        tasks: list[SupportedTask] = []
        if self.model_config.runner_type == "generate":
252
            tasks.extend(self.model_state.get_supported_generation_tasks())
253
254
255
256
257
        if self.is_pooling_model:
            # Do not rely on pooling_runner here, since this information is needed
            # on the first PP rank, while pooling_runner is only initialized
            # on the last PP rank.
            tasks.extend(PoolingRunner.get_supported_tasks(self.model))
258
        return tuple(tasks)
Woosuk Kwon's avatar
Woosuk Kwon committed
259

260
    def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
261
        time_before_load = time.perf_counter()
262
263
264
265
        if load_dummy_weights:
            self.load_config.load_format = "dummy"
        self.eplb.prepare_load()
        eplb_models_added = False
Woosuk Kwon's avatar
Woosuk Kwon committed
266
267
268
269
270
        with DeviceMemoryProfiler() as m:
            model_loader = get_model_loader(self.vllm_config.load_config)
            logger.info("Loading model from scratch...")

            self.model = model_loader.load_model(
271
                vllm_config=self.vllm_config, model_config=self.vllm_config.model_config
Woosuk Kwon's avatar
Woosuk Kwon committed
272
273
274
            )
            if self.lora_config:
                self.model = self.load_lora_model(
275
                    self.model, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
276
                )
277
278
279
280
281

            if self.use_aux_hidden_state_outputs:
                assert self.speculative_config is not None
                set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
            if self.speculator is not None:
282
                self.speculator.load_model(self.model)
283
284
285
                eplb_models_added = self.eplb.maybe_register_speculator(
                    self.speculator, self.speculative_config, load_dummy_weights
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
286
287
288
289
        time_after_load = time.perf_counter()

        self.model_memory_usage = m.consumed_memory
        logger.info(
290
291
            "Model loading took %s GiB and %.6f seconds",
            format_gib(m.consumed_memory),
Woosuk Kwon's avatar
Woosuk Kwon committed
292
293
294
            time_after_load - time_before_load,
        )

295
296
297
298
        if not load_dummy_weights:
            prepare_communication_buffer_for_model(self.model)
            if self.speculator is not None:
                prepare_communication_buffer_for_model(self.speculator.model)
299

300
        # Initialize the components that require the model.
301
        self.model_state = init_model_state(
302
303
            self.vllm_config, self.model, self.encoder_cache, self.device
        )
304
        if self.is_pooling_model and self.is_last_pp_rank:
305
            self.pooling_runner = PoolingRunner(self.model)
306
307
308
309
310
311
        eplb_models_added |= self.eplb.maybe_register_model(
            self.model,
            self.model_config,
            load_dummy_weights,
        )
        self.eplb.maybe_start_async_loop(eplb_models_added)
312

313
314
315
316
317
318
319
320
321
322
323
        if not self.is_first_pp_rank:
            # For non-first PP ranks, create intermediate tensors sized
            # for the max capture size so they can be sliced per batch.
            # Save as persistent member so runtime can copy received data
            # into the same addresses that the CUDA graphs captured.
            self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
                batch_size=self.max_num_tokens,
                dtype=self.model_config.dtype,
                device=self.device,
            )

Woosuk Kwon's avatar
Woosuk Kwon committed
324
325
326
    def get_model(self) -> nn.Module:
        return self.model

327
328
329
330
331
    @functools.cached_property
    def main_stream(self) -> torch.cuda.Stream:
        # Cache the default CUDA stream to avoid lookup overhead.
        return torch.cuda.current_stream(self.device)

Woosuk Kwon's avatar
Woosuk Kwon committed
332
333
334
335
336
337
338
339
340
341
342
    def get_kv_cache_spec(self):
        return get_kv_cache_spec(self.vllm_config)

    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
        kv_cache_config = deepcopy(kv_cache_config)
        self.kv_cache_config = kv_cache_config
        block_sizes = [
            kv_cache_group.kv_cache_spec.block_size
            for kv_cache_group in kv_cache_config.kv_cache_groups
        ]

343
344
345
346
347
348
349
350
351
        block_table_max_model_len = self.max_model_len
        if self.is_encoder_decoder:
            # Cross-attention block tables need to index encoder tokens
            # (e.g., Whisper ~1500), which can exceed decoder max_model_len.
            block_table_max_model_len = max(
                block_table_max_model_len,
                getattr(self.model_config.hf_config, "max_source_positions", 0),
            )

Woosuk Kwon's avatar
Woosuk Kwon committed
352
353
354
355
        self.block_tables = BlockTables(
            block_sizes=block_sizes,
            max_num_reqs=self.max_num_reqs,
            max_num_batched_tokens=self.max_num_tokens,
356
            max_model_len=block_table_max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
357
            device=self.device,
358
359
360
            cp_size=self.dcp_size,
            cp_rank=self.dcp_rank,
            cp_interleave=self.cp_interleave,
Woosuk Kwon's avatar
Woosuk Kwon committed
361
362
        )

363
        self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend(
364
            self.kv_cache_config, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
365
        )
roikoren755's avatar
roikoren755 committed
366
367
368
        initialize_mamba_ssu_backend(
            self.vllm_config.mamba_config, self.kv_cache_config
        )
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
        cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
            attn_cg_support.min_cg_support,
            attn_cg_support.min_cg_attn_backend,
            self.uniform_decode_query_len,
            self.parallel_config.tensor_parallel_size,
            self.kv_cache_config,
            self.max_num_reqs,
        )
        self.cudagraph_manager = ModelCudaGraphManager(
            self.vllm_config,
            self.device,
            cudagraph_mode,
            decode_query_len=self.decode_query_len,
        )
        if self.speculator is not None:
            self.speculator.init_cudagraph_manager(cudagraph_mode)

386
        check_attention_cp_compatibility(self.vllm_config)
387
        if self.speculator is not None:
388
389
            # HACK(woosuk)
            self.speculator.set_attn(
390
                self.model_state,
391
392
393
                self.kv_cache_config,
                self.block_tables,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
394
395

        self.kv_caches: list[torch.Tensor] = []
396
        kv_caches_dict = init_kv_cache(
Woosuk Kwon's avatar
Woosuk Kwon committed
397
398
399
400
401
            self.kv_caches,
            self.compilation_config.static_forward_context,
            self.kv_cache_config,
            self.attn_backends,
            self.device,
402
            self.cache_config.cache_dtype,
Woosuk Kwon's avatar
Woosuk Kwon committed
403
        )
404
405
        self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

Woosuk Kwon's avatar
Woosuk Kwon committed
406
    @torch.inference_mode()
407
    @step_eplb_after(is_dummy=True)
Woosuk Kwon's avatar
Woosuk Kwon committed
408
    def _dummy_run(
409
410
411
        self,
        num_tokens: int,
        *args,
412
        skip_attn: bool = False,
413
        uniform_decode: bool = False,
414
415
        skip_eplb: bool = False,
        is_profile: bool = False,
416
        **kwargs,
417
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
418
419
420
421
422
        if skip_attn and not is_profile:
            raise ValueError(
                "skip_attn must only be True for initial memory profiling."
            )

423
        # Create a dummy scheduler output.
424
        num_reqs = min(num_tokens, self.max_num_reqs)
425
        if uniform_decode:
426
427
428
429
430
431
432
433
434
435
            # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
            # and for spec-decode with MTP we want to make sure the dummy runs use
            # 1+num_speculative_tokens we use max here, this will likely be eventually
            # changed in the worker: https://github.com/vllm-project/vllm/pull/35243
            num_tokens = max(num_tokens, self.decode_query_len)
            num_reqs = num_tokens // self.decode_query_len
            assert num_tokens % self.decode_query_len == 0
        num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
        num_tokens_per_request[-1] += num_tokens % num_reqs

436
437
        assert sum(num_tokens_per_request) == num_tokens
        num_scheduled_tokens = {
438
            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
439
440
441
442
443
        }
        dummy_scheduler_output = SchedulerOutput.make_empty()
        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens

444
445
446
        # Disable any use of KVConnector for dummy runs.
        self.kv_connector.set_disabled(True)

447
        # Get the intermediate tensors for the dummy run.
448
        intermediate_tensors = None
449
        if not self.is_first_pp_rank:
450
451
            assert self.intermediate_tensors is not None
            intermediate_tensors = self.intermediate_tensors[:num_tokens]
452

453
454
        # Execute the model.
        self.execute_model(
455
456
457
458
            dummy_scheduler_output,
            intermediate_tensors=intermediate_tensors,
            dummy_run=True,
            skip_attn_for_dummy_run=skip_attn,
459
            is_profile=is_profile,
460
        )
461
        self.kv_connector.set_disabled(False)
462
463

        # Non-last PP ranks don't produce output for sampling.
464
        if not self.is_last_pp_rank:
465
466
            return None, None

467
        assert self.execute_model_state is not None
468
469
470
471
472
        input_batch = self.execute_model_state.input_batch
        attn_metadata = self.execute_model_state.attn_metadata
        slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
        hidden_states = self.execute_model_state.hidden_states
        aux_hidden_states = self.execute_model_state.aux_hidden_states
473
        self.execute_model_state = None
474
475
476

        # dummy run the eagle speculator's propose to ensure DP/EP sync.
        if self.speculator is not None:
477
            assert self.sampler is not None
478
479
480
481
482
483
484
485
486
487
            mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
            if self.speculator.supports_mm_inputs:
                mm_inputs = (
                    [],
                    torch.zeros(
                        input_batch.num_tokens,
                        dtype=torch.bool,
                        device=self.device,
                    ),
                )
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
            self.speculator.propose(
                input_batch=input_batch,
                attn_metadata=attn_metadata,
                slot_mappings=slot_mappings_by_layer,
                last_hidden_states=hidden_states,
                aux_hidden_states=aux_hidden_states,
                num_sampled=torch.ones(
                    input_batch.num_reqs, dtype=torch.int32, device=self.device
                ),
                num_rejected=torch.zeros(
                    input_batch.num_reqs, dtype=torch.int32, device=self.device
                ),
                last_sampled=self.req_states.last_sampled_tokens,
                next_prefill_tokens=self.req_states.next_prefill_tokens,
                temperature=self.sampler.sampling_states.temperature.gpu,
                seeds=self.sampler.sampling_states.seeds.gpu,
                dummy_run=True,
                skip_attn_for_dummy_run=skip_attn,
506
                mm_inputs=mm_inputs,
507
                is_profile=is_profile,
508
509
            )

510
        assert hidden_states is not None  # Last PP rank always has hidden_states
511
        sample_hidden_states = hidden_states[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
512
513
514
        return hidden_states, sample_hidden_states

    @torch.inference_mode()
515
    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
516
517
        num_reqs = hidden_states.shape[0]
        logits = self.model.compute_logits(hidden_states)
518
519
        dummy_input_batch = InputBatch.make_dummy(
            num_reqs, num_reqs, self.input_buffers
520
        )
521

522
523
524
        # NOTE(woosuk): During the initial memory profiling, the sampler may skip
        # top_k, top_p, and logprobs, using less GPU memory than what is possible
        # during actual execution.
525
526
        assert self.sampler is not None
        self.sampler(logits, dummy_input_batch)
Woosuk Kwon's avatar
Woosuk Kwon committed
527

528
529
530
531
532
    @torch.inference_mode()
    def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
        assert self.pooling_runner is not None
        self.pooling_runner.dummy_pooler_run(hidden_states)

Woosuk Kwon's avatar
Woosuk Kwon committed
533
534
535
    @torch.inference_mode()
    def profile_run(self) -> None:
        hidden_states, sample_hidden_states = self._dummy_run(
536
            self.max_num_tokens, skip_attn=True, is_profile=True
Woosuk Kwon's avatar
Woosuk Kwon committed
537
        )
538

539
        # Only run sampler/pooler on last PP rank (non-last ranks return None).
540
        if self.is_last_pp_rank:
541
            assert sample_hidden_states is not None
542
543
544
545
            if self.pooling_runner is None:
                self._dummy_sampler_run(sample_hidden_states)
            else:
                self._dummy_pooler_run(hidden_states)
546

547
        torch.accelerator.synchronize()
Woosuk Kwon's avatar
Woosuk Kwon committed
548
549
550
551
        del hidden_states, sample_hidden_states
        gc.collect()

    def reset_mm_cache(self) -> None:
552
553
        if self.encoder_cache is not None:
            self.encoder_cache.reset_mm_cache()
554
555

    def reset_encoder_cache(self) -> None:
556
557
        if self.encoder_cache is not None:
            self.encoder_cache.reset_encoder_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
558
559
560
561
562

    def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:
        # SP is not supported yet.
        return num_scheduled_tokens

563
564
565
566
    def profile_cudagraph_memory(self) -> int:
        # NOTE(woosuk): It is TBD whether we keep this API or not.
        return 0

Woosuk Kwon's avatar
Woosuk Kwon committed
567
568
    @torch.inference_mode()
    def capture_model(self) -> int:
569
        assert self.cudagraph_manager is not None
Woosuk Kwon's avatar
Woosuk Kwon committed
570
571
572
573
574
575
576
577
        if not self.cudagraph_manager.needs_capture():
            logger.warning(
                "Skipping CUDA graph capture. To turn on CUDA graph capture, "
                "ensure `cudagraph_mode` was not manually set to `NONE`"
            )
            return 0

        start_time = time.perf_counter()
578
        gc.collect()
579
        torch.accelerator.empty_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
580
581
582
583
        start_free_gpu_memory = torch.cuda.mem_get_info()[0]

        with self.maybe_setup_dummy_loras(self.lora_config):
            self.cudagraph_manager.capture(
584
585
586
                self.model,
                self.model_state,
                self.input_buffers,
587
                self.intermediate_tensors,
588
589
590
                self.block_tables,
                self.attn_groups,
                self.kv_cache_config,
591
                has_lora=self.lora_config is not None,
592
                use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
Woosuk Kwon's avatar
Woosuk Kwon committed
593
            )
594
            if self.speculator is not None:
595
                self.speculator.capture_model()
Woosuk Kwon's avatar
Woosuk Kwon committed
596
597
598
599
600
601
602
603
604
605
606
607
608

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

609
610
611
612
613
614
615
616
617
618
    def _remove_request(self, req_id: str) -> bool:
        if not self.req_states.remove_request(req_id):
            return False
        if self.encoder_cache is not None:
            self.encoder_cache.remove_request(req_id)
        if self.prompt_logprobs_worker is not None:
            self.prompt_logprobs_worker.remove_request(req_id)
        self.lora_state.remove_request(req_id)
        return True

619
    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
620
        finished_req_ids = scheduler_output.finished_req_ids
621
622
623
        preempted_req_ids = scheduler_output.preempted_req_ids
        if preempted_req_ids:
            finished_req_ids = finished_req_ids.union(preempted_req_ids)
624
        for req_id in finished_req_ids:
625
            self._remove_request(req_id)
626

627
    def free_states(self, scheduler_output: SchedulerOutput) -> None:
628
        if self.encoder_cache is not None:
629
            for mm_hash in scheduler_output.free_encoder_mm_hashes:
630
                self.encoder_cache.free_encoder_cache(mm_hash)
Woosuk Kwon's avatar
Woosuk Kwon committed
631

632
    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
633
        for new_req_data in scheduler_output.scheduled_new_reqs:
634
635
            assert new_req_data.prompt_token_ids is not None
            assert new_req_data.prefill_token_ids is not None
Woosuk Kwon's avatar
Woosuk Kwon committed
636
            req_id = new_req_data.req_id
637
638
639
640
641
642

            # Streaming input update: request already exists from a prior
            # chunk. Remove old state so it can be cleanly re-added below
            # with the updated prompt_token_ids and mm_features.
            self._remove_request(req_id)

643
            prompt_len = len(new_req_data.prompt_token_ids)
Woosuk Kwon's avatar
Woosuk Kwon committed
644
645
            self.req_states.add_request(
                req_id=req_id,
646
                prompt_len=prompt_len,
647
                all_token_ids=new_req_data.prefill_token_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
648
649
650
                num_computed_tokens=new_req_data.num_computed_tokens,
            )
            req_index = self.req_states.req_id_to_index[req_id]
651

652
653
            if self.encoder_cache is not None:
                self.encoder_cache.add_request(req_id, new_req_data.mm_features)
654

655
            self.model_state.add_request(req_index, new_req_data)
656
657
658
            self.block_tables.append_block_ids(
                req_index, new_req_data.block_ids, overwrite=True
            )
659
            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
Woosuk Kwon's avatar
Woosuk Kwon committed
660

661
662
            if self.is_last_pp_rank and new_req_data.sampling_params is not None:
                assert self.sampler is not None
663
664
665
                self.sampler.add_request(
                    req_index, prompt_len, new_req_data.sampling_params
                )
666
                assert self.prompt_logprobs_worker is not None
667
668
669
670
                self.prompt_logprobs_worker.add_request(
                    req_id, req_index, new_req_data.sampling_params
                )

671
672
        if scheduler_output.scheduled_new_reqs:
            self.req_states.apply_staged_writes()
673
            self.model_state.apply_staged_writes()
674
675
        if self.sampler is not None:
            self.sampler.apply_staged_writes()
676
677

    def update_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
678
        # Add new blocks for the existing requests.
679
680
        reqs = scheduler_output.scheduled_cached_reqs
        for req_new_block_ids, req_id in zip(reqs.new_block_ids, reqs.req_ids):
Woosuk Kwon's avatar
Woosuk Kwon committed
681
            if req_new_block_ids is not None:
682
                req_index = self.req_states.req_id_to_index[req_id]
683
684
685
                self.block_tables.append_block_ids(
                    req_index, req_new_block_ids, overwrite=False
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
686
687

    def prepare_inputs(
688
        self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor
Woosuk Kwon's avatar
Woosuk Kwon committed
689
690
    ) -> InputBatch:
        num_tokens = scheduler_output.total_num_scheduled_tokens
691
        num_tokens_after_padding = batch_desc.num_tokens
Woosuk Kwon's avatar
Woosuk Kwon committed
692
        assert num_tokens > 0
693
694
        num_tokens_per_req = scheduler_output.num_scheduled_tokens
        num_reqs = len(num_tokens_per_req)
Woosuk Kwon's avatar
Woosuk Kwon committed
695
696
697

        # Decode first, then prefill.
        # batch_idx -> req_id
698
699
700
        req_ids = sorted(num_tokens_per_req, key=num_tokens_per_req.get)  # type: ignore[arg-type]
        numtoks_iter = map(num_tokens_per_req.get, req_ids)
        num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
Woosuk Kwon's avatar
Woosuk Kwon committed
701

702
703
        idx_mapping_iter = map(self.req_states.req_id_to_index.get, req_ids)
        idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.int32, count=num_reqs)
704
        idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
Woosuk Kwon's avatar
Woosuk Kwon committed
705

706
        # Get the number of draft tokens for each request.
707
708
        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
        if not draft_tokens:
709
710
711
            # No draft token scheduled (common case).
            total_num_draft_tokens = 0
            total_num_logits = num_reqs
712
            cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
713
714
715
            cu_num_logits = torch.arange(
                num_reqs + 1, device=self.device, dtype=torch.int32
            )
716
            expanded_idx_mapping = idx_mapping
717
718
719
            expanded_local_pos = torch.zeros(
                num_reqs, dtype=torch.int32, device=self.device
            )
720
        else:
721
722
            num_draft_tokens = np.fromiter(
                (len(draft_tokens.get(req_id, ())) for req_id in req_ids),
723
                dtype=np.int32,
724
                count=num_reqs,
725
726
727
728
            )
            total_num_draft_tokens = int(num_draft_tokens.sum())
            total_num_logits = num_reqs + total_num_draft_tokens

729
730
731
732
            num_logits = num_draft_tokens + 1
            cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
            cu_num_logits_np[0] = 0
            np.cumsum(num_logits, out=cu_num_logits_np[1:])
733
            cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
734

735
            max_expand_len = self.num_speculative_steps + 1
736
            expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
737
                idx_mapping, total_num_logits, cu_num_logits, max_expand_len
738
739
            )

740
        # Get query_start_loc.
741
742
        # num_reqs_padded is None for PIECEWISE graphs (no request padding needed)
        num_reqs_padded = batch_desc.num_reqs or num_reqs
743
744
745
        query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
        query_start_loc_np[0] = 0
        np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1 : num_reqs + 1])
746
747
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
748
        query_start_loc_np[num_reqs + 1 :] = num_tokens
749
        async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
750
751
        query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
        query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1]
752

753
754
755
756
757
758
759
760
761
762
763
        # Get prefill tokens if any.
        if self.req_states.any_prefills(idx_mapping_np):
            prepare_prefill_inputs(
                self.input_buffers.input_ids,
                self.req_states.next_prefill_tokens,
                idx_mapping,
                query_start_loc,
                self.req_states.all_token_ids.gpu,
                self.req_states.prefill_len.gpu,
                self.req_states.num_computed_tokens.gpu,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
764

765
766
767
        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
768
769
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
770
771
772
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
773
        seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
774

775
        dcp_local_seq_lens = None
776
777
        if self.use_dcp:
            # Prepare dcp local seq_lens.
778
779
            prepare_dcp_local_seq_lens(
                self.input_buffers.dcp_local_seq_lens,
780
                self.input_buffers.seq_lens,
781
                num_reqs,
782
783
784
                self.dcp_size,
                self.dcp_rank,
                self.cp_interleave,
785
            )
786
            dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs_padded]
787

788
        # Some input token ids are directly read from the last sampled tokens
789
790
        # and draft tokens. Also, get the logits indices to sample tokens from.
        logits_indices = combine_sampled_and_draft_tokens(
791
            self.input_buffers.input_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
792
793
            idx_mapping,
            self.req_states.last_sampled_tokens,
794
            query_start_loc,
795
796
            seq_lens,
            self.req_states.prefill_len.gpu,
797
798
799
            self.req_states.draft_tokens,
            cu_num_logits,
            total_num_logits,
Woosuk Kwon's avatar
Woosuk Kwon committed
800
801
        )

802
803
804
805
806
807
808
809
810
        # CPU upper bound on seq_lens; padded entries left at zero.
        seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
        np.add(
            self.req_states.num_computed_tokens_np[idx_mapping_np],
            num_scheduled_tokens,
            out=seq_lens_cpu_upper_bound_np[:num_reqs],
        )
        seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)

Woosuk Kwon's avatar
Woosuk Kwon committed
811
812
813
        return InputBatch(
            req_ids=req_ids,
            num_reqs=num_reqs,
814
            num_reqs_after_padding=num_reqs_padded,
Woosuk Kwon's avatar
Woosuk Kwon committed
815
816
            idx_mapping=idx_mapping,
            idx_mapping_np=idx_mapping_np,
817
            expanded_idx_mapping=expanded_idx_mapping,
818
            expanded_local_pos=expanded_local_pos,
Woosuk Kwon's avatar
Woosuk Kwon committed
819
820
821
            num_scheduled_tokens=num_scheduled_tokens,
            num_tokens=num_tokens,
            num_tokens_after_padding=num_tokens_after_padding,
822
            num_draft_tokens=total_num_draft_tokens,
823
            query_start_loc=query_start_loc,
Woosuk Kwon's avatar
Woosuk Kwon committed
824
            query_start_loc_np=query_start_loc_np,
825
            seq_lens=seq_lens,
826
            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
827
828
829
            dcp_local_seq_lens=dcp_local_seq_lens,
            input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
            positions=self.input_buffers.positions[:num_tokens_after_padding],
Woosuk Kwon's avatar
Woosuk Kwon committed
830
            logits_indices=logits_indices,
831
            cu_num_logits=cu_num_logits,
832
            cu_num_logits_np=cu_num_logits_np,
833
            has_structured_output_reqs=scheduler_output.has_structured_output_requests,
Woosuk Kwon's avatar
Woosuk Kwon committed
834
835
        )

836
837
838
    def prepare_attn(
        self, input_batch: InputBatch
    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
839
840
841
842
843
844
845
        # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
        block_tables = self.block_tables.gather_block_tables(
            input_batch.idx_mapping,
            num_reqs_padded=input_batch.num_reqs_after_padding,
        )
        # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
        # Kernel pads beyond num_tokens with PAD_SLOT_ID.
846
847
848
849
        slot_mappings = self.block_tables.compute_slot_mappings(
            input_batch.idx_mapping,
            input_batch.query_start_loc,
            input_batch.positions,
850
            num_tokens_padded=input_batch.num_tokens_after_padding,
851
852
853
854
855
856
857
858
859
860
861
862
        )
        return block_tables, slot_mappings

    def prepare_dummy_attn(
        self, input_batch: InputBatch
    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
        block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
        slot_mappings = self.block_tables.get_dummy_slot_mappings(
            input_batch.num_tokens
        )
        return block_tables, slot_mappings

Woosuk Kwon's avatar
Woosuk Kwon committed
863
864
865
866
867
    def sample(
        self,
        hidden_states: torch.Tensor,
        input_batch: InputBatch,
        grammar_output: GrammarOutput | None,
868
    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
869
870
871
872
        sample_hidden_states = hidden_states[input_batch.logits_indices]
        logits = self.model.compute_logits(sample_hidden_states)
        if grammar_output is not None:
            # Apply grammar bitmask to the logits in-place.
873
            assert self.structured_outputs_worker is not None
874
875
876
877
878
879
            self.structured_outputs_worker.apply_grammar_bitmask(
                logits,
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.grammar_bitmask,
            )
880

881
882
        if input_batch.num_draft_tokens == 0:
            # No draft tokens (common case).
883
884
            assert self.sampler is not None
            sampler_output = self.sampler(logits, input_batch)
885
        else:
886
            # Rejection sampling for spec decoding.
887
            assert self.rejection_sampler is not None
888
            assert self.speculator is not None
889
890
891
892
            sampler_output = self.rejection_sampler(
                logits,
                input_batch,
                # Draft logits are needed for probabilistic rejection sampling.
893
                self.speculator.draft_logits,
894
            )
895
896
897
898

        # Get the number of sampled and rejected tokens.
        # For chunked prefills, num_sampled and num_rejected are both 0.
        num_sampled, num_rejected = get_num_sampled_and_rejected(
899
            sampler_output.num_sampled,
900
901
902
903
904
            input_batch.seq_lens,
            input_batch.cu_num_logits,
            input_batch.idx_mapping,
            self.req_states.prefill_len.gpu,
        )
905
        return sampler_output, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
906
907
908
909

    def postprocess(
        self,
        input_batch: InputBatch,
910
911
        sampled_tokens: torch.Tensor,
        num_sampled: torch.Tensor,
912
        num_rejected: torch.Tensor,
913
914
    ) -> None:
        # Update the number of computed tokens.
915
916
917
918
919
        if self.is_last_pp_rank:
            assert self.sampler is not None
            output_bin_counts = self.sampler.penalties_state.output_bin_counts
        else:
            output_bin_counts = None
920
        post_update(
921
            input_batch.idx_mapping,
922
            self.req_states.num_computed_tokens.gpu,
923
            self.req_states.last_sampled_tokens,
924
            output_bin_counts,
925
926
            sampled_tokens,
            num_sampled,
927
            num_rejected,
928
            input_batch.query_start_loc,
929
930
            self.req_states.all_token_ids.gpu,
            self.req_states.total_len.gpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
931
        )
932
933

        # Update the number of computed prefill tokens.
Woosuk Kwon's avatar
Woosuk Kwon committed
934
        idx_mapping_np = input_batch.idx_mapping_np
935
        computed_prefill = self.req_states.num_computed_prefill_tokens
936
937
938
        computed_prefill[idx_mapping_np] += input_batch.num_scheduled_tokens
        np.minimum(
            computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
Woosuk Kwon's avatar
Woosuk Kwon committed
939
        )
940
941
942
943
        # Advance the CPU mirror optimistically (assume all scheduled accepted).
        self.req_states.num_computed_tokens_np[idx_mapping_np] += (
            input_batch.num_scheduled_tokens
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
944
945
946
947
948

    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: SchedulerOutput,
949
        intermediate_tensors: IntermediateTensors | None = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
950
        dummy_run: bool = False,
951
        skip_attn_for_dummy_run: bool = False,
952
        is_profile: bool = False,
953
    ) -> ModelRunnerOutput | IntermediateTensors | None:
954
955
956
957
958
959
960
961
962
        if not dummy_run:
            # Update the request states.
            self.finish_requests(scheduler_output)
            self.free_states(scheduler_output)
            self.add_requests(scheduler_output)
            self.update_requests(scheduler_output)
            self.block_tables.apply_staged_writes()
            if scheduler_output.total_num_scheduled_tokens == 0:
                # No need to run the model.
963
964
                empty_output = self.kv_connector.no_forward(scheduler_output)
                return empty_output
Woosuk Kwon's avatar
Woosuk Kwon committed
965

966
967
968
969
970
971
        # Get batch descriptor and sync across DP ranks.
        num_reqs = len(scheduler_output.num_scheduled_tokens)
        num_toks = scheduler_output.total_num_scheduled_tokens
        max_query_len = max(scheduler_output.num_scheduled_tokens.values())
        uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len)

972
973
974
975
976
977
978
        skip_compiled = False
        if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
            # Encoder-decoder models such as Whisper should run eager/non-compiled
            # when encoder inputs are scheduled, because this step updates
            # cross-attention cache with dynamic encoder outputs.
            skip_compiled = True

979
980
981
982
983
984
985
986
987
        batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
            self.cudagraph_manager,
            num_reqs,
            num_toks,
            uniform_tok_count,
            self.dp_size,
            self.dp_rank,
            need_eager=is_profile or skip_compiled,
        )
988
989

        if batch_desc.num_tokens == 0:
990
            # All DP ranks have zero tokens to run.
991
992
            empty_output = self.kv_connector.no_forward(scheduler_output)
            return empty_output
993
994
995
996

        if not dummy_run:
            # Common case.
            # Prepare all the inputs and copy to the input buffers.
997
            input_batch = self.prepare_inputs(scheduler_output, batch_desc)
998
999
            block_tables, slot_mappings = self.prepare_attn(input_batch)

1000
1001
            if self.lora_config:
                # Activate LoRA adapters.
1002
                lora_inputs = self.lora_state.make_lora_inputs(
1003
1004
1005
                    input_batch.req_ids,
                    input_batch.idx_mapping_np,
                    input_batch.num_scheduled_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
1006
                )
1007
1008
                self._set_active_loras(*lora_inputs)
        else:
1009
            # No actual tokens to run. A dummy run for DP or memory profiling.
1010
            input_batch = InputBatch.make_dummy(
1011
1012
1013
                batch_desc.num_reqs or num_reqs,
                batch_desc.num_tokens,
                self.input_buffers,
1014
            )
1015
            if not skip_attn_for_dummy_run:
1016
1017
                block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)
            else:
1018
1019
1020
1021
                assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
                    "Attention metadata must be prepared for dummy runs when using "
                    "FULL cudagraph mode."
                )
1022
1023
                block_tables = None
                slot_mappings = None
1024
            # FIXME(woosuk): Fix warmup for LoRA.
Woosuk Kwon's avatar
Woosuk Kwon committed
1025

1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
        attn_metadata = None
        slot_mappings_by_layer = None
        if not (dummy_run and skip_attn_for_dummy_run):
            assert slot_mappings is not None
            slot_mappings_by_layer = build_slot_mappings_by_layer(
                slot_mappings, self.kv_cache_config
            )
            assert block_tables is not None
            attn_metadata = self.model_state.prepare_attn(
                input_batch,
1036
                batch_desc.cg_mode,
1037
1038
1039
1040
1041
1042
                block_tables,
                slot_mappings,
                self.attn_groups,
                self.kv_cache_config,
            )

1043
        inputs_embeds = None
1044
        if self.supports_mm_inputs and self.is_first_pp_rank:
1045
1046
            # Run MM encoder (if needed) and get multimodal embeddings.
            # Only first PP rank prepares multimodal embeddings.
1047
1048
            # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs
            # to obtain inputs_embeds, because the compiled model expects this input.
1049
1050
1051
1052
1053
1054
            inputs_embeds = self.model_state.get_mm_embeddings(
                scheduler_output.scheduled_encoder_inputs,
                input_batch,
                self.req_states,
            )

1055
1056
1057
        model_inputs = {
            "input_ids": input_batch.input_ids,
            "positions": input_batch.positions,
1058
            "inputs_embeds": inputs_embeds,
1059
1060
1061
1062
1063
1064
1065
1066
            # NOTE: Values returned by `prepare_inputs` will override the default
            # values above.
            **self.model_state.prepare_inputs(input_batch, self.req_states),
        }
        if not self.is_first_pp_rank:
            # Update for non-first PP ranks.
            model_inputs["input_ids"] = None
            model_inputs["inputs_embeds"] = None
1067
1068

            # Prepare the intermediate tensors.
1069
            assert intermediate_tensors is not None
1070
1071
            assert self.intermediate_tensors is not None
            n = input_batch.num_tokens_after_padding
1072
            model_inputs["intermediate_tensors"] = IntermediateTensors(
1073
1074
1075
                {
                    k: v[:n].copy_(intermediate_tensors.tensors[k][:n])
                    for k, v in self.intermediate_tensors.tensors.items()
1076
                }
1077
            )
1078
            del intermediate_tensors
1079

Woosuk Kwon's avatar
Woosuk Kwon committed
1080
        # Run model.
1081
        if batch_desc.cg_mode == CUDAGraphMode.FULL:
1082
            # Use explicit cudagraph replay for FULL mode.
Woosuk Kwon's avatar
Woosuk Kwon committed
1083
1084
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
1085
            assert self.cudagraph_manager is not None
1086
            self.kv_connector.pre_forward(scheduler_output)
1087
            model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
Woosuk Kwon's avatar
Woosuk Kwon committed
1088
        else:
1089
1090
1091
1092
1093
1094
            # For piecewise and eager mode, just call model().
            batch_descriptor = BatchDescriptor(
                num_tokens=input_batch.num_tokens_after_padding,
                has_lora=self.lora_config is not None,
            )

Woosuk Kwon's avatar
Woosuk Kwon committed
1095
            with set_forward_context(
1096
                attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
1097
1098
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
1099
                cudagraph_runtime_mode=batch_desc.cg_mode,
Woosuk Kwon's avatar
Woosuk Kwon committed
1100
                num_tokens_across_dp=num_tokens_across_dp,
1101
                batch_descriptor=batch_descriptor,
1102
                slot_mapping=slot_mappings_by_layer,
1103
                skip_compiled=skip_compiled,
Woosuk Kwon's avatar
Woosuk Kwon committed
1104
            ):
1105
                self.kv_connector.pre_forward(scheduler_output)
1106
                model_output = self.model(**model_inputs)
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121

        if self.is_last_pp_rank:
            if self.use_aux_hidden_state_outputs:
                assert isinstance(model_output, tuple)
                hidden_states, aux_hidden_states = model_output
            else:
                assert isinstance(model_output, torch.Tensor)
                hidden_states = model_output
                aux_hidden_states = None
            output_intermediate_tensors = None
        else:
            assert isinstance(model_output, IntermediateTensors)
            hidden_states = None
            aux_hidden_states = None
            output_intermediate_tensors = model_output
Woosuk Kwon's avatar
Woosuk Kwon committed
1122

1123
        kv_connector_output = self.kv_connector.post_forward(scheduler_output)
1124
1125
1126
1127
1128
1129
1130
        self.execute_model_state = ExecuteModelState(
            input_batch=input_batch,
            attn_metadata=attn_metadata,
            slot_mappings_by_layer=slot_mappings_by_layer,
            hidden_states=hidden_states,
            aux_hidden_states=aux_hidden_states,
            kv_connector_output=kv_connector_output,
1131
        )
1132

1133
        if not self.is_last_pp_rank:
1134
            # Non-last PP rank: return IntermediateTensors for sending.
1135
1136
1137
            assert output_intermediate_tensors is not None
            output_intermediate_tensors.kv_connector_output = kv_connector_output
            return output_intermediate_tensors
Woosuk Kwon's avatar
Woosuk Kwon committed
1138
1139
1140
        return None

    @torch.inference_mode()
1141
    @step_eplb_after()
Woosuk Kwon's avatar
Woosuk Kwon committed
1142
    def sample_tokens(
1143
        self, grammar_output: GrammarOutput | None
1144
    ) -> AsyncOutput | ModelRunnerOutput | None:
1145
1146
1147
        if self.execute_model_state is None:
            # The prior execute_model call must have failed.
            return None
1148
1149
1150
1151
1152
1153
1154

        input_batch = self.execute_model_state.input_batch
        attn_metadata = self.execute_model_state.attn_metadata
        slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
        hidden_states = self.execute_model_state.hidden_states
        aux_hidden_states = self.execute_model_state.aux_hidden_states
        kv_connector_output = self.execute_model_state.kv_connector_output
1155
        self.execute_model_state = None
Woosuk Kwon's avatar
Woosuk Kwon committed
1156

1157
        if not self.is_last_pp_rank:
1158
1159
1160
1161
            # Non-last PP rank: hidden_states is None because this rank produced
            # IntermediateTensors instead of final hidden states. Receive the
            # sampled tokens broadcast from the last rank and update local state.
            sampled, num_sampled, num_rejected = pp_receive(
1162
                input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
1163
            )
1164
            self.postprocess(input_batch, sampled, num_sampled, num_rejected)
1165
1166
1167
            return None

        # Last rank: sample tokens
1168
        sampler_output, num_sampled, num_rejected = self.sample(
1169
            hidden_states, input_batch, grammar_output
Woosuk Kwon's avatar
Woosuk Kwon committed
1170
        )
1171
1172

        if self.use_pp:
1173
            # Broadcast to non-last PP ranks (handles spec decode multi-token).
1174
            pp_broadcast(sampler_output.sampled_token_ids, num_sampled, num_rejected)
1175

1176
        assert self.prompt_logprobs_worker is not None
1177
1178
1179
1180
        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
            self.model.compute_logits,
            hidden_states,
            input_batch,
1181
            self.req_states.all_token_ids.gpu,
1182
            self.req_states.num_computed_tokens.gpu,
1183
            self.req_states.prompt_len.np,
1184
1185
1186
            self.req_states.prefill_len.np,
            self.req_states.num_computed_prefill_tokens,
        )
1187
1188
1189
1190
1191
1192
1193
1194

        # Prepare the model runner output.
        model_runner_output = ModelRunnerOutput(
            req_ids=input_batch.req_ids,
            # NOTE(woosuk): req_id_to_index is unused in this model runner.
            # Only for compatibility with the existing model runner and scheduler.
            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
            sampled_token_ids=None,  # type: ignore
1195
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
1196
            kv_connector_output=kv_connector_output,
1197
1198
1199
1200
        )
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,
1201
            num_sampled_tokens=num_sampled,
1202
            main_stream=self.main_stream,
1203
1204
1205
            copy_stream=self.output_copy_stream,
        )

1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
        mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
        if self.speculator is not None and self.speculator.supports_mm_inputs:
            # Get cached multimodal embeddings for draft forward.
            # NOTE: This is done here because postprocess updates
            # num_computed_prefill_tokens.
            prefill_lens = self.req_states.prefill_len.np[input_batch.idx_mapping_np]
            computed_prefill_lens = self.req_states.num_computed_prefill_tokens[
                input_batch.idx_mapping_np
            ]
            mm_inputs = self.model_state.encoder_runner.gather_mm_embeddings(
                input_batch.req_ids,
                input_batch.num_tokens,
                input_batch.num_scheduled_tokens,
                input_batch.query_start_loc_np,
                prefill_lens,
                computed_prefill_lens + 1,  # +1 to consider the skew in eagle
            )

1224
1225
1226
1227
1228
1229
        # Postprocess results and update request states.
        # NOTE: This is intentionally done after creating the AsyncOutput,
        # ensuring that `copy_event` is recorded before calling postprocess.
        # This sequencing may slightly reduce latency as async D2H copy does not
        # need to wait for the postprocess to finish.
        self.postprocess(
1230
            input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
1231
        )
1232

1233
        if self.speculator is not None:
1234
            assert self.sampler is not None
1235
            draft_tokens = self.speculator.propose(
1236
                input_batch,
1237
1238
                attn_metadata,
                slot_mappings_by_layer,
1239
                hidden_states,
1240
                aux_hidden_states,
1241
1242
                num_sampled,
                num_rejected,
1243
1244
1245
1246
                self.req_states.last_sampled_tokens,
                self.req_states.next_prefill_tokens,
                self.sampler.sampling_states.temperature.gpu,
                self.sampler.sampling_states.seeds.gpu,
1247
                mm_inputs=mm_inputs,
1248
            )
1249
            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1250
            self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
1251
1252
1253
1254

        if self.use_async_scheduling:
            return async_output
        return async_output.get_output()
1255
1256
1257

    def take_draft_token_ids(self) -> DraftTokenIds | None:
        return self.draft_tokens_handler.get_draft_tokens()
1258
1259

    @torch.inference_mode()
1260
    @step_eplb_after()
1261
1262
1263
1264
1265
    def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None:
        if self.execute_model_state is None:
            # The prior execute_model call must have failed.
            return None

1266
1267
1268
        input_batch = self.execute_model_state.input_batch
        hidden_states = self.execute_model_state.hidden_states
        kv_connector_output = self.execute_model_state.kv_connector_output
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
        self.execute_model_state = None

        if not self.is_last_pp_rank:
            self.postprocess_pool(input_batch)
            return None

        assert self.pooling_runner is not None
        pooler_output, is_valid = self.pooling_runner.pool(
            hidden_states, input_batch, self.req_states
        )

        # Build the model runner output.
        model_runner_output = ModelRunnerOutput(
            req_ids=input_batch.req_ids,
            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
            kv_connector_output=kv_connector_output,
        )
        async_output = AsyncPoolingOutput(
            model_runner_output=model_runner_output,
            pooler_output=pooler_output,
            is_valid=is_valid,
            main_stream=self.main_stream,
            copy_stream=self.output_copy_stream,
        )
1293
1294

        self.postprocess_pool(input_batch)
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
        if self.use_async_scheduling:
            return async_output
        return async_output.get_output()

    def postprocess_pool(self, input_batch: InputBatch) -> None:
        # Update the number of computed tokens.
        post_update_pool(
            input_batch.idx_mapping,
            self.req_states.num_computed_tokens.gpu,
            input_batch.query_start_loc,
        )

        # Update the number of computed prefill tokens.
        idx_mapping_np = input_batch.idx_mapping_np
        computed_prefill = self.req_states.num_computed_prefill_tokens
        computed_prefill[idx_mapping_np] += input_batch.num_scheduled_tokens
        np.minimum(
            computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
        )
1314
1315
1316
1317
        # Advance the CPU mirror optimistically (assume all scheduled accepted).
        self.req_states.num_computed_tokens_np[idx_mapping_np] += (
            input_batch.num_scheduled_tokens
        )
1318

1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
    ########### EPLB methods start ###########
    @property
    def eplb_state(self):
        return self.eplb.state

    @eplb_state.setter
    def eplb_state(self, state) -> None:
        self.eplb.state = state

    @property
    def eep_eplb_suppressed(self) -> bool:
        return self.eplb.suppressed

    @eep_eplb_suppressed.setter
    def eep_eplb_suppressed(self, suppressed: bool) -> None:
        self.eplb.suppressed = suppressed

    def setup_eplb_from_mapping(
        self,
        expanded_physical_to_logical: torch.Tensor,
        old_num_physical_experts: int,
    ) -> None:
        self.eplb.setup_from_mapping(
            self.model,
            self.model_config,
            expanded_physical_to_logical,
            old_num_physical_experts,
        )

    ########### EPLB methods end ###########

1350
1351
1352
1353
1354

class ExecuteModelState(NamedTuple):
    input_batch: InputBatch
    attn_metadata: dict[str, Any] | None
    slot_mappings_by_layer: dict[str, torch.Tensor] | None
1355
    hidden_states: torch.Tensor | None
1356
1357
    aux_hidden_states: list[torch.Tensor] | None
    kv_connector_output: KVConnectorOutput | None