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

logger = init_logger(__name__)


77
class GPUModelRunner(LoRAModelRunnerMixin):
Woosuk Kwon's avatar
Woosuk Kwon committed
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
105
106
107
    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
108
        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
Woosuk Kwon's avatar
Woosuk Kwon committed
109

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

Woosuk Kwon's avatar
Woosuk Kwon committed
131
132
133
134
        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()

135
136
137
        if self.speculative_config is not None:
            self.do_spec_decode = True
            self.num_speculative_steps = self.speculative_config.num_speculative_tokens
138
            self.speculator = init_speculator(self.vllm_config, self.device)
139
140
141
        else:
            self.do_spec_decode = False
            self.num_speculative_steps = 0
142
            self.speculator = None
Woosuk Kwon's avatar
Woosuk Kwon committed
143
144
145
146
        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,
147
            num_speculative_steps=self.num_speculative_steps,
Woosuk Kwon's avatar
Woosuk Kwon committed
148
149
150
151
152
153
154
155
            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,
        )
156
157
158
159
        self.sampler = Sampler(
            max_num_reqs=self.max_num_reqs,
            vocab_size=self.vocab_size,
            device=self.device,
160
            req_states=self.req_states,
161
            logprobs_mode=self.model_config.logprobs_mode,
162
            num_speculative_tokens=self.num_speculative_steps + 1,
163
        )
164
        self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
Woosuk Kwon's avatar
Woosuk Kwon committed
165
166

        # CUDA graphs.
167
168
169
        self.cudagraph_manager = CudaGraphManager(
            self.vllm_config, self.uses_mrope, self.device
        )
170
171
172
173
        # Structured outputs worker.
        self.structured_outputs_worker = StructuredOutputsWorker(
            max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
            vocab_size=self.vocab_size,
174
            device=self.device,
175
        )
176
177
        # LoRA-related workers.
        self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
178

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

        # KV Connector if configured.
183
184
        self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR

185
186
        # Pipeline parallelism.
        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
187
188
189
190
191
192
        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
193

194
195
196
197
198
199
        # Decode context parallelism.
        self.dcp_size = self.parallel_config.decode_context_parallel_size
        self.use_dcp = self.dcp_size > 1
        self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0
        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size

200
201
202
203
    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

204
205
    @staticmethod
    def get_supported_tasks() -> tuple[str]:
Woosuk Kwon's avatar
Woosuk Kwon committed
206
207
208
209
210
211
212
213
214
215
216
217
218
219
        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(
220
                    self.model, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
221
                )
222
223
            if self.do_spec_decode:
                self.speculator.load_model(self.model)
Woosuk Kwon's avatar
Woosuk Kwon committed
224
225
226
227
        time_after_load = time.perf_counter()

        self.model_memory_usage = m.consumed_memory
        logger.info(
228
229
            "Model loading took %s GiB and %.6f seconds",
            format_gib(m.consumed_memory),
Woosuk Kwon's avatar
Woosuk Kwon committed
230
231
232
            time_after_load - time_before_load,
        )

233
234
235
236
237
238
        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
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
    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,
259
260
261
            cp_size=self.dcp_size,
            cp_rank=self.dcp_rank,
            cp_interleave=self.cp_interleave,
Woosuk Kwon's avatar
Woosuk Kwon committed
262
263
264
        )

        self.attn_backends, self.attn_metadata_builders = init_attn_backend(
265
            self.kv_cache_config, self.vllm_config, self.device
Woosuk Kwon's avatar
Woosuk Kwon committed
266
        )
267
        check_attention_cp_compatibility(self.vllm_config)
268
269
270
271
272
273
274
        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
275
276

        self.kv_caches: list[torch.Tensor] = []
277
        kv_caches_dict = init_kv_cache(
Woosuk Kwon's avatar
Woosuk Kwon committed
278
279
280
281
282
283
            self.kv_caches,
            self.compilation_config.static_forward_context,
            self.kv_cache_config,
            self.attn_backends,
            self.device,
        )
284
285
        self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

Woosuk Kwon's avatar
Woosuk Kwon committed
286
287
288
289
290
291
292
293
        # 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
        )
