From 53b86439fc0dd811c2dc97d8dee00a158ecbf9c6 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 22 Oct 2025 17:16:15 +0800 Subject: [PATCH] Refactor ARModelRunner to streamline hidden state handling - Removed unnecessary imports of numpy and torch within the request scheduling logic. - Simplified the handling of auxiliary hidden states by initializing it to None when not present. - Optimized CPU tensor conversion for text hidden states to improve performance. --- vllm_omni/worker/gpu_ar_model_runner.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 6f32a3c03..7e8ec3135 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -59,8 +59,6 @@ class ARModelRunner(OmniGPUModelRunner): try: new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) if new_reqs: - import numpy as np - import torch for nr in new_reqs: req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) if req_id is None: @@ -272,6 +270,7 @@ class ARModelRunner(OmniGPUModelRunner): hidden_states, _aux_hidden_states = model_output else: hidden_states = model_output + _aux_hidden_states = None text_hidden_states, multimodal_outputs = ( self.extract_multimodal_outputs(hidden_states)) @@ -403,16 +402,17 @@ class ARModelRunner(OmniGPUModelRunner): sampling_metadata, text_hidden_states, sample_hidden_states, - _aux_hidden_states if '_aux_hidden_states' in locals() else None, + _aux_hidden_states, spec_decode_metadata, spec_decode_common_attn_metadata, ) # Convert to per-request tensors on CPU + text_hidden_states_cpu = text_hidden_states.detach().to("cpu").contiguous() pooler_output: list[Optional[torch.Tensor]] = [] prev_logits_index = 0 for logits_index in logits_indices: - pooler_output.append(text_hidden_states[prev_logits_index:logits_index+1].detach().to("cpu").contiguous()) + pooler_output.append(text_hidden_states_cpu[prev_logits_index:logits_index+1]) prev_logits_index = logits_index + 1