【vllm】（四）vLLM v1 Worker — 模块超深度逐行分析之二

四、GPUModelRunner 初始化深度解析

4.1 类定位与 Mixin 组合

python 复制代码

class GPUModelRunner(
    LoRAModelRunnerMixin,
    KVConnectorModelRunnerMixin,
    ECConnectorModelRunnerMixin
):

GPUModelRunner 通过 多重 Mixin 组合 获得额外能力：

LoRAModelRunnerMixin：LoRA 模型加载、动态 LoRA 管理
KVConnectorModelRunnerMixin：KV cache 跨节点传输的 preemption 处理和 metadata 管理
ECConnectorModelRunnerMixin：Encoder-Decoder 连接器的输出管理

这种组合优于多重继承，因为 Mixin 不定义 __init__，不引入菱形继承问题。

4.2 init() 逐行深度解析（行 397-857）

4.2.1 配置提取（行 397-419）

python 复制代码

self.vllm_config = vllm_config
self.model_config = vllm_config.model_config
self.cache_config = vllm_config.cache_config
self.offload_config = vllm_config.offload_config
self.compilation_config = vllm_config.compilation_config
self.lora_config = vllm_config.lora_config
self.load_config = vllm_config.load_config
self.parallel_config = vllm_config.parallel_config
self.scheduler_config = vllm_config.scheduler_config
self.speculative_config = vllm_config.speculative_config
self.observability_config = vllm_config.observability_config

设计目的 ：将全局配置根对象和常用子配置提取为实例属性，避免后续代码中反复写 self.vllm_config.model_config.xxx。这是 Python 中常见的 配置扁平化 模式。

4.2.2 模型配置（行 420-440）

python 复制代码

model_config = self.model_config
self.device = device
self.pin_memory = is_pin_memory_available()
self.dtype = self.model_config.dtype

dtype：模型计算的数据类型（float16/bfloat16/float32），由模型配置决定，影响所有张量分配。

python 复制代码

self.kv_cache_dtype = kv_cache_dtype_str_to_dtype(
    cache_config.cache_dtype, self.model_config
)

kv_cache_dtype：KV cache 的数据类型，可能与模型 dtype 不同。支持：

与模型相同（auto）
FP8 E4M3 / FP8 E5M2（节省 KV cache 显存，cache_dtype="fp8"）
自定义量化格式

kv_cache_dtype_str_to_dtype() 将字符串配置转换为实际 torch.dtype，并处理 FP8 量化相关逻辑。

python 复制代码

self.is_pooling_model = model_config.runner_type == "pooling"
self.enable_prompt_embeds = model_config.enable_prompt_embeds
self.is_multimodal_raw_input_only_model = model_config.is_multimodal_raw_input_only_model
self.is_multimodal_pruning_enabled = False  # 后续在 load_model 中更新
self.requires_sequential_video_encoding = False  # 后续在 load_model 中更新
self.routed_experts_initialized = False  # init_routed_experts_capturer 后置 True

is_pooling_model：是否为 pooling 类模型（Embedding/Reward/Classification），决定输出路径（不走采样，走 pooler）。

enable_prompt_embeds：是否支持用户直接提供 embedding（而非 token IDs），用于 embedding-as-input 场景。

is_multimodal_pruning_enabled / requires_sequential_video_encoding：多模态优化标志，在 load_model() 后根据模型能力更新。

python 复制代码

self.max_model_len = model_config.max_model_len

max_model_len ：模型支持的最大序列长度，决定 block_table、positions 等缓冲区的大小上界。可在运行时通过 update_max_model_len() 缩小（auto-fit 场景）。

4.2.3 调度配置（行 435-445）

python 复制代码

self.calculate_kv_scales = self.cache_config.calculate_kv_scales

calculate_kv_scales：是否需要在首次前向时计算 FP8 KV cache 的缩放因子。默认 True，首次前向后设为 False。为 True 时强制禁用 CUDA Graph（因为计算缩放因子涉及动态操作）。

python 复制代码

self.dcp_world_size = self.parallel_config.decode_context_parallel_size
self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group

DCP (Decode Context Parallel)：

dcp_world_size：DCP 并行度，将长上下文分片到多个 rank
dcp_rank：当前 rank 在 DCP 组内的位置
影响：KV cache 的 seq_len 需要除以 dcp_world_size，attention 需要跨 DCP rank 通信

python 复制代码

self.max_num_tokens = scheduler_config.max_num_batched_tokens
self.max_num_reqs = scheduler_config.max_num_seqs

max_num_tokens：单步最大调度的 token 总数（prefill + decode），决定所有持久缓冲区的第一维度。

max_num_reqs：单步最大并发请求数，决定 block_table、seq_lens 等缓冲区的第一维度。

python 复制代码

self.broadcast_pp_output = (
    self.parallel_config.distributed_executor_backend == "external_launcher"
    and len(get_pp_group().ranks) > 1
)

broadcast_pp_output：使用 torchrun（external_launcher）时，PP 的 last rank 需要通过 broadcast 将 logits 发送给其他 rank。这是因为 torchrun 模式下 rank 间没有直接通信通道。

4.2.4 模型头部配置（行 447-456）

python 复制代码

self.num_query_heads = model_config.get_num_attention_heads(parallel_config)
self.inputs_embeds_size = model_config.get_inputs_embeds_size()
self.use_alibi = model_config.uses_alibi

num_query_heads：当前 TP rank 管理的 query 头数。如果 TP=2 且模型有 32 头，每个 rank 管理 16 头。

inputs_embeds_size ：embedding 层的输出维度（即 hidden_size），用于分配 inputs_embeds 缓冲区。0 表示模型不支持 embedding 输入。

use_alibi：是否使用 ALiBi (Attention with Linear Biases) 位置编码，影响 attention 计算方式。

python 复制代码

self.cascade_attn_enabled = not self.model_config.disable_cascade_attn
self.is_mm_prefix_lm = self.model_config.is_mm_prefix_lm

cascade_attn_enabled：是否启用 Cascade Attention（前缀 KV 共享优化）。当 batch 中多个请求共享公共前缀（如 system prompt）时，前缀部分只需计算一次 attention。

is_mm_prefix_lm：是否为多模态前缀语言模型（如 LLaVA，视觉 token 作为前缀）。

4.2.5 多模态配置（行 458-475）

python 复制代码

self.mm_registry = MULTIMODAL_REGISTRY
self.uses_mrope = model_config.uses_mrope
self.uses_xdrope_dim = model_config.uses_xdrope_dim
self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(model_config)

mm_registry：全局多模态注册表，管理各模型支持的多模态处理器（视觉/音频/视频等）。

uses_mrope：是否使用 Multi-dimensional RoPE (M-RoPE)。

M-RoPE 用于 Qwen2-VL 等模型，将位置编码扩展为 3D（temporal/height/width）
文本输入时三个维度使用相同位置 ID，等价于 1D RoPE
视觉输入时各维度独立编码空间位置

uses_xdrope_dim：XD-RoPE 的维度数（用于 HunYuan-VL 等）。

类似 M-RoPE 但维度数可配置，默认 4D

supports_mm_inputs：模型是否支持多模态输入。影响：

是否初始化 encoder_cache
是否执行 _execute_mm_encoder
是否将 input_ids 转换为 embeddings

python 复制代码

if self.model_config.is_encoder_decoder:
    self.max_encoder_len = scheduler_config.max_num_encoder_input_tokens
else:
    self.max_encoder_len = 0

max_encoder_len：encoder-decoder 模型（如 T5/Whisper）中编码器输入的最大长度。用于分配 cross-attention 的 KV cache。

4.2.6 异步调度配置（行 476-477）

python 复制代码

self.use_async_scheduling = self.scheduler_config.async_scheduling

use_async_scheduling：是否启用异步调度模式。

启用后，execute_model 和 sample_tokens 的执行与下一次调度重叠
GPU→CPU 的 token ID 拷贝在独立 CUDA stream 上异步进行
需要 prepare_inputs_event 协调 CPU tensor 的重用

4.2.7 采样器（行 478-479）

python 复制代码

self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)

Sampler：核心采样器，支持：

贪心采样（argmax）
随机采样（categorical）
Top-k / Top-p / Min-p 过滤
重复惩罚（frequency/presence/repetition penalty）
Temperature 调节
Logprob 计算（可选）

logprobs_mode：控制 logprob 的计算和返回模式，影响性能。

4.2.8 EPLB 状态（行 480-486）

python 复制代码

self.eplb_state: EplbState | None = None
self.eep_eplb_suppressed = False

EPLB (Expert Parallelism Load Balancing)：

eplb_state：惰性初始化，在 load_model() 中根据模型是否为 MoE 创建
eep_eplb_suppressed：Elastic EP 扩缩容时临时抑制 EPLB，避免在 rank 变化期间产生错误的负载均衡决策

4.2.9 惰性初始化声明（行 488-502）

python 复制代码

# self.model: nn.Module  # Set after load_model
self.kv_caches: list[torch.Tensor] = []  # Initialize in initialize_kv_cache
self.cross_layers_kv_cache: torch.Tensor | None = None  # cross-attn 共享 KV
self.cross_layers_attn_backend: type[AttentionBackend] | None = None
self.attn_groups: list[list[AttentionGroup]] = []  # [kv_group_id][attn_group]

kv_caches ：每个 attention 层的 KV cache tensor 列表。在 initialize_kv_cache() 中填充。

cross_layers_kv_cache：encoder-decoder 模型中，多个 decoder 层共享同一份 cross-attention KV cache，以节省显存。

attn_groups：二维列表，第一维是 KV cache group ID，第二维是该 group 内的 attention 子组。每个 AttentionGroup 包含：

该组使用的 attention backend
对应的 KV cache tensor 引用
metadata builder

4.2.10 Encoder Cache（行 503-510）

python 复制代码

self.encoder_cache: dict[str, torch.Tensor] = {}
self.late_interaction_runner = LateInteractionRunner()
self.encoder_cudagraph_manager: EncoderCudaGraphManager | None = None

encoder_cache ：mm_hash → encoder_output 的字典缓存。

视觉编码器（如 ViT）的输出按 mm_hash 缓存
同一图片被多个请求使用时，只计算一次编码器
scheduler 通过 free_encoder_mm_hashes 通知 Worker 释放过期缓存

late_interaction_runner：晚交互模型（如 ColBERT/Jina-ColBERT）的执行器，处理 token-level 交互。

