model_runner.py 37.4 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
23
24
25
26
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
from vllm.v1.outputs import (
    EMPTY_MODEL_RUNNER_OUTPUT,
    ModelRunnerOutput,
)
27
from vllm.v1.worker.gpu.async_utils import AsyncOutput
Woosuk Kwon's avatar
Woosuk Kwon committed
28
29
30
31
32
33
34
from vllm.v1.worker.gpu.attn_utils import (
    build_attn_metadata,
    get_kv_cache_spec,
    init_attn_backend,
    init_kv_cache,
)
from vllm.v1.worker.gpu.block_table import BlockTables
35
from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
Woosuk Kwon's avatar
Woosuk Kwon committed
36
from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
37
from vllm.v1.worker.gpu.dp_utils import (
38
    get_cudagraph_and_dp_padding,
39
40
    make_num_tokens_across_dp,
)
Woosuk Kwon's avatar
Woosuk Kwon committed
41
42
43
from vllm.v1.worker.gpu.input_batch import (
    InputBatch,
    InputBuffers,
44
    combine_sampled_and_draft_tokens,
45
    expand_idx_mapping,
46
    get_num_sampled_and_rejected,
47
    post_update,
48
49
    prepare_pos_seq_lens,
    prepare_prefill_inputs,
Woosuk Kwon's avatar
Woosuk Kwon committed
50
)
51
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
52
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
53
from vllm.v1.worker.gpu.sample.output import SamplerOutput
54
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
55
from vllm.v1.worker.gpu.sample.sampler import Sampler
56
from vllm.v1.worker.gpu.spec_decode import init_speculator
57
from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
58
from vllm.v1.worker.gpu.states import RequestState
59
from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
Woosuk Kwon's avatar
Woosuk Kwon committed
60
61
62
63
64
65
66
67
68
69
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
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

logger = init_logger(__name__)


class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
    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
97
        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
Woosuk Kwon's avatar
Woosuk Kwon committed
98

99
        # Multimodal
100
101
102
103
104
105
106
107
108
109
110
        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,
            )
111
112
113
114
        self.uses_mrope = self.model_config.uses_mrope
        if self.uses_mrope:
            self.mrope_states = MRopeState(
                max_num_reqs=self.max_num_reqs,
115
                max_num_tokens=self.max_num_tokens,
116
117
118
119
                max_model_len=self.max_model_len,
                device=self.device,
            )

Woosuk Kwon's avatar
Woosuk Kwon committed
120
121
122
123
124
125
126
127
128
129
        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()
        if self.use_async_scheduling:
            self.input_prep_event = torch.cuda.Event()
            self.structured_outputs_event = torch.cuda.Event()
        else:
            self.input_prep_event = None
            self.structured_outputs_event = None

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

Woosuk Kwon's avatar
Woosuk Kwon committed
139
140
141
142
        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,
143
            num_speculative_steps=self.num_speculative_steps,
Woosuk Kwon's avatar
Woosuk Kwon committed
144
145
146
147
148
149
150
151
            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,
        )
152
153
154
155
156
157
        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,
        )
158
        self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
Woosuk Kwon's avatar
Woosuk Kwon committed
159
160

        # CUDA graphs.
161
162
163
        self.cudagraph_manager = CudaGraphManager(
            self.vllm_config, self.uses_mrope, self.device
        )
164
165
166
167
168
169
170
171
172
173
        # 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,
        )

        # Buffers for CPU-to-GPU copies.
        self.tmp_idx_mapping = UvaBufferPool(self.max_num_reqs, torch.int32)
        self.tmp_cu_num_logits = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
        self.tmp_query_start_loc = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
Woosuk Kwon's avatar
Woosuk Kwon committed
174

175
176
177
178
    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

Woosuk Kwon's avatar
Woosuk Kwon committed
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
    def get_supported_tasks(self) -> tuple[str]:
        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(
                    self.model,
                    self.vllm_config,
                    self.device,
                )
198
199
            if self.do_spec_decode:
                self.speculator.load_model(self.model)
Woosuk Kwon's avatar
Woosuk Kwon committed
200
201
202
203
        time_after_load = time.perf_counter()

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

209
210
211
212
213
214
        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
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    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(
            self.kv_cache_config,
            self.vllm_config,
            self.device,
        )
