model_runner.py 38 KB
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
2
3
4
5
6
7
8
9
10
11
12
13
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import gc
import time
from copy import deepcopy
from typing import Any

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

from vllm.config import VllmConfig
from vllm.config.compilation import CUDAGraphMode
14
from vllm.distributed.parallel_state import prepare_communication_buffer_for_model
Woosuk Kwon's avatar
Woosuk Kwon committed
15
16
17
from vllm.forward_context import set_forward_context
from vllm.logger import init_logger
from vllm.model_executor.model_loader import get_model_loader
18
from vllm.multimodal import MULTIMODAL_REGISTRY
19
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
Woosuk Kwon's avatar
Woosuk Kwon committed
20
21
22
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
23
from vllm.v1.outputs import ModelRunnerOutput
24
from vllm.v1.worker.gpu.async_utils import AsyncOutput
Woosuk Kwon's avatar
Woosuk Kwon committed
25
26
from vllm.v1.worker.gpu.attn_utils import (
    build_attn_metadata,
27
    build_slot_mappings_by_layer,
Woosuk Kwon's avatar
Woosuk Kwon committed
28
29
30
31
32
    get_kv_cache_spec,
    init_attn_backend,
    init_kv_cache,
)
from vllm.v1.worker.gpu.block_table import BlockTables
33
from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
Woosuk Kwon's avatar
Woosuk Kwon committed
34
from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
35
from vllm.v1.worker.gpu.dp_utils import (
36
    get_cudagraph_and_dp_padding,
37
38
    make_num_tokens_across_dp,
)
Woosuk Kwon's avatar
Woosuk Kwon committed
39
40
41
from vllm.v1.worker.gpu.input_batch import (
    InputBatch,
    InputBuffers,
42
    combine_sampled_and_draft_tokens,
43
    expand_idx_mapping,
44
    get_num_sampled_and_rejected,
45
    post_update,
46
47
    prepare_pos_seq_lens,
    prepare_prefill_inputs,
Woosuk Kwon's avatar
Woosuk Kwon committed
48
)
49
50
51
52
53
from vllm.v1.worker.gpu.kv_connector import (
    NO_OP_KV_CONNECTOR,
    KVConnector,
    get_kv_connector,
)
54
from vllm.v1.worker.gpu.lora_utils import LoraState
55
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
56
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
57
from vllm.v1.worker.gpu.sample.output import SamplerOutput
58
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
59
from vllm.v1.worker.gpu.sample.sampler import Sampler
60
from vllm.v1.worker.gpu.spec_decode import init_speculator
61
from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
62
from vllm.v1.worker.gpu.states import RequestState
63
from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
Woosuk Kwon's avatar
Woosuk Kwon committed
64
65
66
67
68
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

logger = init_logger(__name__)


69
class GPUModelRunner(LoRAModelRunnerMixin):
Woosuk Kwon's avatar
Woosuk Kwon committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
    ):
        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.is_pooling_model = False

        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
100
        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
Woosuk Kwon's avatar
Woosuk Kwon committed
101

102
        # Multimodal
103
104
105
106
107
108
109
110
111
112
113
        self.mm_registry = MULTIMODAL_REGISTRY
        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
            self.model_config
        )
        if self.supports_mm_inputs:
            self.encoder_runner = EncoderRunner(
                max_num_tokens=self.max_num_tokens,
                hidden_size=self.inputs_embeds_size,
                dtype=self.dtype,
                device=self.device,
            )
114
115
116
117
        self.uses_mrope = self.model_config.uses_mrope
        if self.uses_mrope:
            self.mrope_states = MRopeState(
                max_num_reqs=self.max_num_reqs,
118
                max_num_tokens=self.max_num_tokens,
119
120
121
122
                max_model_len=self.max_model_len,
                device=self.device,
            )

Woosuk Kwon's avatar
Woosuk Kwon committed
123
124
125
126
        self.use_async_scheduling = self.scheduler_config.async_scheduling
        self.output_copy_stream = torch.cuda.Stream(self.device)
        self.output_copy_event = torch.cuda.Event()

127
128
129
        if self.speculative_config is not None:
            self.do_spec_decode = True
            self.num_speculative_steps = self.speculative_config.num_speculative_tokens
130
            self.speculator = init_speculator(self.vllm_config, self.device)
131
132
133
        else:
            self.do_spec_decode = False
            self.num_speculative_steps = 0
