[Perf][Model] Add optimized LTX-2.5 DiffVAE operators (#7308)

Signed-off-by: mglyn <1203789601@qq.com>

tests/diffusion/models/ltx2/test_ltx2_ops_platform.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+
+import pytest
+import torch
+from vllm.platforms.interface import DeviceCapability
+
+from vllm_omni.diffusion.models.ltx2 import ops as ltx2_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass(frozen=True)
+class _TensorProbe:
+    is_cuda: bool
+    device: torch.device
+
+
+@pytest.fixture(autouse=True)
+def clear_platform_cache():
+    ltx2_platform._cuda_compute_capability.cache_clear()
+    try:
+        yield
+    finally:
+        ltx2_platform._cuda_compute_capability.cache_clear()
+
+
+@pytest.mark.parametrize(
+    ("capability", "triton", "tilelang", "fusion_expected", "fna_expected"),
+    [
+        (80, True, True, False, False),
+        (90, True, True, True, True),
+        (90, True, False, True, False),
+        (90, False, True, False, False),
+        (100, True, True, False, False),
+        (103, True, True, False, False),
+        (110, True, True, False, False),
+        (120, True, True, False, False),
+        (121, True, True, False, False),
+    ],
+)
+def test_production_resolver_and_tensor_eligibility_agree(
+    monkeypatch, capability, triton, tilelang, fusion_expected, fna_expected
+):
+    monkeypatch.setattr(ltx2_platform, "HAS_TRITON", triton)
+    monkeypatch.setattr(ltx2_platform, "has_tilelang", lambda: tilelang)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True)
+    monkeypatch.setattr(
+        ltx2_platform.current_omni_platform,
+        "get_device_capability",
+        lambda device_id=0: DeviceCapability(major=capability // 10, minor=capability % 10),
+    )
+    # Resolution must be lazy even when FNA is selected.
+    monkeypatch.setitem(sys.modules, "vllm_omni.diffusion.models.ltx2.ops.fna", None)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
+    device = torch.device("cuda:0")
+    operators = ltx2_platform.resolve_ltx2_vae_operators(device)
+    assert (operators is not None) is fusion_expected
+    assert (operators is not None and operators.fna is not None) is fna_expected
+    if operators is not None:
+        from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm, residual_adaln, swiglu
+
+        assert operators.qk_norm_rope is qk_rms_norm.try_qk_rms_norm_scale_rope_3d_exact
+        assert operators.residual_norm is residual_adaln.try_residual_rms_norm_modulate_exact
+        assert operators.residual_add is residual_adaln.try_residual_add3_exact
+        assert operators.swiglu is swiglu.try_swiglu_tiled_exact
+
+    tensor = _TensorProbe(is_cuda=True, device=device)
+    with torch.inference_mode():
+        assert ltx2_platform.is_ltx2_fusion_eligible(tensor) is fusion_expected
+        assert ltx2_platform.is_ltx2_fna_eligible(tensor) is fna_expected
+
+
+@pytest.mark.parametrize(
+    ("device", "is_cuda", "available"),
+    [("cpu", True, True), ("cuda:0", False, True), ("cuda:0", True, False)],
+)
+def test_resolver_rejects_unsupported_platform_before_capability_lookup(monkeypatch, device, is_cuda, available):
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: is_cuda)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: available)
+
+    def unexpected_lookup(_index):
+        raise AssertionError("Unsupported platforms must not query CUDA capability")
+
+    monkeypatch.setattr(ltx2_platform, "_cuda_compute_capability", unexpected_lookup)
+    assert ltx2_platform.resolve_ltx2_vae_operators(torch.device(device)) is None
+
+
+@pytest.mark.parametrize(("compiling", "grad_enabled"), [(True, False), (False, True)])
+def test_runtime_guards_bypass_operator_resolution(monkeypatch, compiling, grad_enabled):
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: compiling)
+    monkeypatch.setattr(torch, "is_grad_enabled", lambda: grad_enabled)
+
+    def unexpected_resolve(_device):
+        raise AssertionError("Ineligible runtime must not select eager operators")
+
+    monkeypatch.setattr(ltx2_platform, "resolve_ltx2_vae_operators", unexpected_resolve)
+    tensor = _TensorProbe(is_cuda=True, device=torch.device("cuda:0"))
+    assert not ltx2_platform.is_ltx2_fusion_eligible(tensor)
+    assert not ltx2_platform.is_ltx2_fna_eligible(tensor)
+
+
+def test_device_capability_is_cached(monkeypatch):
+    calls = 0
+    monkeypatch.setattr(ltx2_platform, "HAS_TRITON", True)
+
+    def get_device_capability(device_id=0):
+        nonlocal calls
+        calls += 1
+        return DeviceCapability(major=9, minor=0)
+
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "get_device_capability", get_device_capability)
+    assert ltx2_platform._cuda_compute_capability(0) == 90
+    assert ltx2_platform._cuda_compute_capability(0) == 90
+    assert calls == 1