encoder_cudagraph_manager ：编码器 CUDA 图管理器，可选启用（cudagraph_mm_encoder=True），在 capture_model() 中初始化。

4.2.11 推测解码配置（行 512-570）

这是 __init__ 中最复杂的配置块之一，根据 speculative_config.method 动态选择 drafter：

python 复制代码

if self.speculative_config and get_pp_group().is_last_rank:
    self.drafter: NgramProposer | NgramProposerGPU | ...

为什么只在 PP last rank 创建 drafter？

在 Pipeline Parallel 中，只有 last rank 执行采样，只有 last rank 有 logits/hidden_states 来驱动推测解码。其他 rank 只做中间层前向，不需要 drafter。

各 Drafter 选择逻辑：

条件	Drafter	说明
`method == "ngram"`	`NgramProposer`	CPU 侧 n-gram 匹配，无需 GPU 资源
`uses_draft_model()`	`DraftModelProposer`	独立小型 draft 模型（如 Llama-68M→Llama-7B）
`use_ngram_gpu()`	`NgramProposerGPU`	GPU 侧 n-gram 匹配，完全在 GPU 上运行
`use_dflash()`	`DFlashProposer`	DFlash 推测解码（需要 aux_hidden_states）
`method == "suffix"`	`SuffixDecodingProposer`	后缀解码（基于 suffix tree 匹配）
`use_eagle()`	`EagleProposer`	EAGLE/EAGLE3 推测解码（基于模型 internal state）
`method == "medusa"`	`MedusaProposer`	Medusa 多头推测解码
`method == "extract_hidden_states"`	`ExtractHiddenStatesProposer`	提取隐藏状态作为 draft

NgramProposerGPU 特殊初始化

python 复制代码

self.num_tokens_no_spec_gpu = torch.zeros(
    self.max_num_reqs, dtype=torch.int32, device=device)
self.token_ids_gpu_tensor = torch.zeros(
    self.max_num_reqs, self.max_model_len, dtype=torch.int32, device=device)
self._ngram_pinned_idx_buf = torch.zeros(
    self.max_num_reqs, dtype=torch.long, pin_memory=True)
self._ngram_pinned_val_buf = torch.zeros(
    self.max_num_reqs, dtype=torch.int32, pin_memory=True)

num_tokens_no_spec_gpu：每个请求的有效 token 数（不含 spec tokens），GPU 上维护
token_ids_gpu_tensor：每个请求的完整 token ID 历史，GPU 上维护
_ngram_pinned_idx_buf/_val_buf：pinned memory 缓冲区，用于 H2D/D2H 数据传输
目的：完全避免 n-gram 匹配的 CPU-GPU 同步

DFlash 特殊标志

python 复制代码

self.use_aux_hidden_state_outputs = True  # DFlash 需要 aux_hidden_states

DFlash 在模型前向中需要中间层的隐藏状态，因此设置 use_aux_hidden_state_outputs=True，使 _model_forward() 返回额外的 aux_hidden_states。

EAGLE3 特殊逻辑

python 复制代码

if self.speculative_config.method == "eagle3":
    self.use_aux_hidden_state_outputs = self.drafter.eagle3_use_aux_hidden_state

EAGLE3 也可能需要 aux_hidden_states，取决于具体配置。

Rejection Sampler

python 复制代码

self.rejection_sampler = RejectionSampler(self.sampler)

所有使用推测解码的场景都需要 RejectionSampler，用于验证 draft tokens 是否与 target model 的分布一致。

4.2.12 推测解码参数（行 572-589）

python 复制代码

self.num_spec_tokens = 0
self.valid_sampled_token_count_gpu: torch.Tensor | None = None
if self.speculative_config:
    self.num_spec_tokens = self.speculative_config.num_speculative_tokens
    draft_config = self.speculative_config.draft_model_config
    if draft_config is not None and draft_config.max_model_len is not None:
        self.effective_drafter_max_model_len = draft_config.max_model_len
    else:
        self.effective_drafter_max_model_len = self.max_model_len

num_spec_tokens：每步推测的 draft token 数量。0 表示不使用推测解码。

effective_drafter_max_model_len：drafter 模型支持的最大序列长度。

如果 draft model 有自己的 max_model_len，使用其值
否则使用 target model 的 max_model_len
当请求的序列长度超过此值时，自动禁用推测解码（避免 drafter OOM）

python 复制代码

self.use_async_spec_decode = (
    self.use_async_scheduling and self.num_spec_tokens > 0
)

use_async_spec_decode：异步调度 + 推测解码的组合优化。启用后，draft token 的验证与下一步调度重叠执行。

4.2.13 请求状态与 InputBatch（行 591-646）

python 复制代码

self.requests: dict[str, CachedRequestState] = {}
self.num_prompt_logprobs: dict[str, int] = {}

requests：请求 ID → CachedRequestState 的字典。存储所有活跃请求的状态（prompt_token_ids、output_token_ids、block_ids、sampling_params 等）。

num_prompt_logprobs：当前处于 prefill 阶段的请求的 prompt logprob 需求。只在 prefill 阶段有效，decode 阶段不需要。

python 复制代码

placeholder_block_size = self.cache_config.block_size or CacheConfig.DEFAULT_BLOCK_SIZE
self._init_block_sizes = [placeholder_block_size]
self._init_kernel_block_sizes = [placeholder_block_size]

self.input_batch = InputBatch(
    max_num_reqs=self.max_num_reqs,
    max_model_len=max(self.max_model_len, self.max_encoder_len),
    max_num_batched_tokens=self.max_num_tokens,
    device=self.device,
    pin_memory=self.pin_memory,
    vocab_size=self.model_config.get_vocab_size(),
    block_sizes=[placeholder_block_size],
    kernel_block_sizes=[placeholder_block_size],
    is_spec_decode=bool(self.vllm_config.speculative_config),
    logitsprocs=build_logitsprocs(...),
    logitsprocs_need_output_token_ids=...,
    is_pooling_model=self.is_pooling_model,
    cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
)

InputBatch 是推理循环中最核心的数据结构，持有：

所有请求的 GPU 端状态（token_ids、block_table、sampling_metadata）
CPU 端镜像（用于快速索引更新）
LogitsProcessor 集合
LoRA mapping

为什么使用 placeholder_block_size？

因为 InputBatch 在 load_model() 之前创建，此时还不知道最终的 block_size（可能被 attention backend 调整）。后续在 initialize_kv_cache() 中通过 may_reinitialize_input_batch() 更新。

logitsprocs ：logits 处理器链，由 build_logitsprocs() 构建，包括：

内置处理器（repetition penalty、frequency penalty 等）
自定义处理器（通过 --logits-processors 指定）

logitsprocs_need_output_token_ids：自定义 logits processor 是否需要已输出的 token IDs。保守设为 True（只要有自定义处理器或 reasoning config）。

cp_kv_cache_interleave_size：Context Parallel 中 KV cache 的交错大小，影响 block_table 的布局。

4.2.14 异步输出流（行 648-660）

python 复制代码

self.async_output_copy_stream: torch.cuda.Stream | None = None
self.prepare_inputs_event: torch.Event | None = None
if self.use_async_scheduling:
    self.async_output_copy_stream = torch.cuda.Stream()
    self.prepare_inputs_event = torch.Event()

async_output_copy_stream：独立 CUDA stream，用于异步 D2H 拷贝采样结果。

主 stream 继续执行下一步计算
拷贝 stream 并行传输数据到 CPU
通过 CUDA event 协调同步

prepare_inputs_event：标记 CPU→GPU 数据传输完成的事件。

异步调度中，上一步的 CPU tensor 在下一步被重用
必须确保上一步的 CPU→GPU 传输完成后才能修改 CPU tensor
synchronize_input_prep() 上下文管理器处理此同步

4.2.15 CUDA Graph 批次大小（行 662-669）

python 复制代码

if (self.compilation_config.cudagraph_capture_sizes
    and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE):
    self.cudagraph_batch_sizes = sorted(
        self.compilation_config.cudagraph_capture_sizes)
else:
    self.cudagraph_batch_sizes = []

cudagraph_batch_sizes：需要捕获 CUDA Graph 的 batch token 数列表（升序排列）。

由 compilation_config 决定
通常包含 decode 阶段的常见 batch size（如 [1, 2, 4, 8, ...]）
CUDAGraphMode.NONE 时为空

4.2.16 持久缓冲区分配（行 674-760）

这是 __init__ 中最庞大的缓冲区分配块。所有这些缓冲区都是 持久分配 的，在 Worker 生命周期内不会释放，支持 CUDA Graph 的静态内存要求：

python 复制代码

self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
self.positions = torch.zeros(self.max_num_tokens, dtype=torch.int64, device=self.device)
self.query_start_loc = self._make_buffer(self.max_num_reqs + 1, dtype=torch.int32)
self.seq_lens = torch.zeros(self.max_num_reqs, dtype=torch.int32, device=self.device)
self.optimistic_seq_lens_cpu = torch.zeros(
    self.max_num_reqs, dtype=torch.int32, pin_memory=self.pin_memory)
self.num_computed_tokens = torch.zeros(
    self.max_num_reqs, dtype=torch.int32, device=self.device)
self.prev_num_draft_tokens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
self.req_indices = self._make_buffer(self.max_num_tokens, dtype=torch.int64)
self.prev_positions = self._make_buffer(self.max_num_reqs, dtype=torch.int64)
self.num_scheduled_tokens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)

各缓冲区详解：

缓冲区	形状	用途
`input_ids`	[max_num_tokens]	输入 token IDs，CpuGpuBuffer（CPU+GPU 双缓冲）
`positions`	[max_num_tokens]	每个 token 的位置 ID（1D RoPE）
`query_start_loc`	[max_num_reqs+1]	请求内 token 的累积起始位置（类似 CSR offset）
`seq_lens`	[max_num_reqs]	每个请求的当前序列长度
`optimistic_seq_lens_cpu`	[max_num_reqs]	乐观估计的序列长度（spec_decode + async）
`num_computed_tokens`	[max_num_reqs]	每个请求已计算 KV cache 的 token 数
`prev_num_draft_tokens`	[max_num_reqs]	上一步每个请求的 draft token 数（用于 async spec decode 修正）
`req_indices`	[max_num_tokens]	每个 token 对应的请求索引
`prev_positions`	[max_num_reqs]	上一步每个请求在 batch 中的位置（用于增量更新）
`num_scheduled_tokens`	[max_num_reqs]	本步每个请求调度的 token 数

CpuGpuBuffer（_make_buffer）：