134
            self.speculator = None
135

Woosuk Kwon's avatar
Woosuk Kwon committed
136
137
138
139
        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,
140
            num_speculative_steps=self.num_speculative_steps,
Woosuk Kwon's avatar
Woosuk Kwon committed
141
142
143
144
145
146
147
148
            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,
        )
149
150
151
152
153
154
        self.sampler = Sampler(
            max_num_reqs=self.max_num_reqs,
            vocab_size=self.vocab_size,
            device=self.device,
            logprobs_mode=self.model_config.logprobs_mode,
        )
155
        self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
Woosuk Kwon's avatar
Woosuk Kwon committed
156
157

        # CUDA graphs.
158
159
160
        self.cudagraph_manager = CudaGraphManager(
            self.vllm_config, self.uses_mrope, self.device
        )
161
162
163
164
        # Structured outputs worker.
        self.structured_outputs_worker = StructuredOutputsWorker(
            max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
            vocab_size=self.vocab_size,
165
            device=self.device,
166
        )
167
168
        # LoRA-related workers.
        self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
169

170
171
        self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR

172
173
174
175
    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

176
177
    @staticmethod
    def get_supported_tasks() -> tuple[str]:
Woosuk Kwon's avatar
Woosuk Kwon committed
178
179
180
181
182
183
184
185
186
187
188
189
190
191
        return ("generate",)

    def load_model(self, *args, **kwargs) -> None:
        time_before_load = time.perf_counter()
        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(
                vllm_config=self.vllm_config,
                model_config=self.vllm_config.model_config,
            )
            if self.lora_config:
                self.model = self.load_lora_model(
192
                    self.model, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
193
                )
194
195
            if self.do_spec_decode:
                self.speculator.load_model(self.model)
Woosuk Kwon's avatar
Woosuk Kwon committed
196
197
198
199
        time_after_load = time.perf_counter()

        self.model_memory_usage = m.consumed_memory
        logger.info(
200
201
            "Model loading took %s GiB and %.6f seconds",
            format_gib(m.consumed_memory),
Woosuk Kwon's avatar
Woosuk Kwon committed
202
203
204
            time_after_load - time_before_load,
        )

205
206
207
208
209
210
        prepare_communication_buffer_for_model(self.model)
        if self.do_spec_decode:
            speculator_model = getattr(self.speculator, "model", None)
            if speculator_model is not None:
                prepare_communication_buffer_for_model(speculator_model)

Woosuk Kwon's avatar
Woosuk Kwon committed
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
    def get_model(self) -> nn.Module:
        return self.model

    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
        ]

        self.block_tables = BlockTables(
            block_sizes=block_sizes,
            max_num_reqs=self.max_num_reqs,
            max_num_batched_tokens=self.max_num_tokens,
            max_model_len=self.max_model_len,
            device=self.device,
        )

        self.attn_backends, self.attn_metadata_builders = init_attn_backend(
234
            self.kv_cache_config, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
235
        )
236
237
238
239
240
241
242
        if self.do_spec_decode:
            # HACK(woosuk)
            self.speculator.set_attn(
                self.kv_cache_config,
                self.attn_metadata_builders,
                self.block_tables,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
243
244

        self.kv_caches: list[torch.Tensor] = []
245
        kv_caches_dict = init_kv_cache(
Woosuk Kwon's avatar
Woosuk Kwon committed
246
247
248
249
250
251
            self.kv_caches,
            self.compilation_config.static_forward_context,
            self.kv_cache_config,
            self.attn_backends,
            self.device,
        )
252
253
        self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

Woosuk Kwon's avatar
Woosuk Kwon committed
254
255
256
257
258
259
260
261
        # Attention groups are not supported.
        self.attn_groups = []  # type: ignore

    def prepare_dummy_attn_metadata(self, input_batch: InputBatch) -> None:
        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
        )
262
263
264
        slot_mappings_by_layer = build_slot_mappings_by_layer(
            slot_mappings, self.kv_cache_config
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
265
266
267
268
        attn_metadata = build_attn_metadata(
            attn_metadata_builders=self.attn_metadata_builders,
            num_reqs=input_batch.num_reqs,
            num_tokens=input_batch.num_tokens,
269
270
271
            query_start_loc_gpu=input_batch.query_start_loc,
            query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np),
            seq_lens=input_batch.seq_lens,
272
            max_seq_len=self.max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
273
274
275
276
277
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
        )
        input_batch.attn_metadata = attn_metadata