242
243
244
245
246
247
248
        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
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269

        self.kv_caches: list[torch.Tensor] = []
        init_kv_cache(
            self.kv_caches,
            self.compilation_config.static_forward_context,
            self.kv_cache_config,
            self.attn_backends,
            self.device,
        )
        # 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
        )
        attn_metadata = build_attn_metadata(
            attn_metadata_builders=self.attn_metadata_builders,
            num_reqs=input_batch.num_reqs,
            num_tokens=input_batch.num_tokens,
270
271
272
            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,
273
            max_seq_len=self.max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
274
275
276
277
278
279
280
281
282
283
284
285
286
287
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
        )
        input_batch.attn_metadata = attn_metadata

    @torch.inference_mode()
    def _dummy_run(
        self,
        num_tokens: int,
        *args,
        skip_attn: bool = True,
        **kwargs,
    ) -> tuple[torch.Tensor, torch.Tensor]:
288
        # Create a dummy scheduler output.
Woosuk Kwon's avatar
Woosuk Kwon committed
289
        num_reqs = min(num_tokens, self.max_num_reqs)
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
        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 = {
            f"_dummy_req_{i}": num_tokens_per_request[i] for i in range(num_reqs)
        }
        dummy_scheduler_output = SchedulerOutput.make_empty()
        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens

        # Execute the model.
        self.execute_model(
            dummy_scheduler_output, dummy_run=True, skip_attn_for_dummy_run=skip_attn
        )
        assert self.execute_model_state is not None
        hidden_states, input_batch = self.execute_model_state
        sample_hidden_states = hidden_states[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
307
308
309
310
311
312
313
314
315
        return hidden_states, sample_hidden_states

    @torch.inference_mode()
    def _dummy_sampler_run(
        self,
        hidden_states: torch.Tensor,
    ) -> None:
        num_reqs = hidden_states.shape[0]
        logits = self.model.compute_logits(hidden_states)
316
317
318
319
320
321
322
        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.
        self.sampler(logits, idx_mapping, idx_mapping_np, pos)
Woosuk Kwon's avatar
Woosuk Kwon committed
323
324
325
326
327
328
329
330

    @torch.inference_mode()
    def profile_run(self) -> None:
        hidden_states, sample_hidden_states = self._dummy_run(
            self.max_num_tokens,
            skip_attn=True,
        )
        self._dummy_sampler_run(sample_hidden_states)
331
        if self.do_spec_decode:
332
            num_tokens_across_dp = make_num_tokens_across_dp(
333
                self.parallel_config.data_parallel_size, self.max_num_tokens
334
335
336
337
338
339
            )
            self.speculator.run_model(
                self.max_num_tokens,
                attn_metadata=None,
                num_tokens_across_dp=num_tokens_across_dp,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
        torch.cuda.synchronize()
        del hidden_states, sample_hidden_states
        gc.collect()

    def reset_mm_cache(self) -> None:
        pass

    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
        if scheduler_output.preempted_req_ids is not None:
            for req_id in scheduler_output.preempted_req_ids:
                self.req_states.remove_request(req_id)
407
408
                if self.supports_mm_inputs:
                    self.encoder_runner.remove_request(req_id)
409
                self.prompt_logprobs_worker.remove_request(req_id)
Woosuk Kwon's avatar
Woosuk Kwon committed
410
411
        for req_id in scheduler_output.finished_req_ids:
            self.req_states.remove_request(req_id)
412
413
            if self.supports_mm_inputs:
                self.encoder_runner.remove_request(req_id)
414
            self.prompt_logprobs_worker.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
435
                prefill_token_ids=new_req_data.prefill_token_ids,
                num_computed_tokens=new_req_data.num_computed_tokens,
                lora_request=new_req_data.lora_request,
            )
            req_index = self.req_states.req_id_to_index[req_id]
436

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

440
441
442
443
444
445
            # 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,
446
                    mm_features=new_req_data.mm_features,
447
448
                )

449
450
451
            self.block_tables.append_block_ids(
                req_index, new_req_data.block_ids, overwrite=True
            )
452
453
454
            self.sampler.add_request(
                req_index, prompt_len, new_req_data.sampling_params
            )
455
456
457
            self.prompt_logprobs_worker.add_request(
                req_id, req_index, new_req_data.sampling_params
            )
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
471
472
473
474
475
        # Add new blocks for the existing requests.
        cached_reqs = scheduler_output.scheduled_cached_reqs
        for i, req_id in enumerate(cached_reqs.req_ids):
            req_index = self.req_states.req_id_to_index[req_id]
            req_new_block_ids = cached_reqs.new_block_ids[i]
            if req_new_block_ids is not None:
476
477
478
                self.block_tables.append_block_ids(
                    req_index, req_new_block_ids, overwrite=False
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
479
480
481
482
483
484
485
486
487
488
489
490
491

    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(
492
493
            scheduler_output.num_scheduled_tokens.keys(),
            key=lambda k: scheduler_output.num_scheduled_tokens[k],
Woosuk Kwon's avatar
Woosuk Kwon committed
494
495
496
497
498
499
500
501
        )
        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
        ]
