model_runner.py 55.3 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
39
from vllm.logger import init_logger
from vllm.model_executor.model_loader import get_model_loader
40
from vllm.multimodal import MULTIMODAL_REGISTRY
41
from vllm.sequence import IntermediateTensors
42
from vllm.tasks import SupportedTask
43
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
Woosuk Kwon's avatar
Woosuk Kwon committed
44
45
46
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
47
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
48
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
49
from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
Woosuk Kwon's avatar
Woosuk Kwon committed
50
from vllm.v1.worker.gpu.attn_utils import (
51
    build_slot_mappings_by_layer,
Woosuk Kwon's avatar
Woosuk Kwon committed
52
53
54
55
56
    get_kv_cache_spec,
    init_attn_backend,
    init_kv_cache,
)
from vllm.v1.worker.gpu.block_table import BlockTables
57
from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
58
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
59
60
61
62
63
from vllm.v1.worker.gpu.cudagraph_utils import (
    BatchExecutionDescriptor,
    ModelCudaGraphManager,
    get_uniform_token_count,
)
64
from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
65
from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after
Woosuk Kwon's avatar
Woosuk Kwon committed
66
67
68
from vllm.v1.worker.gpu.input_batch import (
    InputBatch,
    InputBuffers,
69
    combine_sampled_and_draft_tokens,
70
    expand_idx_mapping,
71
    get_num_sampled_and_rejected,
72
    post_update,
73
    post_update_pool,
74
75
    prepare_pos_seq_lens,
    prepare_prefill_inputs,
Woosuk Kwon's avatar
Woosuk Kwon committed
76
)
77
78
79
80
81
from vllm.v1.worker.gpu.kv_connector import (
    NO_OP_KV_CONNECTOR,
    KVConnector,
    get_kv_connector,
)
82
from vllm.v1.worker.gpu.lora_utils import LoraState
83
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
84
from vllm.v1.worker.gpu.model_states import init_model_state
85
from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
86
from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive
87
from vllm.v1.worker.gpu.sample.output import SamplerOutput
88
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
89
from vllm.v1.worker.gpu.sample.sampler import Sampler
90
from vllm.v1.worker.gpu.spec_decode import init_speculator
91
92
93
from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
    set_eagle3_aux_hidden_state_layers,
)
94
from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler
95
from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
96
from vllm.v1.worker.gpu.states import RequestState
97
from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
Woosuk Kwon's avatar
Woosuk Kwon committed
98
99
100
101
102
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

logger = init_logger(__name__)


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

Woosuk Kwon's avatar
Woosuk Kwon committed
131
132
133
        self.use_async_scheduling = self.scheduler_config.async_scheduling
        self.output_copy_stream = torch.cuda.Stream(self.device)

134
        # Pipeline parallelism.
135
136
137
138
        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

139
140
        # Persistent buffer for intermediate tensors (non-first PP ranks).
        self.intermediate_tensors: IntermediateTensors | None = None
141

142
143
144
145
        # Data parallelism.
        self.dp_size = self.parallel_config.data_parallel_size
        self.dp_rank = self.parallel_config.data_parallel_rank

146
147
148
149
150
151
        # 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

152
153
154
155
156
157
158
159
160
        # 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()

161
        # Speculative decoding.
162
        self.speculator = None
163
        self.num_speculative_steps = 0
164
        self.use_aux_hidden_state_outputs = False
165
166
        if self.speculative_config is not None:
            self.num_speculative_steps = self.speculative_config.num_speculative_tokens
167

168
169
170
171
172
173
            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
174
                if self.use_pp:
175
176
177
178
                    raise ValueError("EAGLE3 with pipeline parallel is not supported.")

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

181
182
183
184
        # Pooling models.
        self.is_pooling_model = self.model_config.runner_type == "pooling"
        self.pooling_runner: PoolingRunner | None = None

185
        # General request states.
Woosuk Kwon's avatar
Woosuk Kwon committed
186
187
188
189
        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,
190
            num_speculative_steps=self.num_speculative_steps,