294
295
296
        slot_mappings_by_layer = build_slot_mappings_by_layer(
            slot_mappings, self.kv_cache_config
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
297
298
299
300
        attn_metadata = build_attn_metadata(
            attn_metadata_builders=self.attn_metadata_builders,
            num_reqs=input_batch.num_reqs,
            num_tokens=input_batch.num_tokens,
301
302
            query_start_loc_gpu=input_batch.query_start_loc,
            query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np),
303
            max_query_len=input_batch.num_scheduled_tokens.max().item(),
304
            seq_lens=input_batch.seq_lens,
305
            max_seq_len=self.max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
306
307
308
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
309
            dcp_local_seq_lens=self.input_buffers.dcp_local_seq_lens,
Woosuk Kwon's avatar
Woosuk Kwon committed
310
311
        )
        input_batch.attn_metadata = attn_metadata
312
        input_batch.slot_mappings = slot_mappings_by_layer
Woosuk Kwon's avatar
Woosuk Kwon committed
313
314
315

    @torch.inference_mode()
    def _dummy_run(
316
        self, num_tokens: int, *args, skip_attn: bool = True, **kwargs
317
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
318
        # Create a dummy scheduler output.
Woosuk Kwon's avatar
Woosuk Kwon committed
319
        num_reqs = min(num_tokens, self.max_num_reqs)
320
321
322
323
        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 = {
324
            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
325
326
327
328
329
        }
        dummy_scheduler_output = SchedulerOutput.make_empty()
        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens

330
331
332
        # Disable any use of KVConnector for dummy runs.
        self.kv_connector.set_disabled(True)

333
334
        # For non-first PP ranks, create dummy intermediate_tensors.
        intermediate_tensors = None
335
        if not self.is_first_pp_rank:
336
337
338
339
340
341
            intermediate_tensors = self.model.make_empty_intermediate_tensors(
                batch_size=num_tokens,
                dtype=self.model_config.dtype,
                device=self.device,
            )

342
343
        # Execute the model.
        self.execute_model(
344
345
346
347
            dummy_scheduler_output,
            intermediate_tensors=intermediate_tensors,
            dummy_run=True,
            skip_attn_for_dummy_run=skip_attn,
348
        )
349
        self.kv_connector.set_disabled(False)
350
351

        # Non-last PP ranks don't produce output for sampling.
352
        if not self.is_last_pp_rank:
353
354
            return None, None

355
        assert self.execute_model_state is not None
356
        hidden_states, input_batch, _ = self.execute_model_state
357
        assert hidden_states is not None  # Last PP rank always has hidden_states
358
        sample_hidden_states = hidden_states[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
359
360
361
        return hidden_states, sample_hidden_states

    @torch.inference_mode()
362
    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
363
364
        num_reqs = hidden_states.shape[0]
        logits = self.model.compute_logits(hidden_states)
365
366
367
        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)
368
369
370
371
        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
        )
372
373
374
        # 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.
375
376
377
378
379
380
381
382
383
        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
384
385
386
387

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

391
        # Only run sampler on last PP rank (non-last ranks return None).
392
        if self.is_last_pp_rank:
393
394
            assert sample_hidden_states is not None
            self._dummy_sampler_run(sample_hidden_states)
395
396
397
398
399
400
401
402
403
404
405
406

            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
407
408
409
410
411
        torch.cuda.synchronize()
        del hidden_states, sample_hidden_states
        gc.collect()

    def reset_mm_cache(self) -> None:
412
413
        if self.supports_mm_inputs:
            self.encoder_runner.reset_mm_cache()
414
415

    def reset_encoder_cache(self) -> None:
416
417
        if self.supports_mm_inputs:
            self.encoder_runner.reset_encoder_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
418
419
420
421
422
423
424
425
426
427
428
429
430
431

    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

432
433
434
435
436
437
438
439
        # 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
440
        start_time = time.perf_counter()
441
        gc.collect()
442
        torch.cuda.empty_cache()
Woosuk Kwon's avatar
Woosuk Kwon committed
443
444
445
        start_free_gpu_memory = torch.cuda.mem_get_info()[0]

        with self.maybe_setup_dummy_loras(self.lora_config):
