Skip to content

[ROCm] [CI] Fix diffusion test #4574

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hsliuustc0106 merged 7 commits into from
Jun 20, 2026
Merged

[ROCm] [CI] Fix diffusion test #4574

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ce963a1
fix voxtral tts test
tjtanaa Jun 12, 2026
50b3efb
increase timeout
tjtanaa Jun 12, 2026
c78924a
revert time out change first
tjtanaa Jun 12, 2026
39f06dc
Merge remote-tracking branch 'origin/main' into fixvoxtral
tjtanaa Jun 15, 2026
afd35a2
fix diffusion test cases
tjtanaa Jun 15, 2026
b08fdf9
fix diffusion worker
tjtanaa Jun 19, 2026
2d960da
fix precommit
tjtanaa Jun 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 17 additions & 8 deletions tests/diffusion/test_diffusion_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,15 @@
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]


def patch_cumem_allocator(mocker: MockerFixture):

@hsliuustc0106 hsliuustc0106 Jun 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good: Helper function centralizes the mock patching pattern.

mock_allocator_class = mocker.Mock()
mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class",
return_value=mock_allocator_class,
)
return mock_allocator_class


@pytest.fixture
def mock_od_config(mocker: MockerFixture):
"""Create a mock OmniDiffusionConfig."""
Expand Down Expand Up @@ -118,15 +127,15 @@ def setup_allocator(self, mocker: MockerFixture):
"""
Unified interception of Allocators, and provision of default security values.
"""
self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
self.mock_allocator_class = patch_cumem_allocator(mocker)
self.mock_allocator = mocker.Mock()
self.mock_allocator_class.get_instance.return_value = self.mock_allocator
self.mock_allocator.get_current_usage.return_value = 4 * 1024**3
self.mock_allocator.sleep = mocker.Mock()

def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
"""Test sleep mode level 1 (offload weights only)."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mock_allocator_class = patch_cumem_allocator(mocker)
mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
mock_platform.get_device_total_memory.return_value = 80 * 1024**3
Expand Down Expand Up @@ -158,7 +167,7 @@ def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):

def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
"""Test sleep mode level 2 (offload all, save buffers)."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mock_allocator_class = patch_cumem_allocator(mocker)
mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3]
mock_platform.get_device_total_memory.return_value = 80 * 1024**3
Expand Down Expand Up @@ -201,7 +210,7 @@ def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):

def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
"""Test that sleep validates memory was actually freed."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mock_allocator_class = patch_cumem_allocator(mocker)
mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_platform.get_free_memory.return_value = 10 * 1024**3
mock_platform.get_device_total_memory.return_value = 80 * 1024**3
Expand All @@ -226,7 +235,7 @@ def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_wor
def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker):
"""Test sleep uses device-scoped fallback when NVML is unavailable."""

mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mock_allocator_class = patch_cumem_allocator(mocker)
mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
mock_get_process_memory.side_effect = [None, None]
Expand All @@ -252,7 +261,7 @@ class TestDiffusionWorkerWakeUp:

def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):
"""Test wake_up without saved buffers (level 1 sleep)."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mock_allocator_class = patch_cumem_allocator(mocker)

# Setup allocator mock
mock_allocator = mocker.Mock()
Expand All @@ -272,7 +281,7 @@ def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):

def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
"""Test wake_up with saved buffers (level 2 sleep)."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mock_allocator_class = patch_cumem_allocator(mocker)

# Setup allocator mock
mock_allocator = mocker.Mock()
Expand Down Expand Up @@ -317,7 +326,7 @@ def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):

def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker):
"""Test wake_up only restores buffers that were saved."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mock_allocator_class = patch_cumem_allocator(mocker)

# Setup allocator mock
mock_allocator = mocker.Mock()
Expand Down
15 changes: 9 additions & 6 deletions vllm_omni/diffusion/worker/diffusion_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,12 @@ def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDi
return vllm_config


def _get_cumem_allocator_class() -> type:

@hsliuustc0106 hsliuustc0106 Jun 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good: Lazy import wrapper prevents hipSetDevice symbol error on ROCm.

from vllm.device_allocator.cumem import CuMemAllocator

return CuMemAllocator


def _resolve_ir_op_priority(od_config: OmniDiffusionConfig, vllm_config: VllmConfig) -> Any:
ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config)
ir_op_priority_func = get_diffusion_ir_op_priority_func(od_config)
Expand Down Expand Up @@ -455,8 +461,7 @@ def sleep(self, level: int = 1) -> bool:
Args:
level: Sleep level. Level 1 offloads weights, level 2 also saves buffers.
"""
from vllm.device_allocator.cumem import CuMemAllocator

CuMemAllocator = _get_cumem_allocator_class()
allocator = CuMemAllocator.get_instance()

usage_before = allocator.get_current_usage()
Expand Down Expand Up @@ -508,8 +513,7 @@ def wake_up(self, tags: list[str] | None = None) -> bool:
tags: List of memory pool tags to re-activate (e.g., ["weights"]
to match Level 1 sleep). If None, all pools are re-activated.
"""
from vllm.device_allocator.cumem import CuMemAllocator

CuMemAllocator = _get_cumem_allocator_class()
allocator = CuMemAllocator.get_instance()
allocator.wake_up(tags)
current_omni_platform.synchronize()
Expand Down Expand Up @@ -631,8 +635,7 @@ def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
if is_sleep_enabled:
current_omni_platform.synchronize()
gc.collect()
from vllm.device_allocator.cumem import CuMemAllocator

CuMemAllocator = _get_cumem_allocator_class()
allocator = CuMemAllocator.get_instance()
if tag == "weights":
assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
Expand Down
Loading