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

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

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

logger = init_logger(__name__)


74
class GPUModelRunner(LoRAModelRunnerMixin):
Woosuk Kwon's avatar
Woosuk Kwon committed
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
100
101
102
103
104
    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
105
        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
Woosuk Kwon's avatar
Woosuk Kwon committed
106

107
        # Multimodal
108
109
110
111
112
113
114
115
116
117
118
        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,
            )
119
120
121
122
        self.uses_mrope = self.model_config.uses_mrope
        if self.uses_mrope:
            self.mrope_states = MRopeState(
                max_num_reqs=self.max_num_reqs,
123
                max_num_tokens=self.max_num_tokens,
124
125
126
127
                max_model_len=self.max_model_len,
                device=self.device,
            )

Woosuk Kwon's avatar
Woosuk Kwon committed
128
129
130
131
        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()

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

Woosuk Kwon's avatar
Woosuk Kwon committed
141
142
143
144
        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,
145
            num_speculative_steps=self.num_speculative_steps,
Woosuk Kwon's avatar
Woosuk Kwon committed
146
147
148
149
150
151
152
153
            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,
        )
154
155
156
157
        self.sampler = Sampler(
            max_num_reqs=self.max_num_reqs,
            vocab_size=self.vocab_size,
            device=self.device,
158
            req_states=self.req_states,
159
            logprobs_mode=self.model_config.logprobs_mode,
160
            num_speculative_tokens=self.num_speculative_steps + 1,
161
        )
162
        self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
Woosuk Kwon's avatar
Woosuk Kwon committed
163
164

        # CUDA graphs.
165
166
167
        self.cudagraph_manager = CudaGraphManager(
            self.vllm_config, self.uses_mrope, self.device
        )
168
169
170
171
        # 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,
172
            device=self.device,
173
        )
174
175
        # LoRA-related workers.
        self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
176

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

        # KV Connector if configured.
181
182
        self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR

183
184
        # Pipeline parallelism.
        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
185
186
187
188
189
190
        if self.use_pp:
            self.is_first_pp_rank = get_pp_group().is_first_rank
            self.is_last_pp_rank = get_pp_group().is_last_rank
        else:
            self.is_first_pp_rank = True
            self.is_last_pp_rank = True
191

192
193
194
195
    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

196
197
    @staticmethod
    def get_supported_tasks() -> tuple[str]:
Woosuk Kwon's avatar
Woosuk Kwon committed
198
199
200
201
202
203
204
205
206
207
208
209
210
211
        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(
212
                    self.model, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
213
                )
214
215
            if self.do_spec_decode:
                self.speculator.load_model(self.model)
Woosuk Kwon's avatar
Woosuk Kwon committed
216
217
218
219
        time_after_load = time.perf_counter()

        self.model_memory_usage = m.consumed_memory
        logger.info(
220
221
            "Model loading took %s GiB and %.6f seconds",
            format_gib(m.consumed_memory),
Woosuk Kwon's avatar
Woosuk Kwon committed
222
223
224
            time_after_load - time_before_load,
        )

225
226
227
228
229
230
        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
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
    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(
254
            self.kv_cache_config, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
255
        )
256
257
258
259
260
261
262
        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
263
264

        self.kv_caches: list[torch.Tensor] = []
265
        kv_caches_dict = init_kv_cache(
Woosuk Kwon's avatar
Woosuk Kwon committed
266
267
268
269
270
271
            self.kv_caches,
            self.compilation_config.static_forward_context,
            self.kv_cache_config,
            self.attn_backends,
            self.device,
        )
272
273
        self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

Woosuk Kwon's avatar
Woosuk Kwon committed
274
275
276
277
278
279
280
281
        # 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
        )