446
447
448
            mrope_positions = None
            if self.uses_mrope:
                mrope_positions = self.mrope_states.mrope_positions
449
450
451
            inputs_embeds = None
            if self.supports_mm_inputs:
                inputs_embeds = self.encoder_runner.inputs_embeds
Woosuk Kwon's avatar
Woosuk Kwon committed
452
453
454
            self.cudagraph_manager.capture(
                model=self.model,
                input_buffers=self.input_buffers,
455
                mrope_positions=mrope_positions,
456
                inputs_embeds=inputs_embeds,
Woosuk Kwon's avatar
Woosuk Kwon committed
457
458
459
                block_tables=self.block_tables,
                attn_metadata_builders=self.attn_metadata_builders,
                kv_cache_config=self.kv_cache_config,
460
                has_lora=self.lora_config is not None,
Woosuk Kwon's avatar
Woosuk Kwon committed
461
            )
462
463
            if self.do_spec_decode:
                self.speculator.capture_model()
Woosuk Kwon's avatar
Woosuk Kwon committed
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483

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

484
    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
485
        finished_req_ids = scheduler_output.finished_req_ids
486
487
488
        preempted_req_ids = scheduler_output.preempted_req_ids
        if preempted_req_ids:
            finished_req_ids = finished_req_ids.union(preempted_req_ids)
489
        for req_id in finished_req_ids:
Woosuk Kwon's avatar
Woosuk Kwon committed
490
            self.req_states.remove_request(req_id)
491
492
            if self.supports_mm_inputs:
                self.encoder_runner.remove_request(req_id)
493
            self.prompt_logprobs_worker.remove_request(req_id)
494
            self.lora_state.remove_request(req_id)
495

496
    def free_states(self, scheduler_output: SchedulerOutput) -> None:
497
498
499
        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
500

501
    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
502
        for new_req_data in scheduler_output.scheduled_new_reqs:
503
504
505
            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
506
            req_id = new_req_data.req_id
507
            prompt_len = len(new_req_data.prompt_token_ids)
Woosuk Kwon's avatar
Woosuk Kwon committed
508
509
            self.req_states.add_request(
                req_id=req_id,
510
                prompt_len=prompt_len,
511
                all_token_ids=new_req_data.prefill_token_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
512
513
514
                num_computed_tokens=new_req_data.num_computed_tokens,
            )
            req_index = self.req_states.req_id_to_index[req_id]
515

516
517
518
            if self.supports_mm_inputs:
                self.encoder_runner.add_request(req_id, new_req_data.mm_features)

519
520
521
522
523
524
            # 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,
525
                    mm_features=new_req_data.mm_features,
526
527
                )

528
529
530
            self.block_tables.append_block_ids(
                req_index, new_req_data.block_ids, overwrite=True
            )
531
532
533
            self.sampler.add_request(
                req_index, prompt_len, new_req_data.sampling_params
            )
534
535
536
            self.prompt_logprobs_worker.add_request(
                req_id, req_index, new_req_data.sampling_params
            )
537
            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
Woosuk Kwon's avatar
Woosuk Kwon committed
538

539
540
        if scheduler_output.scheduled_new_reqs:
            self.req_states.apply_staged_writes()
541
            self.sampler.apply_staged_writes()
542
543
544
545
            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
546
        # Add new blocks for the existing requests.
547
548
        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
549
            if req_new_block_ids is not None:
550
                req_index = self.req_states.req_id_to_index[req_id]
551
552
553
                self.block_tables.append_block_ids(
                    req_index, req_new_block_ids, overwrite=False
                )
Woosuk Kwon's avatar
Woosuk Kwon committed
554
555

    def prepare_inputs(
556
        self, scheduler_output: SchedulerOutput, num_tokens_after_padding: int
Woosuk Kwon's avatar
Woosuk Kwon committed
557
558
559
    ) -> InputBatch:
        num_tokens = scheduler_output.total_num_scheduled_tokens
        assert num_tokens > 0
560
561
        num_tokens_per_req = scheduler_output.num_scheduled_tokens
        num_reqs = len(num_tokens_per_req)