278
        input_batch.slot_mappings = slot_mappings_by_layer
Woosuk Kwon's avatar
Woosuk Kwon committed
279
280
281

    @torch.inference_mode()
    def _dummy_run(
282
        self, num_tokens: int, *args, skip_attn: bool = True, **kwargs
Woosuk Kwon's avatar
Woosuk Kwon committed
283
    ) -> tuple[torch.Tensor, torch.Tensor]:
284
        # Create a dummy scheduler output.
Woosuk Kwon's avatar
Woosuk Kwon committed
285
        num_reqs = min(num_tokens, self.max_num_reqs)
286
287
288
289
        num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
        num_tokens_per_request[-1] += num_tokens % num_reqs
        assert sum(num_tokens_per_request) == num_tokens
        num_scheduled_tokens = {
290
            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
291
292
293
294
295
        }
        dummy_scheduler_output = SchedulerOutput.make_empty()
        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens

296
297
298
        # Disable any use of KVConnector for dummy runs.
        self.kv_connector.set_disabled(True)

299
300
301
302
        # Execute the model.
        self.execute_model(
            dummy_scheduler_output, dummy_run=True, skip_attn_for_dummy_run=skip_attn
        )
303
        self.kv_connector.set_disabled(False)
304
        assert self.execute_model_state is not None
305
        hidden_states, input_batch, _ = self.execute_model_state
306
        sample_hidden_states = hidden_states[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
307
308
309
        return hidden_states, sample_hidden_states

    @torch.inference_mode()
310
    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
311
312
        num_reqs = hidden_states.shape[0]
        logits = self.model.compute_logits(hidden_states)
313
314
315
316
317
318
        idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=self.device)
        idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
        pos = torch.zeros(num_reqs, dtype=torch.int64, device=self.device)
        # 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.
319
        self.sampler(logits, idx_mapping, idx_mapping_np, idx_mapping_np, pos)
Woosuk Kwon's avatar
Woosuk Kwon committed
320
321
322
323

    @torch.inference_mode()
    def profile_run(self) -> None:
        hidden_states, sample_hidden_states = self._dummy_run(
324
            self.max_num_tokens, skip_attn=True
Woosuk Kwon's avatar
Woosuk Kwon committed
325
326
        )
        self._dummy_sampler_run(sample_hidden_states)
327
        if self.do_spec_decode:
