vllm-forward

vllm/vllm/v1/worker/gpu_model_runner.py · ef9339e · L4136–4190 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

4136        with (
4137            set_forward_context(
4138                attn_metadata,
4139                self.vllm_config,
4140                num_tokens=num_tokens_padded,
4141                num_tokens_across_dp=num_tokens_across_dp,
4142                cudagraph_runtime_mode=cudagraph_mode,
4143                batch_descriptor=batch_desc,
4144                ubatch_slices=ubatch_slices_padded,
4145                slot_mapping=slot_mappings,
4146                skip_compiled=has_encoder_input,
4147                is_padding=is_padding,
4148            ),
4149            record_function_or_nullcontext("gpu_model_runner: forward"),
4150            self.maybe_get_kv_connector_output(
4151                scheduler_output,
4152                defer_finalize=defer_kv_connector_finalize,
4153            ) as kv_connector_output,
4154        ):
4155            model_output = self._model_forward(
4156                input_ids=input_ids,
4157                positions=positions,
4158                intermediate_tensors=intermediate_tensors,
4159                inputs_embeds=inputs_embeds,
4160                **model_kwargs,
4161            )
4162
4163        with record_function_or_nullcontext("gpu_model_runner: postprocess"):
4164            if self.use_aux_hidden_state_outputs:
4165                # True when EAGLE 3 is used.
4166                hidden_states, aux_hidden_states = model_output
4167            else:
4168                # Common case.
4169                hidden_states = model_output
4170                aux_hidden_states = None
4171
4172            if not self.broadcast_pp_output:
4173                # Common case.
4174                if not get_pp_group().is_last_rank:
4175                    # Return the intermediate tensors.
4176                    assert isinstance(hidden_states, IntermediateTensors)
4177                    self.kv_connector_output = kv_connector_output
4178                    return hidden_states
4179
4180                if self.is_pooling_model:
4181                    # Return the pooling output.
4182                    return self._pool(
4183                        hidden_states,
4184                        num_scheduled_tokens,
4185                        num_scheduled_tokens_np,
4186                        kv_connector_output,
4187                    )
4188
4189                sample_hidden_states = hidden_states[logits_indices]
4190                logits = self.model.compute_logits(sample_hidden_states)