python 复制代码

def _make_buffer(self, *size, dtype, numpy=True):
    return CpuGpuBuffer(*size, dtype=dtype, device=self.device,
                        pin_memory=self.pin_memory, with_numpy=numpy)

同时在 CPU（pinned memory + numpy）和 GPU 上分配
.gpu 访问 GPU tensor，.np 访问 CPU numpy 数组
CPU→GPU 通过 copy_to_gpu()，GPU→CPU 通过事件同步
pin_memory=True 使 CPU 端使用页锁定内存，加速 DMA 传输

python 复制代码

self.encoder_seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
if self.dcp_world_size > 1:
    self.dcp_local_seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)

encoder_seq_lens：encoder-decoder 模型中编码器输入的长度
dcp_local_seq_lens：DCP 模式下本地 rank 看到的序列长度（= seq_len / dcp_world_size）

python 复制代码

self.inputs_embeds = self._make_buffer(
    self.max_num_tokens, self.inputs_embeds_size, dtype=self.dtype, numpy=False)
self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool)
self.discard_request_mask = self._make_buffer(self.max_num_reqs, dtype=torch.bool)
self.num_decode_draft_tokens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
self.num_accepted_tokens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)

inputs_embeds：[max_num_tokens, embeds_size] 的嵌入缓冲区。numpy=False 因为可能是 bfloat16（numpy 不支持）
is_token_ids：标记每个位置是 token ID 还是 embedding
discard_request_mask：标记需要丢弃采样结果的请求（如被抢占的请求）
num_decode_draft_tokens：decode 阶段每个请求的 draft token 数
num_accepted_tokens：rejection sampling 后每个请求实际接受的 token 数

4.2.17 M-RoPE / XD-RoPE 位置缓冲区（行 762-783）

python 复制代码

if self.uses_mrope:
    self.mrope_positions = self._make_buffer(
        (3, self.max_num_tokens + 1), dtype=torch.int64)

M-RoPE positions：[3, max_num_tokens+1] 的 3D 位置缓冲区。

三个维度分别对应 temporal/height/width
纯文本时三个维度相同，等价于 1D RoPE
视觉输入时各维度独立编码
+1 的巧妙设计 ：额外的 dummy position 使张量变为非连续（non-contiguous），以兼容 torch.compile。详见 PR #12128

python 复制代码

if self.uses_xdrope_dim > 0:
    self.xdrope_positions = self._make_buffer(
        (self.uses_xdrope_dim, self.max_num_tokens + 1), dtype=torch.int64)

XD-RoPE positions：类似 M-RoPE 但维度数可配置，用于 HunYuan-VL 等模型。

4.2.18 Pipeline Parallel 中间张量（行 785-786）

python 复制代码

self.intermediate_tensors: IntermediateTensors | None = None

非 PP 首 rank 的中间张量引用，在 _preprocess() 中被赋值。首 rank 保持 None。

4.2.19 Arange 缓存（行 788-795）

python 复制代码

arange_size = max(self.max_num_reqs + 1, self.max_num_tokens)
self.arange_np = np.arange(arange_size, dtype=np.int64)
self.query_pos = self._make_buffer(arange_size, dtype=torch.int64)
self._arange_scratch = np.empty(arange_size, dtype=np.int64)

arange_np：不可变的 [0, 1, 2, ...] 数组，作为 batch 计算的源
query_pos：CpuGpuBuffer，存储计算后的 batched arange 结果
_arange_scratch：numpy 临时缓冲区，避免每步分配

用途示例：将 [2, 5, 3]（每个请求的 token 数）展开为 [0,1,0,1,2,3,4,0,1,2]（每个 token 的请求内位置）。

4.2.20 跨层 KV 共享（行 797-815）

python 复制代码

self.shared_kv_cache_layers: dict[str, str] = {}
self.kv_sharing_fast_prefill_eligible_layers: set[str] = set()
self.kv_sharing_fast_prefill_logits_indices = None
if self.cache_config.kv_sharing_fast_prefill:
    self.kv_sharing_fast_prefill_logits_indices = torch.zeros(
        self.max_num_tokens, dtype=torch.int32, device=self.device)

shared_kv_cache_layers ：{layer_name: target_layer_name} 映射。

key 层使用 value 层的 KV cache
如 You-Only-Cache-Once 架构中，多个 attention 层共享同一份 KV cache
在 get_kv_cache_spec() 中扫描 kv_sharing_target_layer_name 属性构建

kv_sharing_fast_prefill：快速 prefill 优化。

共享 KV 的层中，只有"源层"参与 prefill 计算，"消费层"提前退出
logits_indices 用于在 fast prefill 模式下正确索引 logits
限制：不能与 prompt logprob 同时使用（因为中间层被跳过，无法计算完整 logprob）

4.2.21 Uniform Decode 配置（行 817）

python 复制代码

self.uniform_decode_query_len = 1 + self.num_spec_tokens

uniform_decode_query_len：统一解码时每个请求的 query 长度。

不使用 spec decode 时 = 1（每个请求只生成 1 个 token）
使用 spec decode 时 = 1 + num_spec_tokens（包含 draft tokens）
用于判断当前 batch 是否为 "uniform decode"（所有请求 query 长度相同）

4.2.22 CudagraphDispatcher（行 819）

python 复制代码

self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)

CudagraphDispatcher：运行时 CUDA Graph 分发器。

根据当前 batch 的 token 数、LoRA 状态、是否 uniform decode，选择合适的 CUDA Graph
如果没有匹配的图，回退到 eager 模式
支持多种 CUDAGraphMode：NONE / FULL（包含 attention）/ PARTIAL（不含 attention）

4.2.23 多模态预算（行 821-823）

python 复制代码

self.mm_budget = (
    MultiModalBudget(self.vllm_config, self.mm_registry)
    if self.supports_mm_inputs
    else None
)

MultiModalBudget：多模态显存预算计算器。

根据模型配置和注册表，计算每种模态的最大输入 token 数
用于 profile_run 中确定多模态编码器的显存需求
只在支持多模态的模型上创建

4.2.24 其他配置（行 825-830）

python 复制代码

self.reorder_batch_threshold: int | None = None
self.runner_only_attn_layers: set[str] = set()

reorder_batch_threshold ：batch 重排序阈值。某些 attention backend（如 MLA）希望将 decode 和 prefill 请求分开排列以提高效率。在 calculate_reorder_batch_threshold() 中设置。

runner_only_attn_layers：只存在于 ModelRunner 的 KVCacheConfig 中、不存在于 Scheduler 的 KVCacheConfig 中的 attention 层（如 KV sharing 消费层、encoder-only attention）。

4.2.25 Draft Token 缓存（行 832-857）

python 复制代码

self._draft_token_ids: list[list[int]] | torch.Tensor | None = None
self._draft_token_req_ids: list[str] | None = None
self.transfer_event = torch.Event()
self.sampled_token_ids_pinned_cpu = torch.empty(
    (self.max_num_reqs, 1), dtype=torch.int64, device="cpu",
    pin_memory=self.pin_memory)

_draft_token_ids ：上一步生成的 draft token IDs，供下一步 scheduler 使用。
_draft_token_req_ids ：对应请求 ID 列表。
transfer_event ：CUDA event，用于同步 D2H 拷贝。
sampled_token_ids_pinned_cpu：pinned memory 缓冲区，用于高效 D2H 拷贝采样结果。

N-gram GPU 异步缓冲区

python 复制代码

if self.speculative_config is not None and self.speculative_config.use_ngram_gpu():
    self._num_valid_draft_tokens_cpu = torch.empty(
        self.max_num_reqs, dtype=torch.int32, pin_memory=True)
    self._num_valid_draft_tokens_event = torch.cuda.Event()
    self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream()

ngram_gpu 路径的专用异步缓冲区：

_num_valid_draft_tokens_cpu：pinned memory，用于接收 D2H 的有效 draft token 数
_num_valid_draft_tokens_event/copy_stream：专用 CUDA event 和 stream

Spec Decode 异步缓冲区

python 复制代码

if self.num_spec_tokens:
    self.draft_token_ids_event = torch.Event()
    self.num_accepted_tokens_event = torch.Event()
    self.draft_token_ids_copy_stream = torch.cuda.Stream()
    self.draft_token_ids_cpu = torch.empty(
        (self.max_num_reqs, self.num_spec_tokens),
        dtype=torch.int64, device="cpu", pin_memory=True)
    if self.use_async_scheduling:
        self.valid_sampled_token_count_event = torch.Event()
        self.valid_sampled_token_count_copy_stream = torch.cuda.Stream()
        self.valid_sampled_token_count_cpu = torch.empty(
            self.max_num_reqs, dtype=torch.int32, device="cpu",
            pin_memory=True)

启用推测解码时的额外缓冲区：

draft_token_ids_cpu：[max_num_reqs, num_spec_tokens]，pinned memory，用于异步 D2H 拷贝 draft tokens
valid_sampled_token_count_cpu：异步调度下每个请求的有效采样 token 数
每个缓冲区都有专属 CUDA event 和 stream，实现真正的异步传输重叠

4.2.26 模型权重卸载器（行 853-854）

python 复制代码

set_offloader(create_offloader(self.offload_config))

在 __init__ 最后阶段设置全局权重卸载器。必须在任何 get_offloader 调用之前执行。支持：

CPU offloading（权重默认在 CPU，需要时加载到 GPU）
Disk offloading（权重默认在磁盘，需要时加载到 GPU）
None（权重常驻 GPU）

4.2.27 临时执行状态（行 856-857）

python 复制代码

self.execute_model_state: ExecuteModelState | None = None
self.kv_connector_output: KVConnectorOutput | None = None
self.mamba_state_idx: dict[str, int] = {}
self._mamba_copy_bufs: mamba_utils.MambaCopyBuffers | None = None
self.layerwise_nvtx_hooks_registered = False

execute_model_state：execute_model() 和 sample_tokens() 之间的临时状态传递。

execute_model 结束时保存（logits、hidden_states、spec_decode_metadata 等）
sample_tokens 开始时恢复
使用后立即清空（防止泄漏）

kv_connector_output：KV 连接器的输出引用，跨方法传递。

mamba_state_idx / _mamba_copy_bufs：Mamba SSM 模型的状态管理。

mamba_state_idx：层名 → 在 mamba cache 中的索引
_mamba_copy_bufs：Mamba 状态拷贝缓冲区（惰性初始化）

layerwise_nvtx_hooks_registered：是否已注册逐层 NVTX profiling hooks。