Woosuk Kwon's avatar
Woosuk Kwon committed
191
192
193
194
195
196
197
198
            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,
        )
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215

        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,
            )
216
217
218
219
220
            if self.speculative_config is not None:
                self.rejection_sampler = RejectionSampler(
                    self.sampler,
                    self.speculative_config,
                )
221
222
223
224
225
226
            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
227

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

236
        # For transferring state from execute_model to subsequent sample_tokens call.
237
        self.execute_model_state: ExecuteModelState | None = None
238

239
240
241
        # Expert parallelism load balancer.
        self.eplb = EPLBController(self.parallel_config, self.device)

242
243
244
245
    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

246
247
248
    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
        tasks: list[SupportedTask] = []
        if self.model_config.runner_type == "generate":
249
            tasks.extend(self.model_state.get_supported_generation_tasks())
250
251
252
253
254
        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))
255
        return tuple(tasks)
Woosuk Kwon's avatar
Woosuk Kwon committed
256

257
    def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
258
        time_before_load = time.perf_counter()
259
260
261
262
        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
263
264
265
266
267
        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(
268
                vllm_config=self.vllm_config, model_config=self.vllm_config.model_config
Woosuk Kwon's avatar
Woosuk Kwon committed
269
270
271
            )
            if self.lora_config:
                self.model = self.load_lora_model(
272
                    self.model, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
273
                )
274
275
276
277
278

            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:
279
                self.speculator.load_model(self.model)
280
281
282
                eplb_models_added = self.eplb.maybe_register_speculator(
                    self.speculator, self.speculative_config, load_dummy_weights
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
283
284
285
286
        time_after_load = time.perf_counter()

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

292
293
294
295
        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)
296

297
        # Initialize the components that require the model.
298
        self.model_state = init_model_state(
299
300
            self.vllm_config, self.model, self.encoder_cache, self.device
        )
301
        if self.is_pooling_model and self.is_last_pp_rank:
302
            self.pooling_runner = PoolingRunner(self.model)
303
304
305
306
307
308
        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)
309

310
311
312
313
314
315
316
317
318
319
320
        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
321
322
323
    def get_model(self) -> nn.Module:
        return self.model

324
325
326
327
328
    @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
329
330
331
332
333
334
335
336
337
338
339
    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
        ]

340
341
342
343
344
345
346
347
348
        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
349
350
351
352
        self.block_tables = BlockTables(
            block_sizes=block_sizes,
            max_num_reqs=self.max_num_reqs,
            max_num_batched_tokens=self.max_num_tokens,
353
            max_model_len=block_table_max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
354
            device=self.device,
355
356
357
            cp_size=self.dcp_size,
            cp_rank=self.dcp_rank,
            cp_interleave=self.cp_interleave,
Woosuk Kwon's avatar
Woosuk Kwon committed
358
359
        )

360
        self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend(
361
            self.kv_cache_config, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
362
        )
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
        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)

380
        check_attention_cp_compatibility(self.vllm_config)
381
        if self.speculator is not None:
382
383
            # HACK(woosuk)
            self.speculator.set_attn(
384
                self.model_state,
385
386
387
                self.kv_cache_config,
                self.block_tables,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
388
389

        self.kv_caches: list[torch.Tensor] = []
390
        kv_caches_dict = init_kv_cache(
Woosuk Kwon's avatar
Woosuk Kwon committed
391
392
393
394
395
            self.kv_caches,
            self.compilation_config.static_forward_context,
            self.kv_cache_config,
            self.attn_backends,
            self.device,
396
            self.cache_config.cache_dtype,
Woosuk Kwon's avatar
Woosuk Kwon committed
397
        )
398
399
        self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

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

417
        # Create a dummy scheduler output.
418
        num_reqs = min(num_tokens, self.max_num_reqs)
419
        if uniform_decode:
420
421
422
423
424
425
426
427
428
429
            # 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

430
431
        assert sum(num_tokens_per_request) == num_tokens
        num_scheduled_tokens = {
432
            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
433
434
435
436
437
        }
        dummy_scheduler_output = SchedulerOutput.make_empty()
        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens

438
439
440
        # Disable any use of KVConnector for dummy runs.
        self.kv_connector.set_disabled(True)

441
        # Get the intermediate tensors for the dummy run.
442
        intermediate_tensors = None
443
        if not self.is_first_pp_rank:
444
445
            assert self.intermediate_tensors is not None
            intermediate_tensors = self.intermediate_tensors[:num_tokens]
446

447
448
        # Execute the model.
        self.execute_model(
449
450
451
452
            dummy_scheduler_output,
            intermediate_tensors=intermediate_tensors,
            dummy_run=True,
            skip_attn_for_dummy_run=skip_attn,
453
            is_profile=is_profile,
454
        )
455
        self.kv_connector.set_disabled(False)
456
457

        # Non-last PP ranks don't produce output for sampling.
458
        if not self.is_last_pp_rank:
459
460
            return None, None

461
        assert self.execute_model_state is not None
462
463
464
465
466
467
        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
        num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp
468
        self.execute_model_state = None
469
470
471

        # dummy run the eagle speculator's propose to ensure DP/EP sync.
        if self.speculator is not None:
472
            assert self.sampler is not None
473
474
475
476
477
478
479
480
481
482
            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,
                    ),
                )
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
            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,
                num_tokens_across_dp=num_tokens_across_dp,
                dummy_run=True,
                skip_attn_for_dummy_run=skip_attn,
502
                mm_inputs=mm_inputs,
503
                is_profile=is_profile,
504
505
            )

506
        assert hidden_states is not None  # Last PP rank always has hidden_states
507
        sample_hidden_states = hidden_states[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
508
509
510
        return hidden_states, sample_hidden_states

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

518
519
520
        # 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.
521
522
        assert self.sampler is not None
        self.sampler(logits, dummy_input_batch)
Woosuk Kwon's avatar
Woosuk Kwon committed
523

524
525
526
527
528
    @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
529
530
531
    @torch.inference_mode()
    def profile_run(self) -> None:
        hidden_states, sample_hidden_states = self._dummy_run(
532
            self.max_num_tokens, skip_attn=True, is_profile=True
Woosuk Kwon's avatar
Woosuk Kwon committed
533
        )
534

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

543
        torch.accelerator.synchronize()
Woosuk Kwon's avatar
Woosuk Kwon committed
544
545
546
547
        del hidden_states, sample_hidden_states
        gc.collect()

    def reset_mm_cache(self) -> None:
548
549
        if self.encoder_cache is not None:
            self.encoder_cache.reset_mm_cache()
550
551

    def reset_encoder_cache(self) -> None:
552
553
        if self.encoder_cache is not None:
            self.encoder_cache.reset_encoder_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
554
555
556
557
558

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

559
560
561
562
    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
563
564
    @torch.inference_mode()
    def capture_model(self) -> int:
565
        assert self.cudagraph_manager is not None
Woosuk Kwon's avatar
Woosuk Kwon committed
566
567
568
569
570
571
572
573
        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()
574
        gc.collect()
575
        torch.accelerator.empty_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
576
577
578
579
        start_free_gpu_memory = torch.cuda.mem_get_info()[0]

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

        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

605
606
607
608
609
610
611
612
613
614
    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

615
    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
616
        finished_req_ids = scheduler_output.finished_req_ids
617
618
619
        preempted_req_ids = scheduler_output.preempted_req_ids
        if preempted_req_ids:
            finished_req_ids = finished_req_ids.union(preempted_req_ids)
620
        for req_id in finished_req_ids:
621
            self._remove_request(req_id)
622

623
    def free_states(self, scheduler_output: SchedulerOutput) -> None:
624
        if self.encoder_cache is not None:
625
            for mm_hash in scheduler_output.free_encoder_mm_hashes:
626
                self.encoder_cache.free_encoder_cache(mm_hash)
Woosuk Kwon's avatar
Woosuk Kwon committed
627

628
    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
629
        for new_req_data in scheduler_output.scheduled_new_reqs:
630
631
            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
632
            req_id = new_req_data.req_id
633
634
635
636
637
638

            # 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)

639
            prompt_len = len(new_req_data.prompt_token_ids)
Woosuk Kwon's avatar
Woosuk Kwon committed
640
641
            self.req_states.add_request(
                req_id=req_id,
642
                prompt_len=prompt_len,
643
                all_token_ids=new_req_data.prefill_token_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
644
645
646
                num_computed_tokens=new_req_data.num_computed_tokens,
            )
            req_index = self.req_states.req_id_to_index[req_id]
647

648
649
            if self.encoder_cache is not None:
                self.encoder_cache.add_request(req_id, new_req_data.mm_features)
650

651
            self.model_state.add_request(req_index, new_req_data)
652
653
654
            self.block_tables.append_block_ids(
                req_index, new_req_data.block_ids, overwrite=True
            )
655
            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
Woosuk Kwon's avatar
Woosuk Kwon committed
656

657
658
            if self.is_last_pp_rank and new_req_data.sampling_params is not None:
                assert self.sampler is not None
659
660
661
                self.sampler.add_request(
                    req_index, prompt_len, new_req_data.sampling_params
                )
662
                assert self.prompt_logprobs_worker is not None
663
664
665
666
                self.prompt_logprobs_worker.add_request(
                    req_id, req_index, new_req_data.sampling_params
                )

667
668
        if scheduler_output.scheduled_new_reqs:
            self.req_states.apply_staged_writes()
669
            self.model_state.apply_staged_writes()
670
671
        if self.sampler is not None:
            self.sampler.apply_staged_writes()
672
673

    def update_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
674
        # Add new blocks for the existing requests.
675
676
        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
677
            if req_new_block_ids is not None:
678
                req_index = self.req_states.req_id_to_index[req_id]
679
680
681
                self.block_tables.append_block_ids(
                    req_index, req_new_block_ids, overwrite=False
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
682
683

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

        # Decode first, then prefill.
        # batch_idx -> req_id
694
695
696
        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
697

698
699
        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)
700
        idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
Woosuk Kwon's avatar
Woosuk Kwon committed
701

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

725
726
727
728
            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:])