502
503
        idx_mapping_np = np.array(idx_mapping_list, dtype=np.int32)
        idx_mapping = self.tmp_idx_mapping.copy_to_gpu(idx_mapping_np)
Woosuk Kwon's avatar
Woosuk Kwon committed
504

505
506
507
508
509
        # Get the number of draft tokens for each request.
        if not scheduler_output.scheduled_spec_decode_tokens:
            # 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
517
518
519
520
521
522
523
524
525
526
        else:
            draft_tokens = scheduler_output.scheduled_spec_decode_tokens
            num_draft_tokens = np.array(
                [
                    len(draft_tokens[req_id]) if req_id in draft_tokens else 0
                    for req_id in req_ids
                ],
                dtype=np.int32,
            )
            total_num_draft_tokens = int(num_draft_tokens.sum())
            total_num_logits = num_reqs + total_num_draft_tokens

527
528
529
530
531
532
533
534
535
536
537
            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:])
            cu_num_logits = self.tmp_cu_num_logits.copy_to_gpu(cu_num_logits_np)

            expanded_idx_mapping = expand_idx_mapping(
                idx_mapping,
                total_num_logits,
                cu_num_logits,
                max_expand_len=self.num_speculative_steps + 1,
538
539
            )

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

543
        # Get query_start_loc.
544
545
546
        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])
547
548
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
549
550
551
552
553
554
555
556
        query_start_loc_np[num_reqs + 1 :] = num_tokens
        self.tmp_query_start_loc.copy_to_gpu(
            query_start_loc_np,
            out=self.input_buffers.query_start_loc,
        )
        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]
557

558
        # Get prefill tokens.
559
        prepare_prefill_inputs(
560
561
562
            self.input_buffers.input_ids,
            self.req_states.next_prefill_tokens,
            idx_mapping,
563
            query_start_loc,
564
565
            self.req_states.prefill_token_ids.gpu,
            self.req_states.prefill_len.gpu,
566
            self.req_states.num_computed_tokens.gpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
567
568
        )

569
570
571
        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
572
573
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
574
575
576
577
578
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
        seq_lens = self.input_buffers.seq_lens[:num_reqs]

579
580
581
582
583
584
585
586
587
        # 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,
            )

588
        # Some input token ids are directly read from the last sampled tokens
589
590
        # and draft tokens. Also, get the logits indices to sample tokens from.
        logits_indices = combine_sampled_and_draft_tokens(
591
            self.input_buffers.input_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
592
593
            idx_mapping,
            self.req_states.last_sampled_tokens,
594
            query_start_loc,
595
596
            seq_lens,
            self.req_states.prefill_len.gpu,
597
598
599
            self.req_states.draft_tokens,
            cu_num_logits,
            total_num_logits,
Woosuk Kwon's avatar
Woosuk Kwon committed
600
601
602
603
        )

        # Compute slot mappings: [num_kv_cache_groups, num_tokens]
        slot_mappings = self.block_tables.compute_slot_mappings(
604
605
606
            idx_mapping,
            query_start_loc,
            self.input_buffers.positions[:num_tokens],
Woosuk Kwon's avatar
Woosuk Kwon committed
607
608
609
610
611
612
613
        )

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

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

653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
    @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
