From 4563999ec0e4376e81df21168542b1ccd1ae2fad Mon Sep 17 00:00:00 2001 From: akshatvishu Date: Thu, 30 Jul 2026 20:39:48 +0530 Subject: [PATCH] [BugFix][Diffusion] Add CPU LAPACK fallback for FlowUniPC (#5329) Signed-off-by: akshatvishu --- .../models/cosmos3/test_cosmos3_pipeline.py | 5 -- tests/diffusion/utils/test_flow_matching.py | 69 +++++++++++++++++++ .../scheduling_flow_unipc_multistep.py | 5 +- vllm_omni/diffusion/utils/flow_matching.py | 12 ++++ 4 files changed, 84 insertions(+), 7 deletions(-) create mode 100644 tests/diffusion/utils/test_flow_matching.py diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 994acd5a2..3093e6ceb 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -16,7 +16,6 @@ from PIL import Image from torch import nn from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -858,10 +857,6 @@ def test_flow_unipc_reuses_scheduler_and_forwards_each_request_shift(make_cosmos assert all(call["sigmas"] is None for call in scheduler.set_timesteps_calls) -@pytest.mark.skipif( - current_omni_platform.is_rocm(), - reason="Pytorch shipped with vLLM does not built with LAPACK support", -) def test_flow_unipc_reproducible_with_same_seed(make_cosmos3_pipeline) -> None: from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler, diff --git a/tests/diffusion/utils/test_flow_matching.py b/tests/diffusion/utils/test_flow_matching.py new file mode 100644 index 000000000..7f08bb624 --- /dev/null +++ b/tests/diffusion/utils/test_flow_matching.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.utils import flow_matching + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_safe_linalg_solve_uses_torch_native_path(monkeypatch: pytest.MonkeyPatch): + matrix = torch.tensor([[3.0, 1.0], [1.0, 2.0]], dtype=torch.float32) + rhs = torch.tensor([9.0, 8.0], dtype=torch.float32) + expected = torch.tensor([2.0, 3.0], dtype=torch.float32) + + def fake_torch_solve(solve_matrix: torch.Tensor, solve_rhs: torch.Tensor) -> torch.Tensor: + assert torch.equal(solve_matrix, matrix) + assert torch.equal(solve_rhs, rhs) + return expected + + def fail_numpy(*_args, **_kwargs): + raise AssertionError("NumPy fallback should not run") + + monkeypatch.setattr(flow_matching.torch.linalg, "solve", fake_torch_solve) + monkeypatch.setattr(flow_matching.np.linalg, "solve", fail_numpy) + + result = flow_matching.safe_linalg_solve(matrix, rhs) + + assert result is expected + + +def test_safe_linalg_solve_falls_back_on_missing_cpu_lapack(monkeypatch: pytest.MonkeyPatch): + matrix = torch.tensor([[3.0, 1.0], [1.0, 2.0]], dtype=torch.float32) + rhs = torch.tensor([9.0, 8.0], dtype=torch.float32) + expected = np.array([2.0, 3.0], dtype=np.float32) + + def fail_with_missing_lapack(*_args, **_kwargs): + raise RuntimeError( + "Calling torch.linalg.lu_factor on a CPU tensor requires compiling PyTorch with LAPACK. " + "Please use PyTorch built with LAPACK support." + ) + + monkeypatch.setattr(flow_matching.torch.linalg, "solve", fail_with_missing_lapack) + monkeypatch.setattr(flow_matching.np.linalg, "solve", lambda *_args, **_kwargs: expected) + + result = flow_matching.safe_linalg_solve(matrix, rhs) + + assert torch.equal(result, torch.from_numpy(expected)) + assert result.dtype == matrix.dtype + assert result.device == matrix.device + + +def test_safe_linalg_solve_reraises_unrelated_runtime_error(monkeypatch: pytest.MonkeyPatch): + matrix = torch.tensor([[3.0, 1.0], [1.0, 2.0]], dtype=torch.float32) + rhs = torch.tensor([9.0, 8.0], dtype=torch.float32) + + def fail_with_unrelated_error(*_args, **_kwargs): + raise RuntimeError("solver exploded for another reason") + + def fail_numpy(*_args, **_kwargs): + raise AssertionError("NumPy fallback should not run") + + monkeypatch.setattr(flow_matching.torch.linalg, "solve", fail_with_unrelated_error) + monkeypatch.setattr(flow_matching.np.linalg, "solve", fail_numpy) + + with pytest.raises(RuntimeError, match="solver exploded for another reason"): + flow_matching.safe_linalg_solve(matrix, rhs) diff --git a/vllm_omni/diffusion/models/schedulers/scheduling_flow_unipc_multistep.py b/vllm_omni/diffusion/models/schedulers/scheduling_flow_unipc_multistep.py index 3efe564bc..895e12827 100644 --- a/vllm_omni/diffusion/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/vllm_omni/diffusion/models/schedulers/scheduling_flow_unipc_multistep.py @@ -23,6 +23,7 @@ from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, Sch from diffusers.utils import deprecate from vllm_omni.diffusion.models.schedulers.base import BaseScheduler +from vllm_omni.diffusion.utils.flow_matching import safe_linalg_solve class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): @@ -431,7 +432,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) else: assert isinstance(R, torch.Tensor) - rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) + rhos_p = safe_linalg_solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) else: D1s = None @@ -565,7 +566,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if order == 1: rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) else: - rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) + rhos_c = safe_linalg_solve(R, b).to(device).to(x.dtype) if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 diff --git a/vllm_omni/diffusion/utils/flow_matching.py b/vllm_omni/diffusion/utils/flow_matching.py index e81d405ec..ec643644b 100644 --- a/vllm_omni/diffusion/utils/flow_matching.py +++ b/vllm_omni/diffusion/utils/flow_matching.py @@ -3,10 +3,22 @@ import math +import numpy as np import torch from torch import nn +def safe_linalg_solve(matrix: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor: + try: + return torch.linalg.solve(matrix, rhs) + except RuntimeError as exc: + if matrix.device.type != "cpu" or "requires compiling PyTorch with LAPACK" not in str(exc): + raise + # Some ROCm wheels still fail here at runtime even when the LAPACK probe reports support. + solved = np.linalg.solve(matrix.detach().cpu().numpy(), rhs.detach().cpu().numpy()) + return torch.from_numpy(solved).to(device=matrix.device, dtype=matrix.dtype) + + def swish(x: torch.Tensor) -> torch.Tensor: return x * torch.sigmoid(x)