328
            num_tokens_across_dp = make_num_tokens_across_dp(
329
                self.parallel_config.data_parallel_size, self.max_num_tokens
330
331
332
333
            )
            self.speculator.run_model(
                self.max_num_tokens,
                attn_metadata=None,
334
                slot_mappings=None,
335
336
                num_tokens_across_dp=num_tokens_across_dp,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
337
338
339
340
341
        torch.cuda.synchronize()
        del hidden_states, sample_hidden_states
        gc.collect()

    def reset_mm_cache(self) -> None:
342
343
344
345
        self.encoder_runner.reset_mm_cache()

    def reset_encoder_cache(self) -> None:
        self.encoder_runner.reset_encoder_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360

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

    @torch.inference_mode()
    def capture_model(self) -> int:
        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()
361
        gc.collect()
362
        torch.cuda.empty_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
363
364
365
        start_free_gpu_memory = torch.cuda.mem_get_info()[0]

        with self.maybe_setup_dummy_loras(self.lora_config):
366
367
368
            mrope_positions = None
            if self.uses_mrope:
                mrope_positions = self.mrope_states.mrope_positions
369
370
371
            inputs_embeds = None
            if self.supports_mm_inputs:
                inputs_embeds = self.encoder_runner.inputs_embeds
Woosuk Kwon's avatar
Woosuk Kwon committed
372
373
374
            self.cudagraph_manager.capture(
                model=self.model,
                input_buffers=self.input_buffers,
375
                mrope_positions=mrope_positions,
376
                inputs_embeds=inputs_embeds,
Woosuk Kwon's avatar
Woosuk Kwon committed
377
378
379
380
                block_tables=self.block_tables,
                attn_metadata_builders=self.attn_metadata_builders,
                kv_cache_config=self.kv_cache_config,
            )
381
382
            if self.do_spec_decode:
                self.speculator.capture_model()
Woosuk Kwon's avatar
Woosuk Kwon committed
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402

        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

    def warmup_for_prefill(self) -> None:
        # For FlashInfer, we would like to execute a dummy prefill run
        # to trigger JIT compilation.
        if all("FLASHINFER" in b.get_name() for b in self.attn_backends.values()):
            self._dummy_run(self.max_num_tokens, skip_attn=False)
            torch.cuda.synchronize()

403
    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
404
405
406
407
408
409
        finished_req_ids = scheduler_output.finished_req_ids
        if scheduler_output.preempted_req_ids:
            finished_req_ids = finished_req_ids.union(
                scheduler_output.preempted_req_ids
            )
        for req_id in finished_req_ids:
Woosuk Kwon's avatar
Woosuk Kwon committed
410
            self.req_states.remove_request(req_id)
411
412
            if self.supports_mm_inputs:
                self.encoder_runner.remove_request(req_id)
413
            self.prompt_logprobs_worker.remove_request(req_id)
414
            self.lora_state.remove_request(req_id)
415

416
    def free_states(self, scheduler_output: SchedulerOutput) -> None:
417
418
419
        if self.supports_mm_inputs:
            for mm_hash in scheduler_output.free_encoder_mm_hashes:
                self.encoder_runner.free_encoder_cache(mm_hash)
Woosuk Kwon's avatar
Woosuk Kwon committed
420

421
    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
422
        for new_req_data in scheduler_output.scheduled_new_reqs:
423
424
425
            assert new_req_data.prompt_token_ids is not None
            assert new_req_data.prefill_token_ids is not None
            assert new_req_data.sampling_params is not None
Woosuk Kwon's avatar
Woosuk Kwon committed
426
            req_id = new_req_data.req_id
427
            prompt_len = len(new_req_data.prompt_token_ids)
Woosuk Kwon's avatar
Woosuk Kwon committed
428
429
            self.req_states.add_request(
                req_id=req_id,
430
                prompt_len=prompt_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
431
432
433
434
                prefill_token_ids=new_req_data.prefill_token_ids,
                num_computed_tokens=new_req_data.num_computed_tokens,
            )
            req_index = self.req_states.req_id_to_index[req_id]
435

436
437
438
            if self.supports_mm_inputs:
                self.encoder_runner.add_request(req_id, new_req_data.mm_features)

439
440
441
442
443
444
            # Pre-compute M-RoPE positions for prefill.
            if self.uses_mrope:
                self.mrope_states.init_prefill_mrope_positions(
                    req_index,
                    self.model,  # type: ignore
                    new_req_data.prefill_token_ids,
445
                    mm_features=new_req_data.mm_features,
446
447
                )

448
449
450
            self.block_tables.append_block_ids(
                req_index, new_req_data.block_ids, overwrite=True
            )
451
452
453
            self.sampler.add_request(
                req_index, prompt_len, new_req_data.sampling_params
            )
454
455
456
            self.prompt_logprobs_worker.add_request(
                req_id, req_index, new_req_data.sampling_params
            )
457
            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
Woosuk Kwon's avatar
Woosuk Kwon committed
458

459
460
461
462
463
464
465
466
467
468
469
        if scheduler_output.scheduled_new_reqs:
            self.req_states.apply_staged_writes()
            self.sampler.apply_staged_writes(
                self.req_states.prefill_token_ids.gpu,
                self.req_states.prefill_len.np,
                self.req_states.prompt_len,
            )
            if self.uses_mrope:
                self.mrope_states.apply_staged_writes()

    def update_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
470
        # Add new blocks for the existing requests.
471
472
        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
473
            if req_new_block_ids is not None:
474
                req_index = self.req_states.req_id_to_index[req_id]
475
476
477
                self.block_tables.append_block_ids(
                    req_index, req_new_block_ids, overwrite=False
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
478
479
480
481
482
483
484
485
486
487
488
489
490

    def prepare_inputs(
        self,
        scheduler_output: SchedulerOutput,
        num_tokens_after_padding: int,
    ) -> InputBatch:
        num_tokens = scheduler_output.total_num_scheduled_tokens
        assert num_tokens > 0
        num_reqs = len(scheduler_output.num_scheduled_tokens)

        # Decode first, then prefill.
        # batch_idx -> req_id
        req_ids = sorted(
491
492
            scheduler_output.num_scheduled_tokens.keys(),
            key=lambda k: scheduler_output.num_scheduled_tokens[k],
Woosuk Kwon's avatar
Woosuk Kwon committed
493
494
495
496
497
498
499
500
        )
        num_scheduled_tokens = np.array(
            [scheduler_output.num_scheduled_tokens[i] for i in req_ids], dtype=np.int32
        )

        idx_mapping_list = [
            self.req_states.req_id_to_index[req_id] for req_id in req_ids
        ]
501
        idx_mapping_np = np.array(idx_mapping_list, dtype=np.int32)
502
        idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
Woosuk Kwon's avatar
Woosuk Kwon committed
503

504
        # Get the number of draft tokens for each request.
505
506
        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
        if not draft_tokens:
507
508
509
            # No draft token scheduled (common case).
            total_num_draft_tokens = 0
            total_num_logits = num_reqs
510
            cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
511
512
513
            cu_num_logits = torch.arange(
                num_reqs + 1, device=self.device, dtype=torch.int32
            )
514
            expanded_idx_mapping = idx_mapping
515
516
        else:
            num_draft_tokens = np.array(
517
                [len(draft_tokens.get(req_id, ())) for req_id in req_ids],
518
519
520
521
522
                dtype=np.int32,
            )
            total_num_draft_tokens = int(num_draft_tokens.sum())
            total_num_logits = num_reqs + total_num_draft_tokens

523
524
525
526
            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:])