总结

vLLM v1 Worker 模块的四层架构（WorkerBase → Worker → GPUModelRunner → 子模块）体现了清晰的职责分离：

WorkerBase：定义接口契约，不触碰硬件
Worker：管理设备生命周期（初始化→加载→显存管理→休眠/恢复），是"运维层"
GPUModelRunner：编排推理流程（状态更新→输入准备→预处理→前向→采样→后处理），是"执行层"
子模块（Sampler/AttentionBackend/InputBatch/SpecDecode/...）：各司其职的"工具层"

GPUModelRunner 的 __init__ 方法（460+ 行）是整个模块最密集的配置点，每个变量、每个分支都服务于特定的业务场景或性能优化。其核心设计哲学是 预分配一切------所有 GPU/CPU 缓冲区在初始化时就按最大尺寸分配，运行时零分配，从而：

支持 CUDA Graph 的静态内存要求
避免运行时内存碎片化
消除动态分配的延迟开销
使推理延迟可预测

这种"重初始化、轻运行"的设计模式，是高性能推理系统的标准实践。

Phase 0：前置检查（行 3771--3840）

5.0.1 方法签名与防重入检查

python 复制代码

def execute_model(
    self,
    scheduler_output: "SchedulerOutput",
    intermediate_tensors: IntermediateTensors | None = None,
) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None:

参数说明 ：
- scheduler_output：调度器本轮的完整输出，包含已调度请求、新请求、已结束请求、block分配等信息。
- intermediate_tensors：Pipeline Parallelism (PP) 场景下，从前一 stage 传递过来的中间张量；非 PP 首阶为 None。
返回值 ：多态返回------普通同步结果 ModelRunnerOutput、异步结果 AsyncModelRunnerOutput、PP 中间张量 IntermediateTensors，或空结果 None。

python 复制代码

if self.execute_model_state is not None:
    raise RuntimeError(
        "State error: sample_tokens() must be called "
        "after execute_model() returns None."
    )

防重入机制 ：execute_model_state 是一个临时状态变量，在 execute_model() 结束时被设置为 ExecuteModelState(...)，在后续的 sample_tokens() 中被消费并清空。若在 sample_tokens() 未被调用前再次进入 execute_model()，说明上层调用违反了协议------execute_model 和 sample_tokens 必须成对交替调用。这种设计将 forward 执行 与 采样/后处理 解耦为两步，为异步调度（async scheduling）提供了基础：forward 完成后可以先返回 None（不阻塞等采样），让调度器在 GPU 执行采样的同时开始下一轮调度。

5.0.2 RoutedExperts 缓冲清理

python 复制代码

if self.routed_experts_initialized:
    capturer = RoutedExpertsCapturer.get_instance()
    if capturer is not None:
        capturer.clear_buffer()
    else:
        logger.error("RoutedExpertsCapturer not initialized.")

当 MoE（Mixture of Experts）模型的 routed experts 捕获器已初始化时，每步执行前需要清空其缓冲区。RoutedExpertsCapturer 使用单例模式（get_instance()），用于捕获当前 step 的专家路由信息。清空缓冲确保不会混入上一步的陈旧路由数据。若单例不存在则记录错误------这通常意味着初始化时序有问题。

5.0.3 ngram_gpu scheduler_output 深拷贝

python 复制代码

if (
    self.speculative_config is not None
    and self.speculative_config.use_ngram_gpu()
):
    num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy()
    spec_decode_tokens_copy = (
        scheduler_output.scheduled_spec_decode_tokens.copy()
    )
    scheduler_output = replace(
        scheduler_output,
        num_scheduled_tokens=num_scheduled_tokens_copy,
        scheduled_spec_decode_tokens=spec_decode_tokens_copy,
    )

设计目的 ：当使用 ngram GPU 推测解码时，worker 会在 _update_states() 中修改 scheduler_output.num_scheduled_tokens 和 scheduled_spec_decode_tokens（裁剪无效 draft tokens）。但 scheduler_output 是在 engine core 进程和 worker 进程之间共享的（通过共享内存），直接修改会影响调度器侧的状态。因此使用 dataclasses.replace() 创建一个浅替换的新对象，其中可能被修改的两个字段使用 .copy() 深拷贝。replace() 比 deepcopy() 快得多，因为大部分字段是引用复制。

5.0.4 KV Transfer 抢占处理

python 复制代码

if has_kv_transfer_group():
    kv_connector_metadata = scheduler_output.kv_connector_metadata
    assert kv_connector_metadata is not None
    get_kv_transfer_group().handle_preemptions(kv_connector_metadata)

KV Transfer（跨节点 KV 缓存迁移，用于 disaggregated prefill 等场景）需要处理请求被抢占的情况。当请求被抢占时，其 KV 缓存可能正在传输中，handle_preemptions() 负责取消正在进行的传输、释放传输资源，确保不会出现悬空的传输操作。has_kv_transfer_group() 是一个轻量级检查，避免在非 KV transfer 场景下的不必要开销。

5.0.5 num_scheduled_tokens=0 早返路径

python 复制代码

num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
with (
    record_function_or_nullcontext("gpu_model_runner: preprocess"),
    self.synchronize_input_prep(),
):
    deferred_state_corrections_fn = self._update_states(scheduler_output)

关键观察 ：即使 num_scheduled_tokens == 0，_update_states() 仍然必须执行！因为调度器可能发送了"有请求结束/抢占但没有新token调度"的信息，这些状态更新不能跳过。

python 复制代码

if not num_scheduled_tokens:
    if (
        self.parallel_config.distributed_executor_backend
        == "external_launcher"
        and self.parallel_config.data_parallel_size > 1
    ):
        self._dummy_run(1)
    if not has_kv_transfer_group():
        return EMPTY_MODEL_RUNNER_OUTPUT
    return self.kv_connector_no_forward(scheduler_output, self.vllm_config)

早返逻辑的三条路径：

External Launcher + DP ：当使用外部启动器（如 torchrun）且数据并行 > 1 时，即使本 worker 没有 token 需要处理，也需要调用 _dummy_run(1) 执行一次 dummy 前向，确保 coordinate_batch_across_dp 同步点被正确触达。否则其他 DP rank 可能永远阻塞在同步屏障上。
普通场景 ：返回 EMPTY_MODEL_RUNNER_OUTPUT（预分配的空结果常量，避免重复分配）。
KV Transfer 场景 ：即使没有 token 调度，KV transfer 可能仍有正在进行的保存/加载操作需要推进，因此调用 kv_connector_no_forward() 处理传输流程但不执行模型前向。

5.0.6 EC Transfer Consumer 特殊路径

python 复制代码

if has_ec_transfer() and not get_ec_transfer().is_consumer:
    with self.maybe_get_ec_connector_output(
        scheduler_output,
        encoder_cache=self.encoder_cache,
    ) as ec_connector_output:
        self._execute_mm_encoder(scheduler_output)
        return make_empty_encoder_model_runner_output(scheduler_output)

Encoder-Consumer (EC) Transfer ：在 disaggregated 架构中，可能有一个节点专门执行 encoder（prefill），另一个节点作为 consumer 只做 decode。当当前节点不是 consumer（即是 producer）时，只需执行 MM encoder 并返回空结果。Consumer 节点会在后续正常流程中使用 encoder 的输出。maybe_get_ec_connector_output 是一个 context manager，管理 EC 连接器的生命周期。

5.0.7 kv_sharing_fast_prefill 断言

python 复制代码

if self.cache_config.kv_sharing_fast_prefill:
    assert not self.num_prompt_logprobs, (
        "--kv-sharing-fast-prefill produces incorrect "
        "logprobs for prompt tokens, tokens, please disable "
        "it when the requests need prompt logprobs"
    )

KV sharing fast prefill 是一种优化：当多个请求共享相同前缀时，跳过重复前缀的 prefill 计算。但这种优化会导致 prompt token 的 logprobs 不正确（因为跳过了部分计算），因此需要断言确保不会同时启用。

5.0.8 批次元数据准备

python 复制代码

num_reqs = self.input_batch.num_reqs
req_ids = self.input_batch.req_ids
tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens

从持久化批次中提取当前活跃请求的 ID 列表，根据调度器为每个请求分配的 token 数构建 numpy 数组。max_num_scheduled_tokens 用于确定 CUDA graph 填充大小和 attention metadata 的 max_query_len。num_tokens_unpadded 是未填充的总 token 数，与填充后的 num_tokens_padded 区分------后者用于实际 GPU 执行。

Phase 1：_update_states()（行 1059--1412）

这是状态管理的核心方法，负责将调度器的决策反映到 worker 的内部状态中。

5.1.1 清除已完成请求

python 复制代码

for req_id in scheduler_output.finished_req_ids:
    self.requests.pop(req_id, None)
    self.num_prompt_logprobs.pop(req_id, None)
self.late_interaction_runner.on_requests_finished(
    scheduler_output.finished_req_ids
)

两层清理：

从 self.requests（dict[str, CachedRequestState]）中移除请求的缓存状态。
从 self.num_prompt_logprobs 中移除 prompt logprobs 计数。
通知 late_interaction_runner（用于 ColPali 等交互式模型）请求已完成。

python 复制代码

for req_id in scheduler_output.finished_req_ids:
    self.input_batch.remove_request(req_id)

从持久化批次中移除请求，释放其 slot。remove_request 不会立即触发 condense，而是在后续统一执行。

5.1.2 新分配 block 的零初始化

python 复制代码

if scheduler_output.new_block_ids_to_zero:
    self._zero_block_ids(scheduler_output.new_block_ids_to_zero)

新分配的 KV 缓存 block 可能包含陈旧的 NaN 或垃圾数据（GPU 内存复用），如果不清零，attention 或 SSM（State Space Model）计算会被污染。_zero_block_ids 将这些 block 对应的 GPU 内存区域置零。

5.1.3 释放 Encoder 缓存

python 复制代码

for mm_hash in scheduler_output.free_encoder_mm_hashes:
    self.encoder_cache.pop(mm_hash, None)

多模态 encoder 的输出（如视觉 embedding）缓存在 self.encoder_cache 中，key 为 mm_hash。当请求完成后，对应的 encoder 输出不再需要，释放以回收 GPU 内存。

5.1.4 移除未调度请求

python 复制代码

scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
cached_req_ids = self.input_batch.req_id_to_index.keys()
resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids)

集合运算推导：