Woosuk Kwon's avatar
Woosuk Kwon committed
562
563
564

        # Decode first, then prefill.
        # batch_idx -> req_id
565
566
567
        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
568

569
570
        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)
571
        idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
Woosuk Kwon's avatar
Woosuk Kwon committed
572

573
        # Get the number of draft tokens for each request.
574
575
        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
        if not draft_tokens:
576
577
578
            # No draft token scheduled (common case).
            total_num_draft_tokens = 0
            total_num_logits = num_reqs
579
            cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
580
581
582
            cu_num_logits = torch.arange(
                num_reqs + 1, device=self.device, dtype=torch.int32
            )
583
            expanded_idx_mapping = idx_mapping
584
585
586
            expanded_local_pos = torch.zeros(
                num_reqs, dtype=torch.int32, device=self.device
            )
587
588
        else:
            num_draft_tokens = np.array(
589
                [len(draft_tokens.get(req_id, ())) for req_id in req_ids],
590
591
592
593
594
                dtype=np.int32,
            )
            total_num_draft_tokens = int(num_draft_tokens.sum())
            total_num_logits = num_reqs + total_num_draft_tokens

595
596
597
598
            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:])
599
            cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
600

601
            max_expand_len = self.num_speculative_steps + 1
602
            expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
603
                idx_mapping, total_num_logits, cu_num_logits, max_expand_len
604
605
            )

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

609
        # Get query_start_loc.
610
611
612
        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])
613
614
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
615
        query_start_loc_np[num_reqs + 1 :] = num_tokens
616
617
        async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)

618
619
620
        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]
621
        max_query_len = num_scheduled_tokens.max().item()
622

623
624
625
626
627
628
629
630
631
632
633
        # Get prefill tokens if any.
        if self.req_states.any_prefills(idx_mapping_np):
            prepare_prefill_inputs(
                self.input_buffers.input_ids,
                self.req_states.next_prefill_tokens,
                idx_mapping,
                query_start_loc,
                self.req_states.all_token_ids.gpu,
                self.req_states.prefill_len.gpu,
                self.req_states.num_computed_tokens.gpu,
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
634

635
636
637
        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
638
639
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
640
641
642
643
644
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
        seq_lens = self.input_buffers.seq_lens[:num_reqs]

645
646
        if self.use_dcp:
            # Prepare dcp local seq_lens.
647
648
            prepare_dcp_local_seq_lens(
                self.input_buffers.dcp_local_seq_lens,
649
                self.input_buffers.seq_lens,
650
                num_reqs,
651
652
653
                self.dcp_size,
                self.dcp_rank,
                self.cp_interleave,
654
            )
655
        dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs]
656

657
658
659
660
661
662
663
664
665
        # 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,
            )

666
        # Some input token ids are directly read from the last sampled tokens
667
668
        # and draft tokens. Also, get the logits indices to sample tokens from.
        logits_indices = combine_sampled_and_draft_tokens(
669
            self.input_buffers.input_ids,
Woosuk Kwon's avatar
Woosuk Kwon committed
670
671
            idx_mapping,
            self.req_states.last_sampled_tokens,
672
            query_start_loc,
673
674
            seq_lens,
            self.req_states.prefill_len.gpu,
675
676
677
            self.req_states.draft_tokens,
            cu_num_logits,
            total_num_logits,
Woosuk Kwon's avatar
Woosuk Kwon committed
678
679
680
681
        )

        # Compute slot mappings: [num_kv_cache_groups, num_tokens]
        slot_mappings = self.block_tables.compute_slot_mappings(
682
683
684
            idx_mapping,
            query_start_loc,
            self.input_buffers.positions[:num_tokens],
Woosuk Kwon's avatar
Woosuk Kwon committed
685
        )
686
687
688
689
        # 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