527
            cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
528

529
            max_expand_len = self.num_speculative_steps + 1
530
            expanded_idx_mapping = expand_idx_mapping(
531
                idx_mapping, total_num_logits, cu_num_logits, max_expand_len
532
533
            )

Woosuk Kwon's avatar
Woosuk Kwon committed
534
535
536
        # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
        block_tables = self.block_tables.gather_block_tables(idx_mapping)

537
        # Get query_start_loc.
538
539
540
        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])
541
542
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
543
        query_start_loc_np[num_reqs + 1 :] = num_tokens
544
545
        async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)

546
547
548
        query_start_loc_np = query_start_loc_np[: num_reqs + 1]
        query_start_loc_cpu = torch.from_numpy(query_start_loc_np)
        query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
549

550
        # Get prefill tokens.
551
        prepare_prefill_inputs(
552
553
554
            self.input_buffers.input_ids,
            self.req_states.next_prefill_tokens,
            idx_mapping,
555
            query_start_loc,
556
557
            self.req_states.prefill_token_ids.gpu,
            self.req_states.prefill_len.gpu,
558
            self.req_states.num_computed_tokens.gpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
559
560
        )

561
562
563
        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
564
565
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
566
567
568
569
570
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
        seq_lens = self.input_buffers.seq_lens[:num_reqs]

571
572
573
574
575
576
577
578
579
        # Prepare M-RoPE positions.
        if self.uses_mrope:
            self.mrope_states.prepare_mrope_positions(
                idx_mapping,
                query_start_loc,
                self.req_states.prefill_len.gpu,
                self.req_states.num_computed_tokens.gpu,
            )

580
        # Some input token ids are directly read from the last sampled tokens
581
582
        # and draft tokens. Also, get the logits indices to sample tokens from.
        logits_indices = combine_sampled_and_draft_tokens(
583
            self.input_buffers.input_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
584
585
            idx_mapping,
            self.req_states.last_sampled_tokens,
586
            query_start_loc,
587
588
            seq_lens,
            self.req_states.prefill_len.gpu,
589
590
591
            self.req_states.draft_tokens,
            cu_num_logits,
            total_num_logits,
Woosuk Kwon's avatar
Woosuk Kwon committed
592
593
594
595
        )

        # Compute slot mappings: [num_kv_cache_groups, num_tokens]
        slot_mappings = self.block_tables.compute_slot_mappings(
596
597
598
            idx_mapping,
            query_start_loc,
            self.input_buffers.positions[:num_tokens],
Woosuk Kwon's avatar
Woosuk Kwon committed
599
        )