282
283
284
        slot_mappings_by_layer = build_slot_mappings_by_layer(
            slot_mappings, self.kv_cache_config
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
285
286
287
288
        attn_metadata = build_attn_metadata(
            attn_metadata_builders=self.attn_metadata_builders,
            num_reqs=input_batch.num_reqs,
            num_tokens=input_batch.num_tokens,
289
290
            query_start_loc_gpu=input_batch.query_start_loc,
            query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np),
291
            max_query_len=input_batch.num_scheduled_tokens.max().item(),
292
            seq_lens=input_batch.seq_lens,
293
            max_seq_len=self.max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
294
295
296
297
298
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
        )
        input_batch.attn_metadata = attn_metadata
299
        input_batch.slot_mappings = slot_mappings_by_layer
Woosuk Kwon's avatar
Woosuk Kwon committed
300
301
302

    @torch.inference_mode()
    def _dummy_run(
303
        self, num_tokens: int, *args, skip_attn: bool = True, **kwargs
304
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
305
        # Create a dummy scheduler output.
Woosuk Kwon's avatar
Woosuk Kwon committed
306
        num_reqs = min(num_tokens, self.max_num_reqs)
307
308
309
310
        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 = {
311
            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
312
313
314
315
316
        }
        dummy_scheduler_output = SchedulerOutput.make_empty()
        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens

317
318
319
        # Disable any use of KVConnector for dummy runs.
        self.kv_connector.set_disabled(True)

320
321
        # For non-first PP ranks, create dummy intermediate_tensors.
        intermediate_tensors = None
322
        if not self.is_first_pp_rank:
323
324
325
326
327
328
            intermediate_tensors = self.model.make_empty_intermediate_tensors(
                batch_size=num_tokens,
                dtype=self.model_config.dtype,
                device=self.device,
            )

329
330
        # Execute the model.
        self.execute_model(
331
332
333
334
            dummy_scheduler_output,
            intermediate_tensors=intermediate_tensors,
            dummy_run=True,
            skip_attn_for_dummy_run=skip_attn,
335
        )
336
        self.kv_connector.set_disabled(False)
337
338

        # Non-last PP ranks don't produce output for sampling.
339
        if not self.is_last_pp_rank:
340
341
            return None, None

342
        assert self.execute_model_state is not None
343
        hidden_states, input_batch, _ = self.execute_model_state
344
        assert hidden_states is not None  # Last PP rank always has hidden_states
345
        sample_hidden_states = hidden_states[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
346
347
348
        return hidden_states, sample_hidden_states

    @torch.inference_mode()
349
    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
350
351
        num_reqs = hidden_states.shape[0]
        logits = self.model.compute_logits(hidden_states)
352
353
354
        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)
355
356
357
358
        dummy_input_ids = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)
        expanded_local_pos = torch.zeros(
            num_reqs, dtype=torch.int32, device=self.device
        )
359
360
361
        # 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.