729
            cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
730

731
            max_expand_len = self.num_speculative_steps + 1
732
            expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
733
                idx_mapping, total_num_logits, cu_num_logits, max_expand_len
734
735
            )

736
        # Get query_start_loc.
737
738
        # num_reqs_padded is None for PIECEWISE graphs (no request padding needed)
        num_reqs_padded = batch_desc.num_reqs or num_reqs
739
740
741
        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])
742
743
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
744
        query_start_loc_np[num_reqs + 1 :] = num_tokens
745
        async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
746
747
        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]
748

749
750
751
752
753
754
755
756
757
758
759
        # 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
760

761
762
763
        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
764
765
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
766
767
768
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
769
        seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
770

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

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

        return InputBatch(
            req_ids=req_ids,
            num_reqs=num_reqs,
801
            num_reqs_after_padding=num_reqs_padded,
Woosuk Kwon's avatar
Woosuk Kwon committed
802
803
            idx_mapping=idx_mapping,
            idx_mapping_np=idx_mapping_np,
804
            expanded_idx_mapping=expanded_idx_mapping,
805
            expanded_local_pos=expanded_local_pos,
Woosuk Kwon's avatar
Woosuk Kwon committed
806
807
808
            num_scheduled_tokens=num_scheduled_tokens,
            num_tokens=num_tokens,
            num_tokens_after_padding=num_tokens_after_padding,
809
            num_draft_tokens=total_num_draft_tokens,
810
            query_start_loc=query_start_loc,
Woosuk Kwon's avatar
Woosuk Kwon committed
811
            query_start_loc_np=query_start_loc_np,
812
            seq_lens=seq_lens,
813
814
815
            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
816
            logits_indices=logits_indices,
817
            cu_num_logits=cu_num_logits,
818
            cu_num_logits_np=cu_num_logits_np,
819
            has_structured_output_reqs=scheduler_output.has_structured_output_requests,
Woosuk Kwon's avatar
Woosuk Kwon committed
820
821
        )