scheduled_req_ids：本轮被调度的请求 ID 集合。
cached_req_ids：当前持久化批次中的请求 ID 集合。
resumed_req_ids：被抢占后重新恢复的请求。

差集 (scheduled_req_ids - resumed_req_ids) 表示"本轮正常调度（非恢复）的请求"。再用 cached_req_ids 减去它，得到"在批次中但本轮未被调度的请求"------即被抢占的请求或跳过的请求。

特殊处理 resumed 请求 ：resumed_req_ids 同时出现在 scheduled_req_ids 和 cached_req_ids 中。在 forced preemption（如 reset_prefix_cache）场景下，同一个请求可能先被抢占再立即恢复。此时需要先将它从批次中清除（重置其 slot 状态），然后在后续的"新增/恢复请求"路径中重新添加。

python 复制代码

for req_id in unscheduled_req_ids:
    self.input_batch.remove_request(req_id)

关键设计 ：移除请求时保留 self.requests 中的 CachedRequestState，因为被抢占的请求将来会被重新调度。只从持久化批次中移除------因为批次 slot 是稀缺资源，未调度的请求不需要占用。

5.1.5 新增请求注册

python 复制代码

for new_req_data in scheduler_output.scheduled_new_reqs:
    req_id = new_req_data.req_id
    if req_id in self.requests:
        req_state = self._update_streaming_request(req_id, new_req_data)
        reqs_to_add.append(req_state)
        continue

流式请求复用 ：某些场景下（如 streaming），同一个 req_id 可能在 self.requests 中已存在。此时不创建新的 CachedRequestState，而是更新现有状态（如新的 sampling params）。

python 复制代码

sampling_params = new_req_data.sampling_params
pooling_params = new_req_data.pooling_params

if (
    sampling_params
    and sampling_params.sampling_type == SamplingType.RANDOM_SEED
):
    generator = torch.Generator(device=self.device)
    generator.manual_seed(sampling_params.seed)
else:
    generator = None

随机数生成器 ：当采样类型为 RANDOM_SEED（用户指定了固定种子）时，创建一个独立的 torch.Generator 并设置种子。这确保了确定性采样------同一个种子总是产生相同的输出序列。其他采样类型（greedy、random 无种子）使用全局随机状态。

python 复制代码

if self.is_pooling_model:
    assert pooling_params is not None
    task = pooling_params.task
    model = cast(VllmModelForPooling, self.get_model())
    to_update = model.pooler.get_pooling_updates(task)
    to_update.apply(pooling_params)

Pooling 模型处理 ：对于嵌入/池化模型（如 ColPali、E5），需要根据 task 类型更新 pooling 参数（如池化策略、归一化方式等）。get_pooling_updates() 返回需要应用的更新列表。

python 复制代码

req_state = CachedRequestState(
    req_id=req_id,
    prompt_token_ids=new_req_data.prompt_token_ids,
    prompt_embeds=new_req_data.prompt_embeds,
    mm_features=new_req_data.mm_features,
    sampling_params=sampling_params,
    pooling_params=pooling_params,
    generator=generator,
    block_ids=new_req_data.block_ids,
    num_computed_tokens=new_req_data.num_computed_tokens,
    output_token_ids=[],
    lora_request=new_req_data.lora_request,
)
self.requests[req_id] = req_state
self.late_interaction_runner.register_request(req_id, pooling_params)

CachedRequestState 初始化：

prompt_token_ids：prompt 的 token ID 列表（文本输入时使用）。
prompt_embeds：prompt 的 embedding 张量（软 token / embedding 输入时使用，与 token_ids 互斥）。
mm_features：多模态特征列表（图像位置信息等）。
block_ids：KV 缓存 block 分配元组（支持多 KV 缓存组，每个组一个列表）。
num_computed_tokens：已计算的 token 数（对于新 prefill 请求，通常为 0）。
output_token_ids：输出 token ID 列表，初始为空。
lora_request：LoRA 适配器请求（如果有）。

python 复制代码

if sampling_params and sampling_params.prompt_logprobs is not None:
    self.num_prompt_logprobs[req_id] = (
        self.input_batch.vocab_size
        if sampling_params.prompt_logprobs == -1
        else sampling_params.prompt_logprobs
    )

Prompt logprobs 请求记录：-1 表示计算全部 token 的 logprobs，存储为 vocab_size（最大可能值）；否则存储用户指定的 top-k 数。

python 复制代码

if self.uses_mrope:
    self._init_mrope_positions(req_state)
if self.uses_xdrope_dim > 0:
    self._init_xdrope_positions(req_state)

M-RoPE (Multi-dimensional RoPE) ：Qwen2-VL 等模型使用多维位置编码（时间/高度/宽度各一维），需要预初始化位置信息。XD-RoPE：HunYuan-VL 使用交叉维度 RoPE，同样需要初始化。

python 复制代码

reqs_to_add.append(req_state)
if is_ngram_gpu:
    ngram_gpu_new_reqs.append(req_state)

新请求标记为待添加。ngram_gpu 场景下额外追踪，因为新增请求需要完整初始化其 ngram GPU 张量（而非增量更新）。

5.1.6 已调度缓存请求状态更新

python 复制代码

is_last_rank = get_pp_group().is_last_rank
req_data = scheduler_output.scheduled_cached_reqs
scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens

获取 PP 信息和 spec decode token 数据。is_last_rank 决定了谁负责维护 output_token_ids------只有最后一阶（负责采样）才拥有真实的采样结果。

异步 Spec Decode 的乐观扩展机制

python 复制代码

original_num_spec_per_req: dict[str, int] = {}
if (
    self.speculative_config is not None
    and self.speculative_config.use_ngram_gpu()
):
    for req_id, toks in scheduled_spec_tokens.items():
        original_num_spec_per_req[req_id] = len(toks)
    update_scheduler_for_invalid_drafts(
        self._num_valid_draft_tokens_event,
        self._num_valid_draft_tokens_cpu,
        scheduler_output,
        self.input_batch.req_id_to_index,
    )

ngram GPU 的 draft token 裁剪 ：ngram GPU 在上一步生成了 draft tokens，但调度器可能发现其中一些是无效的（如违反结构化输出约束）。update_scheduler_for_invalid_drafts 使用上一步 rejection sampling 的结果来裁剪无效 draft tokens，更新 scheduler_output 中的 scheduled_spec_decode_tokens 和 num_scheduled_tokens。同时保存原始 draft 长度用于后续恢复。

python 复制代码

if self.use_async_spec_decode:
    self.prev_num_draft_tokens.np.fill(0)

异步 spec decode 模式下，每步开始时将"上一轮 draft tokens 数"归零，后续会在循环中按请求更新。

逐请求状态更新循环

python 复制代码

for i, req_id in enumerate(req_data.req_ids):
    req_state = self.requests[req_id]
    num_computed_tokens = req_data.num_computed_tokens[i]
    new_block_ids = req_data.new_block_ids[i]
    resumed_from_preemption = req_id in req_data.resumed_req_ids
    num_output_tokens = req_data.num_output_tokens[i]
    req_index = self.input_batch.req_id_to_index.get(req_id)

变量含义：

num_computed_tokens：调度器记录的已计算 token 数。对于 decode 请求，这是上一步的已计算数；对于恢复的请求，这是抢占前的进度。
new_block_ids：本轮新分配的 block ID 列表（按 KV 缓存组分组）。
resumed_from_preemption：标识该请求是否从抢占中恢复。
num_output_tokens：调度器确认的有效输出 token 数。
req_index：请求在持久化批次中的当前 slot 索引。None 表示请求不在批次中（被抢占后等待恢复）。

python 复制代码

if req_state.prev_num_draft_len and self.use_async_scheduling:
    if req_index is None:
        req_state.prev_num_draft_len = 0
    else:
        optimistic_num_accepted = req_state.prev_num_draft_len
        req_state.output_token_ids.extend([-1] * optimistic_num_accepted)

        deferred_spec_decode_corrections.append(
            (req_id, optimistic_num_accepted, req_state)
        )

异步调度的乐观占位机制 ：在异步调度模式 + spec decode 场景下，当上一轮有 draft tokens 时，我们乐观地假设所有 draft tokens 都被接受 。在 output_token_ids 中用 -1 占位（-1 是无效 token ID，表示待修正）。这允许 _prepare_inputs 基于乐观假设构建正确的 positions 和 token_ids，而不需要等待 rejection sampling 结果。真正的接受数量在 forward 完成后通过 deferred_spec_decode_corrections 回调修正。

python 复制代码

prev_req_index = (
    self.input_batch.prev_req_id_to_index.get(req_id)
    if self.input_batch.prev_req_id_to_index
    else None
)
if prev_req_index is not None:
    self.prev_num_draft_tokens.np[prev_req_index] = optimistic_num_accepted

记录上一轮批次中该请求的 draft token 数，用于后续 GPU 端的 num_computed_tokens 修正内核。

output_token_ids 更新

python 复制代码

if not is_last_rank:
    if not req_data.new_token_ids:
        new_token_ids: list[int] = []
    else:
        new_token_ids = req_data.new_token_ids[i]
        num_new_tokens = (
            num_computed_tokens + len(new_token_ids) - req_state.num_tokens
        )
        if num_new_tokens == 1:
            req_state.output_token_ids.append(new_token_ids[-1])
        elif num_new_tokens > 0:
            req_state.output_token_ids.extend(
                new_token_ids[-num_new_tokens:]
            )

非最后 PP 阶的处理 ：中间 PP 阶没有采样能力，需要调度器发送上一轮的采样结果 new_token_ids。num_new_tokens 的计算：num_computed_tokens + len(new_token_ids) 是当前已知的总 token 数，减去 req_state.num_tokens（之前记录的总 token 数）得到新增 token 数。

优化：num_new_tokens == 1 是最常见的 decode 场景，使用 append 比 extend 切片更快。

python 复制代码

elif num_output_tokens < len(req_state.output_token_ids):
    del req_state.output_token_ids[num_output_tokens:]
    if req_index is not None:
        end_idx = (
            self.input_batch.num_prompt_tokens[req_index]
            + num_output_tokens
        )
        self.input_batch.num_tokens_no_spec[req_index] = end_idx

最后 PP 阶的回退处理 ：当 num_output_tokens < len(output_token_ids) 时，说明某些输出 token 被丢弃------可能是 KV load 失败或异步 spec decode 的乐观占位需要回退。截断 output_token_ids 并同步更新 num_tokens_no_spec。

Block IDs 更新

python 复制代码

