recipes/docker-images/glm53-kindling/experimental/draft-trunc/model_runner.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NOTE: Coding style guide for this file:
This model runner is shared by all models: text and multimodal, generative
and embedding, public and private. As a result, this file must only contain
code that is common to every model. Model-specific behavior belongs in the
appropriate model-specific files.

In other words:
* Be paranoid about changing this file. It should remain stable.
* Be even more paranoid about adding new lines. It should remain minimal.

Even for shared features (for example, different parallelism modes), keep the
complexity out of this path. The less common the feature, the more it should be
hidden. Prefer utility functions defined elsewhere and call them from here,
instead of embedding feature-specific logic directly.
"""

import functools
import gc
import time
from contextlib import AbstractContextManager
from copy import deepcopy
from typing import Any, NamedTuple

import numpy as np
import torch

from vllm.v1.worker.gpu.spec_decode import draft_trunc
import torch.nn as nn

import vllm.envs as envs
from vllm.compilation.counter import compilation_counter
from vllm.compilation.cuda_graph import CUDAGraphStat
from vllm.config import VllmConfig
from vllm.config.compilation import CUDAGraphMode
from vllm.distributed.aux_output_connector.worker import (
    AuxOutputWorkerConnector,
    get_aux_output_connector,
)
from vllm.distributed.parallel_state import get_dcp_group, get_pp_group
from vllm.forward_context import BatchDescriptor, set_forward_context
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
    initialize_mamba_ssu_backend,
)
from vllm.model_executor.model_loader import get_model_loader
from vllm.model_executor.models.interfaces import requires_raw_input_tokens
from vllm.model_executor.offloader import (
    create_offloader,
    get_offloader,
    set_offloader,
)
from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.encoder_budget import (
    MultiModalBudget,
)
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.tasks import SupportedTask
from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.kv_cache_interface import (
    KVCacheConfig,
    MambaSpec,
    UniformTypeKVCacheSpecs,
)
from vllm.v1.outputs import (
    DraftTokenIds,
    ECConnectorOutput,
    ModelRunnerOutput,
)
from vllm.v1.watermarking import create_watermarker
from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
from vllm.v1.watermarking.spec_decode import (
    create_speculative_target_watermarker,
    speculative_target_watermark_key,
)
from vllm.v1.worker.block_table import get_block_table_width
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
from vllm.v1.worker.gpu import pcp_manager as pcp
from vllm.v1.worker.gpu.async_utils import (
    AsyncOutput,
    AsyncPoolingOutput,
    StepTimingCollector,
)
from vllm.v1.worker.gpu.attn_utils import (
    FastPrefillHelper,
    build_slot_mappings_by_layer,
    get_kv_cache_spec,
    init_attn_backend,
    init_kv_cache,
)
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cudagraph_utils import (
    BatchExecutionDescriptor,
    ModelCudaGraphManager,
    has_compiled_submodule,
    make_cudagraph_stats,
)
from vllm.v1.worker.gpu.cudagraph_utils import (
    profile_cudagraph_memory as _profile_cudagraph_memory,
)
from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp
from vllm.v1.worker.gpu.ec_connector import get_ec_connector
from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after
from vllm.v1.worker.gpu.input_batch import (
    InputBatch,
    InputBuffers,
    combine_sampled_and_draft_tokens,
    expand_idx_mapping,
    post_update,
    post_update_num_computed_tokens,
    prepare_pos_seq_lens,
    prepare_prefill_inputs,
    set_dummy_context,
)
from vllm.v1.worker.gpu.kv_connector import (
    NO_OP_KV_CONNECTOR,
    KVConnector,
    get_kv_connector,
)
from vllm.v1.worker.gpu.lora_utils import (
    LoraState,
    create_lora_capture_hook,
    get_lora_capture_cases,
    get_num_active_loras_for_dispatch,
)
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
from vllm.v1.worker.gpu.model_states import init_model_state
from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
from vllm.v1.worker.gpu.pp_utils import PPHandler
from vllm.v1.worker.gpu.sample.batch_shard import (
    BatchSharder,
    all_to_all_logits,
    gather_sampler_output,
)
from vllm.v1.worker.gpu.sample.logits_processor import build_custom_logits_processors
from vllm.v1.worker.gpu.sample.output import SamplerOutput
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
from vllm.v1.worker.gpu.sample.sampler import Sampler
from vllm.v1.worker.gpu.shutdown import free_before_shutdown
from vllm.v1.worker.gpu.spec_decode import init_speculator
from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
    AdaptiveVerificationManager,
    maybe_create_adaptive_verification_manager,
    resolve_adaptive_cudagraph_mode,
)
from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
    set_eagle3_aux_hidden_state_layers,
    verify_supports_aux_hidden_states_over_pp,
)
from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
    RejectionSampler,
    get_max_chunk_logits,
)
from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
from vllm.v1.worker.gpu.states import RequestState
from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
from vllm.v1.worker.gpu.ubatch_utils import (
    UBatchRunner,
    UBatchState,
    maybe_build_ubatch_runner,
)
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
from vllm.v1.worker.utils import (
    KVBlockZeroer,
    clear_layer_kv_caches,
    copy_kv_cache_blocks_inplace,
    get_uniform_decode_token_count,
)
from vllm.v1.worker.workspace import lock_workspace, use_workspace_lane

logger = init_logger(__name__)


class GPUModelRunner(LoRAModelRunnerMixin):
    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._draft_workspace_lane = int(
            self.speculative_config is not None and self.speculative_config.use_dspark()
        )
        self.observability_config = vllm_config.observability_config
        self.jit_warmup_registry = JitWarmupRegistry(vllm_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
            ]

        # Lazily initialized in _init_kv_zero_meta() when the KV cache needs
        # zeroing (e.g. hybrid models with fp8 KV cache).
        self.kv_block_zeroer: KVBlockZeroer | None = None

        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
        self.is_encoder_decoder = self.model_config.is_encoder_decoder

        self.output_copy_stream = torch.cuda.Stream(self.device)

        # Pipeline parallelism.
        self.use_pp = self.parallel_config.pipeline_parallel_size > 1
        self.is_first_pp_rank = get_pp_group().is_first_rank
        self.is_last_pp_rank = get_pp_group().is_last_rank

        # Size the UVA buffer pools to the max number of concurrent in-flight
        # steps. Must run before any pooled buffer is constructed
        set_default_max_concurrency(vllm_config.max_concurrent_batches)

        # PP broadcast/recv helper. Runs the collective on a side stream.
        self.pp_handler: PPHandler | None = None

        # Persistent buffer for intermediate tensors (non-first PP ranks).
        self.intermediate_tensors: IntermediateTensors | None = None

        # Data parallelism.
        self.dp_size = self.parallel_config.data_parallel_size
        self.dp_rank = self.parallel_config.data_parallel_rank

        # Dual batch overlap. Created in initialize_kv_cache(), once everything
        # it runs the microbatched forward with exists.
        self.ubatch_runner: UBatchRunner | None = None

        # Detect EP all2all peer faults to prevent emitting corrupted output.
        # Only meaningful for MoE + DP with an FT-capable all2all backend.
        self.check_ep_fault = False
        if self.dp_size > 1 and self.model_config.is_moe:
            self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance

        # 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

        # Multimodal
        self.mm_registry = MULTIMODAL_REGISTRY
        self.supports_mm_inputs = self.model_config.supports_multimodal_inputs
        self.uses_inputs_embeds = (
            self.supports_mm_inputs or self.model_config.enable_prompt_embeds
        )
        self.encoder_cache = None
        if self.supports_mm_inputs and self.is_first_pp_rank:
            self.encoder_cache = EncoderCache()
        self.ec_connector = get_ec_connector(vllm_config, self.encoder_cache)

        # Speculative decoding.
        self.speculator = None
        self.use_aux_hidden_state_outputs = False
        self.num_speculative_steps = vllm_config.num_speculative_tokens
        if self.speculative_config is not None:
            if self.is_last_pp_rank:
                self.speculator = init_speculator(self.vllm_config, self.device)

            if self.speculative_config.method in (
                "eagle3",
                "dflash",
                "dspark",
                "extract_hidden_states",
            ):
                # Drafting may require auxiliary hidden states from target model outputs
                self.use_aux_hidden_state_outputs = True

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

        self.pcp_manager: pcp.PCPManager | None = None

        # Pooling models.
        self.is_pooling_model = self.model_config.runner_type == "pooling"
        self.pooling_runner: PoolingRunner | None = None

        num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens)

        self.step_timing = StepTimingCollector()

        # General request states.
        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,
            num_speculative_steps=self.num_speculative_steps,
            vocab_size=self.vocab_size,
            device=self.device,
            num_prefill_lookahead=num_prefill_lookahead,
        )
        self.adaptive_verification: AdaptiveVerificationManager | None = None
        self.input_buffers = InputBuffers(
            max_num_reqs=self.max_num_reqs,
            max_num_tokens=self.max_num_tokens,
            device=self.device,
        )
        self.fast_prefill: FastPrefillHelper | None = None
        if self.use_pp:
            self.pp_handler = PPHandler(
                max_num_reqs=self.max_num_reqs,
                num_speculative_steps=self.num_speculative_steps,
                device=self.device,
            )

        # Samplers and decode_query_len created in load_model() after
        # model_state exists (num_new_sampled_tokens_per_step from ModelState).
        self.sampler: Sampler | None = None
        self.rejection_sampler: RejectionSampler | None = None
        self.batch_sharder: BatchSharder | None = None
        self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
        self.structured_outputs_worker: StructuredOutputsWorker | None = None
        self.cudagraph_manager: ModelCudaGraphManager | None = None

        # LoRA-related workers.
        self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
        self.lora_capture_cases = [0]
        if self.lora_config:
            self.lora_capture_cases = get_lora_capture_cases(
                self.lora_config, self.compilation_config
            )

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

        # For transferring state from execute_model to subsequent sample_tokens call.
        self.execute_model_state: ExecuteModelState | None = None

        # Expert parallelism load balancer.
        self.eplb = EPLBController(self.parallel_config, self.device)
        # The AuxOutput Connector owns R3 capture, copying, and storage.
        self.aux_output_connector: AuxOutputWorkerConnector | None = None

        set_offloader(create_offloader(self.vllm_config.offload_config))

    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

    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
        tasks: list[SupportedTask] = []
        if self.model_config.runner_type == "generate":
            tasks.extend(self.model_state.get_supported_generation_tasks())
        if self.is_pooling_model:
            # Do not rely on pooling_runner here, since this information is needed
            # on the first PP rank, while pooling_runner is only initialized
            # on the last PP rank.
            tasks.extend(PoolingRunner.get_supported_tasks(self.model))
        return tuple(tasks)

    def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
        time_before_load = time.perf_counter()
        if load_dummy_weights:
            self.load_config.load_format = "dummy"
        self.eplb.prepare_load()
        eplb_models_added = False
        with DeviceMemoryProfiler() as m:
            model_loader = get_model_loader(self.vllm_config.load_config)
            logger.info_once("Loading model from scratch...")

            # Capture warmup providers selected while constructing the model.
            with self.jit_warmup_registry.activate():
                self.model = model_loader.load_model(
                    vllm_config=self.vllm_config,
                    model_config=self.vllm_config.model_config,
                )
            if self.lora_config:
                self.model = self.load_lora_model(
                    self.model, self.vllm_config, self.device
                )

            if self.use_aux_hidden_state_outputs:
                assert self.speculative_config is not None
                set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config)
                if self.use_pp:
                    assert self.speculative_config.method is not None
                    verify_supports_aux_hidden_states_over_pp(
                        self.model, self.speculative_config.method
                    )
                    assert self.pp_handler is not None
                    self.pp_handler.configure_aux_hidden_state_relay(self.model)
            if isinstance(self.speculator, DraftModelSpeculator):
                with use_workspace_lane(self._draft_workspace_lane):
                    self.speculator.load_model(self.model)
                    eplb_models_added = self.eplb.maybe_register_speculator(
                        self.speculator, self.speculative_config, load_dummy_weights
                    )
        time_after_load = time.perf_counter()

        self.model_memory_usage = m.consumed_memory
        logger.info(
            "Model loading took %s GiB memory and %.6f seconds",
            format_gib(m.consumed_memory),
            time_after_load - time_before_load,
        )

        # Initialize the components that require the model.
        self.model_state = init_model_state(
            self.vllm_config, self.model, self.encoder_cache, self.device
        )

        self.decode_query_len = (
            self.num_speculative_steps
            + self.model_state.num_new_sampled_tokens_per_step
        )

        if self.parallel_config.enable_batch_sharded_sampling:
            if hasattr(self.model, "compute_logits_local"):
                self.batch_sharder = BatchSharder(
                    max_num_reqs=self.max_num_reqs,
                    max_num_logits_per_req=self.decode_query_len,
                    device=self.device,
                )
                logger.info("Batch-sharded sampling enabled.")
            else:
                logger.warning_once(
                    "Disabling batch-sharded sampling: %s does not implement "
                    "compute_logits_local",
                    type(self.model).__name__,
                )

        # Initialize samplers. Model states may override via custom_sampler().
        config_processors = self.model_config.logits_processors or ()
        custom_logits_processors = build_custom_logits_processors(
            self.vllm_config, self.req_states, self.is_pooling_model, config_processors
        )
        if self.is_last_pp_rank and not self.is_pooling_model:
            # V2 bypasses TopKTopPSampler, which registers native warmups.
            # ROCm only: on CUDA these warmups add ~2 min to every engine start.
            if current_platform.is_rocm():
                from vllm.v1.sample.ops.topk_topp_sampler import (
                    register_top_k_top_p_warmups,
                )

                with self.jit_warmup_registry.activate():
                    register_top_k_top_p_warmups()

            sampler_kwargs: dict[str, Any] = {
                "vllm_config": self.vllm_config,
                "max_num_reqs": self.max_num_reqs,
                "vocab_size": self.vocab_size,
                "device": self.device,
                "req_states": self.req_states,
                "logprobs_mode": self.model_config.logprobs_mode,
                "num_speculative_tokens": self.decode_query_len,
                "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
                "enable_trace_replay": self.model_config.enable_trace_replay,
                "return_sampling_mask": self.model_config.return_sampling_mask,
                "custom_logits_processors": custom_logits_processors,
            }
            if self.vllm_config.watermark_config is None:
                self.sampler = Sampler(**sampler_kwargs)
            else:
                wm_config = self.vllm_config.watermark_config
                watermarker = create_watermarker(wm_config)
                if self.speculative_config is not None:
                    watermarker = create_speculative_target_watermarker(watermarker)
                self.sampler = GPUWatermarkSampler(
                    watermarker,
                    deduplicate_contexts=wm_config.deduplicate_contexts,
                    deduplicate_contexts_max_history=(
                        wm_config.deduplicate_contexts_max_history
                    ),
                    **sampler_kwargs,
                )
            custom = self.model_state.custom_sampler(self.sampler)

            if custom:
                self.vllm_config._check_watermarking_unsupported(custom_sampler=True)
                self.sampler, self.rejection_sampler = custom
            elif self.speculative_config is not None:
                self.rejection_sampler = RejectionSampler(
                    self.sampler,
                    self.speculative_config,
                    self.device,
                    watermark_key=speculative_target_watermark_key(
                        self.vllm_config.watermark_config
                    ),
                )
            self.prompt_logprobs_worker = PromptLogprobsWorker(
                self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode
            )
            self.structured_outputs_worker = StructuredOutputsWorker(
                max_num_logits=self.max_num_reqs * self.decode_query_len,
                vocab_size=self.vocab_size,
                device=self.device,
                mask_stride=self.decode_query_len,
                num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
            )

        if self.is_pooling_model and self.is_last_pp_rank:
            self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
        eplb_models_added |= self.eplb.maybe_register_model(
            self.model, self.model_config, load_dummy_weights
        )
        self.eplb.maybe_start_async_loop(eplb_models_added)

        if not self.is_first_pp_rank:
            # For non-first PP ranks, create intermediate tensors sized
            # for the max capture size so they can be sliced per batch.
            # Save as persistent member so runtime can copy received data
            # into the same addresses that the CUDA graphs captured.
            self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
                batch_size=self.max_num_tokens,
                dtype=self.model_config.dtype,
                device=self.device,
            )

        get_offloader().post_init()

    def get_model(self) -> nn.Module:
        return self.model

    def get_draft_model(self) -> nn.Module | None:
        speculator = self.speculator
        if not isinstance(speculator, DraftModelSpeculator):
            return None
        return speculator.model

    def reload_weights(self, *args, **kwargs) -> None:
        # TODO(Wentao): Use full version instead of import when fully migrated to v2
        from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1

        GPUModelRunnerV1.reload_weights(self, *args, **kwargs)  # type: ignore[arg-type]

    def update_config(self, *args, **kwargs) -> None:
        # TODO(Wentao): Use full version instead of import when fully migrated to v2
        from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1

        GPUModelRunnerV1.update_config(self, *args, **kwargs)  # type: ignore[arg-type]

        # v2 reads config via self.vllm_config (e.g. in load_model), so keep it
        # in sync with the attributes the v1 helper just replaced.
        self.vllm_config.model_config = self.model_config
        self.vllm_config.load_config = self.load_config

    @functools.cached_property
    def main_stream(self) -> torch.cuda.Stream:
        # Cache the default CUDA stream to avoid lookup overhead.
        return torch.cuda.current_stream(self.device)

    def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
        encoder_runner = getattr(self.model_state, "encoder_runner", None)
        if encoder_runner is None:
            return {}
        return encoder_runner.get_encoder_timing_stats()

    def get_kv_cache_spec(self):
        return get_kv_cache_spec(self.vllm_config)

    def initialize_kv_cache(
        self,
        kv_cache_config: KVCacheConfig,
        is_profiling: bool = False,
        kv_cache_allocation_context: AbstractContextManager | None = None,
    ) -> None:
        # GPUWorker finalizes the PD interleave before KV cache initialization.
        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
        kv_cache_config = deepcopy(kv_cache_config)
        self.kv_cache_config = kv_cache_config

        block_table_max_model_len = self.max_model_len
        if self.is_encoder_decoder:
            # Cross-attention block tables need to index encoder tokens, which
            # can exceed the decoder's max_model_len.
            block_table_max_model_len = max(
                block_table_max_model_len,
                self.scheduler_config.max_num_encoder_input_tokens,
                getattr(self.model_config.hf_config, "max_source_positions", 0),
            )

        block_sizes = []
        max_num_blocks_per_group = []
        slot_mapping_enabled = []
        for kv_cache_group in kv_cache_config.kv_cache_groups:
            spec = kv_cache_group.kv_cache_spec
            block_sizes.append(spec.block_size)
            layer_spec = (
                spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
            )
            slot_mapping_enabled.append(layer_spec.uses_slot_mapping)
            # Let each cache type account for CP. Attention KV is DCP-sharded,
            # while Mamba/GDN recurrent state is replicated across DCP ranks.
            max_num_blocks = spec.max_num_blocks_per_req(
                self.vllm_config, block_table_max_model_len
            )
            # Preserve each cache type's alignment requirements after applying
            # its topology-aware block-table width.
            max_num_blocks = get_block_table_width(
                max_num_blocks,
                spec.block_size,
                token_alignment=spec.block_table_token_alignment,
            )
            max_num_blocks_per_group.append(max_num_blocks)

        target_attn_layer_names = None
        if isinstance(self.speculator, DraftModelSpeculator):
            # Adaptive verification validates target attention separately.
            target_attn_layer_names = {
                layer_name
                for group in self.kv_cache_config.kv_cache_groups
                for layer_name in group.layer_names
            } - self.speculator.draft_attn_layer_names
        draft_attn_layer_names = None
        if isinstance(self.speculator, DraftModelSpeculator):
            draft_attn_layer_names = self.speculator.draft_attn_layer_names
        # Metadata builders select attention kernels that need JIT warmup.
        with self.jit_warmup_registry.activate():
            (
                self.attn_groups,
                attn_cg_support,
                self.kernel_block_sizes,
            ) = init_attn_backend(
                self.kv_cache_config,
                self.vllm_config,
                self.device,
                draft_layer_names=draft_attn_layer_names,
            )
        additional_attn_cg_support = self.model_state.get_additional_cg_support()
        attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
        # The speculator clears the flag at load time when the checkpoint has
        # no confidence head, so it holds the effective value.
        self.adaptive_verification = maybe_create_adaptive_verification_manager(
            enable_adaptive_verification=getattr(
                self.speculator, "enable_adaptive_verification", False
            ),
            attn_groups=self.attn_groups,
            attn_cg_support=attn_cg_support,
            req_states=self.req_states,
            query_start_loc=self.input_buffers.query_start_loc,
            num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
            max_total_logits=get_max_chunk_logits(self.vocab_size),
            vllm_config=self.vllm_config,
            target_layer_names=target_attn_layer_names,
            additional_attn_cg_support=additional_attn_cg_support,
        )

        self.block_tables = BlockTables(
            block_sizes=block_sizes,
            max_num_reqs=self.max_num_reqs,
            max_num_batched_tokens=self.max_num_tokens,
            max_num_blocks_per_group=max_num_blocks_per_group,
            device=self.device,
            kernel_block_sizes=self.kernel_block_sizes,
            slot_mapping_enabled=slot_mapping_enabled,
            cp_size=self.dcp_size,
            cp_rank=self.dcp_rank,
            cp_interleave=self.cp_interleave,
        )
        self.pcp_manager = pcp.maybe_build_pcp_manager(
            self.vllm_config,
            self.device,
            self.supports_mm_inputs,
            self.block_tables,
            cls=self.pcp_manager_cls,
        )
        self.ubatch_runner = maybe_build_ubatch_runner(
            self.vllm_config,
            self.device,
            self.model_state,
            self.attn_groups,
            self.kv_cache_config,
            self.max_num_reqs,
        )
        if self.speculator is not None:
            self.speculator.pcp_manager = self.pcp_manager
        initialize_mamba_ssu_backend(
            self.vllm_config.mamba_config,
            self.kv_cache_config,
            use_replayssm=self.cache_config.use_replayssm,
        )
        piecewise_capture_available = bool(
            envs.VLLM_USE_BREAKABLE_CUDAGRAPH or has_compiled_submodule(self.model)
        )
        if self.adaptive_verification is not None:
            self.compilation_config.cudagraph_mode = resolve_adaptive_cudagraph_mode(
                self.compilation_config.cudagraph_mode,
                piecewise_capture_available=piecewise_capture_available,
            )
        cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
            attn_cg_support.min_cg_support,
            attn_cg_support.min_cg_attn_backend,
            self.decode_query_len,
            use_v2_model_runner=True,
            tensor_parallel_size=self.parallel_config.tensor_parallel_size,
            kv_cache_config=self.kv_cache_config,
            max_num_reqs=self.max_num_reqs,
            is_profiling=is_profiling,
            piecewise_capture_available=piecewise_capture_available,
        )
        self.cudagraph_manager = ModelCudaGraphManager(
            self.vllm_config,
            self.device,
            cudagraph_mode,
            decode_query_len=self.decode_query_len,
            lora_capture_cases=self.lora_capture_cases,
            varlen_decode=self.adaptive_verification is not None,
            ubatch_runner=self.ubatch_runner,
        )
        if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None:
            self.fast_prefill = FastPrefillHelper(
                self.cudagraph_manager, self.max_num_tokens
            )
        check_attention_cp_compatibility(self.vllm_config, target_attn_layer_names)
        if isinstance(self.speculator, DraftModelSpeculator):
            # HACK(woosuk)
            self.speculator.set_attn(
                self.model_state,
                self.kv_cache_config,
                self.block_tables,
                self.input_buffers,
                self.attn_groups,
            )
        if self.speculator is not None:
            # After set_attn, so the speculator can size its cudagraph mode
            # to its own attention support.
            self.speculator.init_cudagraph_manager(cudagraph_mode)

        # Capture warmup providers that depend on allocated KV-cache strides.
        with self.jit_warmup_registry.activate():
            kv_caches_dict = init_kv_cache(
                self.compilation_config.static_forward_context,
                self.kv_cache_config,
                self.device,
                self.kernel_block_sizes,
                self.vllm_config,
                kv_cache_allocation_context=kv_cache_allocation_context,
                block_tables=self.block_tables,
            )
        self.kv_caches = [
            cache for cache in kv_caches_dict.values() if cache.device == self.device
        ]
        if is_profiling:
            self.kv_connector = NO_OP_KV_CONNECTOR
        else:
            self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

            # AuxOutput connector requires resolved kv_cache_config.
            if self.vllm_config.aux_output_config.enabled:
                self.aux_output_connector = get_aux_output_connector(
                    self.model, self.vllm_config, kv_cache_config
                )

    def _init_kv_zero_meta(self) -> None:
        """Build KV-block zeroing metadata; invoked from gpu_worker."""
        self.kv_block_zeroer = KVBlockZeroer(
            self.device,
            attn_groups_iter=(g for groups in self.attn_groups for g in groups),
            kernel_block_sizes=self.kernel_block_sizes,
            static_forward_context=self.compilation_config.static_forward_context,
            num_blocks=self.kv_cache_config.num_blocks,
        )

    @torch.inference_mode()
    @step_eplb_after(is_dummy=True)
    def _dummy_run(
        self,
        num_tokens: int,
        *args,
        skip_attn: bool = False,
        uniform_decode: bool = False,
        context_len: int = 0,
        skip_eplb: bool = False,
        is_profile: bool = False,
        valid_dummy_state_slots: bool = False,
        **kwargs,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        if skip_attn and not is_profile:
            raise ValueError(
                "skip_attn must only be True for initial memory profiling."
            )

        # Create a dummy scheduler output.
        num_reqs = min(num_tokens, self.max_num_reqs)
        if uniform_decode:
            # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
            # and for spec-decode with MTP we want to make sure the dummy runs use
            # 1+num_speculative_tokens we use max here, this will likely be eventually
            # changed in the worker: https://github.com/vllm-project/vllm/pull/35243
            num_tokens = max(num_tokens, self.decode_query_len)
            num_reqs = num_tokens // self.decode_query_len
            assert num_tokens % self.decode_query_len == 0
        elif self.speculator is not None:
            num_reqs = min(
                num_reqs, self.max_num_tokens // self.speculator.num_query_per_req
            )
        # Distribute the remainder evenly so no dummy request exceeds
        # ceil(num_tokens / num_reqs) <= max_model_len tokens.
        num_tokens_per_request = [
            num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
            for i in range(num_reqs)
        ]

        assert sum(num_tokens_per_request) == num_tokens
        num_scheduled_tokens = {
            f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
        }
        dummy_scheduler_output = SchedulerOutput.make_empty()
        dummy_scheduler_output.total_num_scheduled_tokens = num_tokens
        dummy_scheduler_output.num_scheduled_tokens = num_scheduled_tokens

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

        # Get the intermediate tensors for the dummy run.
        intermediate_tensors = None
        if not self.is_first_pp_rank:
            assert self.intermediate_tensors is not None
            intermediate_tensors = self.intermediate_tensors[:num_tokens]

        max_loras = self.lora_config.max_loras if self.lora_config is not None else 0
        with self.maybe_dummy_run_with_lora(
            self.lora_config,
            num_scheduled_tokens=np.array(num_tokens_per_request, dtype=np.int32),
            num_sampled_tokens=None,
            remove_lora=True,
            num_active_loras=max_loras,
        ):
            # Execute the model.
            self.execute_model(
                dummy_scheduler_output,
                intermediate_tensors=intermediate_tensors,
                dummy_run=True,
                skip_attn_for_dummy_run=skip_attn,
                is_profile=is_profile,
                context_len=context_len,
                valid_dummy_state_slots=valid_dummy_state_slots,
            )
        self.kv_connector.set_disabled(False)

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

        assert self.execute_model_state is not None
        input_batch = self.execute_model_state.input_batch
        attn_metadata = self.execute_model_state.attn_metadata
        slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
        hidden_states = self.execute_model_state.hidden_states
        aux_hidden_states = self.execute_model_state.aux_hidden_states
        dp_sync = self.execute_model_state.dp_sync
        self.execute_model_state = None

        self.step_timing.forward_end()

        # dummy run the eagle speculator's propose to ensure DP/EP sync.
        if self.speculator is not None:
            assert self.sampler is not None
            self.step_timing.drafter_start()
            mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
            if self.speculator.supports_mm_inputs:
                all_false = torch.zeros(
                    input_batch.num_tokens, dtype=torch.bool, device="cpu"
                )
                mm_inputs = [], all_false

            # Let the target override the hidden state fed to the drafter
            # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
            # target returns a persistent buffer sized at max_num_batched_tokens;
            # slice to the active token count that propose() expects.
            spec_hidden_states = hidden_states
            if hasattr(self.model, "get_mtp_target_hidden_states"):
                pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
                spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]]  # type: ignore[union-attr]
            if isinstance(self.sampler, GPUWatermarkSampler):
                self.speculator.prepare_watermarking(
                    self.sampler._get_contexts(input_batch.idx_mapping),
                    self.sampler.watermarking.gpu[input_batch.idx_mapping],
                )
            with use_workspace_lane(self._draft_workspace_lane):
                self.speculator.propose(
                    input_batch=input_batch,
                    attn_metadata=attn_metadata,
                    slot_mappings=slot_mappings_by_layer,
                    last_hidden_states=spec_hidden_states,
                    aux_hidden_states=aux_hidden_states,
                    num_sampled=torch.ones(
                        input_batch.num_reqs, dtype=torch.int32, device=self.device
                    ),
                    num_rejected=torch.zeros(
                        input_batch.num_reqs, dtype=torch.int32, device=self.device
                    ),
                    last_sampled=self.req_states.last_sampled_tokens,
                    next_prefill_tokens=self.req_states.next_prefill_tokens,
                    temperature=self.sampler.sampling_states.temperature.gpu,
                    seeds=self.sampler.sampling_states.seeds.gpu,
                    dp_sync=dp_sync,
                    dummy_run=True,
                    skip_attn_for_dummy_run=skip_attn,
                    mm_inputs=mm_inputs,
                    is_profile=is_profile,
                )
            self.step_timing.drafter_end()

        assert hidden_states is not None  # Last PP rank always has hidden_states
        sample_hidden_states = hidden_states[input_batch.logits_indices]
        return hidden_states, sample_hidden_states

    @torch.inference_mode()
    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
        num_reqs = hidden_states.shape[0]
        logits = self.model.compute_logits(hidden_states)
        dummy_input_batch = InputBatch.make_dummy(
            num_reqs, num_reqs, self.input_buffers
        )

        # 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.
        assert self.sampler is not None
        self.sampler(logits, dummy_input_batch)

    @torch.inference_mode()
    def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
        assert self.pooling_runner is not None
        self.pooling_runner.dummy_pooler_run(hidden_states)

    @torch.inference_mode()
    def profile_run(self) -> None:
        if self.supports_mm_inputs and self.is_first_pp_rank:
            mm_config = self.model_config.multimodal_config
            if mm_config is not None and not mm_config.skip_mm_profiling:
                mm_budget = MultiModalBudget(
                    self.vllm_config, self.mm_registry, enable_cache=False
                )
                dummy_mm_inputs = mm_budget.get_dummy_encoder_profile_inputs()
                self.model_state.encoder_runner.profile_encoder_cache(
                    dummy_mm_inputs, mm_budget
                )

        hidden_states, sample_hidden_states = self._dummy_run(
            self.max_num_tokens, skip_attn=True, is_profile=True
        )

        # Only run sampler/pooler on last PP rank (non-last ranks return None).
        if self.is_last_pp_rank:
            assert sample_hidden_states is not None
            if self.pooling_runner is None:
                self._dummy_sampler_run(sample_hidden_states)
            else:
                self._dummy_pooler_run(hidden_states)

        torch.accelerator.synchronize()
        del hidden_states, sample_hidden_states
        self.reset_encoder_cache()
        gc.collect()

    def reset_mm_cache(self) -> None:
        if self.encoder_cache is not None:
            self.encoder_cache.reset_mm_cache()

    def reset_encoder_cache(self) -> None:
        if self.encoder_cache is not None:
            self.encoder_cache.reset_encoder_cache()
        if self.pooling_runner is not None:
            self.pooling_runner.clear()

    @torch.inference_mode()
    def profile_cudagraph_memory(self) -> int:
        """Estimate the GPU memory required to capture CUDA graphs."""
        return _profile_cudagraph_memory(self)

    def needs_cudagraph_capture(self) -> bool:
        """Whether capture_model() has any CUDA graphs left to capture."""
        return (
            self.model_state.supports_mm_inputs
            and self.model_state.encoder_runner.has_cudagraph()
        ) or (
            self.cudagraph_manager is not None
            and self.cudagraph_manager.needs_capture()
        )

    @torch.inference_mode()
    def capture_model(self, *, profile_only: bool = False) -> int:
        assert self.cudagraph_manager is not None
        capture_encoder = (
            self.model_state.supports_mm_inputs
            and self.model_state.encoder_runner.has_cudagraph()
        )
        capture_decoder = self.cudagraph_manager.needs_capture()
        if not capture_encoder and not capture_decoder:
            logger.warning(
                "Skipping encoder and decoder CUDA graph capture. To enable "
                "encoder capture, ensure `cudagraph_mm_encoder` is enabled; "
                "to enable decoder capture, ensure `cudagraph_mode` is not `NONE`."
            )
            return 0

        compilation_counter.num_gpu_runner_capture_triggers += 1

        start_time = time.perf_counter()
        with freeze_gc_for_cudagraph_capture():
            torch.accelerator.empty_cache()
            start_free_gpu_memory = torch.accelerator.get_memory_info()[0]

            with self.maybe_setup_dummy_loras(self.lora_config):
                if capture_encoder:
                    self.model_state.encoder_runner.capture()

                if capture_decoder:
                    input_buffers = self.input_buffers
                    if self.pcp_manager is not None:
                        input_buffers = self.pcp_manager.input_buffers
                    self.cudagraph_manager.capture(
                        self.model,
                        self.model_state,
                        input_buffers,
                        self.intermediate_tensors,
                        self.block_tables,
                        self.attn_groups,
                        self.kv_cache_config,
                        pcp_manager=self.pcp_manager,
                        has_lora=self.lora_config is not None,
                        use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
                        lora_capture_hook=create_lora_capture_hook(
                            self.lora_config, self
                        ),
                    )
                    if self.speculator is not None:
                        with use_workspace_lane(self._draft_workspace_lane):
                            self.speculator.capture()
                    if self.adaptive_verification is not None:
                        with self.step_timing.collect() as timings:
                            for batch in self.adaptive_verification.batches_to_profile(
                                self.cudagraph_manager.captured_token_counts()
                            ):
                                self._dummy_run(**batch)
                        self.adaptive_verification.set_initial_cost_curves(timings)
                    self.kv_connector.reset_capture_state()

            end_free_gpu_memory = torch.accelerator.get_memory_info()[0]

        if not profile_only:
            # Lock workspace to prevent resizing during execution. A resize after
            # capture frees the static cuda graph buffer.
            lock_workspace()

        end_time = time.perf_counter()
        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 _remove_request(self, req_id: str) -> bool:
        # Call model_state.remove_request *before* req_states.remove_request
        # so the model_state can still look up the slot index.
        self.model_state.remove_request(req_id)
        req_idx = self.req_states.remove_request(req_id)
        if req_idx is None:
            return False
        if self.pooling_runner is not None:
            self.pooling_runner.remove_request(req_idx)
        if self.pp_handler is not None:
            self.pp_handler.on_req_idx_freed(req_idx)
        if self.encoder_cache is not None:
            self.encoder_cache.remove_request(req_id)
        if self.prompt_logprobs_worker is not None:
            self.prompt_logprobs_worker.remove_request(req_id)
        self.lora_state.remove_request(req_id)
        return True

    def finish_requests(self, scheduler_output: SchedulerOutput) -> None:
        finished_req_ids = scheduler_output.finished_req_ids
        if self.pooling_runner is not None:
            # Preempted docs keep their query-use reservation until rescheduled.
            self.pooling_runner.on_requests_finished(finished_req_ids)
        preempted_req_ids = scheduler_output.preempted_req_ids
        if preempted_req_ids:
            finished_req_ids = finished_req_ids.union(preempted_req_ids)
        # Sorted so every TP rank frees request slots in the same order.
        # Features like batch-sharded sampling derive rank request ownership
        # from the slot index.
        for req_id in sorted(finished_req_ids):
            self._remove_request(req_id)

    def free_states(self, scheduler_output: SchedulerOutput) -> None:
        if self.encoder_cache is not None:
            for mm_hash in scheduler_output.free_encoder_mm_hashes:
                self.encoder_cache.free_encoder_cache(mm_hash)

    def update_pp_decode_requests(self):
        # For non-last PP ranks, update decode requests with sampler output from
        # the prior step in which they were scheduled (pp_size steps ago).
        if self.pp_handler is not None:
            outputs = self.pp_handler.get_prev_sampled_outputs(
                self.req_states.draft_tokens
            )
            if outputs is not None:
                self.postprocess_sampled(**outputs)

    def add_requests(self, scheduler_output: SchedulerOutput) -> None:
        for new_req_data in scheduler_output.scheduled_new_reqs:
            assert new_req_data.prefill_token_ids is not None
            req_id = new_req_data.req_id

            # Streaming input update: request already exists from a prior
            # chunk. Remove old state so it can be cleanly re-added below
            # with the updated prompt_token_ids and mm_features.
            self._remove_request(req_id)

            prompt_len = new_req_data.prompt_len
            sampling_params = new_req_data.sampling_params
            self.req_states.add_request(
                req_id=req_id,
                prompt_len=prompt_len,
                all_token_ids=new_req_data.prefill_token_ids,
                num_computed_tokens=new_req_data.num_computed_tokens,
                max_tokens=sampling_params.max_tokens if sampling_params else 1,  # type: ignore[arg-type]
            )
            req_index = self.req_states.req_id_to_index[req_id]
            if self.adaptive_verification is not None:
                self.adaptive_verification.add_request(req_index)

            if self.pooling_runner is not None:
                assert new_req_data.pooling_params is not None
                self.pooling_runner.add_request(
                    req_id,
                    req_index,
                    new_req_data.pooling_params,
                    new_req_data.prompt_token_ids,
                )

            if self.encoder_cache is not None:
                self.encoder_cache.add_request(req_id, new_req_data.mm_features)

            self.model_state.add_request(req_index, new_req_data)
            self.block_tables.append_block_ids(
                req_index, new_req_data.block_ids, overwrite=True
            )
            self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)

            if self.is_last_pp_rank and new_req_data.sampling_params is not None:
                assert self.sampler is not None
                self.sampler.add_request(req_index, new_req_data.sampling_params)
                assert self.prompt_logprobs_worker is not None
                self.prompt_logprobs_worker.add_request(
                    req_id, req_index, new_req_data.sampling_params
                )

        if scheduler_output.scheduled_new_reqs:
            self.req_states.apply_staged_writes()
            self.model_state.apply_staged_writes()
        if self.sampler is not None:
            self.sampler.apply_staged_writes()

    def update_requests(self, scheduler_output: SchedulerOutput) -> None:
        # Add new blocks and update num_computed_tokens for the existing requests.
        reqs = scheduler_output.scheduled_cached_reqs
        table_updates = scheduler_output.block_table_updates or {}
        for req_id, block_ids in table_updates.items():
            if (idx := self.req_states.req_id_to_index.get(req_id)) is not None:
                self.block_tables.append_block_ids(idx, block_ids, overwrite=True)
        num_computed_tokens_np = self.req_states.num_computed_tokens_np
        for req_id, num_computed_tokens, req_new_block_ids in zip(
            reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids
        ):
            req_index = self.req_states.req_id_to_index[req_id]
            num_computed_tokens_np[req_index] = num_computed_tokens
            if req_new_block_ids is not None and req_id not in table_updates:
                self.block_tables.append_block_ids(
                    req_index, req_new_block_ids, overwrite=False
                )

        # Update CPU num_computed_prefill_tokens.
        np.minimum(
            self.req_states.num_computed_tokens_np,
            self.req_states.prefill_len.np,
            out=self.req_states.num_computed_prefill_tokens,
        )

        # Zero GPU memory for freshly allocated cache blocks to prevent
        # stale NaN/data from corrupting attention or SSM computation.
        if scheduler_output.new_block_ids_to_zero:
            assert self.kv_block_zeroer is not None
            self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero)

        # Apply copy-on-write block copies for partial prefix-cache hits, after
        # zeroing new blocks and before the forward pass reads them.
        if scheduler_output.kv_cache_block_copies:
            copy_kv_cache_blocks_inplace(
                self.kv_caches,
                self.kv_cache_config.num_blocks,
                scheduler_output.kv_cache_block_copies,
            )

    def gather_batch_req_state(
        self, scheduler_output: SchedulerOutput, dummy_run: bool
    ) -> tuple["BatchReqState | None", int | None]:
        """Gather CPU request state for the scheduled batch, in batch order.
        Returns (batch_state, uniform_decode_token_count)
        """
        num_tokens_per_req = scheduler_output.num_scheduled_tokens
        num_reqs = len(num_tokens_per_req)
        num_toks = scheduler_output.total_num_scheduled_tokens
        max_query_len = max(scheduler_output.num_scheduled_tokens.values())

        if dummy_run:
            # Dummy batches are uniform by construction.
            return None, get_uniform_decode_token_count(
                num_reqs, num_toks, max_query_len, has_prefill=False
            )

        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
        # batch_idx -> req_id
        req_ids = sort_batch_req_ids(
            num_tokens_per_req, draft_tokens, self.decode_query_len
        )

        numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids)
        num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)

        idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids)
        idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs)
        prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np]
        num_computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens[
            idx_mapping_np
        ]
        is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np

        if self.adaptive_verification is not None and draft_tokens:
            num_toks = self.adaptive_verification.get_num_tokens(
                num_tokens_per_req, draft_tokens
            )

        batch_state = BatchReqState(
            req_ids=req_ids,
            num_scheduled_tokens=num_scheduled_tokens,
            num_tokens=num_toks,
            idx_mapping_np=idx_mapping_np,
            prefill_len_np=prefill_len_np,
            num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
            is_prefilling_np=is_prefilling_np,
            has_prefill=bool(is_prefilling_np.any()),
        )
        return batch_state, get_uniform_decode_token_count(
            num_reqs, num_toks, max_query_len, batch_state.has_prefill
        )

    def prepare_inputs(
        self,
        scheduler_output: SchedulerOutput,
        batch_req_state: "BatchReqState",
        batch_desc: BatchExecutionDescriptor,
        num_active_loras: int,
    ) -> InputBatch:
        num_tokens = batch_req_state.num_tokens
        num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
        assert num_tokens > 0
        is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
        if envs.VLLM_MOE_SKIP_PADDING:
            is_padding[:num_tokens].fill_(False)
            is_padding[num_tokens:].fill_(True)

        req_ids = batch_req_state.req_ids
        num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
        idx_mapping_np = batch_req_state.idx_mapping_np
        idx_mapping = async_tensor_h2d(
            idx_mapping_np, device=self.device, dtype=torch.int32
        )
        num_reqs = len(req_ids)

        # Get the number of draft tokens for each request.
        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
        num_draft_tokens_per_req = None
        if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1:
            # No draft token scheduled (common case).
            total_num_draft_tokens = 0
            total_num_logits = num_reqs
            cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
            cu_num_logits = torch.arange(
                num_reqs + 1, device=self.device, dtype=torch.int32
            )
            expanded_idx_mapping = idx_mapping
            expanded_local_pos = torch.zeros(
                num_reqs, dtype=torch.int32, device=self.device
            )
        else:
            num_draft_tokens_per_req = np.fromiter(
                (len(draft_tokens.get(req_id, ())) for req_id in req_ids),
                dtype=np.int32,
                count=num_reqs,
            )
            num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step
            total_num_draft_tokens = int(num_draft_tokens_per_req.sum())
            total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
            num_logits = num_draft_tokens_per_req + num_bonus_tokens
            # combine_sampled_and_draft_tokens places a request's logits rows
            # at [query_end - num_logits, query_end). Fewer query rows than
            # that would silently select the preceding request's hidden states.
            assert (num_scheduled_tokens_np >= num_logits).all()
            cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
            cu_num_logits_np[0] = 0
            np.cumsum(num_logits, out=cu_num_logits_np[1:])
            cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device)

        # The general branch also serves a no-draft batch with k != 1, and
        # compact_batch's budget is only primed when drafts are scheduled.
        adaptive_verification = self.adaptive_verification if draft_tokens else None
        num_scheduled_tokens_upper_bound = num_scheduled_tokens_np
        if adaptive_verification is not None:
            # num_scheduled_tokens represents the draft budget evenly distributed across
            # all verification requests, `reallocate_drafts` will unevenly assign the
            # draft budget to requests on the GPU side only.
            num_scheduled_tokens_np, cu_num_logits_np = (
                adaptive_verification.compact_batch(
                    num_draft_tokens_per_req,
                    num_scheduled_tokens_np,
                    cu_num_logits_np,
                )
            )

        # Get query_start_loc.
        # num_reqs_padded is None for PIECEWISE graphs (no request padding needed)
        num_reqs_padded = batch_desc.num_reqs or num_reqs
        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_np, out=query_start_loc_np[1 : num_reqs + 1])
        # Pad for full CUDA graph mode.
        # Some attention backends like FA3 require query_start_loc to be non-decreasing.
        query_start_loc_np[num_reqs + 1 :] = num_tokens
        query_start_loc = self.input_buffers.query_start_loc
        async_tensor_h2d(query_start_loc_np, out=query_start_loc)
        if adaptive_verification is not None:
            cu_num_logits, query_start_loc, total_num_draft_tokens = (
                adaptive_verification.reallocate_drafts(req_ids, idx_mapping)
            )
            total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
        if num_draft_tokens_per_req is not None:
            expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
                idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len
            )
        query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
        query_start_loc = query_start_loc[: num_reqs_padded + 1]

        # Get prefill tokens if any.
        if batch_req_state.has_prefill:
            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,
            )

        # Prepare positions and seq_lens.
        prepare_pos_seq_lens(
            idx_mapping,
            query_start_loc,
            self.req_states.num_computed_tokens.gpu,
            self.input_buffers.positions,
            self.input_buffers.seq_lens,
        )
        seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]

        # Some input token ids are directly read from the last sampled tokens
        # and draft tokens. Also, get the logits indices to sample tokens from.
        logits_indices = combine_sampled_and_draft_tokens(
            self.input_buffers.input_ids,
            idx_mapping,
            self.req_states.last_sampled_tokens,
            query_start_loc,
            seq_lens,
            self.req_states.prefill_len.gpu,
            self.req_states.draft_tokens,
            cu_num_logits,
            total_num_logits,
            self.model_state.num_new_sampled_tokens_per_step,
        )
        if draft_trunc.ENABLED and envs.VLLM_MOE_SKIP_PADDING:
            est = getattr(self.speculator, "acceptance_estimator", None)
            draft_trunc.mark_dead(is_padding, query_start_loc, cu_num_logits, num_reqs,
                                  self.req_states.draft_tokens.shape[-1],
                                  est.predictions if est is not None else None, idx_mapping)

        fast_prefill = None
        if self.fast_prefill is not None:
            fast_prefill = self.fast_prefill.prepare(
                logits_indices,
                num_reqs,
                cu_num_logits_np,
                batch_req_state.has_prefill,
                batch_desc,
                num_active_loras=num_active_loras,
            )

        # CPU upper bound on seq_lens; padded entries left at zero.
        num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
        seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
        np.add(
            num_computed_tokens_np,
            num_scheduled_tokens_upper_bound,
            out=seq_lens_cpu_upper_bound_np[:num_reqs],
        )
        seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)

        prompt_lens = None
        if self.model_config.rswa_window is not None:
            # prompt_lens is only used in R-SWA case.
            prompt_lens = self.req_states.prompt_len.gpu[idx_mapping]

        input_batch = InputBatch(
            req_ids=req_ids,
            num_reqs=num_reqs,
            num_reqs_after_padding=num_reqs_padded,
            idx_mapping=idx_mapping,
            idx_mapping_np=idx_mapping_np,
            expanded_idx_mapping=expanded_idx_mapping,
            expanded_local_pos=expanded_local_pos,
            num_scheduled_tokens=num_scheduled_tokens_upper_bound,
            num_tokens=num_tokens,
            num_tokens_after_padding=num_tokens_after_padding,
            num_draft_tokens=total_num_draft_tokens,
            num_draft_tokens_per_req=num_draft_tokens_per_req,
            query_start_loc=query_start_loc,
            query_start_loc_np=query_start_loc_np,
            seq_lens=seq_lens,
            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
            dcp_local_seq_lens=None,
            num_computed_tokens_np=num_computed_tokens_np,
            prefill_len_np=batch_req_state.prefill_len_np,
            num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,
            is_prefilling_np=batch_req_state.is_prefilling_np,
            has_prefill=batch_req_state.has_prefill,
            input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
            positions=self.input_buffers.positions[:num_tokens_after_padding],
            is_padding=is_padding,
            logits_indices=logits_indices,
            cu_num_logits=cu_num_logits,
            cu_num_logits_np=cu_num_logits_np,
            has_structured_output_reqs=scheduler_output.has_structured_output_requests,
            prompt_lens=prompt_lens,
            fast_prefill=fast_prefill,
            max_query_len=(
                int(num_scheduled_tokens_upper_bound.max())
                if adaptive_verification is not None
                else None
            ),
        )
        if self.pcp_manager is not None:
            input_batch = self.pcp_manager.partition_batch(input_batch, batch_desc)
        return input_batch

    def prepare_attn(
        self, input_batch: InputBatch
    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
        if self.pcp_manager is not None:
            return self.pcp_manager.prepare_attn(input_batch)

        # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
        block_tables = self.block_tables.gather_block_tables(
            input_batch.idx_mapping, num_reqs_padded=input_batch.num_reqs_after_padding
        )
        # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
        # Kernel pads beyond num_tokens with PAD_SLOT_ID.
        slot_mappings = self.block_tables.compute_slot_mappings(
            input_batch.idx_mapping,
            input_batch.query_start_loc,
            input_batch.positions,
            num_tokens_padded=input_batch.num_tokens_after_padding,
        )
        return block_tables, slot_mappings

    def prepare_dummy_attn(
        self, input_batch: InputBatch, valid_state_slots: bool = False
    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
        block_table_provider = self.pcp_manager or self.block_tables
        block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs)
        if valid_state_slots:
            state_slots = torch.arange(
                1, input_batch.num_reqs + 1, dtype=torch.int32, device=self.device
            )
            for block_table in block_tables:
                block_table[:, 0].copy_(state_slots)
        slot_mappings = block_table_provider.get_dummy_slot_mappings(
            input_batch.num_tokens
        )
        return block_tables, slot_mappings

    def sample(
        self,
        hidden_states: torch.Tensor,
        input_batch: InputBatch,
        grammar_output: GrammarOutput | None,
    ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
        shard_metadata = None
        global_input_batch = input_batch
        if self.batch_sharder is not None:
            # Shard the inputs along the batch dimension to sample in parallel
            # across TP ranks.
            input_batch, sorted_logits_indices, grammar_output, shard_metadata = (
                self.batch_sharder.shard_sampler_inputs(input_batch, grammar_output)
            )
            # The hidden states must be gathered in rank-owner-sorted order
            # before computing the partial-vocab logits, so that the all-to-all
            # produces full-vocab logits for just the locally-owned requests.
            sample_hidden_states = hidden_states[sorted_logits_indices]
            local_logits = self.model.compute_logits_local(sample_hidden_states)
            logits = all_to_all_logits(local_logits, shard_metadata)
            logits = logits[:, : self.vocab_size]
        else:
            sample_hidden_states = hidden_states[input_batch.logits_indices]
            logits = self.model.compute_logits(sample_hidden_states)

        # A diffusion prefill has no logit rows even when a bitmask row
        # arrived for it.
        if grammar_output is not None and logits.shape[0] > 0:
            # Apply grammar bitmask to the logits in-place.
            assert self.structured_outputs_worker is not None
            self.structured_outputs_worker.apply_grammar_bitmask(
                logits,
                input_batch,
                grammar_output.structured_output_request_ids,
                grammar_output.grammar_bitmask,
            )

        sampler_output: SamplerOutput | None
        if input_batch.num_reqs == 0:
            # This rank owns no requests this step. It contributes an
            # all-padding block to the gather below.
            sampler_output = None
        elif input_batch.num_draft_tokens == 0 or self.rejection_sampler is None:
            assert self.sampler is not None
            sampler_output = self.sampler(logits, input_batch)
        else:
            # Rejection sampling for spec decoding.
            assert self.rejection_sampler is not None
            assert self.speculator is not None
            sampler_output = self.rejection_sampler(
                logits,
                input_batch,
                # Draft logits are needed for probabilistic rejection sampling.
                self.speculator.draft_logits,
            )

        if shard_metadata is not None:
            # Gather the sharded sampler outputs from the TP ranks into a single
            # sampler output.
            assert self.sampler is not None
            sampler_output = gather_sampler_output(
                sampler_output,
                shard_metadata,
                device=self.device,
                global_batch=global_input_batch,
                local_batch=input_batch,
                gather_num_nans=self.sampler.compute_nans,
                logprobs_dims=self.sampler.get_logprobs_dims(
                    global_input_batch.idx_mapping_np,
                    # Rejection sampler does not return logprob token ids.
                    include_token_ids=(
                        global_input_batch.num_draft_tokens == 0
                        or self.rejection_sampler is None
                    ),
                ),
            )

        assert sampler_output is not None
        return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected

    def postprocess_sampled(
        self,
        idx_mapping: torch.Tensor,  # May include -1 for masked entries
        sampled_tokens: torch.Tensor,
        num_sampled: torch.Tensor,
        num_rejected: torch.Tensor,
        query_start_loc: torch.Tensor | None = None,
    ) -> None:
        # Update the number of computed tokens.
        output_bin_counts = None
        if self.is_last_pp_rank:
            assert self.sampler is not None
            output_bin_counts = self.sampler.penalties_state.output_bin_counts
        post_update(
            idx_mapping,
            self.req_states.num_computed_tokens.gpu,
            self.req_states.last_sampled_tokens,
            output_bin_counts,
            sampled_tokens,
            num_sampled,
            num_rejected,
            query_start_loc,
            self.req_states.all_token_ids.gpu,
            self.req_states.total_len.gpu,
        )

        self.model_state.postprocess_state(
            idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu
        )

    def _merge_ec_connector_no_forward(
        self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput
    ) -> ModelRunnerOutput:
        """Let the EC connector send/recv on a step with no work to run."""
        return ModelRunnerOutput.with_ec_conn_output(
            output,
            self.ec_connector.no_forward(scheduler_output).ec_connector_output,
        )

    @torch.inference_mode()
    def execute_model(
        self,
        scheduler_output: SchedulerOutput,
        intermediate_tensors: IntermediateTensors | None = None,
        dummy_run: bool = False,
        skip_attn_for_dummy_run: bool = False,
        is_profile: bool = False,
        context_len: int = 0,
        valid_dummy_state_slots: bool = False,
    ) -> ModelRunnerOutput | IntermediateTensors | None:
        if not dummy_run:
            # Update the request states.
            self.update_pp_decode_requests()
            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 self.aux_output_connector is not None:
                # Register this step before the GPU forward.
                self.aux_output_connector.begin_step(
                    scheduler_output.aux_output_connector_metadata
                )
            if scheduler_output.total_num_scheduled_tokens == 0:
                # No need to run the model.
                empty_output = self.kv_connector.no_forward(scheduler_output)
                return self._merge_ec_connector_no_forward(
                    scheduler_output, empty_output
                )

        # Get batch descriptor and sync across DP ranks.
        num_reqs = len(scheduler_output.num_scheduled_tokens)
        num_toks = scheduler_output.total_num_scheduled_tokens
        max_query_len = max(scheduler_output.num_scheduled_tokens.values())
        batch_req_state, uniform_tok_count = self.gather_batch_req_state(
            scheduler_output, dummy_run
        )
        if batch_req_state is not None:
            num_toks = batch_req_state.num_tokens
            if self.pcp_manager is not None:
                num_toks = self.pcp_manager.get_num_tokens_for_dispatch(
                    batch_req_state.num_scheduled_tokens,
                    batch_req_state.is_prefilling_np,
                )

        num_active_loras = 0
        if self.lora_config:
            req_ids = list(scheduler_output.num_scheduled_tokens.keys())
            num_active_loras = get_num_active_loras_for_dispatch(
                self.lora_config, self.lora_state, req_ids, dummy_run
            )

        skip_compiled = False
        if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
            # Encoder-decoder models such as Whisper should run eager/non-compiled
            # when encoder inputs are scheduled, because this step updates
            # cross-attention cache with dynamic encoder outputs.
            skip_compiled = True

        batch_desc, dp_sync = dispatch_cg_and_sync_dp(
            self.cudagraph_manager,
            num_reqs,
            num_toks,
            uniform_tok_count,
            self.dp_size,
            self.dp_rank,
            max_query_len=max_query_len,
            need_eager=is_profile or skip_compiled,
            num_active_loras=num_active_loras,
            parallel_config=self.parallel_config,
            allow_ubatching=(
                self.ubatch_runner is not None and not skip_attn_for_dummy_run
            ),
            uniform_decode=uniform_tok_count == self.decode_query_len,
        )

        if batch_desc.num_tokens == 0:
            # All DP ranks have zero tokens to run.
            empty_output = self.kv_connector.no_forward(scheduler_output)
            return self._merge_ec_connector_no_forward(scheduler_output, empty_output)

        cudagraph_stats = None
        if not dummy_run:
            # Common case.
            # Prepare all the inputs and copy to the input buffers.
            if self.observability_config.cudagraph_metrics:
                cudagraph_stats = make_cudagraph_stats(batch_desc, num_toks)
            assert batch_req_state is not None
            input_batch = self.prepare_inputs(
                scheduler_output, batch_req_state, batch_desc, num_active_loras
            )
            block_tables, slot_mappings = self.prepare_attn(input_batch)
            # Mamba "align" pre-copy: migrate recurrent state across block
            # boundaries before the forward. Runs only on real batches, and
            # before model_state.prepare_attn gathers num_accepted_tokens so the
            # boundary reset is visible to the attention metadata.
            self.model_state.preprocess_state(
                input_batch,
                block_tables,
                self.kv_cache_config,
                self.req_states.num_computed_tokens.gpu,
            )

            if self.lora_config:
                # Activate LoRA adapters.
                lora_inputs = self.lora_state.make_lora_inputs(
                    input_batch.req_ids,
                    input_batch.idx_mapping_np,
                    input_batch.num_scheduled_tokens,
                )
                self._set_active_loras(*lora_inputs)
        else:
            # No actual tokens to run. A dummy run for DP or memory profiling.
            dummy_num_reqs = batch_desc.num_reqs or num_reqs
            input_batch = InputBatch.make_dummy(
                dummy_num_reqs,
                batch_desc.num_tokens,
                self.input_buffers,
                max_query_len=batch_desc.max_query_len,
                # Profiling and warmup must route the dummy tokens to experts
                # so MoE memory is measured and MoE kernels are exercised.
                is_padding=not is_profile,
            )
            if self.pcp_manager is not None:
                input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)
            if skip_attn_for_dummy_run:
                assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
                    "Attention metadata must be prepared for dummy runs when using "
                    "FULL cudagraph mode."
                )
                block_tables, slot_mappings = None, None
            else:
                block_tables, slot_mappings = self.prepare_dummy_attn(
                    input_batch, valid_dummy_state_slots
                )
                if context_len:
                    set_dummy_context(
                        input_batch,
                        self.block_tables,
                        context_len,
                        self.kv_cache_config.num_blocks,
                        self.max_model_len,
                        input_block_tables=block_tables,
                    )

        if self.dcp_size > 1 and not (dummy_run and skip_attn_for_dummy_run):
            input_buffers = self.input_buffers
            if self.pcp_manager is not None:
                input_buffers = self.pcp_manager.input_buffers
            input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
                input_buffers.dcp_local_seq_lens,
                input_batch.seq_lens,
                input_batch.num_reqs,
                self.dcp_size,
                self.dcp_rank,
                self.cp_interleave,
                num_reqs_padded=input_batch.num_reqs_after_padding,
            )

        attn_metadata = None
        slot_mappings_by_layer = None
        ubatch_state: UBatchState | None = None
        if batch_desc.num_ubatches > 1:
            assert self.ubatch_runner is not None
            assert block_tables is not None and slot_mappings is not None
            ubatch_state = self.ubatch_runner.prepare(
                input_batch,
                block_tables,
                slot_mappings,
                cg_mode=batch_desc.cg_mode,
                for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
            )
        elif not (dummy_run and skip_attn_for_dummy_run):
            assert slot_mappings is not None
            slot_mappings_by_layer = build_slot_mappings_by_layer(
                slot_mappings, self.kv_cache_config
            )
            assert block_tables is not None
            attn_groups = self.attn_groups
            if dummy_run and is_profile:
                # Mamba layers take a cheap warmup path with no metadata;
                # attention metadata is still built so those kernels tune.
                attn_groups = [
                    [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)]
                    for groups in attn_groups
                ]
            attn_metadata = self.model_state.prepare_attn(
                input_batch,
                batch_desc.cg_mode,
                block_tables,
                slot_mappings,
                attn_groups,
                self.kv_cache_config,
                # FULL replay reads capture-time metadata buffers. Re-stage them
                # from the zeroed dummy block tables instead of retaining state
                # indices from the previous real batch.
                for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL,
            )

        input_ids = input_batch.input_ids
        inputs_embeds = None
        ec_connector_output = None
        if self.uses_inputs_embeds and self.is_first_pp_rank:
            # Prepare inputs_embeds (MM encoder outputs and/or prompt_embeds
            # overlay). Only first PP rank prepares them.
            if dummy_run:
                # Obtain embeddings of correct shape for compiled model.
                inputs_embeds = self.model_state.dummy_inputs_embeds(
                    input_batch.num_tokens_after_padding
                )
            else:
                scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
                if self.supports_mm_inputs and self.lora_config is not None:
                    set_active_mm_loras(
                        model=self.model,
                        lora_manager=self.lora_manager,
                        encoder_cache=self.encoder_cache,
                        req_id_to_index=self.req_states.req_id_to_index,
                        lora_state=self.lora_state,
                        scheduled_encoder_inputs=scheduled_encoder_inputs,
                    )
                with self.ec_connector.maybe_get_output(
                    scheduler_output
                ) as ec_connector_output:
                    inputs_embeds = self.model_state.prepare_inputs_embeds(
                        scheduled_encoder_inputs, input_batch, self.req_states
                    )
            if inputs_embeds is not None and not requires_raw_input_tokens(self.model):
                input_ids = None

        model_inputs = {
            "input_ids": input_ids,
            "positions": input_batch.positions,
            "inputs_embeds": inputs_embeds,
            "intermediate_tensors": None,
            # NOTE: Values returned by `prepare_inputs` will override the default
            # values above.
            **self.model_state.prepare_inputs(input_batch, self.req_states),
        }
        if not self.is_first_pp_rank:
            # Update for non-first PP ranks.
            model_inputs["input_ids"] = None
            model_inputs["inputs_embeds"] = None

            # Prepare the intermediate tensors.
            assert intermediate_tensors is not None
            assert self.intermediate_tensors is not None
            n = input_batch.num_tokens_after_padding
            new_tensors = {
                k: v[:n]
                if dummy_run
                else v[:n].copy_(intermediate_tensors.tensors[k][:n])
                for k, v in self.intermediate_tensors.tensors.items()
            }
            model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors)
            del intermediate_tensors

        # Update the EPLB meta.
        ubatch_slices = ubatch_state.slices if ubatch_state is not None else None
        self.eplb.prepare_forward(
            self.model_config, input_batch.num_tokens, ubatch_slices
        )

        self.step_timing.record_batch(
            input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL
        )
        self.step_timing.forward_start()

        connector_kwargs = dict(
            scheduler_output=scheduler_output,
            request_state_indices=input_batch.idx_mapping,
            request_ids=input_batch.req_ids,
            num_tokens=input_batch.num_tokens,
        )

        # Run model.
        if batch_desc.cg_mode == CUDAGraphMode.FULL:
            # Use explicit cudagraph replay for FULL mode.
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
            assert self.cudagraph_manager is not None
            self.kv_connector.pre_forward(
                **connector_kwargs, attn_metadata=attn_metadata
            )
            model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
        else:
            # For piecewise and eager mode, just call model().
            batch_descriptor = BatchDescriptor(
                num_tokens=input_batch.num_tokens_after_padding,
                has_lora=self.lora_config is not None,
                num_active_loras=batch_desc.num_active_loras,
            )

            with set_forward_context(
                attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
                cudagraph_runtime_mode=batch_desc.cg_mode,
                num_tokens_across_dp=(
                    dp_sync.num_tokens_across_dp if dp_sync is not None else None
                ),
                batch_descriptor=batch_descriptor,
                ubatch_slices=ubatch_slices,
                slot_mapping=slot_mappings_by_layer,
                skip_compiled=skip_compiled,
                is_padding=input_batch.is_padding,
            ):
                self.kv_connector.pre_forward(**connector_kwargs)
                if ubatch_state is not None:
                    assert self.ubatch_runner is not None
                    model_output = self.ubatch_runner.run(
                        self.model, model_inputs, ubatch_state
                    )
                elif batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
                    # Run the PIECEWISE graph (compiled PW cudagraph or breakable
                    # cudagraph, chosen inside run_pw_graph). cg_mode is only
                    # PIECEWISE after the cudagraph manager exists.
                    assert self.cudagraph_manager is not None
                    model_output = self.cudagraph_manager.run_pw_graph(
                        self.model, model_inputs
                    )
                else:
                    # Eager (NONE): call the raw model directly.
                    model_output = self.model(**model_inputs)

        self.kv_connector.finish_forward()

        if self.is_last_pp_rank:
            if self.use_aux_hidden_state_outputs:
                assert isinstance(model_output, tuple)
                hidden_states, aux_hidden_states = model_output
            else:
                assert isinstance(model_output, torch.Tensor)
                hidden_states = model_output
                aux_hidden_states = None
            output_intermediate_tensors = None
        else:
            assert isinstance(model_output, IntermediateTensors)
            hidden_states = None
            aux_hidden_states = None
            output_intermediate_tensors = model_output

        finished_req_ids = scheduler_output.finished_req_ids
        self.execute_model_state = ExecuteModelState(
            input_batch=input_batch,
            attn_metadata=attn_metadata,
            slot_mappings_by_layer=slot_mappings_by_layer,
            hidden_states=hidden_states,
            aux_hidden_states=aux_hidden_states,
            dp_sync=dp_sync,
            finished_req_ids=finished_req_ids,
            ec_connector_output=ec_connector_output,
            cudagraph_stats=cudagraph_stats,
        )

        if not self.is_last_pp_rank:
            # Non-last PP rank: return IntermediateTensors for sending.
            assert output_intermediate_tensors is not None
            assert self.pp_handler is not None
            return self.pp_handler.relay_aux_hidden_states(
                model_inputs["intermediate_tensors"], output_intermediate_tensors
            )
        return None

    @torch.inference_mode()
    @step_eplb_after()
    def sample_tokens(
        self, grammar_output: GrammarOutput | None
    ) -> AsyncOutput | ModelRunnerOutput | None:
        if self.execute_model_state is None:
            # The prior execute_model call must have failed.
            return None

        input_batch = self.execute_model_state.input_batch
        attn_metadata = self.execute_model_state.attn_metadata
        slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
        hidden_states = self.execute_model_state.hidden_states
        aux_hidden_states = self.execute_model_state.aux_hidden_states
        dp_sync = self.execute_model_state.dp_sync
        finished_req_ids = self.execute_model_state.finished_req_ids
        ec_connector_output = self.execute_model_state.ec_connector_output
        cudagraph_stats = self.execute_model_state.cudagraph_stats
        self.execute_model_state = None

        if not self.is_last_pp_rank:
            # 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.
            assert self.pp_handler is not None
            all_decode_next = self.pp_handler.receive(input_batch)
            # Optimistically update num_computed_tokens for entire batch here.
            # Will be adjusted for rejections if necessary in update_requests.
            self.postprocess_num_computed_tokens(input_batch)
            if not all_decode_next:
                # Might contain non-final prefill chunks, which will be scheduled
                # in the immediate next step (rather than in pp_size steps).
                self.model_state.postprocess_state(input_batch.idx_mapping, 0)

            # Post-step KV connector related operations.
            kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
            # The first PP rank holds the encoder cache, so pass its EC output on.
            output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
            return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)

        # Last rank: sample tokens
        assert hidden_states is not None
        draft_hidden_states = hidden_states
        if self.pcp_manager is not None:
            hidden_states, aux_hidden_states, input_batch = (
                self.pcp_manager.restore_for_sampling(hidden_states, aux_hidden_states)
            )

        sampler_output, num_sampled, num_rejected = self.sample(
            hidden_states, input_batch, grammar_output
        )

        if self.pp_handler is not None:
            # Broadcast to non-last PP ranks (handles spec decode multi-token).
            self.pp_handler.broadcast(
                sampler_output.sampled_token_ids,
                num_sampled,
                num_rejected,
                input_batch,
            )

        assert self.prompt_logprobs_worker is not None
        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
            self.model.compute_logits,
            hidden_states,
            input_batch,
            self.req_states.all_token_ids.gpu,
            self.req_states.num_computed_tokens.gpu,
            self.req_states.prompt_len.np,
        )

        # 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
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
            cudagraph_stats=cudagraph_stats,
        )
        pending_aux_output = None
        if self.aux_output_connector is not None:
            pending_aux_output = self.aux_output_connector.prepare_output(input_batch)

        # Start async output copy here so that it can overlap with speculator proposal.
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,
            num_sampled_tokens=num_sampled,
            main_stream=self.main_stream,
            copy_stream=self.output_copy_stream,
            check_ep_fault=self.check_ep_fault,
            pending_aux_output=pending_aux_output,
        )

        mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
        # The encoder runner exists only on the first PP rank, so later ranks
        # have no cached embeddings to gather.
        if (
            self.speculator is not None
            and self.speculator.supports_mm_inputs
            and self.model_state.supports_mm_inputs
        ):
            # Get cached multimodal embeddings for draft forward.
            # NOTE: This is done here because postprocess updates
            # num_computed_prefill_tokens.
            # The EAGLE/MTP drafter reads one position ahead of the target.
            # TODO(TheEpicDolphin): Gather MM embeddings for all speculative
            # steps during multi-module MTP.
            mm_inputs = self.model_state.gather_mm_embeddings(
                input_batch, draft_lookahead=1
            )

        # 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_sampled(
            input_batch.idx_mapping,
            sampler_output.sampled_token_ids,
            num_sampled,
            num_rejected,
            input_batch.query_start_loc,
        )

        if self.speculator is not None:
            assert self.sampler is not None
            if isinstance(self.speculator, DraftModelSpeculator) and draft_trunc.graded():
                self.speculator.observe_verification(
                    input_batch.idx_mapping, num_sampled, num_rejected
                )
            if draft_trunc.LOG:
                draft_trunc.log_step(
                    self.speculator.draft_token_confidence_probs[input_batch.idx_mapping],
                    input_batch.cu_num_logits, num_sampled,
                    getattr(self.speculator, "acceptance_estimator", None), input_batch.idx_mapping,
                )
            # Let the target override the hidden state fed to the drafter
            # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
            # target returns a persistent buffer sized at max_num_batched_tokens;
            # slice to the active token count that propose() expects.
            spec_hidden_states = draft_hidden_states
            if hasattr(self.model, "get_mtp_target_hidden_states"):
                pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()
                spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]
            if isinstance(self.sampler, GPUWatermarkSampler):
                self.speculator.prepare_watermarking(
                    self.sampler._get_contexts(input_batch.idx_mapping),
                    self.sampler.watermarking.gpu[input_batch.idx_mapping],
                )
            with use_workspace_lane(self._draft_workspace_lane):
                draft_tokens = self.speculator.propose(
                    input_batch,
                    attn_metadata,
                    slot_mappings_by_layer,
                    spec_hidden_states,
                    aux_hidden_states,
                    num_sampled,
                    num_rejected,
                    self.req_states.last_sampled_tokens,
                    self.req_states.next_prefill_tokens,
                    self.sampler.sampling_states.temperature.gpu,
                    self.sampler.sampling_states.seeds.gpu,
                    dp_sync=dp_sync,
                    mm_inputs=mm_inputs,
                )
            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
            if self.adaptive_verification is not None:
                self.adaptive_verification.record_confidences(
                    self.speculator.draft_token_confidence_probs, input_batch
                )

        if self.num_speculative_steps > 0:
            # Spec-decode and diffusion LLMs both use draft tokens but the latter does
            # not have a speculator (i.e. self.speculator is None)
            self.draft_tokens_handler.set_draft_tokens(
                input_batch,
                self.req_states.draft_tokens[input_batch.idx_mapping],
            )
            if self.pp_handler is not None:
                self.pp_handler.broadcast_drafts(
                    self.req_states.draft_tokens, input_batch
                )

        # Post-step KV connector related operations.
        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
        model_runner_output.kv_connector_output = kv_connector_output
        model_runner_output.ec_connector_output = ec_connector_output

        return async_output

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

    @torch.inference_mode()
    @step_eplb_after()
    def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None:
        if self.execute_model_state is None:
            # The prior execute_model call must have failed.
            return None

        input_batch = self.execute_model_state.input_batch
        hidden_states = self.execute_model_state.hidden_states
        finished_req_ids = self.execute_model_state.finished_req_ids
        ec_connector_output = self.execute_model_state.ec_connector_output
        self.execute_model_state = None

        # Post-step KV connector related operations.
        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)

        if not self.is_last_pp_rank:
            self.postprocess_num_computed_tokens(input_batch)
            output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
            return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)

        assert self.pooling_runner is not None
        pooler_output, finished_mask = self.pooling_runner.pool(
            hidden_states, input_batch, self.req_states
        )

        # Build the model runner output.
        model_runner_output = ModelRunnerOutput(
            req_ids=input_batch.req_ids,
            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
            kv_connector_output=kv_connector_output,
            ec_connector_output=ec_connector_output,
        )
        async_output = AsyncPoolingOutput(
            model_runner_output=model_runner_output,
            pooler_output=pooler_output,
            finished_mask=finished_mask,
            main_stream=self.main_stream,
            copy_stream=self.output_copy_stream,
        )

        self.postprocess_num_computed_tokens(input_batch)
        return async_output

    def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None:
        # Update the number of computed tokens.
        post_update_num_computed_tokens(
            input_batch.idx_mapping,
            self.req_states.num_computed_tokens.gpu,
            input_batch.query_start_loc,
        )

    def shutdown(self) -> None:
        """Release GPU tensors (model weights, KV caches, workspace) so that
        memory is reclaimable when running in the same process."""
        torch.accelerator.synchronize()
        if self.aux_output_connector is not None:
            self.aux_output_connector.close()
        set_offloader(None)
        self.cudagraph_manager = None
        self.fast_prefill = None
        self.pooling_runner = None
        if hasattr(self, "kv_caches"):
            self.kv_caches.clear()
        if hasattr(self, "attn_groups"):
            self.attn_groups.clear()
        if hasattr(self, "kv_cache_config"):
            del self.kv_cache_config
        if hasattr(self, "model_state") and self.model_state.supports_mm_inputs:
            self.model_state.encoder_runner.clear()
        free_before_shutdown(self.vllm_config)
        if hasattr(self, "model_state"):
            del self.model_state
        # Detach the layer-level KV/state cache tensors before dropping the
        # models; the model objects can outlive this runner.
        speculator = getattr(self, "speculator", None)
        if speculator is not None:
            if draft_model := getattr(speculator, "model", None):
                clear_layer_kv_caches(draft_model.modules())
            self.speculator = None
        if hasattr(self, "model"):
            clear_layer_kv_caches(self.model.modules())
            del self.model

        gc.collect()
        torch.accelerator.empty_cache()
        logger.debug("Cleaned up model weights, KV caches, and workspace")

    ########### EPLB methods start ###########
    @property
    def eplb_state(self):
        return self.eplb.state

    @eplb_state.setter
    def eplb_state(self, state) -> None:
        self.eplb.state = state

    @property
    def eep_eplb_suppressed(self) -> bool:
        return self.eplb.suppressed

    @eep_eplb_suppressed.setter
    def eep_eplb_suppressed(self, suppressed: bool) -> None:
        self.eplb.suppressed = suppressed

    def setup_eplb_from_mapping(
        self, expanded_physical_to_logical: torch.Tensor
    ) -> None:
        self.eplb.setup_from_mapping(
            self.model, self.model_config, expanded_physical_to_logical
        )

    ########### EPLB methods end ###########

    # Out-of-tree hardware runners can select a PCP manager class.
    @property
    def pcp_manager_cls(self) -> type[pcp.PCPManager]:
        return pcp.PCPManager


class ExecuteModelState(NamedTuple):
    input_batch: InputBatch
    attn_metadata: dict[str, Any] | None
    slot_mappings_by_layer: dict[str, torch.Tensor] | None
    hidden_states: torch.Tensor | None
    aux_hidden_states: list[torch.Tensor] | None
    dp_sync: DPSyncState | None
    finished_req_ids: set[str]
    ec_connector_output: ECConnectorOutput | None
    cudagraph_stats: CUDAGraphStat | None


class BatchReqState(NamedTuple):
    """CPU request state for a scheduled batch, in batch (sorted) order."""

    req_ids: list[str]
    num_scheduled_tokens: np.ndarray  # [num_reqs]
    # May be less than scheduler_output.total_num_scheduled_tokens:
    # adaptive verification trims the draft budget before running.
    num_tokens: int
    idx_mapping_np: np.ndarray  # [num_reqs]
    prefill_len_np: np.ndarray  # [num_reqs]
    num_computed_prefill_tokens_np: np.ndarray  # [num_reqs]
    is_prefilling_np: np.ndarray  # [num_reqs]
    has_prefill: bool


def sort_batch_req_ids(
    num_tokens_per_req: dict[str, int],
    draft_tokens: dict[str, list[int]],
    decode_query_len: int,
) -> list[str]:
    # Order verification/decode -> short_extend -> prefill;
    # split_decodes_and_prefills relies on decode-like requests leading.
    key = lambda r: (
        not draft_tokens.get(r),
        (num := num_tokens_per_req[r]) != decode_query_len,
        num,
    )
    return sorted(num_tokens_per_req, key=key)