GPUModelRunner是真正执行模型前向传播的组件，主要的功能：
- load_model：模型加载与初始化。
- execute_model ：模型执行，负责驱动模型的执行并输出结果。
- capture_model ：图捕获。
- LORA：主要靠继承的LoRAModelRunnerMixin实现。
load_model
load_model完成模型的加载：
- 通过get_model_loader获取model_loader，再调用model_loader.load_model加载模型。
- 调用self.load_lora_model加载LORA部分，包括：模型转换，权重注入，调度器适配。
- 调用self.drafter.load_model(self.model)加载drafter模型，用于Speculative Decoding。
- 如果是MOE模型，则创建EplbState，用于后续专家并行负载均衡。
py
def load_model(self, eep_scale_up: bool = False) -> None:
...
with DeviceMemoryProfiler() as m: # noqa: SIM117
time_before_load = time.perf_counter()
model_loader = get_model_loader(self.load_config)
if not hasattr(self, "model"):
self.model = model_loader.load_model(
vllm_config=self.vllm_config,
model_config=self.model_config)
else:
model_loader.load_weights(self.model,
model_config=self.model_config)
if self.lora_config:
self.model = self.load_lora_model(self.model,
self.model_config,
self.scheduler_config,
self.lora_config,
self.device)
if hasattr(self, "drafter"):
self.drafter.load_model(self.model)
if self.use_aux_hidden_state_outputs:
self.model.set_aux_hidden_state_layers(
self.model.get_eagle3_aux_hidden_state_layers())
prepare_communication_buffer_for_model(self.model)
if is_mixture_of_experts(
self.model) and self.parallel_config.enable_eplb:
logger.info("EPLB is enabled for model %s.",
self.model_config.model)
self.eplb_state = EplbState.build(
self.model,
self.device,
self.parallel_config,
global_expert_load,
old_global_expert_indices,
rank_mapping,
)execute_model
execute_model完成模型的执行：
- 状态更新：调用self._update_states更新状态。
- 输入准备：调用self._prepare_inputs根据调度器的输出准备模型的输入数据和相关元数据，包括：
1）attn_metadata: dict[str, Any]: 映射每一层到其对应的注意力元数据。
2）attention_cuda_graphs: bool: attention是否使用cugraph。
3）logits_indices: torch.Tensor: logits 的索引。
4）spec_decode_metadata: Optional[SpecDecodeMetadata]: 推测解码的元数据。
5）num_scheduled_tokens: np.ndarray: 每个request的调度token数。
6）spec_decode_common_attn_metadata: Optional[CommonAttentionMetadata]: 推测解码的通用注意力元数据。 - Padding num_input_tokens：
1）如使用CUDA图，则padding到对应分档的Graph Size。
2）如使能序列并行，需要padding到tp_size的整数倍（vllm中sp_size == tp_size）。
3）DP Padding：DP Group内的所有RANK的num_input_tokens对齐到最大值。 - 多模态Encode：如果是多模态模型，调用self._execute_mm_encoder生成多模态的tokens：mm_embeds。
- 多模态Embedding：如果是多模态模型，调用self.model.get_input_embeddings生成输入的embedding（多模态需要处理不同来源的输入，且图像和视频适合预计算）。
- PP并行intermediate_tensors：PP并行下，非第一个RANK调用self.sync_and_slice_intermediate_tensors处理前级RANK输入的intermediate_tensors。
- KV Cache加载：调用self.maybe_setup_kv_connector(scheduler_output)加载KV Cache。
- 模型前向传播：调用self.model执行一次前向传播。
- KV Cache存储：调用self.maybe_wait_for_kv_save()等待新生成的KV Cache写入完成。
- 计算Logits：调用self.model.compute_logits(sample_hidden_states, None)计算新生成token的Logits。
- 文法约束：调用self.apply_grammar_bitmask清零的不符合文化的token Logits。
- Sampling：调用self.sampler进行token采样，生成新token id（考虑temperature，top-p，top-k，repetition penalty等）。
- Reject Sampling：调用self.rejection_sampler对投机推理的结果进行拒绝采样，最后生成的output tokens = accepted tokens + recovered tokens + bonus tokens。
- Discard Tokens：在Partial Prefill等场景下，可能存在Prompt的token还未处理完的情况，所以这时要丢弃生成的token，并回退torch.Generator的状态。
- Parse Valid Tokens：如果不存在投机推理，则所有token都有效，否则则调用self.rejection_sampler.parse_output获取输出中有效的tokens。
- 投机推理：调用self.propose_draft_token_ids进行下一轮的投机推理。
- EPLB：调用self.eplb_step()处理EP的Load Balance步。
- 返回结果：ModelRunnerOutput
py
@torch.inference_mode()
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[ModelRunnerOutput, IntermediateTensors]:
self._update_states(scheduler_output)
...
(attn_metadata, attention_cuda_graphs, logits_indices,
spec_decode_metadata, num_scheduled_tokens_np,
spec_decode_common_attn_metadata) = (
self._prepare_inputs(scheduler_output))
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
if (self.use_cuda_graph
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
num_input_tokens = self.vllm_config.pad_for_cudagraph(
num_scheduled_tokens)
else:
tp_size = self.vllm_config.parallel_config.tensor_parallel_size
if self.compilation_config.pass_config. \
enable_sequence_parallelism and tp_size > 1:
num_input_tokens = round_up(num_scheduled_tokens, tp_size)
else:
num_input_tokens = num_scheduled_tokens
num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
num_input_tokens += num_pad
if self.is_multimodal_model:
self._execute_mm_encoder(scheduler_output)
mm_embeds = self._gather_mm_embeddings(scheduler_output)
else:
mm_embeds = []
if self.is_multimodal_model and get_pp_group().is_first_rank:
input_ids = self.input_ids[:num_scheduled_tokens]
inputs_embeds = self.model.get_input_embeddings(
input_ids=input_ids,
multimodal_embeddings=mm_embeds or None,
)
self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds)
inputs_embeds = self.inputs_embeds[:num_input_tokens]
input_ids = None
else:
input_ids = self.input_ids[:num_input_tokens]
inputs_embeds = None
...
if get_pp_group().is_first_rank:
intermediate_tensors = None
else:
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
num_input_tokens, intermediate_tensors, True)
...
skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
with set_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
skip_cuda_graphs=skip_cuda_graphs,
):
self.maybe_setup_kv_connector(scheduler_output)
model_output = self.model(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
)
self.maybe_wait_for_kv_save()
finished_sending, finished_recving = (
self.get_finished_kv_transfers(scheduler_output))
if self.use_aux_hidden_state_outputs:
hidden_states, aux_hidden_states = model_output
else:
hidden_states = model_output
aux_hidden_states = None
broadcast_pp_output = \
self.parallel_config.distributed_executor_backend \
== "external_launcher" and len(get_pp_group().ranks) > 0
if not get_pp_group().is_last_rank:
...
else:
...
sample_hidden_states = hidden_states[logits_indices]
logits = self.model.compute_logits(sample_hidden_states, None)
...
if scheduler_output.grammar_bitmask is not None:
self.apply_grammar_bitmask(scheduler_output, logits)
sampling_metadata = self.input_batch.sampling_metadata
if spec_decode_metadata is None:
sampler_output = self.sampler(
logits=logits,
sampling_metadata=sampling_metadata,
)
else:
bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
sampler_output = self.sampler(
logits=bonus_logits,
sampling_metadata=sampling_metadata,
)
bonus_token_ids = sampler_output.sampled_token_ids
target_logits = logits[spec_decode_metadata.target_logits_indices]
output_token_ids = self.rejection_sampler(
spec_decode_metadata,
None, # draft_probs
target_logits,
bonus_token_ids,
sampling_metadata,
)
sampler_output.sampled_token_ids = output_token_ids
discard_sampled_tokens_req_indices = []
for i, req_id in enumerate(self.input_batch.req_ids):
req_state = self.requests[req_id]
seq_len = (req_state.num_computed_tokens +
scheduler_output.num_scheduled_tokens[req_id])
if seq_len < req_state.num_tokens:
generator = self.input_batch.generators.get(i)
if generator is not None:
generator.set_offset(generator.get_offset() - 4)
discard_sampled_tokens_req_indices.append(i)
...
sampled_token_ids = sampler_output.sampled_token_ids
max_gen_len = sampled_token_ids.shape[-1]
if max_gen_len == 1:
valid_sampled_token_ids = sampled_token_ids.tolist()
else:
valid_sampled_token_ids = self.rejection_sampler.parse_output(
sampled_token_ids,
self.input_batch.vocab_size,
)
...
if not self.speculative_config:
# Speculative decoding is not enabled.
spec_token_ids = None
else:
assert spec_decode_common_attn_metadata is not None
spec_token_ids = self.propose_draft_token_ids(
scheduler_output,
valid_sampled_token_ids,
sampling_metadata,
hidden_states,
sample_hidden_states,
aux_hidden_states,
spec_decode_metadata,
spec_decode_common_attn_metadata,
)
self.eplb_step()
return ModelRunnerOutput(
req_ids=self.input_batch.req_ids,
req_id_to_index=self.input_batch.req_id_to_index,
sampled_token_ids=valid_sampled_token_ids,
spec_token_ids=spec_token_ids,
logprobs=logprobs_lists,
prompt_logprobs_dict=prompt_logprobs_dict,
pooler_output=[],
finished_sending=finished_sending,
finished_recving=finished_recving,
num_nans_in_logits=num_nans_in_logits,
)_calc_spec_decode_metadata
在execute_model中，会调用_calc_spec_decode_metadata计算投机推理的SpecDecodeMetadata。
以代码中的注释为例：
py
Inputs:
cu_num_scheduled_tokens: [ 4, 104, 107, 207, 209]
num_draft_tokens: [ 3, 0, 2, 0, 1]
Outputs:
cu_num_draft_tokens: [ 3, 3, 5, 5, 6]
logits_indices: [ 0, 1, 2, 3, 103, 104, 105, 106, 206, 207, 208]
target_logits_indices: [ 0, 1, 2, 5, 6, 9]
bonus_logits_indices: [ 3, 4, 7, 8, 10]输入：
累积的调度token数：cu_num_scheduled_tokens = [4, 104, 107, 207, 209]
在上一个Step每个Request已经通过draft model投机推理产生的token数：[3, 0, 2, 0, 1]
所以输入的5个Request：
- Request 0：从位置4开始，有3个draft token，本轮推理要Sample 4个(3+1)token。
- Request 1：从位置104开始，有0个draft token，本轮推理要Sample 1个token。
- Request 2：从位置107开始，有2个draft token，本轮推理要Sample 3个token。
- Request 3：从位置207开始，有0个draft token，本轮推理要Sample 1个token。
- Request 4：从位置209开始，有1个draft token，本轮推理要Sample 2个token。
输出：
累积的草稿token数cu_num_draft_tokens: [3, 3, 5, 5, 6]
需要Sample的token在输出的logits中的位置logits_indices：
- Request 0：[0(4-3-1), 1, 2, 3]
- Request 1：[103(104-0-1)]
- Request 2：[104(107-2-1), 105, 106]
- Request 3：[206(207-0-1)]
- Request 4：[207(209-1-1), 208]
target model生成的用于验证draft token在logits_indices中的位置target_logits_indices：
- Request 0：[0， 1， 2]
- Request 1：[]
- Request 2：[5， 6]
- Request 3：[]
- Request 4：[9]
所有的draft token都被接受后，next token在logits_indices中的位置bonus_logits_indices：
- Request 0：[3]
- Request 1：[4]
- Request 2：[7]
- Request 3：[8]
- Request 4：[10]