362
363
364
365
366
367
368
369
370
        self.sampler(
            logits,
            idx_mapping,
            idx_mapping_np,
            idx_mapping_np,
            pos,
            dummy_input_ids,
            expanded_local_pos,
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
371
372
373
374

    @torch.inference_mode()
    def profile_run(self) -> None:
        hidden_states, sample_hidden_states = self._dummy_run(
375
            self.max_num_tokens, skip_attn=True
Woosuk Kwon's avatar
Woosuk Kwon committed
376
        )
377

378
        # Only run sampler on last PP rank (non-last ranks return None).
379
        if self.is_last_pp_rank:
380
381
            assert sample_hidden_states is not None
            self._dummy_sampler_run(sample_hidden_states)
382
383
384
385
386
387
388
389
390
391
392
393

            if self.do_spec_decode:
                num_tokens_across_dp = make_num_tokens_across_dp(
                    self.parallel_config.data_parallel_size, self.max_num_tokens
                )
                self.speculator.run_model(
                    self.max_num_tokens,
                    attn_metadata=None,
                    slot_mappings=None,
                    num_tokens_across_dp=num_tokens_across_dp,
                )

Woosuk Kwon's avatar
Woosuk Kwon committed
394
395
396
397
398
        torch.cuda.synchronize()
        del hidden_states, sample_hidden_states
        gc.collect()

    def reset_mm_cache(self) -> None:
399
400
        if self.supports_mm_inputs:
            self.encoder_runner.reset_mm_cache()
401
402

    def reset_encoder_cache(self) -> None:
403
404
        if self.supports_mm_inputs:
            self.encoder_runner.reset_encoder_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
405
406
407
408
409
410
411
412
413
414
415
416
417
418

    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

419
420
421
422
423
424
425
426
        # TODO (zhanqiu): support CUDA graph for PP.
        if self.use_pp:
            logger.warning_once(
                "Skipping CUDA graph capture because pipeline parallel is "
                "enabled. Pipeline parallel is currently eager-only.",
            )
            return 0

Woosuk Kwon's avatar
Woosuk Kwon committed
427
        start_time = time.perf_counter()
428
        gc.collect()
429
        torch.cuda.empty_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
430
431
432
        start_free_gpu_memory = torch.cuda.mem_get_info()[0]

        with self.maybe_setup_dummy_loras(self.lora_config):
433
434
435
            mrope_positions = None
            if self.uses_mrope:
                mrope_positions = self.mrope_states.mrope_positions
436
437
438
            inputs_embeds = None
            if self.supports_mm_inputs:
                inputs_embeds = self.encoder_runner.inputs_embeds
Woosuk Kwon's avatar
Woosuk Kwon committed
439
440
441
            self.cudagraph_manager.capture(
                model=self.model,
                input_buffers=self.input_buffers,
442
                mrope_positions=mrope_positions,
443
                inputs_embeds=inputs_embeds,
Woosuk Kwon's avatar
Woosuk Kwon committed
444
445
446
447
                block_tables=self.block_tables,
                attn_metadata_builders=self.attn_metadata_builders,
                kv_cache_config=self.kv_cache_config,
            )
448
449
            if self.do_spec_decode:
                self.speculator.capture_model()
Woosuk Kwon's avatar
Woosuk Kwon committed
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469

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

470
    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
471
        finished_req_ids = scheduler_output.finished_req_ids
472
473
474
        preempted_req_ids = scheduler_output.preempted_req_ids
        if preempted_req_ids:
            finished_req_ids = finished_req_ids.union(preempted_req_ids)
475
        for req_id in finished_req_ids:
Woosuk Kwon's avatar
Woosuk Kwon committed
476
            self.req_states.remove_request(req_id)
477
478
            if self.supports_mm_inputs:
                self.encoder_runner.remove_request(req_id)
479
            self.prompt_logprobs_worker.remove_request(req_id)
480
            self.lora_state.remove_request(req_id)
481

482
    def free_states(self, scheduler_output: SchedulerOutput) -> None:
483
484
485
        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
486

487
    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
488
        for new_req_data in scheduler_output.scheduled_new_reqs:
489
490
491
            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
492
            req_id = new_req_data.req_id
493
            prompt_len = len(new_req_data.prompt_token_ids)
Woosuk Kwon's avatar
Woosuk Kwon committed
494
495
            self.req_states.add_request(
                req_id=req_id,
496
                prompt_len=prompt_len,
497
                all_token_ids=new_req_data.prefill_token_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
498
499
500
                num_computed_tokens=new_req_data.num_computed_tokens,
            )
            req_index = self.req_states.req_id_to_index[req_id]
501

502
503
504
            if self.supports_mm_inputs:
                self.encoder_runner.add_request(req_id, new_req_data.mm_features)

505
506
507
508
509
510
            # 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,
511
                    mm_features=new_req_data.mm_features,
512
513
                )

514
515
516
            self.block_tables.append_block_ids(
                req_index, new_req_data.block_ids, overwrite=True
            )
517
518
519
            self.sampler.add_request(
                req_index, prompt_len, new_req_data.sampling_params
            )
520
521
522
            self.prompt_logprobs_worker.add_request(
                req_id, req_index, new_req_data.sampling_params
            )
523
            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
Woosuk Kwon's avatar
Woosuk Kwon committed
524

525
526
        if scheduler_output.scheduled_new_reqs:
            self.req_states.apply_staged_writes()
527
            self.sampler.apply_staged_writes()
528
529
530
531
            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