if not resumed_from_preemption:
    if new_block_ids is not None:
        for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
            block_ids.extend(new_ids)
else:
    assert req_index is None
    assert new_block_ids is not None
    req_state.block_ids = new_block_ids

正常调度 vs 恢复：

正常调度 ：新 block 追加到已有 block 列表末尾（extend），因为 KV 缓存是增量分配的。
恢复调度 ：抢占后恢复时，旧 block 已被回收，调度器分配了全新的 block 集合，因此直接替换（= 赋值）而非追加。断言 req_index is None 确保恢复的请求不在批次中（已被移除）。

不在批次中的请求处理

python 复制代码

if req_index is None:
    if self.use_async_scheduling and num_output_tokens > 0:
        resumed_token_ids = req_data.all_token_ids[req_id]
        req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]
    reqs_to_add.append(req_state)
    if is_ngram_gpu:
        ngram_gpu_new_reqs.append(req_state)
    continue

请求不在持久化批次中（被抢占或之前未调度），需要重新添加。异步调度场景下还需要从 all_token_ids 恢复输出 token ID（因为调度器知道完整的历史 token 序列）。

在批次中的请求增量更新

python 复制代码

self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
if new_block_ids is not None:
    self.input_batch.block_table.append_row(new_block_ids, req_index)

if not is_last_rank:
    start_token_index = num_computed_tokens
    end_token_index = num_computed_tokens + len(new_token_ids)
    self.input_batch.token_ids_cpu[
        req_index, start_token_index:end_token_index
    ] = new_token_ids
    self.input_batch.num_tokens_no_spec[req_index] = end_token_index

增量同步到 InputBatch：

更新已计算 token 数。
追加新分配的 block 到 block table。
非 PP 最后阶写入新的 token IDs 到 token_ids_cpu 的正确位置。
更新 num_tokens_no_spec（不含 spec tokens 的总 token 数）。

python 复制代码

self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens)
if original_num_spec_per_req:
    orig = original_num_spec_per_req.get(req_id, 0)
    if orig != req_state.prev_num_draft_len:
        req_state.prev_num_draft_len = orig

Spec token 更新 ：将调度器确认的 draft token IDs 写入 token_ids_cpu 和 spec_token_ids。之后恢复 prev_num_draft_len 为原始值（ngram 裁剪前），以保持与调度器的一致性。

5.1.7 批次整理

python 复制代码

for request in reqs_to_add:
    self.input_batch.add_request(request)
    self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens)

self.input_batch.condense()
self._may_reorder_batch(scheduler_output)
self.input_batch.refresh_metadata()

执行顺序至关重要：

添加请求：将新请求和恢复请求填入空 slot。
Condense：紧凑排列，消除中间的空 slot（详见第六章）。
May reorder：某些 attention backend（如 GDN）可能需要重排请求顺序以优化计算。
Refresh metadata ：根据批次变化重建 SamplingMetadata。

5.1.8 ngram GPU 增量更新

python 复制代码

if is_ngram_gpu:
    update_ngram_gpu_tensors_incremental(
        self.input_batch,
        self.token_ids_gpu_tensor,
        self.num_tokens_no_spec_gpu,
        ngram_gpu_new_reqs,
        self.device,
        _pinned_idx_buf=self._ngram_pinned_idx_buf,
        _pinned_val_buf=self._ngram_pinned_val_buf,
    )

只更新新增/恢复请求的 GPU 侧 ngram 张量，而非全量重建。ngram_gpu_new_reqs 包含了所有需要完整初始化的请求列表。使用 pinned memory buffer 进行高效的 CPU→GPU 传输。

5.1.9 Deferred Spec Decode Corrections

python 复制代码

if deferred_spec_decode_corrections:

    def correct_spec_decode_token_counts():
        valid_sampled_token_count = self._get_valid_sampled_token_count()
        if not valid_sampled_token_count:
            return
        prev_req_id_to_index = self.input_batch.prev_req_id_to_index
        if not prev_req_id_to_index:
            return
        for (
            req_id,
            optimistic_num_accepted,
            req_state,
        ) in deferred_spec_decode_corrections:
            prev_req_index = prev_req_id_to_index.get(req_id)
            if prev_req_index is None:
                continue
            num_accepted = valid_sampled_token_count[prev_req_index] - 1
            correction = optimistic_num_accepted - num_accepted
            req_state.num_computed_tokens -= correction
            cur_req_index = self.input_batch.req_id_to_index.get(req_id)
            if cur_req_index is None:
                continue
            self.input_batch.num_computed_tokens_cpu[cur_req_index] -= correction
            if is_ngram_gpu and correction > 0:
                self.input_batch.num_tokens_no_spec[cur_req_index] -= correction
                self.num_tokens_no_spec_gpu[cur_req_index] -= correction

    return correct_spec_decode_token_counts
else:
    return None

延迟修正的核心设计 ：_update_states 返回一个闭包函数而非立即执行修正。这个闭包在以下时机被调用：

Mamba 模型 ：在 Phase 2 中 mamba preprocess 之前（因为 mamba 依赖正确的 num_computed_tokens）。
异步调度 ：在 forward 完成后（因为 valid_sampled_token_count 来自上一步的 GPU 采样结果）。

修正逻辑：

valid_sampled_token_count[prev_req_index] - 1：上一步实际接受的 token 数（减1是因为包含1个真实采样token）。
correction = optimistic - actual：需要回退的差值。
修正 num_computed_tokens（CPU 和请求状态）。
若 ngram GPU 模式，还修正 num_tokens_no_spec。

Phase 2：_prepare_inputs()（行 1774--2085）

这个阶段将 CPU 侧的状态数据转化为 GPU 模型执行所需的输入张量。

5.2.1 Block Table 异步预提交

python 复制代码

self.input_batch.block_table.commit_block_table(num_reqs)

关键优化 ：最先启动 block table 的 CPU→GPU 拷贝（异步 non-blocking），让传输与后续的 CPU 计算重叠。这是一个经典的计算-通信重叠优化。

5.2.2 请求索引展开

python 复制代码

req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)

将请求索引按各请求的 token 数展开。例如 3 个请求分别有 2、5、3 个 token：[0,0,1,1,1,1,1,2,2,2]。这个数组用于将 per-token 数据映射到 per-request 数据。

python 复制代码

cu_num_tokens = self._get_cumsum_and_arange(
    num_scheduled_tokens, self.query_pos.np
)

cu_num_tokens 是累积和：[2,7,10]。query_pos 是每个 token 在其请求内的位置：[0,1,0,1,2,3,4,0,1,2]。这两个数组是 attention 计算的核心输入。

5.2.3 Positions 计算

python 复制代码

positions_np = (
    self.input_batch.num_computed_tokens_cpu[req_indices]
    + self.query_pos.np[: cu_num_tokens[-1]]
)

Position = 已计算 token 数 + 请求内偏移 。例如请求有 5 个已计算 token，本轮调度 3 个新 token，则 positions 为 [5, 6, 7]。这是 RoPE 等位置编码的输入。

5.2.4 Token IDs 收集

python 复制代码

token_indices = (
    positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]
)
token_indices_tensor = torch.from_numpy(token_indices)

torch.index_select(
    self.input_batch.token_ids_cpu_tensor.flatten(),
    0,
    token_indices_tensor,
    out=self.input_ids.cpu[:total_num_scheduled_tokens],
)

二维→一维索引映射 ：token_ids_cpu 的形状是 [max_num_reqs, max_model_len]。将 positions 和请求索引组合成一维索引 position + req_idx * max_model_len，然后使用 torch.index_select 高效收集。之所以用 torch 而非 numpy，是因为 torch 在大张量上更快。

5.2.5 Attention 元数据准备

python 复制代码

self.query_start_loc.np[0] = 0
self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens
self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1])
self.query_start_loc.copy_to_gpu()

query_start_loc 是 FlashAttention 等内核的必要输入------每个请求的 query token 起始位置。填充尾部的非递减值是因为某些内核（如 FlashAttention）要求 query_start_loc 非递减。

5.2.6 乐观序列长度与 Discard Mask

python 复制代码

torch.add(
    self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],
    torch.from_numpy(num_scheduled_tokens),
    out=self.optimistic_seq_lens_cpu[:num_reqs],
)

self.discard_request_mask.np[:num_reqs] = (
    self.optimistic_seq_lens_cpu[:num_reqs].numpy() < num_tokens_np
)
self.discard_request_mask.copy_to_gpu(num_reqs)

乐观 seq_len = 已计算 + 已调度 。对于 spec decode，这假设所有 draft tokens 都被接受。discard mask：如果乐观 seq_len < 实际总 token 数，说明该请求的某些 token 是 draft token，采样的结果应该被丢弃（因为它们对应的是 chunked prefill 的中间 token 或无效 draft token）。

5.2.7 GPU 侧状态同步

python 复制代码

if (
    self.use_async_spec_decode
    and self.valid_sampled_token_count_gpu is not None
    and prev_req_id_to_index
):
    self.prev_positions.copy_to_gpu(num_reqs)
    self.prev_num_draft_tokens.copy_to_gpu()
    cpu_values = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs].to(
        device=self.device, non_blocking=True
    )
    update_num_computed_tokens_for_batch_change(
        self.num_computed_tokens,
        self.num_accepted_tokens.gpu[:num_reqs],
        self.prev_positions.gpu[:num_reqs],
        self.valid_sampled_token_count_gpu,
        self.prev_num_draft_tokens.gpu,
        cpu_values,
    )
else:
    self.num_computed_tokens[:num_reqs].copy_(
        self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],
        non_blocking=True,
    )

异步 spec decode 的 GPU 修正 ：在异步模式下，CPU 侧的 num_computed_tokens 是乐观值，真正的修正需要在 GPU 上完成。update_num_computed_tokens_for_batch_change 是一个 CUDA kernel，使用上一步的 valid_sampled_token_count_gpu 来修正当前步的 num_computed_tokens。

同步模式：直接将 CPU 值拷贝到 GPU。

5.2.8 Slot Mapping 生成

python 复制代码

self.input_batch.block_table.compute_slot_mapping(
    num_reqs,
    self.query_start_loc.gpu[: num_reqs + 1],
    self.positions[:total_num_scheduled_tokens],
)

compute_slot_mapping 是 BlockTable 的核心方法，将每个 token 的 position 映射到 KV 缓存中的物理 slot ID。详见第七章。

5.2.9 logits_indices 与 spec_decode_metadata

python 复制代码