822
823
824
    def prepare_attn(
        self, input_batch: InputBatch
    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
825
826
827
828
829
830
831
        # 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.
832
833
834
835
        slot_mappings = self.block_tables.compute_slot_mappings(
            input_batch.idx_mapping,
            input_batch.query_start_loc,
            input_batch.positions,
836
            num_tokens_padded=input_batch.num_tokens_after_padding,
837
838
839
840
841
842
843
844
845
846
847
848
        )
        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
849
850
851
852
853
    def sample(
        self,
        hidden_states: torch.Tensor,
        input_batch: InputBatch,
        grammar_output: GrammarOutput | None,
854
    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
855
856
857
858
        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.
859
            assert self.structured_outputs_worker is not None
860
861
862
863
864
865
            self.structured_outputs_worker.apply_grammar_bitmask(
                logits,
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.grammar_bitmask,
            )
866

867
868
        if input_batch.num_draft_tokens == 0:
            # No draft tokens (common case).
869
870
            assert self.sampler is not None
            sampler_output = self.sampler(logits, input_batch)
871
        else:
872
            # Rejection sampling for spec decoding.
873
            assert self.rejection_sampler is not None
874
            assert self.speculator is not None
875
876
877
878
            sampler_output = self.rejection_sampler(
                logits,
                input_batch,
                # Draft logits are needed for probabilistic rejection sampling.
879
                self.speculator.draft_logits,
880
            )
881
882
883
884

        # 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(
885
            sampler_output.num_sampled,
886
887
888
889
890
            input_batch.seq_lens,
            input_batch.cu_num_logits,
            input_batch.idx_mapping,
            self.req_states.prefill_len.gpu,
        )
891
        return sampler_output, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
892
893
894
895

    def postprocess(
        self,
        input_batch: InputBatch,
896
897
        sampled_tokens: torch.Tensor,
        num_sampled: torch.Tensor,
898
        num_rejected: torch.Tensor,
899
900
    ) -> None:
        # Update the number of computed tokens.
901
902
903
904
905
        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
906
        post_update(
907
            input_batch.idx_mapping,
908
            self.req_states.num_computed_tokens.gpu,
909
            self.req_states.last_sampled_tokens,
910
            output_bin_counts,
911
912
            sampled_tokens,
            num_sampled,
913
            num_rejected,
914
            input_batch.query_start_loc,
915
916
            self.req_states.all_token_ids.gpu,
            self.req_states.total_len.gpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
917
        )
918
919

        # Update the number of computed prefill tokens.
Woosuk Kwon's avatar
Woosuk Kwon committed
920
        idx_mapping_np = input_batch.idx_mapping_np
921
        computed_prefill = self.req_states.num_computed_prefill_tokens
922
923
924
        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
925
926
927
928
929
930
        )

    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: SchedulerOutput,
931
        intermediate_tensors: IntermediateTensors | None = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
932
        dummy_run: bool = False,
933
        skip_attn_for_dummy_run: bool = False,
934
        is_profile: bool = False,
935
    ) -> ModelRunnerOutput | IntermediateTensors | None:
936
937
938
939
940
941
942
943
944
        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.
945
946
                empty_output = self.kv_connector.no_forward(scheduler_output)
                return empty_output
Woosuk Kwon's avatar
Woosuk Kwon committed
947

948
949
950
951
952
953
        # 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)

954
955
956
957
958
959
960
        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

961
962
963
964
965
966
967
968
969
        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,
        )
970
971

        if batch_desc.num_tokens == 0:
972
            # All DP ranks have zero tokens to run.
973
974
            empty_output = self.kv_connector.no_forward(scheduler_output)
            return empty_output
975
976
977
978

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

982
983
            if self.lora_config:
                # Activate LoRA adapters.