532
        # Add new blocks for the existing requests.
533
534
        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
535
            if req_new_block_ids is not None:
536
                req_index = self.req_states.req_id_to_index[req_id]
537
538
539
                self.block_tables.append_block_ids(
                    req_index, req_new_block_ids, overwrite=False
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
540
541

    def prepare_inputs(
542
        self, scheduler_output: SchedulerOutput, num_tokens_after_padding: int
Woosuk Kwon's avatar
Woosuk Kwon committed
543
544
545
    ) -> InputBatch:
        num_tokens = scheduler_output.total_num_scheduled_tokens
        assert num_tokens > 0
546
547
        num_tokens_per_req = scheduler_output.num_scheduled_tokens
        num_reqs = len(num_tokens_per_req)
Woosuk Kwon's avatar
Woosuk Kwon committed
548
549
550

        # Decode first, then prefill.
        # batch_idx -> req_id
551
552
553
        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
554

555
556
        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)
557
        idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
Woosuk Kwon's avatar
Woosuk Kwon committed
558

559
        # Get the number of draft tokens for each request.
560
561
        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
        if not draft_tokens:
562
563
564
            # No draft token scheduled (common case).
            total_num_draft_tokens = 0
            total_num_logits = num_reqs
565
            cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
566
567
568
            cu_num_logits = torch.arange(
                num_reqs + 1, device=self.device, dtype=torch.int32
            )
569
            expanded_idx_mapping = idx_mapping
570
571
572
            expanded_local_pos = torch.zeros(
                num_reqs, dtype=torch.int32, device=self.device
            )
573
574
        else:
            num_draft_tokens = np.array(
575
                [len(draft_tokens.get(req_id, ())) for req_id in req_ids],
576
577
578
579
580
                dtype=np.int32,
            )
            total_num_draft_tokens = int(num_draft_tokens.sum())
            total_num_logits = num_reqs + total_num_draft_tokens

581
582
583
584
            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:])
585
            cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
586

587
            max_expand_len = self.num_speculative_steps + 1
588
            expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
589
                idx_mapping, total_num_logits, cu_num_logits, max_expand_len
590
591
            )

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

595
        # Get query_start_loc.
596
597
598
        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])
599
600
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
601
        query_start_loc_np[num_reqs + 1 :] = num_tokens
602
603
        async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)

604
605
606
        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]
607
        max_query_len = num_scheduled_tokens.max().item()
608

609
        # Get prefill tokens.
610
        prepare_prefill_inputs(
611
612
613
            self.input_buffers.input_ids,
            self.req_states.next_prefill_tokens,
            idx_mapping,
614
            query_start_loc,
615
            self.req_states.all_token_ids.gpu,
616
            self.req_states.prefill_len.gpu,
617
            self.req_states.num_computed_tokens.gpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
618
619
        )

620
621
622
        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
623
624
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
625
626
627
628
629
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
        seq_lens = self.input_buffers.seq_lens[:num_reqs]

630
631
632
633
634
635
636
637
638
        # 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,
            )

639
        # Some input token ids are directly read from the last sampled tokens
640
641
        # and draft tokens. Also, get the logits indices to sample tokens from.
        logits_indices = combine_sampled_and_draft_tokens(
642
            self.input_buffers.input_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
643
644
            idx_mapping,
            self.req_states.last_sampled_tokens,
645
            query_start_loc,
646
647
            seq_lens,
            self.req_states.prefill_len.gpu,
648
649
650
            self.req_states.draft_tokens,
            cu_num_logits,
            total_num_logits,
Woosuk Kwon's avatar
Woosuk Kwon committed
651
652
653
654
        )

        # Compute slot mappings: [num_kv_cache_groups, num_tokens]
        slot_mappings = self.block_tables.compute_slot_mappings(
655
656
657
            idx_mapping,
            query_start_loc,
            self.input_buffers.positions[:num_tokens],
Woosuk Kwon's avatar
Woosuk Kwon committed
658
        )
659
660
661
662
        # 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