600
601
602
603
        # Layer name -> slot mapping.
        slot_mappings_by_layer = build_slot_mappings_by_layer(
            slot_mappings, self.kv_cache_config
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
604
605
606
607
608
609

        # Layer name -> attention metadata.
        attn_metadata = build_attn_metadata(
            attn_metadata_builders=self.attn_metadata_builders,
            num_reqs=num_reqs,
            num_tokens=num_tokens,
610
            query_start_loc_gpu=query_start_loc,
611
            query_start_loc_cpu=query_start_loc_cpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
612
            seq_lens=self.input_buffers.seq_lens,
613
            max_seq_len=self.max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
614
615
616
617
618
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
        )

619
        input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
620
        positions = self.input_buffers.positions[:num_tokens_after_padding]
621
622
        mrope_positions = None
        if self.uses_mrope:
623
624
            mrope_positions = self.mrope_states.mrope_positions
            mrope_positions = mrope_positions[:, :num_tokens_after_padding]
Woosuk Kwon's avatar
Woosuk Kwon committed
625
626
627
628
629
        return InputBatch(
            req_ids=req_ids,
            num_reqs=num_reqs,
            idx_mapping=idx_mapping,
            idx_mapping_np=idx_mapping_np,
630
            expanded_idx_mapping=expanded_idx_mapping,
Woosuk Kwon's avatar
Woosuk Kwon committed
631
632
633
            num_scheduled_tokens=num_scheduled_tokens,
            num_tokens=num_tokens,
            num_tokens_after_padding=num_tokens_after_padding,
634
            num_draft_tokens=total_num_draft_tokens,
635
            query_start_loc=query_start_loc,
Woosuk Kwon's avatar
Woosuk Kwon committed
636
            query_start_loc_np=query_start_loc_np,
637
            seq_lens=seq_lens,
Woosuk Kwon's avatar
Woosuk Kwon committed
638
639
            input_ids=input_ids,
            positions=positions,
640
            mrope_positions=mrope_positions,
641
            inputs_embeds=None,
Woosuk Kwon's avatar
Woosuk Kwon committed
642
            attn_metadata=attn_metadata,
643
            slot_mappings=slot_mappings_by_layer,
Woosuk Kwon's avatar
Woosuk Kwon committed
644
            logits_indices=logits_indices,
645
            cu_num_logits=cu_num_logits,
646
            cu_num_logits_np=cu_num_logits_np,
Woosuk Kwon's avatar
Woosuk Kwon committed
647
648
        )

649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
    @torch.inference_mode()
    def get_mm_embeddings(
        self,
        scheduled_encoder_inputs: dict[str, list[int]],
        input_batch: InputBatch,
    ) -> tuple[list[torch.Tensor], torch.Tensor]:
        mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
            scheduled_encoder_inputs
        )
        self.encoder_runner.execute_mm_encoder(self.model, mm_hashes, mm_kwargs)
        mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings(
            input_batch.req_ids,
            input_batch.num_tokens,
            input_batch.num_scheduled_tokens,
            input_batch.query_start_loc_np,
            self.req_states.prefill_len.np[input_batch.idx_mapping_np],
            self.req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np],
        )
        return mm_embeds, is_mm_embed

Woosuk Kwon's avatar
Woosuk Kwon committed
669
670
671
672
673
    def sample(
        self,
        hidden_states: torch.Tensor,
        input_batch: InputBatch,
        grammar_output: GrammarOutput | None,
674
    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
675
        sample_hidden_states = hidden_states[input_batch.logits_indices]
676
        sample_pos = input_batch.positions[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
677
678
679
        logits = self.model.compute_logits(sample_hidden_states)
        if grammar_output is not None:
            # Apply grammar bitmask to the logits in-place.
680
681
682
683
684
685
            self.structured_outputs_worker.apply_grammar_bitmask(
                logits,
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.grammar_bitmask,
            )
686

687
        # Sample tokens and compute logprobs (if needed).
688
689
690
691
        sampler_output = self.sampler(
            logits,
            input_batch.expanded_idx_mapping,
            input_batch.idx_mapping_np,
692
            input_batch.cu_num_logits_np,
693
694
            sample_pos,
        )
695
696
697

        if input_batch.num_draft_tokens == 0:
            # No draft tokens (common case).
698
699
700
            num_sampled = torch.ones(
                input_batch.num_reqs, dtype=torch.int32, device=self.device
            )
701
        else:
702
            # Rejection sampling for spec decoding.
703
704
705
706
707
708
709
710
            input_ids = input_batch.input_ids[input_batch.logits_indices]
            sampled_tokens, num_sampled = rejection_sample(
                sampler_output.sampled_token_ids,
                input_ids,
                input_batch.cu_num_logits,
                self.num_speculative_steps,
            )
            sampler_output.sampled_token_ids = sampled_tokens
711
712
713
714
715
716
717
718
719
720

        # 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(
            num_sampled,
            input_batch.seq_lens,
            input_batch.cu_num_logits,
            input_batch.idx_mapping,
            self.req_states.prefill_len.gpu,
        )
721
        return sampler_output, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
722
723
724
725

    def postprocess(
        self,
        input_batch: InputBatch,
726
727
        sampled_tokens: torch.Tensor,
        num_sampled: torch.Tensor,
728
        num_rejected: torch.Tensor,
729
730
    ) -> None:
        # Update the number of computed tokens.
731
        post_update(
732
            input_batch.idx_mapping,
733
            self.req_states.num_computed_tokens.gpu,
734
            self.req_states.last_sampled_tokens,
735
            self.sampler.penalties_state.output_bin_counts,
736
737
            sampled_tokens,
            num_sampled,
738
            num_rejected,
739
            input_batch.query_start_loc,
Woosuk Kwon's avatar
Woosuk Kwon committed
740
        )
741
742

        # Update the number of computed prefill tokens.
Woosuk Kwon's avatar
Woosuk Kwon committed
743
        idx_mapping_np = input_batch.idx_mapping_np
744
        computed_prefill = self.req_states.num_computed_prefill_tokens
745
746
747
        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
748
749
        )

