Fixed memory leak and Remove dead code (#3312)

Signed-off-by: Dan250124 <416947747@qq.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

tests/engine/test_orchestrator.py

@@ -93,6 +93,7 @@ class FakeStageClient:
 class FakeOutputProcessor:
     def __init__(self, *, request_outputs: list[object] | None = None) -> None:
         self.request_outputs = list(request_outputs or [])
+        self.abort_calls: list[list[str]] = []
 
     def add_request(self, *_args, **_kwargs) -> None:
         return None
@@ -103,6 +104,10 @@ class FakeOutputProcessor:
             reqs_to_abort=[],
         )
 
+    def abort_requests(self, request_ids, internal: bool = False):
+        self.abort_calls.append(request_ids)
+        return request_ids
+
     def update_scheduler_stats(self, _scheduler_stats) -> None:
         return None
 
@@ -599,6 +604,8 @@ async def test_run_abort(orchestrator_factory) -> None:
 
         for stage in stages:
             assert stage.abort_calls == [["req-abort"]]
+        for proc in processors:
+            assert proc.abort_calls == [["req-abort"]]
         assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)

vllm_omni/engine/orchestrator.py

@@ -873,6 +873,13 @@ class Orchestrator:
 
         if processed.reqs_to_abort:
             await self.stage_clients[stage_id].abort_requests_async(processed.reqs_to_abort)
+            # Same cleanup as _handle_abort — engine-driven aborts (stop
+            # token, max tokens, etc.) also skip the normal process_outputs()
+            # completion path, so accumulated CPU tensors in OmniRequestState
+            # and OutputProcessor would leak.
+            self.output_processors[stage_id].abort_requests(processed.reqs_to_abort, internal=True)
+            for req_id in processed.reqs_to_abort:
+                self.request_states.pop(req_id, None)
 
         if raw_outputs.scheduler_stats is not None:
             processor.update_scheduler_stats(raw_outputs.scheduler_stats)
@@ -1082,6 +1089,11 @@ class Orchestrator:
         all_ids_to_abort = self._cfg_tracker.abort_parents(request_ids)
         for stage_id in range(self.num_stages):
             await self.stage_clients[stage_id].abort_requests_async(all_ids_to_abort)
+            # Clean up OutputProcessor state (e.g. mm_accumulated tensors) that
+            # would otherwise leak — normal completion cleans up via
+            # process_outputs(), but aborted requests never produce an
+            # EngineCoreOutput, so we must purge them explicitly.
+            self.output_processors[stage_id].abort_requests(all_ids_to_abort, internal=True)
         for req_id in all_ids_to_abort:
             self.request_states.pop(req_id, None)
         logger.info("[Orchestrator] Aborted request(s) %s", request_ids)

vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py

@@ -175,8 +175,6 @@ class _RequestState:
     # Rolling tail of previously-decoded latents used as VAE receptive-field context.
     # Shape (n_pad_frames, feat_dim) on GPU. None before first decode.
     decode_pad: torch.Tensor | None = None
-    # Audio chunks already emitted (CPU float32), concatenated for cumulative output.
-    audio_chunks: list[torch.Tensor] = dataclasses.field(default_factory=list)
     decode_step_count: int = 0
     request_start_time: float = 0.0
     prefill_completed: bool = False
@@ -184,7 +182,6 @@ class _RequestState:
     prompt_cache: dict | None = None
     prefill_masks: tuple | None = None
     is_stopping: bool = False
-    last_decoded_audio: torch.Tensor | None = None
 
 
 @dataclasses.dataclass
@@ -1034,8 +1031,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         all_latents = vae_input  # [pad + new]
         state.decode_pad = all_latents[-self._n_decode_pad_frames :].detach()
 
-        state.audio_chunks.append(new_audio)
-        state.last_decoded_audio = new_audio
         self._perf.stop("vae_decode")
         return new_audio
 
@@ -1146,7 +1141,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             state = self._get_or_create_state(req_id)
             state.prefill_text = ""
             state.decode_pad = None
-            state.audio_chunks = []
             state.prefill_completed = False
             state.decode_step_count = 0
             state.precomputed_stop_logits = None
@@ -1155,7 +1149,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             state.prev_feat_embed = None
             state.curr_prefix_feat_cond = None
             state.is_stopping = False
-            state.last_decoded_audio = None
 
             # Voice clone / continuation
             ref_audio = info_dict.get("reference_audio") or info_dict.get("ref_audio")