663
664
665
666
667
668

        # Layer name -> attention metadata.
        attn_metadata = build_attn_metadata(
            attn_metadata_builders=self.attn_metadata_builders,
            num_reqs=num_reqs,
            num_tokens=num_tokens,
669
            query_start_loc_gpu=query_start_loc,
670
            query_start_loc_cpu=query_start_loc_cpu,
671
            max_query_len=max_query_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
672
            seq_lens=self.input_buffers.seq_lens,
673
            max_seq_len=self.max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
674
675
676
677
678
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
        )

679
        input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
680
        positions = self.input_buffers.positions[:num_tokens_after_padding]
681
682
        mrope_positions = None
        if self.uses_mrope:
683
684
            mrope_positions = self.mrope_states.mrope_positions
            mrope_positions = mrope_positions[:, :num_tokens_after_padding]
Woosuk Kwon's avatar
Woosuk Kwon committed
685
686
687
688
689
        return InputBatch(
            req_ids=req_ids,
            num_reqs=num_reqs,
            idx_mapping=idx_mapping,
            idx_mapping_np=idx_mapping_np,
690
            expanded_idx_mapping=expanded_idx_mapping,
691
            expanded_local_pos=expanded_local_pos,
Woosuk Kwon's avatar
Woosuk Kwon committed
692
693
694
            num_scheduled_tokens=num_scheduled_tokens,
            num_tokens=num_tokens,
            num_tokens_after_padding=num_tokens_after_padding,
695
            num_draft_tokens=total_num_draft_tokens,
696
            query_start_loc=query_start_loc,
Woosuk Kwon's avatar
Woosuk Kwon committed
697
            query_start_loc_np=query_start_loc_np,
698
            seq_lens=seq_lens,
Woosuk Kwon's avatar
Woosuk Kwon committed
699
700
            input_ids=input_ids,
            positions=positions,
701
            mrope_positions=mrope_positions,
702
            inputs_embeds=None,
Woosuk Kwon's avatar
Woosuk Kwon committed
703
            attn_metadata=attn_metadata,
704
            slot_mappings=slot_mappings_by_layer,
Woosuk Kwon's avatar
Woosuk Kwon committed
705
            logits_indices=logits_indices,
706
            cu_num_logits=cu_num_logits,
707
            cu_num_logits_np=cu_num_logits_np,
708
            has_structured_output_reqs=scheduler_output.has_structured_output_requests,
Woosuk Kwon's avatar
Woosuk Kwon committed
709
710
        )

711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
    @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
731
732
733
734
735
    def sample(
        self,
        hidden_states: torch.Tensor,
        input_batch: InputBatch,
        grammar_output: GrammarOutput | None,
736
    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
737
        sample_hidden_states = hidden_states[input_batch.logits_indices]
738
        sample_pos = input_batch.positions[input_batch.logits_indices]
739
        input_ids = input_batch.input_ids[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
740
741
742
        logits = self.model.compute_logits(sample_hidden_states)
        if grammar_output is not None:
            # Apply grammar bitmask to the logits in-place.
743
744
745
746
747
748
            self.structured_outputs_worker.apply_grammar_bitmask(
                logits,
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.grammar_bitmask,
            )
749

750
        # Sample tokens and compute logprobs (if needed).
751
752
753
754
        sampler_output = self.sampler(
            logits,
            input_batch.expanded_idx_mapping,
            input_batch.idx_mapping_np,
755
            input_batch.cu_num_logits_np,
756
            sample_pos,
757
758
            input_ids,
            input_batch.expanded_local_pos,
759
        )
760
761
762

        if input_batch.num_draft_tokens == 0:
            # No draft tokens (common case).
763
764
765
            num_sampled = torch.ones(
                input_batch.num_reqs, dtype=torch.int32, device=self.device
            )
766
        else:
767
            # Rejection sampling for spec decoding.
768
769
770
771
772
773
774
            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
775
776
777
778
779
780
781
782
783
784

        # 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,
        )
785
        return sampler_output, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
786
787
788
789

    def postprocess(
        self,
        input_batch: InputBatch,
790
791
        sampled_tokens: torch.Tensor,
        num_sampled: torch.Tensor,
792
        num_rejected: torch.Tensor,
793
794
    ) -> None:
        # Update the number of computed tokens.
795
        post_update(
796
            input_batch.idx_mapping,
797
            self.req_states.num_computed_tokens.gpu,
798
            self.req_states.last_sampled_tokens,
799
            self.sampler.penalties_state.output_bin_counts,
800
801
            sampled_tokens,
            num_sampled,
802
            num_rejected,
803
            input_batch.query_start_loc,
804
805
            self.req_states.all_token_ids.gpu,
            self.req_states.total_len.gpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
806
        )
807
808

        # Update the number of computed prefill tokens.
Woosuk Kwon's avatar
Woosuk Kwon committed
809
        idx_mapping_np = input_batch.idx_mapping_np
810
        computed_prefill = self.req_states.num_computed_prefill_tokens
811
812
813
        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
814
815
        )