if not use_spec_decode:
    logits_indices = query_start_loc[1:] - 1
    spec_decode_metadata = None

无 spec decode ：logits 只需要从每个请求的最后一个 token 计算（query_start_loc[i+1] - 1）。这是 chunked prefill + decode 的通用路径。

python 复制代码

else:
    num_draft_tokens = np.zeros(num_reqs, dtype=np.int32)
    num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32)
    for req_id, draft_token_ids in scheduled_spec_decode_tokens.items():
        req_idx = self.input_batch.req_id_to_index[req_id]
        draft_len = len(draft_token_ids)
        num_draft_tokens[req_idx] = draft_len
        if (
            self.input_batch.num_computed_tokens_cpu[req_idx]
            >= self.input_batch.num_prompt_tokens[req_idx]
        ):
            num_decode_draft_tokens[req_idx] = draft_len
    spec_decode_metadata = self._calc_spec_decode_metadata(
        num_draft_tokens, cu_num_tokens
    )
    logits_indices = spec_decode_metadata.logits_indices

Spec decode 路径 ：需要为每个 draft token 位置计算 logits（用于 rejection sampling）。num_decode_draft_tokens 区分 prefill 和 decode 阶段------prefill 阶段的 draft tokens 不参与 rejection（用 -1 标记），因为 guided decoding 可能回滚这些 tokens。

5.2.10 返回值

返回 (logits_indices, spec_decode_metadata)，供后续 Phase 3/4/5 使用。

Phase 3：_preprocess()（行 3193--3311）

5.3.1 MM Encoder 执行（首阶 + 非编解码模型）

python 复制代码

if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder:
    with self.maybe_get_ec_connector_output(
        scheduler_output,
        encoder_cache=self.encoder_cache,
    ) as ec_connector_output:
        self._execute_mm_encoder(scheduler_output)
        mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output)

MM Encoder 执行流程：

在 EC connector context 中执行视觉/音频编码器。
_gather_mm_embeddings：将编码器输出放置到正确的 token 位置上（根据 mm_position 信息），生成 mm_embeds（多模态 embedding）和 is_mm_embed（布尔掩码，标识哪些位置是多模态 token）。

python 复制代码

inputs_embeds_scheduled = self.model.embed_input_ids(
    self.input_ids.gpu[:num_scheduled_tokens],
    multimodal_embeddings=mm_embeds,
    is_multimodal=is_mm_embed,
)
self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens)

统一嵌入化 ：多模态模型始终使用 embeddings 作为输入（而非 token IDs），即使某些 token 是文本 token。embed_input_ids 将文本 token ID 转换为 embedding，然后与视觉 embedding 合并。_prepare_mm_inputs 处理 padding（cudagraph 需要固定大小的输入）。

5.3.2 Prompt Embeds 路径

python 复制代码

elif self.enable_prompt_embeds and is_first_rank:
    token_ids_idx = (
        self.is_token_ids.gpu[:num_scheduled_tokens]
        .nonzero(as_tuple=False)
        .squeeze(1)
    )
    if token_ids_idx.numel() > 0:
        token_ids = self.input_ids.gpu[token_ids_idx]
        tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids)
        self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
    inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
    model_kwargs = self._init_model_kwargs()
    input_ids = None

混合输入 ：当启用 prompt embeds 时，同一批次中可能既有 token IDs 又有 embeddings。is_token_ids 掩码标识哪些位置是 token IDs，需要转换为 embeddings。其余位置已经是 embeddings。input_ids = None 告知模型使用 embeddings 路径。

5.3.3 纯文本路径

python 复制代码

else:
    input_ids = self.input_ids.gpu[:num_input_tokens]
    inputs_embeds = None
    model_kwargs = self._init_model_kwargs()

性能关键路径：纯文本模型直接使用 token IDs，不做 embedding 转换。这样 embedding 层被包含在 CUDA graph 中，避免了额外内核启动和内存传输。

5.3.4 Positions 处理

python 复制代码

if self.uses_mrope:
    positions = self.mrope_positions.gpu[:, :num_input_tokens]
elif self.uses_xdrope_dim > 0:
    positions = self.xdrope_positions.gpu[:, :num_input_tokens]
else:
    positions = self.positions[:num_input_tokens]
    if num_input_tokens > num_scheduled_tokens:
        self.positions[num_scheduled_tokens:num_input_tokens].zero_()

三种位置编码：

M-RoPE ：多维位置编码（Qwen2-VL），形状 [dim, num_tokens]。
XD-RoPE ：交叉维度位置编码（HunYuan-VL），形状 [dim, num_tokens]。
标准 RoPE/ALiBi ：一维位置，形状 [num_tokens]。Padding 区域置零（避免无效位置影响 CUDA graph 执行）。

5.3.5 Pipeline Parallel 中间张量

python 复制代码

if is_first_rank:
    intermediate_tensors = None
else:
    intermediate_tensors = self.sync_and_slice_intermediate_tensors(
        num_input_tokens, intermediate_tensors, True
    )

非首阶 PP rank 从前一阶接收中间张量并切片到当前批次大小。

5.3.6 Encoder-Decoder 特殊处理

python 复制代码

if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
    encoder_outputs = self._execute_mm_encoder(scheduler_output)
    model_kwargs.update({"encoder_outputs": encoder_outputs})

Encoder-Decoder 模型（如 T5、BART）的 encoder 输出直接传递给 decoder 作为 encoder_outputs kwarg。

Phase 4：Forward（行 3900--4080 + 关联上下文）

5.4.1 CUDA Graph 模式调整

python 复制代码

if self.calculate_kv_scales:
    cudagraph_mode = CUDAGraphMode.NONE
    self.calculate_kv_scales = False

KV scales 计算涉及动态操作（如 max 操作），与 CUDA graph 的静态捕获不兼容。首次执行时降级为 eager 模式。

python 复制代码

has_encoder_input = (
    self.model_config.is_encoder_decoder and num_encoder_reqs > 0
)

Encoder-Decoder 模型在有 encoder 输入的步不能使用 CUDA graph（因为 encoder 输出是动态的）。

5.4.2 set_forward_context

python 复制代码

with (
    set_forward_context(
        attn_metadata,
        self.vllm_config,
        num_tokens=num_tokens_padded,
        num_tokens_across_dp=num_tokens_across_dp,
        cudagraph_runtime_mode=cudagraph_mode,
        batch_descriptor=batch_desc,
        ubatch_slices=ubatch_slices_padded,
        slot_mapping=slot_mappings,
        skip_compiled=has_encoder_input,
    ),
    record_function_or_nullcontext("gpu_model_runner: forward"),
    self.maybe_get_kv_connector_output(
        scheduler_output,
        defer_finalize=defer_kv_connector_finalize,
    ) as kv_connector_output,
):

三个 context manager 嵌套：

set_forward_context ：将当前步的 attention metadata、CUDA graph 模式等推入全局上下文（ForwardContext），使模型的各层可以通过全局访问获取这些信息，而无需层层传递参数。这是 vLLM 的核心设计模式------通过 thread-local / 全局变量避免函数签名膨胀。
record_function_or_nullcontext：性能分析标记（PyTorch profiler 可见），生产环境中可能为空 context。
maybe_get_kv_connector_output ：KV Transfer 的连接器输出管理。defer_finalize=True 时（spec decode 场景），连接器的 wait_for_save 和 clear_metadata 推迟到 draft model 执行后，以便 draft model 也能保存其 KV 缓存。

5.4.3 _model_forward

python 复制代码

model_output = self._model_forward(
    input_ids=input_ids,
    positions=positions,
    intermediate_tensors=intermediate_tensors,
    inputs_embeds=inputs_embeds,
    **model_kwargs,
)

_model_forward 是对 self.model() 的薄包装，便于子类覆盖（如测试时 mock）。实际执行可能是：

CUDA graph 路径：重放预捕获的 CUDA graph，零开销。
Eager 路径：正常的 PyTorch 前向执行。

5.4.4 后处理与 Logits 计算

python 复制代码

if self.use_aux_hidden_state_outputs:
    hidden_states, aux_hidden_states = model_output
else:
    hidden_states = model_output
    aux_hidden_states = None

EAGLE 3 等 spec decode 方法需要中间层的 hidden states 作为 drafter 的输入，通过 aux_hidden_states 传递。

python 复制代码

if not self.broadcast_pp_output:
    if not get_pp_group().is_last_rank:
        assert isinstance(hidden_states, IntermediateTensors)
        hidden_states.kv_connector_output = kv_connector_output
        self.kv_connector_output = kv_connector_output
        return hidden_states

PP 中间阶 ：非最后阶直接返回 IntermediateTensors 给 APD（All-Parallel Dispatch）层进行跨阶通信。同时保存 kv_connector_output 引用。

python 复制代码

if self.is_pooling_model:
    return self._pool(
        hidden_states,
        num_scheduled_tokens,
        num_scheduled_tokens_np,
        kv_connector_output,
    )

sample_hidden_states = hidden_states[logits_indices]
logits = self.model.compute_logits(sample_hidden_states)

Pooling 模型 ：直接进入 pooling 流程。生成模型 ：只从 logits_indices 位置取 hidden states 计算 logits，避免为所有 token 计算 logits（节省计算量）。logits_indices 在 Phase 2 中确定------decode 时只取最后一个 token，spec decode 时取所有 draft token + 最后一个 token。

5.4.5 PP 广播路径

python 复制代码

else:
    sample_hidden_states = hidden_states[logits_indices]
    if not get_pp_group().is_last_rank:
        get_pp_group().send_tensor_dict(
            hidden_states.tensors,
            all_gather_group=get_tp_group(),
            all_gather_tensors=all_gather_tensors,
        )
        logits = None
    else:
        logits = self.model.compute_logits(sample_hidden_states)

    model_output_broadcast_data: dict[str, Any] = {}
    if logits is not None:
        model_output_broadcast_data["logits"] = logits.contiguous()

    broadcasted = get_pp_group().broadcast_tensor_dict(
        model_output_broadcast_data, src=len(get_pp_group().ranks) - 1
    )
    assert broadcasted is not None
    logits = broadcasted["logits"]

PP + broadcast_pp_output 模式：所有 PP 阶都需要 logits（用于 logits processor），因此最后阶计算 logits 后广播给其他阶。中间阶先发送 hidden_states，然后接收广播的 logits。

5.4.6 保存执行状态

python 复制代码

