From d58beb13d95549937f9bb3b63f729793106f8ee9 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 8 Jul 2026 07:01:57 +0800 Subject: [PATCH] rebase: align vllm-omni with vLLM 65a7b4628409 Signed-off-by: tzhouam --- .../voxcpm2/test_talker_state_eviction.py | 24 +++++-- vllm_omni/patch.py | 67 +++++++++++++++++++ 2 files changed, 85 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py index 28da0b1bc..f5a94376a 100644 --- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py +++ b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py @@ -52,6 +52,10 @@ def _make_bare_talker(): talker._cfm_buffers = None talker._last_audio_output_req_ids = [] talker._batched_fsq_fusion_max_batch = 32 + # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to + # default (False) to avoid AttributeError when test exercises forward + # path that checks this flag. + talker._enable_unified_decode_graph = False return talker @@ -329,23 +333,27 @@ class TestDecodeBatchContract: for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True): torch.testing.assert_close(actual_res, expected_res) - torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"]) + torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden) assert [state.decode_step_count for state in batch_states] == [1, 1] def test_batch_decode_matches_per_request_state_updates(self) -> None: + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + _DecodeResidualMeta, + ) + seq_talker = self._make_decode_talker() batch_talker = self._make_decode_talker() seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] metas = [ - {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])}, - {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])}, + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])), + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])), ] batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]]) for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True): - seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu")) + seq_talker._finish_decode(state, meta, res_out.unsqueeze(0)) batch_talker._finish_decode_batch( [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)], @@ -360,12 +368,16 @@ class TestDecodeBatchContract: torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu) def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None: + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + _DecodeResidualMeta, + ) + talker = self._make_decode_talker() talker._enable_batched_cfm = False states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] metas = [ - {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])}, - {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])}, + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])), + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])), ] batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]]) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 50eb4d991..4f2441294 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -417,3 +417,70 @@ def _patch_fp8_use_quack_fused_bias(): _patch_fp8_use_quack_fused_bias() + + +# ============================================================================= +# Patch CuMemAllocator._python_free_callback to fix CUDA double-free on shutdown +# ============================================================================= +# WHY: CuMemAllocator._python_free_callback guards the asleep-entry double-free +# skip with ``if data.is_asleep and current_platform.is_rocm():`` — only ROCm +# gets the safe empty-handle return. On CUDA, the callback falls through and +# returns the original handle, causing cuMemRelease on already-freed memory +# (CUDA_ERROR_INVALID_VALUE) during EngineCore subprocess atexit cleanup. +# +# This happens because ``sleep()`` calls ``unmap_and_release()`` on ALL +# platforms, then sets ``is_asleep = True``. When the atexit handler +# (``_shutdown_singleton`` -> ``release_pools()`` -> GC) triggers the free +# callback, the asleep entries must return an empty chunk list so the C +# extension skips ``cuMemRelease`` — exactly the same logic already used by +# the ROCm guard. +# +# The fix removes the ``current_platform.is_rocm()`` condition so the guard +# applies to CUDA (and any future platform that implements cumem). +# +# FRAGILITY: Relies on ``_python_free_callback`` being a regular method +# (not a slot or C extension). If CuMemAllocator is rewritten in C/Cython, +# this monkey-patch will silently become a no-op. +def _patch_cumem_free_callback_cuda() -> None: + try: + from vllm.device_allocator.cumem import CuMemAllocator + except ImportError: + _PATCH_LOGGER.debug("[cumem-cuda] CuMemAllocator not available; skipping patch") + return + + _original_free_callback = CuMemAllocator._python_free_callback + + if getattr(_original_free_callback, "_omni_cumem_cuda_patched", False): + return + + # The upstream bug: `_python_free_callback` only skips the double-free + # for ROCm (line ~206). We wrap the method to extend the guard to all + # platforms. + def _patched_free_callback(self, ptr: int) -> tuple: + data = self.pointer_to_data.pop(ptr) + if data.cpu_backup_tensor is not None: + data.cpu_backup_tensor = None + if data.is_asleep: + # sleep() already called unmap_and_release() on this allocation. + # Return an empty chunk list so the C extension skips + # cuMemRelease, avoiding a double-free. Same logic as the + # existing ROCm guard, but applied to all platforms. + device, size, d_mem, _ = data.handle + result = (device, size, d_mem, []) + _PATCH_LOGGER.debug( + "[cumem-cuda] Free callback: asleep entry %s -> empty handle", + ptr, + ) + return result + # Drain pending kernels before the C extension's cuMemUnmap. + torch.accelerator.synchronize(data.handle[0]) + return data.handle + + _patched_free_callback._omni_cumem_cuda_patched = True + CuMemAllocator._python_free_callback = _patched_free_callback + _PATCH_LOGGER.info( + "[cumem-cuda] CuMemAllocator._python_free_callback patched: asleep guard extended to all platforms." + ) + + +_patch_cumem_free_callback_cuda()