984
                lora_inputs = self.lora_state.make_lora_inputs(
985
986
987
                    input_batch.req_ids,
                    input_batch.idx_mapping_np,
                    input_batch.num_scheduled_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
988
                )
989
990
                self._set_active_loras(*lora_inputs)
        else:
991
            # No actual tokens to run. A dummy run for DP or memory profiling.
992
            input_batch = InputBatch.make_dummy(
993
994
995
                batch_desc.num_reqs or num_reqs,
                batch_desc.num_tokens,
                self.input_buffers,
996
            )
997
            if not skip_attn_for_dummy_run:
998
999
                block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)
            else:
1000
1001
1002
1003
                assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
                    "Attention metadata must be prepared for dummy runs when using "
                    "FULL cudagraph mode."
                )
1004
1005
                block_tables = None
                slot_mappings = None
1006
            # FIXME(woosuk): Fix warmup for LoRA.
Woosuk Kwon's avatar
Woosuk Kwon committed
1007

1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
        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,
1018
                batch_desc.cg_mode,
1019
1020
1021
1022
1023
1024
                block_tables,
                slot_mappings,
                self.attn_groups,
                self.kv_cache_config,
            )

1025
        inputs_embeds = None
1026
        if self.supports_mm_inputs and self.is_first_pp_rank:
1027
1028
            # Run MM encoder (if needed) and get multimodal embeddings.
            # Only first PP rank prepares multimodal embeddings.
1029
1030
            # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs
            # to obtain inputs_embeds, because the compiled model expects this input.
1031
1032
1033
1034
1035
1036
            inputs_embeds = self.model_state.get_mm_embeddings(
                scheduler_output.scheduled_encoder_inputs,
                input_batch,
                self.req_states,
            )

1037
1038
1039
        model_inputs = {
            "input_ids": input_batch.input_ids,
            "positions": input_batch.positions,
1040
            "inputs_embeds": inputs_embeds,
1041
1042
1043
1044
1045
1046
1047
1048
            # 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
1049
1050

            # Prepare the intermediate tensors.
1051
            assert intermediate_tensors is not None
1052
1053
            assert self.intermediate_tensors is not None
            n = input_batch.num_tokens_after_padding
1054
            model_inputs["intermediate_tensors"] = IntermediateTensors(
1055
1056
1057
                {
                    k: v[:n].copy_(intermediate_tensors.tensors[k][:n])
                    for k, v in self.intermediate_tensors.tensors.items()
1058
                }
1059
            )
1060
            del intermediate_tensors
1061

Woosuk Kwon's avatar
Woosuk Kwon committed
1062
        # Run model.
1063
        if batch_desc.cg_mode == CUDAGraphMode.FULL:
1064
            # Use explicit cudagraph replay for FULL mode.
Woosuk Kwon's avatar
Woosuk Kwon committed
1065
1066
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
1067
            assert self.cudagraph_manager is not None
1068
            self.kv_connector.pre_forward(scheduler_output)
1069
            model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
Woosuk Kwon's avatar
Woosuk Kwon committed
1070
        else:
1071
1072
1073
1074
1075
1076
            # 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
1077
            with set_forward_context(
1078
                attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
1079
1080
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
1081
                cudagraph_runtime_mode=batch_desc.cg_mode,
Woosuk Kwon's avatar
Woosuk Kwon committed
1082
                num_tokens_across_dp=num_tokens_across_dp,
1083
                batch_descriptor=batch_descriptor,
1084
                slot_mapping=slot_mappings_by_layer,
1085
                skip_compiled=skip_compiled,
Woosuk Kwon's avatar
Woosuk Kwon committed
1086
            ):
1087
                self.kv_connector.pre_forward(scheduler_output)
1088
                model_output = self.model(**model_inputs)
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103

        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
1104

1105
        kv_connector_output = self.kv_connector.post_forward(scheduler_output)
1106
1107
1108
1109
1110
1111
1112
1113
        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,
            num_tokens_across_dp=num_tokens_across_dp,
1114
        )
1115

1116
        if not self.is_last_pp_rank:
1117
            # Non-last PP rank: return IntermediateTensors for sending.
1118
1119
1120
            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
1121
1122
1123
        return None

    @torch.inference_mode()
1124
    @step_eplb_after()
Woosuk Kwon's avatar
Woosuk Kwon committed
1125
    def sample_tokens(
1126
        self, grammar_output: GrammarOutput | None
1127
    ) -> AsyncOutput | ModelRunnerOutput | None:
1128
1129
1130
        if self.execute_model_state is None:
            # The prior execute_model call must have failed.
            return None
1131
1132
1133
1134
1135
1136
1137
1138

        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
        num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp
1139
        self.execute_model_state = None
Woosuk Kwon's avatar
Woosuk Kwon committed
1140

1141
        if not self.is_last_pp_rank:
1142
1143
1144
1145
            # 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(
1146
                input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
1147
            )
1148
            self.postprocess(input_batch, sampled, num_sampled, num_rejected)
1149
1150
1151
            return None

        # Last rank: sample tokens
1152
        sampler_output, num_sampled, num_rejected = self.sample(
1153
            hidden_states, input_batch, grammar_output
Woosuk Kwon's avatar
Woosuk Kwon committed
1154
        )
1155
1156

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

1160
        assert self.prompt_logprobs_worker is not None
1161
1162
1163
1164
        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
            self.model.compute_logits,
            hidden_states,
            input_batch,
1165
            self.req_states.all_token_ids.gpu,
1166
            self.req_states.num_computed_tokens.gpu,
1167
            self.req_states.prompt_len.np,
1168
1169
1170
            self.req_states.prefill_len.np,
            self.req_states.num_computed_prefill_tokens,
        )
1171
1172
1173
1174
1175
1176
1177
1178

        # 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
1179
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
1180
            kv_connector_output=kv_connector_output,
1181
1182
1183
1184
        )
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,
1185
            num_sampled_tokens=num_sampled,
1186
            main_stream=self.main_stream,
1187
1188
1189
            copy_stream=self.output_copy_stream,
        )

1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
        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
            )

1208
1209
1210
1211
1212
1213
        # 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(
1214
            input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
1215
        )
1216

1217
        if self.speculator is not None:
1218
            assert self.sampler is not None
1219
            draft_tokens = self.speculator.propose(
1220
                input_batch,
1221
1222
                attn_metadata,
                slot_mappings_by_layer,
1223
                hidden_states,
1224
                aux_hidden_states,
1225
1226
                num_sampled,
                num_rejected,
1227
1228
1229
1230
                self.req_states.last_sampled_tokens,
                self.req_states.next_prefill_tokens,
                self.sampler.sampling_states.temperature.gpu,
                self.sampler.sampling_states.seeds.gpu,
1231
                num_tokens_across_dp=num_tokens_across_dp,
1232
                mm_inputs=mm_inputs,
1233
            )
1234
            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1235
            self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
1236
1237
1238
1239

        if self.use_async_scheduling:
            return async_output
        return async_output.get_output()
1240
1241
1242

    def take_draft_token_ids(self) -> DraftTokenIds | None:
        return self.draft_tokens_handler.get_draft_tokens()
1243
1244

    @torch.inference_mode()
1245
    @step_eplb_after()
1246
1247
1248
1249
1250
    def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None:
        if self.execute_model_state is None:
            # The prior execute_model call must have failed.
            return None

1251
1252
1253
        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
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
        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,
        )
1278
1279

        self.postprocess_pool(input_batch)
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
        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
        )
1299

1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
    ########### 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 ###########

1331
1332
1333
1334
1335

class ExecuteModelState(NamedTuple):
    input_batch: InputBatch
    attn_metadata: dict[str, Any] | None
    slot_mappings_by_layer: dict[str, torch.Tensor] | None
1336
    hidden_states: torch.Tensor | None
1337
1338
1339
    aux_hidden_states: list[torch.Tensor] | None
    kv_connector_output: KVConnectorOutput | None
    num_tokens_across_dp: torch.Tensor | None