From 26a37abfaf14eda9d429931c593ec5d1c19e8eb2 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 8 Sep 2026 12:04:17 +0800 Subject: [PATCH] test: align SANA SP CPU fixtures with vLLM 0.29 GEMM dispatch Signed-off-by: tzhouam --- .../sana_video/test_transformer_sana_video_sp.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py b/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py index 3f4c62075..6a0f0548e 100644 --- a/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py @@ -41,8 +41,8 @@ def tp1_group(): os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29501") - init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") - initialize_model_parallel() + init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://", backend="gloo") + initialize_model_parallel(backend="gloo") yield cleanup_dist_env_and_memory() @@ -54,7 +54,7 @@ def force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) @@ -482,8 +482,12 @@ class _LockstepP2PDist: def test_tiny_transformer_sp2_lockstep_matches_dense(tp1_group, force_default_gemm, mocker, task) -> None: """Two lockstep rank threads over uneven frame shards must reproduce the dense tiny-transformer output, and must actually have sharded the work.""" + from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + # These CPU tensor tests must not probe a CUDA device during construction. + mocker.patch("vllm_omni.diffusion.attention.selector._cached_get_backend_cls", return_value=SDPABackend) + torch.manual_seed(3) model = SanaVideoTransformer3DModel(**_TINY_SP_CONFIG).eval() for _, param in sorted(model.named_parameters()):