673
674
675
676
677
    def sample(
        self,
        hidden_states: torch.Tensor,
        input_batch: InputBatch,
        grammar_output: GrammarOutput | None,
678
    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
679
        sample_hidden_states = hidden_states[input_batch.logits_indices]
680
        sample_pos = input_batch.positions[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
681
682
683
        logits = self.model.compute_logits(sample_hidden_states)
        if grammar_output is not None:
            # Apply grammar bitmask to the logits in-place.
684
685
686
687
688
689
            self.structured_outputs_worker.apply_grammar_bitmask(
                logits,
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.grammar_bitmask,
            )
690

691
        # Sample tokens and compute logprobs (if needed).
692
693
694
695
696
697
        sampler_output = self.sampler(
            logits,
            input_batch.expanded_idx_mapping,
            input_batch.idx_mapping_np,
            sample_pos,
        )
698
699
700

        if input_batch.num_draft_tokens == 0:
            # No draft tokens (common case).
701
702
703
            num_sampled = torch.ones(
                input_batch.num_reqs, dtype=torch.int32, device=self.device
            )
704
        else:
705
            # Rejection sampling for spec decoding.
706
707
708
709
710
711
712
713
            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
714
715
716
717
718
719
720
721
722
723

        # 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,
        )
724
        return sampler_output, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
725
726
727
728

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

        # Update the number of computed prefill tokens.
Woosuk Kwon's avatar
Woosuk Kwon committed
746
        idx_mapping_np = input_batch.idx_mapping_np
747
748
749
750
751
        computed_prefill = self.req_states.num_computed_prefill_tokens
        # TODO(woosuk): Simplify this.
        computed_prefill[idx_mapping_np] = np.minimum(
            computed_prefill[idx_mapping_np] + input_batch.num_scheduled_tokens,
            self.req_states.prefill_len.np[idx_mapping_np],
Woosuk Kwon's avatar
Woosuk Kwon committed
752
753
        )

754
755
756
757
758
759
760
    @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,
761
        num_rejected: torch.Tensor,
762
763
764
765
766
767
768
    ) -> torch.Tensor:
        assert self.speculator is not None
        draft_tokens = self.speculator.propose(
            input_batch,
            last_hidden_states,
            aux_hidden_states,
            num_sampled,
769
            num_rejected,
770
771
            self.req_states.last_sampled_tokens,
            self.req_states.next_prefill_tokens,
772
773
            self.sampler.sampling_states.temperature.gpu,
            self.sampler.sampling_states.seeds.gpu,
774
775
776
        )
        return draft_tokens

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

797
798
799
800
801
802
803
        # 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(
804
                scheduler_output.total_num_scheduled_tokens,
805
806
807
                cudagraph_size,
                self.parallel_config.data_parallel_size,
                self.parallel_config.data_parallel_rank,
808
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
809
        )
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
        if num_tokens_after_padding == 0:
            # All DP ranks have zero tokens to run.
            return EMPTY_MODEL_RUNNER_OUTPUT

        if not dummy_run:
            # Common case.
            # Prepare all the inputs and copy to the input buffers.
            input_batch = self.prepare_inputs(
                scheduler_output,
                num_tokens_after_padding,
            )
            if self.lora_config:
                # Activate LoRA adapters.
                lora_inputs = self.req_states.make_lora_inputs(
                    input_batch.req_ids,
                    input_batch.idx_mapping_np,
                    input_batch.num_scheduled_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
827
                )
828
                self._set_active_loras(*lora_inputs)
829
830
831
832
833
834
835
836
837
838
839
840

            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
                ]
841
        else:
842
            # No actual tokens to run. A dummy run for DP or memory profiling.
843
844
845
846
847
848
849
            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,
            )
850
851
852
853
            if self.uses_mrope:
                input_batch.mrope_positions = self.mrope_states.mrope_positions[
                    :, :num_tokens_after_padding
                ]
854
855
856
            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
857
858

        # Run model.
859
        if use_cudagraph:
Woosuk Kwon's avatar
Woosuk Kwon committed
860
861
862
863
864
865
866
867
            # 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.
            hidden_states = self.cudagraph_manager.run(
                input_batch.num_tokens_after_padding
            )
        else:
            # Run PyTorch model in eager mode.
868
869
870
            positions = input_batch.positions
            if self.uses_mrope:
                assert input_batch.mrope_positions is not None
871
                positions = input_batch.mrope_positions
Woosuk Kwon's avatar
Woosuk Kwon committed
872
873
874
875
            with set_forward_context(
                input_batch.attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
876
877
                # TODO(woosuk): Support piecewise CUDA graph.
                cudagraph_runtime_mode=CUDAGraphMode.NONE,
Woosuk Kwon's avatar
Woosuk Kwon committed
878
879
880
881
                num_tokens_across_dp=num_tokens_across_dp,
            ):
                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
        self.execute_model_state = hidden_states, input_batch
Woosuk Kwon's avatar
Woosuk Kwon committed
887
888
889
890
891
892
893
894
        return None

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

898
        sampler_output, num_sampled, num_rejected = self.sample(
899
            hidden_states, input_batch, grammar_output
Woosuk Kwon's avatar
Woosuk Kwon committed
900
        )
901
902
903
904
905
906
907
908
909
910
        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,
        )
911
912
913
914
915
916
917
918

        # 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
919
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
920
921
922
923
        )
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,
924
            num_sampled_tokens=num_sampled,
925
926
927
928
929
930
931
932
933
934
            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(
935
            input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
936
        )
937
        if self.do_spec_decode:
938
            draft_tokens = self.propose_draft(
939
940
941
                input_batch,
                hidden_states,
                None,  # aux_hidden_states
942
943
                num_sampled,
                num_rejected,
944
            )
945
            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
946
947
948
949

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