From 21a3035c101976a17fde1f8485659945acc13f9b Mon Sep 17 00:00:00 2001 From: Dan250124 <416947747@qq.com> Date: Sun, 3 May 2026 22:29:47 +0800 Subject: [PATCH] Fixed memory leak and Remove dead code (#3312) Signed-off-by: Dan250124 <416947747@qq.com> Co-authored-by: Claude Opus 4.6 --- tests/engine/test_orchestrator.py | 7 +++++++ vllm_omni/engine/orchestrator.py | 12 ++++++++++++ .../model_executor/models/voxcpm2/voxcpm2_talker.py | 7 ------- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index c07762ad5..c33e23b05 100644 --- a/tests/engine/test_orchestrator.py +++ b/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) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e1bcf39a5..b53bdae84 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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) diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 15d125b53..40cf68181 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/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")