rebase: align vllm-omni with vLLM 65a7b4628409

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

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]])
 

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()