self.execute_model_state = ExecuteModelState(
    scheduler_output,
    logits,
    spec_decode_metadata,
    spec_decode_common_attn_metadata,
    hidden_states,
    sample_hidden_states,
    aux_hidden_states,
    ec_connector_output,
    cudagraph_stats,
    slot_mappings,
)
self.kv_connector_output = kv_connector_output

if deferred_state_corrections_fn:
    deferred_state_corrections_fn()

将 forward 的所有产物打包为 ExecuteModelState，供后续 sample_tokens() 使用。这是 execute_model → sample_tokens 两步协议的桥梁。最后执行延迟修正------此时 forward 已完成，GPU 上的采样结果可用了。

Phase 5：Sample + Bookkeeping（行 3311--3469 + 4124--4460）

5.5.1 sample_tokens() 入口

python 复制代码

@torch.inference_mode
def sample_tokens(
    self, grammar_output: "GrammarOutput | None"
) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:

sample_tokens 是 execute_model 的配对方法，在 execute_model 返回 None（表示"需要等采样"）后被调度器调用。它从 execute_model_state 中恢复 forward 的产物，执行采样和后处理。

python 复制代码

if self.execute_model_state is None:
    kv_connector_output = self.kv_connector_output
    self.kv_connector_output = None
    if self.use_async_scheduling and get_pp_group().world_size > 1:
        self._pp_receive_prev_sampled_token_ids_to_input_batch()
    if not kv_connector_output:
        return None
    if kv_connector_output.is_empty():
        return EMPTY_MODEL_RUNNER_OUTPUT
    output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
    output.kv_connector_output = kv_connector_output
    return output

KV transfer 场景下的特殊处理 ：当 execute_model_state 为 None（说明 execute_model 返回的是中间结果而非需要采样的结果），但仍有 KV transfer 输出需要传递。

5.5.2 Grammar Bitmask 应用

python 复制代码

if grammar_output is not None:
    apply_grammar_bitmask(
        scheduler_output, grammar_output, self.input_batch, logits
    )

结构化输出（JSON/regex 约束）通过 grammar bitmask 实现：在 logits 上应用布尔掩码，将不合规的 token 的 logits 设为 -inf。grammar_output 是由 grammar processor 在 CPU 上生成的 bitmask，apply_grammar_bitmask 将其应用到 GPU logits 上。

5.5.3 _sample()

python 复制代码

def _sample(self, logits, spec_decode_metadata):
    sampling_metadata = self.input_batch.sampling_metadata
    self.input_batch.update_async_output_token_ids()
    
    if spec_decode_metadata is None:
        return self.sampler(
            logits=logits,
            sampling_metadata=sampling_metadata,
        )

    if self.use_async_scheduling and self._draft_token_req_ids is not None:
        draft_token_ids_cpu, _ = self._get_draft_token_ids_cpu()
        self.input_batch.update_async_spec_token_ids(draft_token_ids_cpu)

    sampler_output = self.rejection_sampler(
        spec_decode_metadata, None, logits, sampling_metadata,
    )
    return sampler_output

两条采样路径：

无 spec decode ：直接调用 self.sampler，执行 temperature scaling、top-k/top-p 过滤、随机采样等标准采样流程。
有 spec decode ：使用 rejection_sampler 进行 rejection sampling------将 target model 的 logits 与 draft tokens 比较，接受匹配的 draft tokens，拒绝不匹配的并从 target model 重新采样。

异步调度的 output_token_ids 更新 ：在采样前，需要从上一步的 GPU→CPU 异步拷贝结果中更新 output_token_ids（用于 frequency/presence penalty 计算）。update_async_spec_token_ids 类似地更新 draft token IDs。

5.5.4 _bookkeeping_sync()

python 复制代码

def _bookkeeping_sync(self, scheduler_output, sampler_output, logits, hidden_states, ...):

账本同步的核心方法，将采样结果写回持久化状态。

python 复制代码

num_nans_in_logits = {}
if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
    num_nans_in_logits = self._get_nans_in_logits(logits)

NaN 检测（调试用，默认关闭）。

python 复制代码

discard_sampled_tokens_req_indices = np.nonzero(
    self.discard_request_mask.np[:num_reqs]
)[0]
for i in discard_sampled_tokens_req_indices:
    gen = self.input_batch.generators.get(int(i))
    if gen is not None:
        gen.set_offset(gen.get_offset() - 4)

Discard 处理 ：对于被 discard mask 标记的请求（chunked prefill 的中间 token），回退随机数生成器的 offset 4 字节（torch.Generator 的内部状态单位），确保这些请求的随机状态不受影响------因为它们虽然被采样了（统一处理更简单），但结果会被丢弃。

python 复制代码

req_ids_output_copy = self.input_batch.req_ids.copy()
req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()

防御性拷贝 ：异步调度场景下，返回结果后批次可能被修改（下一轮 _update_states），因此必须拷贝 ID 映射。同步场景下这有一定开销但确保安全。

python 复制代码

if not self.use_async_scheduling:
    if max_gen_len == 1:
        valid_sampled_token_ids = self._to_list(sampled_token_ids)
        for i in discard_sampled_tokens_req_indices:
            valid_sampled_token_ids[int(i)].clear()
    else:
        valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output(
            sampled_token_ids, self.input_batch.vocab_size,
            discard_sampled_tokens_req_indices, ...
        )

同步模式 ：直接将 GPU 采样结果转为 CPU 列表。Spec decode 时使用 RejectionSampler.parse_output 解析 rejection sampling 的输出格式（接受/拒绝的 token 序列）。

python 复制代码

else:
    valid_sampled_token_ids = []
    invalid_req_indices = discard_sampled_tokens_req_indices.tolist()
    
    if self.input_batch.prev_sampled_token_ids is None:
        assert sampled_token_ids.shape[-1] == 1
        self.input_batch.prev_sampled_token_ids = sampled_token_ids

异步模式 ：不等待 GPU→CPU 拷贝，而是将 sampled_token_ids 保留在 GPU 上（prev_sampled_token_ids），下一轮 _prepare_inputs 中直接在 GPU 上将其写入 input_ids。只记录哪些请求的采样应被丢弃（invalid_req_indices）。

5.5.5 采样结果写回 token_ids_cpu

python 复制代码

for req_idx in range(num_sampled_tokens):
    if self.use_async_scheduling:
        sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None
    else:
        sampled_ids = valid_sampled_token_ids[req_idx]

    num_sampled_ids: int = len(sampled_ids) if sampled_ids else 0
    if not sampled_ids:
        continue

    start_idx = self.input_batch.num_tokens_no_spec[req_idx]
    end_idx = start_idx + num_sampled_ids
    assert end_idx <= self.max_model_len

    self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids
    self.input_batch.is_token_ids[req_idx, start_idx:end_idx] = True
    self.input_batch.num_tokens_no_spec[req_idx] = end_idx

    req_state.output_token_ids.extend(sampled_ids)

关键逻辑：

同步模式 ：将真实的采样 token IDs 写入 token_ids_cpu。
异步模式 ：写入 -1 占位符（真实 IDs 还在 GPU 上），下一轮会被 GPU 侧覆盖。
max_model_len 溢出检查：确保 token 数不超过模型最大长度。
双重写入 ：同时更新 token_ids_cpu（InputBatch 的持久化状态）和 output_token_ids（CachedRequestState 的请求级状态）。

5.5.6 异步输出封装

python 复制代码

if not self.use_async_scheduling:
    return output

async_output = AsyncGPUModelRunnerOutput(
    model_runner_output=output,
    sampled_token_ids=sampler_output.sampled_token_ids,
    logprobs_tensors=sampler_output.logprobs_tensors,
    invalid_req_indices=invalid_req_indices,
    async_output_copy_stream=self._get_or_create_async_output_copy_stream(),
    vocab_size=self.input_batch.vocab_size,
)

self.input_batch.set_async_sampled_token_ids(
    async_output.sampled_token_ids_cpu,
    async_output.async_copy_ready_event,
)

AsyncGPUModelRunnerOutput ：封装了 GPU→CPU 异步拷贝的结果。它使用独立的 CUDA stream 进行拷贝，不阻塞主 stream。sampled_token_ids_cpu 是拷贝目标，async_copy_ready_event 是同步事件------CPU 侧需要结果时等待该事件。

set_async_sampled_token_ids 将 CPU 侧张量和事件存入 InputBatch，下一轮 _update_states 中 update_async_output_token_ids 会等待事件完成后用真实 IDs 替换占位符。

5.5.7 Spec Decode Draft Token 提议

python 复制代码

spec_config = self.speculative_config
if spec_config is not None:
    input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
        spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
        <= self.effective_drafter_max_model_len
    )
    use_gpu_toks = (
        spec_config.use_eagle()
        or spec_config.uses_draft_model()
        or spec_config.uses_extract_hidden_states()
    ) and not spec_config.disable_padded_drafter_batch

Drafter 执行决策树：

输入是否适配 drafter ：max_seq_len + spec_tokens <= effective_drafter_max_model_len。如果序列已经接近 drafter 的最大长度，则跳过 draft 生成（避免 OOM 或精度下降）。
是否使用 GPU sampled tokens：EAGLE/DraftModel 等方法可以直接使用 GPU 上的采样结果作为 drafter 输入，无需等 CPU 拷贝。ngram CPU 方法需要等 CPU 结果。

python 复制代码

if use_gpu_toks:
    if input_fits_in_drafter:
        propose_draft_token_ids(sampled_token_ids)
    elif self.valid_sampled_token_count_event is not None:
        next_token_ids, valid_sampled_tokens_count = (
            self.drafter.prepare_next_token_ids_padded(...)
        )
        self._copy_valid_sampled_token_count(...)
elif spec_config.use_ngram_gpu() and not spec_config.disable_padded_drafter_batch:
    if input_fits_in_drafter:
        propose_draft_token_ids(sampled_token_ids)
else:
    propose_drafts_after_bookkeeping = input_fits_in_drafter

三级降级：

EAGLE/DraftModel + 输入适配：直接用 GPU sampled tokens 运行 drafter，与 bookkeeping 并行。
输入不适配 ：零化 draft tokens（避免调度器使用陈旧的 draft tokens），但通过 prepare_next_token_ids_padded 准备下一步的 padded 输入。
ngram CPU：必须在 bookkeeping 后运行（需要 CPU 侧的 sampled token IDs）。

5.5.8 EPLB Step

python 复制代码

self.eplb_step()

Expert Parallelism Load Balancing step------在 MoE 模型中平衡专家间的负载。这是在采样/后处理之后执行的，不阻塞推理关键路径。