750
751
752
753
754
755
756
    @torch.inference_mode()
    def propose_draft(
        self,
        input_batch: InputBatch,
        last_hidden_states: torch.Tensor,
        aux_hidden_states: list[torch.Tensor] | None,
        num_sampled: torch.Tensor,
757
        num_rejected: torch.Tensor,
758
759
760
761
762
763
764
    ) -> torch.Tensor:
        assert self.speculator is not None
        draft_tokens = self.speculator.propose(
            input_batch,
            last_hidden_states,
            aux_hidden_states,
            num_sampled,
765
            num_rejected,
766
767
            self.req_states.last_sampled_tokens,
            self.req_states.next_prefill_tokens,
768
769
            self.sampler.sampling_states.temperature.gpu,
            self.sampler.sampling_states.seeds.gpu,
770
771
772
        )
        return draft_tokens

Woosuk Kwon's avatar
Woosuk Kwon committed
773
774
775
776
777
778
    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: SchedulerOutput,
        intermediate_tensors: Any | None = None,
        dummy_run: bool = False,
779
        skip_attn_for_dummy_run: bool = False,
Woosuk Kwon's avatar
Woosuk Kwon committed
780
781
    ) -> ModelRunnerOutput | None:
        assert intermediate_tensors is None
782
783
784
785
786
787
788
789
790
        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.
791
792
                empty_output = self.kv_connector.no_forward(scheduler_output)
                return empty_output
Woosuk Kwon's avatar
Woosuk Kwon committed
793

