228 lines
9.7 KiB
Python
228 lines
9.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
|
from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
|
|
def _clear_hub_module_cache(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from vllm_omni.diffusion.attention.backends import flash_attn_hub
|
|
|
|
monkeypatch.setattr(flash_attn_hub, "_hub_modules", {})
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
|
|
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
|
|
|
# Temporarily hide/remove kernels module if present
|
|
monkeypatch.setitem(sys.modules, "kernels", None)
|
|
|
|
# Use monkeypatch to mock CudaOmniPlatform capability and package checker to allow FLASH_ATTN
|
|
from vllm.platforms.interface import DeviceCapability
|
|
|
|
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
|
|
|
monkeypatch.setattr(
|
|
CudaOmniPlatform,
|
|
"get_device_capability",
|
|
classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
|
|
)
|
|
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
|
|
|
|
with pytest.raises(ImportError, match="explicitly selected"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
|
|
with pytest.raises(ImportError, match="explicitly selected"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
|
|
|
kernels_module = types.ModuleType("kernels")
|
|
monkeypatch.setitem(sys.modules, "kernels", kernels_module)
|
|
with pytest.raises(ValueError, match="require.*Hopper GPU"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
|
|
|
monkeypatch.setattr(
|
|
CudaOmniPlatform,
|
|
"get_device_capability",
|
|
classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
|
|
)
|
|
with pytest.raises(ValueError, match="current FA2 kernels require"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
|
|
with pytest.raises(ValueError, match="require.*Hopper GPU"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_explicit_kernels_hub_variant_miss_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
|
|
"""Explicit Hub backends must fail loudly when no build variant resolves (#6971)."""
|
|
from vllm.platforms.interface import DeviceCapability
|
|
|
|
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
|
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
|
|
|
monkeypatch.setattr(
|
|
CudaOmniPlatform,
|
|
"get_device_capability",
|
|
classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
|
|
)
|
|
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
|
|
_clear_hub_module_cache(monkeypatch)
|
|
|
|
kernels_module = types.ModuleType("kernels")
|
|
|
|
def _missing_variant(repo_id, version=None, **kwargs):
|
|
raise FileNotFoundError(f"Cannot find a build variant for this system in {repo_id}")
|
|
|
|
kernels_module.get_kernel = _missing_variant # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "kernels", kernels_module)
|
|
|
|
with pytest.raises(RuntimeError, match="no compatible build"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
|
with pytest.raises(RuntimeError, match="no compatible build"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_explicit_kernels_hub_preflight_uses_v2_when_v1_misses(monkeypatch: pytest.MonkeyPatch):
|
|
"""Preflight must share _load_hub_module's (1, 2) policy, not pin version=1."""
|
|
from vllm.platforms.interface import DeviceCapability
|
|
|
|
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
|
|
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
|
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
|
|
|
monkeypatch.setattr(
|
|
CudaOmniPlatform,
|
|
"get_device_capability",
|
|
classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
|
|
)
|
|
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
|
|
_clear_hub_module_cache(monkeypatch)
|
|
|
|
kernels_module = types.ModuleType("kernels")
|
|
versions_tried: list[int | None] = []
|
|
|
|
def _v1_miss_v2_hit(repo_id, version=None, **kwargs):
|
|
versions_tried.append(version)
|
|
if version == 1:
|
|
raise FileNotFoundError(f"Cannot find a build variant for this system in {repo_id}")
|
|
return types.SimpleNamespace(name=f"{repo_id}@v{version}")
|
|
|
|
kernels_module.get_kernel = _v1_miss_v2_hit # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "kernels", kernels_module)
|
|
|
|
fa3_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
|
assert fa3_path == DiffusionAttentionBackendEnum.FLASH_ATTN_3_HUB.get_path()
|
|
fa2_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
|
|
assert fa2_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
|
|
assert versions_tried == [1, 2, 1, 2]
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_explicit_flash_attention_unavailable_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
|
|
from vllm.platforms.interface import DeviceCapability
|
|
|
|
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
|
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
|
|
|
monkeypatch.setattr(
|
|
CudaOmniPlatform,
|
|
"get_device_capability",
|
|
classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
|
|
)
|
|
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
|
|
|
|
with pytest.raises(ValueError, match="explicitly selected but is unsupported"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
|
|
|
|
monkeypatch.setattr(
|
|
CudaOmniPlatform,
|
|
"get_device_capability",
|
|
classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
|
|
)
|
|
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
|
|
monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: False))
|
|
with pytest.raises(ValueError, match="requires CuTe FlashAttention-4"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_explicit_sage_attention_unsupported_arch_does_not_reach_kernel(monkeypatch: pytest.MonkeyPatch):
|
|
from vllm.platforms.interface import DeviceCapability
|
|
|
|
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
|
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
|
|
|
monkeypatch.setattr(
|
|
CudaOmniPlatform,
|
|
"get_device_capability",
|
|
classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
|
|
)
|
|
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
|
|
|
|
with pytest.raises(ValueError, match="does not provide a kernel for sm_103"):
|
|
CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN", head_size=64)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
|
def test_kernels_hub_execution():
|
|
"""Verify basic forward of flash_attn_hub and flash_attn_3_hub, comparing with SDPA reference."""
|
|
device = torch.device(current_omni_platform.device_type)
|
|
dtype = torch.bfloat16
|
|
|
|
num_heads = 8
|
|
head_dim = 64
|
|
seq_len = 32
|
|
batch_size = 1
|
|
|
|
torch.manual_seed(42)
|
|
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
|
|
k = q.clone()
|
|
v = q.clone()
|
|
|
|
# 1. Test PyTorch SDPA reference
|
|
sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
|
|
attn_metadata_sdpa = AttentionMetadata(attn_mask=None)
|
|
output_ref = sdpa_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
|
|
|
|
# 2. Test FlashAttentionHubBackend (FlashAttention 2)
|
|
from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttentionHubImpl
|
|
|
|
fa_hub_impl = FlashAttentionHubImpl(
|
|
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
|
|
)
|
|
output_fa_hub = fa_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
|
|
assert output_fa_hub.shape == q.shape
|
|
assert not torch.isnan(output_fa_hub).any()
|
|
max_diff = torch.max(torch.abs(output_ref - output_fa_hub)).item()
|
|
assert max_diff < 1e-2, f"FlashAttentionHub output differs too much from SDPA reference: {max_diff}"
|
|
|
|
# 3. Test FlashAttention3HubBackend (FlashAttention 3, Hopper+ only)
|
|
major, _minor = torch.cuda.get_device_capability()
|
|
if major < 9:
|
|
pytest.skip("FLASH_ATTN_3_HUB execution requires Hopper-class GPU (compute capability >= 9.0)")
|
|
|
|
from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttention3HubImpl
|
|
|
|
fa3_hub_impl = FlashAttention3HubImpl(
|
|
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
|
|
)
|
|
output_fa3_hub = fa3_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
|
|
assert output_fa3_hub.shape == q.shape
|
|
assert not torch.isnan(output_fa3_hub).any()
|
|
max_diff = torch.max(torch.abs(output_ref - output_fa3_hub)).item()
|
|
assert max_diff < 1e-2, f"FlashAttention3Hub output differs too much from SDPA reference: {max_diff}"
|