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
        if self.supports_mm_inputs:
            self.encoder_runner.reset_mm_cache()
344
345

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

    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()
363
        gc.collect()
364
        torch.cuda.empty_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
365
366
367
        start_free_gpu_memory = torch.cuda.mem_get_info()[0]

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

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

405
    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
406
        finished_req_ids = scheduler_output.finished_req_ids
407
408
409
        preempted_req_ids = scheduler_output.preempted_req_ids
        if preempted_req_ids:
            finished_req_ids = finished_req_ids.union(preempted_req_ids)
410
        for req_id in finished_req_ids:
Woosuk Kwon's avatar
Woosuk Kwon committed
411
            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
            self.lora_state.remove_request(req_id)
416

417
    def free_states(self, scheduler_output: SchedulerOutput) -> None:
418
419
420
        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
421

422
    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
423
        for new_req_data in scheduler_output.scheduled_new_reqs:
424
425
426
            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
427
            req_id = new_req_data.req_id
428
            prompt_len = len(new_req_data.prompt_token_ids)
Woosuk Kwon's avatar
Woosuk Kwon committed
429
430
            self.req_states.add_request(
                req_id=req_id,
431
                prompt_len=prompt_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
432
433
434
435
                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]
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
            )
458
            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
Woosuk Kwon's avatar
Woosuk Kwon committed
459

460
461
462
463
464
465
466
467
468
469
470
        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
471
        # Add new blocks for the existing requests.
472
473
        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
474
            if req_new_block_ids is not None:
475
                req_index = self.req_states.req_id_to_index[req_id]
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

    def prepare_inputs(
481
        self, scheduler_output: SchedulerOutput, num_tokens_after_padding: int
Woosuk Kwon's avatar
Woosuk Kwon committed
482
483
484
    ) -> InputBatch:
        num_tokens = scheduler_output.total_num_scheduled_tokens
        assert num_tokens > 0
485
486
        num_tokens_per_req = scheduler_output.num_scheduled_tokens
        num_reqs = len(num_tokens_per_req)
Woosuk Kwon's avatar
Woosuk Kwon committed
487
488
489

        # Decode first, then prefill.
        # batch_idx -> req_id
490
491
492
        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
493

494
495
        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)
496
        idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
Woosuk Kwon's avatar
Woosuk Kwon committed
497

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

517
518
519
520
            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:])
521
            cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
522

523
            max_expand_len = self.num_speculative_steps + 1
524
            expanded_idx_mapping = expand_idx_mapping(
525
                idx_mapping, total_num_logits, cu_num_logits, max_expand_len
526
527
            )

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

531
        # Get query_start_loc.
532
533
534
        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])
535
536
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
537
        query_start_loc_np[num_reqs + 1 :] = num_tokens
538
539
        async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)

540
541
542
        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]
543

544
        # Get prefill tokens.
545
        prepare_prefill_inputs(
546
547
548
            self.input_buffers.input_ids,
            self.req_states.next_prefill_tokens,
            idx_mapping,
549
            query_start_loc,
550
551
            self.req_states.prefill_token_ids.gpu,
            self.req_states.prefill_len.gpu,
552
            self.req_states.num_computed_tokens.gpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
553
554
        )

555
556
557
        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
558
559
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
560
561
562
563
564
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
        seq_lens = self.input_buffers.seq_lens[:num_reqs]

565
566
567
568
569
570
571
572
573
        # 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,
            )

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

        # Compute slot mappings: [num_kv_cache_groups, num_tokens]
        slot_mappings = self.block_tables.compute_slot_mappings(
590
591
592
            idx_mapping,
            query_start_loc,
            self.input_buffers.positions[:num_tokens],
Woosuk Kwon's avatar
Woosuk Kwon committed
593
        )
594
595
596
597
        # 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
598
599
600
601
602
603

        # Layer name -> attention metadata.
        attn_metadata = build_attn_metadata(
            attn_metadata_builders=self.attn_metadata_builders,
            num_reqs=num_reqs,
            num_tokens=num_tokens,
604
            query_start_loc_gpu=query_start_loc,
605
            query_start_loc_cpu=query_start_loc_cpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
606
            seq_lens=self.input_buffers.seq_lens,
607
            max_seq_len=self.max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
608
609
610
611
612
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
        )

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

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

681
        # Sample tokens and compute logprobs (if needed).