794
795
796
797
798
799
800
        # Get the CUDA graph size. None means no CUDA graph is used.
        cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
            scheduler_output.total_num_scheduled_tokens,
            scheduler_output.num_scheduled_tokens.values(),
        )
        use_cudagraph, num_tokens_after_padding, num_tokens_across_dp = (
            get_cudagraph_and_dp_padding(
801
                scheduler_output.total_num_scheduled_tokens,
802
803
804
                cudagraph_size,
                self.parallel_config.data_parallel_size,
                self.parallel_config.data_parallel_rank,
805
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
806
        )
807
808
        if num_tokens_after_padding == 0:
            # All DP ranks have zero tokens to run.
809
810
            empty_output = self.kv_connector.no_forward(scheduler_output)
            return empty_output
811
812
813
814
815

        if not dummy_run:
            # Common case.
            # Prepare all the inputs and copy to the input buffers.
            input_batch = self.prepare_inputs(
816
                scheduler_output, num_tokens_after_padding
817
818
819
            )
            if self.lora_config:
                # Activate LoRA adapters.
820
                lora_inputs = self.lora_state.make_lora_inputs(
821
822
823
                    input_batch.req_ids,
                    input_batch.idx_mapping_np,
                    input_batch.num_scheduled_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
824
                )
825
                self._set_active_loras(*lora_inputs)
826
827
828
829
830
831
832
833
834
835
836
837

            if self.supports_mm_inputs:
                # Execute the multimodal encoder.
                mm_embeds, is_mm_embed = self.get_mm_embeddings(
                    scheduler_output.scheduled_encoder_inputs, input_batch
                )
                inputs_embeds = self.encoder_runner.get_inputs_embeds(
                    self.model, input_batch.input_ids, mm_embeds, is_mm_embed
                )
                input_batch.inputs_embeds = inputs_embeds[
                    : input_batch.num_tokens_after_padding
                ]
838
        else:
839
            # No actual tokens to run. A dummy run for DP or memory profiling.
840
841
842
843
844
845
846
            num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
            input_batch = InputBatch.make_dummy(
                num_reqs=num_reqs,
                num_tokens=num_tokens_after_padding,
                input_buffers=self.input_buffers,
                device=self.device,
            )
847
848
849
850
            if self.uses_mrope:
                input_batch.mrope_positions = self.mrope_states.mrope_positions[
                    :, :num_tokens_after_padding
                ]
851
852
853
            if not skip_attn_for_dummy_run:
                self.prepare_dummy_attn_metadata(input_batch)
            # FIXME(woosuk): Fix warmup for LoRA.
Woosuk Kwon's avatar
Woosuk Kwon committed
854
855

        # Run model.
856
        if use_cudagraph:
Woosuk Kwon's avatar
Woosuk Kwon committed
857
858
859
            # Run CUDA graph.
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
860
            self.kv_connector.pre_forward(scheduler_output)
Woosuk Kwon's avatar
Woosuk Kwon committed
861
862
863
864
865
            hidden_states = self.cudagraph_manager.run(
                input_batch.num_tokens_after_padding
            )
        else:
            # Run PyTorch model in eager mode.
866
867
868
            positions = input_batch.positions
            if self.uses_mrope:
                assert input_batch.mrope_positions is not None
869
                positions = input_batch.mrope_positions
Woosuk Kwon's avatar
Woosuk Kwon committed
870
871
872
873
            with set_forward_context(
                input_batch.attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
874
875
                # TODO(woosuk): Support piecewise CUDA graph.
                cudagraph_runtime_mode=CUDAGraphMode.NONE,
Woosuk Kwon's avatar
Woosuk Kwon committed
876
                num_tokens_across_dp=num_tokens_across_dp,
877
                slot_mapping=input_batch.slot_mappings,
Woosuk Kwon's avatar
Woosuk Kwon committed
878
            ):
879
                self.kv_connector.pre_forward(scheduler_output)
Woosuk Kwon's avatar
Woosuk Kwon committed
880
881
                hidden_states = self.model(
                    input_ids=input_batch.input_ids,
882
                    positions=positions,
883
                    inputs_embeds=input_batch.inputs_embeds,
Woosuk Kwon's avatar
Woosuk Kwon committed
884
885
                )

886
887
        kv_connector_output = self.kv_connector.post_forward(scheduler_output)
        self.execute_model_state = hidden_states, input_batch, kv_connector_output
Woosuk Kwon's avatar
Woosuk Kwon committed
888
889
890
891
892
893
894
895
        return None

    @torch.inference_mode()
    def sample_tokens(
        self,
        grammar_output: GrammarOutput | None,
    ) -> AsyncOutput | ModelRunnerOutput:
        assert self.execute_model_state is not None
896
        hidden_states, input_batch, kv_connector_output = self.execute_model_state
Woosuk Kwon's avatar
Woosuk Kwon committed
897
898
        self.execute_model_state = None  # type: ignore

899
        sampler_output, num_sampled, num_rejected = self.sample(
900
            hidden_states, input_batch, grammar_output
Woosuk Kwon's avatar
Woosuk Kwon committed
901
        )
902
903
904
905
906
907
908
909
910
911
        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
            self.model.compute_logits,
            hidden_states,
            input_batch,
            self.req_states.prefill_token_ids.gpu,
            self.req_states.num_computed_tokens.gpu,
            self.req_states.prompt_len,
            self.req_states.prefill_len.np,
            self.req_states.num_computed_prefill_tokens,
        )
912
913
914
915
916
917
918
919

        # 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
920
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
921
            kv_connector_output=kv_connector_output,
922
923
924
925
        )
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,
926
            num_sampled_tokens=num_sampled,
927
928
929
930
931
932
933
934
935
936
            copy_stream=self.output_copy_stream,
            copy_event=self.output_copy_event,
        )

        # 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(
937
            input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
938
        )
939
        if self.do_spec_decode:
940
            draft_tokens = self.propose_draft(
941
942
943
                input_batch,
                hidden_states,
                None,  # aux_hidden_states
944
945
                num_sampled,
                num_rejected,
946
            )
947
            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
948
949
950
951

        if self.use_async_scheduling:
            return async_output
        return async_output.get_output()