690
691
692
693
694
695

        # Layer name -> attention metadata.
        attn_metadata = build_attn_metadata(
            attn_metadata_builders=self.attn_metadata_builders,
            num_reqs=num_reqs,
            num_tokens=num_tokens,
696
            query_start_loc_gpu=query_start_loc,
697
            query_start_loc_cpu=query_start_loc_cpu,
698
            max_query_len=max_query_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
699
            seq_lens=self.input_buffers.seq_lens,
700
            max_seq_len=self.max_model_len,
Woosuk Kwon's avatar
Woosuk Kwon committed
701
702
703
            block_tables=block_tables,
            slot_mappings=slot_mappings,
            kv_cache_config=self.kv_cache_config,
704
            dcp_local_seq_lens=dcp_local_seq_lens,
Woosuk Kwon's avatar
Woosuk Kwon committed
705
706
        )

707
        input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
708
        positions = self.input_buffers.positions[:num_tokens_after_padding]
709
710
        mrope_positions = None
        if self.uses_mrope:
711
712
            mrope_positions = self.mrope_states.mrope_positions
            mrope_positions = mrope_positions[:, :num_tokens_after_padding]
Woosuk Kwon's avatar
Woosuk Kwon committed
713
714
715
716
717
        return InputBatch(
            req_ids=req_ids,
            num_reqs=num_reqs,
            idx_mapping=idx_mapping,
            idx_mapping_np=idx_mapping_np,
718
            expanded_idx_mapping=expanded_idx_mapping,
719
            expanded_local_pos=expanded_local_pos,
Woosuk Kwon's avatar
Woosuk Kwon committed
720
721
722
            num_scheduled_tokens=num_scheduled_tokens,
            num_tokens=num_tokens,
            num_tokens_after_padding=num_tokens_after_padding,
723
            num_draft_tokens=total_num_draft_tokens,
724
            query_start_loc=query_start_loc,
Woosuk Kwon's avatar
Woosuk Kwon committed
725
            query_start_loc_np=query_start_loc_np,
726
            seq_lens=seq_lens,
Woosuk Kwon's avatar
Woosuk Kwon committed
727
728
            input_ids=input_ids,
            positions=positions,
729
            mrope_positions=mrope_positions,
730
            inputs_embeds=None,
Woosuk Kwon's avatar
Woosuk Kwon committed
731
            attn_metadata=attn_metadata,
732
            slot_mappings=slot_mappings_by_layer,
Woosuk Kwon's avatar
Woosuk Kwon committed
733
            logits_indices=logits_indices,
734
            cu_num_logits=cu_num_logits,
735
            cu_num_logits_np=cu_num_logits_np,
736
            has_structured_output_reqs=scheduler_output.has_structured_output_requests,
Woosuk Kwon's avatar
Woosuk Kwon committed
737
738
        )

739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
    @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
759
760
761
762
763
    def sample(
        self,
        hidden_states: torch.Tensor,
        input_batch: InputBatch,
        grammar_output: GrammarOutput | None,
764
    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
765
        sample_hidden_states = hidden_states[input_batch.logits_indices]
766
        sample_pos = input_batch.positions[input_batch.logits_indices]
767
        input_ids = input_batch.input_ids[input_batch.logits_indices]
Woosuk Kwon's avatar
Woosuk Kwon committed
768
769
770
        logits = self.model.compute_logits(sample_hidden_states)
        if grammar_output is not None:
            # Apply grammar bitmask to the logits in-place.
771
772
773
774
775
776
            self.structured_outputs_worker.apply_grammar_bitmask(
                logits,
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.grammar_bitmask,
            )
777

778
        # Sample tokens and compute logprobs (if needed).
779
780
781
782
        sampler_output = self.sampler(
            logits,
            input_batch.expanded_idx_mapping,
            input_batch.idx_mapping_np,
783
            input_batch.cu_num_logits_np,
784
            sample_pos,
785
786
            input_ids,
            input_batch.expanded_local_pos,
787
        )
788
789
790

        if input_batch.num_draft_tokens == 0:
            # No draft tokens (common case).
791
792
793
            num_sampled = torch.ones(
                input_batch.num_reqs, dtype=torch.int32, device=self.device
            )
794
        else:
795
            # Rejection sampling for spec decoding.
796
797
798
799
800
801
802
            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
803
804
805
806
807
808
809
810
811
812

        # 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,
        )
813
        return sampler_output, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