682
683
684
685
        sampler_output = self.sampler(
            logits,
            input_batch.expanded_idx_mapping,
            input_batch.idx_mapping_np,
686
            input_batch.cu_num_logits_np,
687
688
            sample_pos,
        )
689
690
691

        if input_batch.num_draft_tokens == 0:
            # No draft tokens (common case).
692
693
694
            num_sampled = torch.ones(
                input_batch.num_reqs, dtype=torch.int32, device=self.device
            )
695
        else:
696
            # Rejection sampling for spec decoding.
697
698
699
700
701
702
703
704
            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
705
706
707
708
709
710
711
712
713
714

        # 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,
        )
715
        return sampler_output, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
716
717
718
719

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

        # Update the number of computed prefill tokens.
Woosuk Kwon's avatar
Woosuk Kwon committed
737
        idx_mapping_np = input_batch.idx_mapping_np
738
        computed_prefill = self.req_states.num_computed_prefill_tokens
739
740
741
        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
742
743
        )

744
745
746
747
748
749
750
    @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,
751
        num_rejected: torch.Tensor,
752
753
754
755
756
757
758
    ) -> torch.Tensor:
        assert self.speculator is not None
        draft_tokens = self.speculator.propose(
            input_batch,
            last_hidden_states,
            aux_hidden_states,
            num_sampled,
759
            num_rejected,
760
761
            self.req_states.last_sampled_tokens,
            self.req_states.next_prefill_tokens,
762
763
            self.sampler.sampling_states.temperature.gpu,
            self.sampler.sampling_states.seeds.gpu,
764
765
766
        )
        return draft_tokens

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

788
789
790
791
792
793
794
        # 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(
795
                scheduler_output.total_num_scheduled_tokens,
796
797
798
                cudagraph_size,
                self.parallel_config.data_parallel_size,
                self.parallel_config.data_parallel_rank,
799
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
800
        )
801
802
        if num_tokens_after_padding == 0:
            # All DP ranks have zero tokens to run.
803
804
            empty_output = self.kv_connector.no_forward(scheduler_output)
            return empty_output
805
806
807
808
809

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

            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
                ]
832
        else:
833
            # No actual tokens to run. A dummy run for DP or memory profiling.
834
835
836
837
838
839
840
            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,
            )
841
842
843
844
            if self.uses_mrope:
                input_batch.mrope_positions = self.mrope_states.mrope_positions[
                    :, :num_tokens_after_padding
                ]
845
846
847
            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
848
849

        # Run model.
850
        if use_cudagraph:
Woosuk Kwon's avatar
Woosuk Kwon committed
851
852
853
            # 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.
854
            self.kv_connector.pre_forward(scheduler_output)
Woosuk Kwon's avatar
Woosuk Kwon committed
855
856
857
858
859
            hidden_states = self.cudagraph_manager.run(
                input_batch.num_tokens_after_padding
            )
        else:
            # Run PyTorch model in eager mode.
860
861
862
            positions = input_batch.positions
            if self.uses_mrope:
                assert input_batch.mrope_positions is not None
863
                positions = input_batch.mrope_positions
Woosuk Kwon's avatar
Woosuk Kwon committed
864
865
866
867
            with set_forward_context(
                input_batch.attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
868
869
                # TODO(woosuk): Support piecewise CUDA graph.
                cudagraph_runtime_mode=CUDAGraphMode.NONE,
Woosuk Kwon's avatar
Woosuk Kwon committed
870
                num_tokens_across_dp=num_tokens_across_dp,
871
                slot_mapping=input_batch.slot_mappings,
Woosuk Kwon's avatar
Woosuk Kwon committed
872
            ):
873
                self.kv_connector.pre_forward(scheduler_output)
Woosuk Kwon's avatar
Woosuk Kwon committed
874
875
                hidden_states = self.model(
                    input_ids=input_batch.input_ids,
876
                    positions=positions,
877
                    inputs_embeds=input_batch.inputs_embeds,
Woosuk Kwon's avatar
Woosuk Kwon committed
878
879
                )

880
881
        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
882
883
884
885
        return None

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

892
        sampler_output, num_sampled, num_rejected = self.sample(
893
            hidden_states, input_batch, grammar_output
Woosuk Kwon's avatar
Woosuk Kwon committed
894
        )
895
896
897
898
899
900
901
902
903
904
        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,
        )
905
906
907
908
909
910
911
912

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

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