816
817
818
819
820
821
822
    @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,
823
        num_rejected: torch.Tensor,
824
825
826
827
828
829
830
    ) -> torch.Tensor:
        assert self.speculator is not None
        draft_tokens = self.speculator.propose(
            input_batch,
            last_hidden_states,
            aux_hidden_states,
            num_sampled,
831
            num_rejected,
832
833
            self.req_states.last_sampled_tokens,
            self.req_states.next_prefill_tokens,
834
835
            self.sampler.sampling_states.temperature.gpu,
            self.sampler.sampling_states.seeds.gpu,
836
837
838
        )
        return draft_tokens

Woosuk Kwon's avatar
Woosuk Kwon committed
839
840
841
842
    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: SchedulerOutput,
843
        intermediate_tensors: IntermediateTensors | None = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
844
        dummy_run: bool = False,
845
        skip_attn_for_dummy_run: bool = False,
846
    ) -> ModelRunnerOutput | IntermediateTensors | None:
847
848
849
850
851
852
853
854
855
        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.
856
857
                empty_output = self.kv_connector.no_forward(scheduler_output)
                return empty_output
Woosuk Kwon's avatar
Woosuk Kwon committed
858

859
860
861
862
863
864
865
        # 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(
866
                scheduler_output.total_num_scheduled_tokens,
867
868
869
                cudagraph_size,
                self.parallel_config.data_parallel_size,
                self.parallel_config.data_parallel_rank,
870
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
871
        )
872
873
        if num_tokens_after_padding == 0:
            # All DP ranks have zero tokens to run.
874
875
            empty_output = self.kv_connector.no_forward(scheduler_output)
            return empty_output
876
877
878
879
880

        if not dummy_run:
            # Common case.
            # Prepare all the inputs and copy to the input buffers.
            input_batch = self.prepare_inputs(
881
                scheduler_output, num_tokens_after_padding
882
883
884
            )
            if self.lora_config:
                # Activate LoRA adapters.
885
                lora_inputs = self.lora_state.make_lora_inputs(
886
887
888
                    input_batch.req_ids,
                    input_batch.idx_mapping_np,
                    input_batch.num_scheduled_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
889
                )
890
                self._set_active_loras(*lora_inputs)
891

892
            # Only first PP rank prepares multimodal embeddings.
893
            if self.supports_mm_inputs and self.is_first_pp_rank:
894
895
896
897
898
899
900
901
902
                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
                ]
903
        else:
904
            # No actual tokens to run. A dummy run for DP or memory profiling.
905
906
907
908
909
910
911
            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,
            )
912
913
914
915
            if self.uses_mrope:
                input_batch.mrope_positions = self.mrope_states.mrope_positions[
                    :, :num_tokens_after_padding
                ]
916
917
918
            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
919
920

        # Run model.
921
        if use_cudagraph:
Woosuk Kwon's avatar
Woosuk Kwon committed
922
923
924
            # 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.
925
            self.kv_connector.pre_forward(scheduler_output)