814
815
816
817

    def postprocess(
        self,
        input_batch: InputBatch,
818
819
        sampled_tokens: torch.Tensor,
        num_sampled: torch.Tensor,
820
        num_rejected: torch.Tensor,
821
822
    ) -> None:
        # Update the number of computed tokens.
823
        post_update(
824
            input_batch.idx_mapping,
825
            self.req_states.num_computed_tokens.gpu,
826
            self.req_states.last_sampled_tokens,
827
            self.sampler.penalties_state.output_bin_counts,
828
829
            sampled_tokens,
            num_sampled,
830
            num_rejected,
831
            input_batch.query_start_loc,
832
833
            self.req_states.all_token_ids.gpu,
            self.req_states.total_len.gpu,
Woosuk Kwon's avatar
Woosuk Kwon committed
834
        )
835
836

        # Update the number of computed prefill tokens.
Woosuk Kwon's avatar
Woosuk Kwon committed
837
        idx_mapping_np = input_batch.idx_mapping_np
838
        computed_prefill = self.req_states.num_computed_prefill_tokens
839
840
841
        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
842
843
        )

844
845
846
847
848
849
850
    @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,
851
        num_rejected: torch.Tensor,
852
853
854
855
856
857
858
    ) -> torch.Tensor:
        assert self.speculator is not None
        draft_tokens = self.speculator.propose(
            input_batch,
            last_hidden_states,
            aux_hidden_states,
            num_sampled,
859
            num_rejected,
860
861
            self.req_states.last_sampled_tokens,
            self.req_states.next_prefill_tokens,
862
863
            self.sampler.sampling_states.temperature.gpu,
            self.sampler.sampling_states.seeds.gpu,
864
865
866
        )
        return draft_tokens

Woosuk Kwon's avatar
Woosuk Kwon committed
867
868
869
870
    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: SchedulerOutput,
871
        intermediate_tensors: IntermediateTensors | None = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
872
        dummy_run: bool = False,
873
        skip_attn_for_dummy_run: bool = False,
874
    ) -> ModelRunnerOutput | IntermediateTensors | None:
875
876
877
878
879
880
881
882
883
        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.
884
885
                empty_output = self.kv_connector.no_forward(scheduler_output)
                return empty_output
Woosuk Kwon's avatar
Woosuk Kwon committed
886

887
888
889
890
891
892
893
        # Get local cudagraph mode and size.
        local_cudagraph_mode, local_cudagraph_size = (
            self.cudagraph_manager.get_cudagraph_runtime_mode(
                num_reqs=len(scheduler_output.num_scheduled_tokens),
                num_tokens=scheduler_output.total_num_scheduled_tokens,
                max_query_len=max(scheduler_output.num_scheduled_tokens.values()),
            )
894
        )
895
896
897

        # DP sync: num_tokens + cudagraph_size + cudagraph_mode
        num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode = (
898
            get_cudagraph_and_dp_padding(
899
                scheduler_output.total_num_scheduled_tokens,
900
901
                local_cudagraph_size,
                local_cudagraph_mode.value,
902
903
                self.parallel_config.data_parallel_size,
                self.parallel_config.data_parallel_rank,
904
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
905
        )
906
        cudagraph_runtime_mode = CUDAGraphMode(synced_cudagraph_mode)
907
908
        if num_tokens_after_padding == 0:
            # All DP ranks have zero tokens to run.
909
910
            empty_output = self.kv_connector.no_forward(scheduler_output)
            return empty_output
911
912
913
914
915

        if not dummy_run:
            # Common case.
            # Prepare all the inputs and copy to the input buffers.
            input_batch = self.prepare_inputs(
916
                scheduler_output, num_tokens_after_padding
917
918
919
            )
            if self.lora_config:
                # Activate LoRA adapters.
920
                lora_inputs = self.lora_state.make_lora_inputs(
921
922
923
                    input_batch.req_ids,
                    input_batch.idx_mapping_np,
                    input_batch.num_scheduled_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
924
                )
925
                self._set_active_loras(*lora_inputs)
926

927
            # Only first PP rank prepares multimodal embeddings.
928
            if self.supports_mm_inputs and self.is_first_pp_rank:
929
930
931
932
933
934
935
936
937
                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
                ]
