[BugFix][Diffusion] Add CPU LAPACK fallback for FlowUniPC (#5329)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

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,

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)

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

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)