Woosuk Kwon's avatar
Woosuk Kwon committed
926
927
928
929
930
            hidden_states = self.cudagraph_manager.run(
                input_batch.num_tokens_after_padding
            )
        else:
            # Run PyTorch model in eager mode.
931
932
933
            positions = input_batch.positions
            if self.uses_mrope:
                assert input_batch.mrope_positions is not None
934
                positions = input_batch.mrope_positions
935

936
937
938
939
940
941
942
943
944
            if self.is_first_pp_rank:
                input_ids = input_batch.input_ids
                inputs_embeds = input_batch.inputs_embeds
                assert intermediate_tensors is None
            else:
                input_ids = None
                inputs_embeds = None
                assert intermediate_tensors is not None

Woosuk Kwon's avatar
Woosuk Kwon committed
945
946
947
948
            with set_forward_context(
                input_batch.attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
949
950
                # TODO(woosuk): Support piecewise CUDA graph.
                cudagraph_runtime_mode=CUDAGraphMode.NONE,
Woosuk Kwon's avatar
Woosuk Kwon committed
951
                num_tokens_across_dp=num_tokens_across_dp,
952
                slot_mapping=input_batch.slot_mappings,
Woosuk Kwon's avatar
Woosuk Kwon committed
953
            ):
954
                self.kv_connector.pre_forward(scheduler_output)
955
956
957
958
959
960
                hidden_states = self.model(
                    input_ids=input_ids,
                    positions=positions,
                    inputs_embeds=inputs_embeds,
                    intermediate_tensors=intermediate_tensors,
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
961

962
        kv_connector_output = self.kv_connector.post_forward(scheduler_output)
963

964
        if not self.is_last_pp_rank:
965
966
967
968
969
970
971
972
973
            # Non-last PP rank: return IntermediateTensors for sending.
            assert isinstance(hidden_states, IntermediateTensors)
            hidden_states.kv_connector_output = kv_connector_output
            self.execute_model_state = (None, input_batch, kv_connector_output)
            return hidden_states

        assert isinstance(hidden_states, torch.Tensor)
        # Last rank (or no PP): hidden_states is a tensor for sampling.
        self.execute_model_state = (hidden_states, input_batch, kv_connector_output)
Woosuk Kwon's avatar
Woosuk Kwon committed
974
975
976
977
        return None

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

984
985
986
        # Non-last PP rank: hidden_states is None because this rank produced
        # IntermediateTensors instead of final hidden states. Receive the
        # sampled tokens broadcast by the last rank and update local state.
987
        if not self.is_last_pp_rank:
988
            received = pp_receive(
989
                input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
990
            )
991
992
993
            assert received is not None
            sampled, num_sampled, num_rejected = received
            self.postprocess(input_batch, sampled, num_sampled, num_rejected)
994
995
996
            return None

        # Last rank: sample tokens
997
        sampler_output, num_sampled, num_rejected = self.sample(
998
            hidden_states, input_batch, grammar_output
Woosuk Kwon's avatar
Woosuk Kwon committed
999
        )
1000
1001
1002

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

1005
1006
1007
1008
        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
            self.model.compute_logits,
            hidden_states,
            input_batch,
1009
            self.req_states.all_token_ids.gpu,
1010
            self.req_states.num_computed_tokens.gpu,
1011
            self.req_states.prompt_len.np,
1012
1013
1014
            self.req_states.prefill_len.np,
            self.req_states.num_computed_prefill_tokens,
        )
1015
1016
1017
1018
1019
1020
1021
1022

        # 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
1023
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
1024
            kv_connector_output=kv_connector_output,
1025
1026
1027
1028
        )
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,
1029
            num_sampled_tokens=num_sampled,
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
            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(
1040
            input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
1041
        )
1042
        if self.do_spec_decode:
1043
            draft_tokens = self.propose_draft(
1044
1045
1046
                input_batch,
                hidden_states,
                None,  # aux_hidden_states
1047
1048
                num_sampled,
                num_rejected,
1049
            )
1050
            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1051
            self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
1052
1053
1054
1055

        if self.use_async_scheduling:
            return async_output
        return async_output.get_output()
1056
1057
1058

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