938
        else:
939
            # No actual tokens to run. A dummy run for DP or memory profiling.
940
941
942
943
944
945
946
            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,
            )
947
948
949
950
            if self.uses_mrope:
                input_batch.mrope_positions = self.mrope_states.mrope_positions[
                    :, :num_tokens_after_padding
                ]
951
952
953
            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
954
955

        # Run model.
956
957
        if cudagraph_runtime_mode == CUDAGraphMode.FULL:
            # Use explicit cudagraph replay for FULL mode.
Woosuk Kwon's avatar
Woosuk Kwon committed
958
959
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
960
            self.kv_connector.pre_forward(scheduler_output)
961
            hidden_states = self.cudagraph_manager.run_fullgraph(
Woosuk Kwon's avatar
Woosuk Kwon committed
962
963
964
                input_batch.num_tokens_after_padding
            )
        else:
965
            # For piecewise and eager mode, just call model().
966
967
968
            positions = input_batch.positions
            if self.uses_mrope:
                assert input_batch.mrope_positions is not None
969
                positions = input_batch.mrope_positions
970

971
972
973
974
975
976
977
978
979
            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

980
981
982
983
984
            batch_descriptor = BatchDescriptor(
                num_tokens=input_batch.num_tokens_after_padding,
                has_lora=self.lora_config is not None,
            )

Woosuk Kwon's avatar
Woosuk Kwon committed
985
986
987
988
            with set_forward_context(
                input_batch.attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
989
                cudagraph_runtime_mode=cudagraph_runtime_mode,
Woosuk Kwon's avatar
Woosuk Kwon committed
990
                num_tokens_across_dp=num_tokens_across_dp,
991
                batch_descriptor=batch_descriptor,
992
                slot_mapping=input_batch.slot_mappings,
Woosuk Kwon's avatar
Woosuk Kwon committed
993
            ):
994
                self.kv_connector.pre_forward(scheduler_output)
995
996
997
998
999
1000
                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
1001

1002
        kv_connector_output = self.kv_connector.post_forward(scheduler_output)
1003

1004
        if not self.is_last_pp_rank:
1005
1006
1007
1008
1009
1010
1011
1012
1013
            # 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
1014
1015
1016
1017
        return None

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

1024
        if not self.is_last_pp_rank:
1025
1026
1027
1028
            # Non-last PP rank: hidden_states is None because this rank produced
            # IntermediateTensors instead of final hidden states. Receive the
            # sampled tokens broadcast from the last rank and update local state.
            sampled, num_sampled, num_rejected = pp_receive(
1029
                input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
1030
            )
1031
            self.postprocess(input_batch, sampled, num_sampled, num_rejected)
1032
1033
1034
            return None

        # Last rank: sample tokens
1035
        sampler_output, num_sampled, num_rejected = self.sample(
1036
            hidden_states, input_batch, grammar_output
Woosuk Kwon's avatar
Woosuk Kwon committed
1037
        )
1038
1039

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

1043
1044
1045
1046
        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
            self.model.compute_logits,
            hidden_states,
            input_batch,
1047
            self.req_states.all_token_ids.gpu,
1048
            self.req_states.num_computed_tokens.gpu,
1049
            self.req_states.prompt_len.np,
1050
1051
1052
            self.req_states.prefill_len.np,
            self.req_states.num_computed_prefill_tokens,
        )
1053
1054
1055
1056
1057
1058
1059
1060

        # 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
1061
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
1062
            kv_connector_output=kv_connector_output,
1063
1064
1065
1066
        )
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,
1067
            num_sampled_tokens=num_sampled,
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
            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(
1078
            input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
Woosuk Kwon's avatar
Woosuk Kwon committed
1079
        )
1080
        if self.do_spec_decode:
1081
            draft_tokens = self.propose_draft(
1082
1083
1084
                input_batch,
                hidden_states,
                None,  # aux_hidden_states
1085
1086
                num_sampled,
                num_rejected,
1087
            )
1088
            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1089
            self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
1090
1091
1092
1093

        if self.use_async_scheduling:
            return async_output
        return async_output.get_output()
1094
1095
1096

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