36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
|
|
||
|
|
import importlib
|
||
|
|
import sys
|
||
|
|
import types
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||
|
|
|
||
|
|
|
||
|
|
def test_sage_attention_rejects_mask_instead_of_ignoring_it(monkeypatch):
|
||
|
|
fake_package = types.ModuleType("sageattention")
|
||
|
|
fake_package.sageattn = lambda *args, **kwargs: pytest.fail("unexpected Sage kernel call")
|
||
|
|
monkeypatch.setitem(sys.modules, "sageattention", fake_package)
|
||
|
|
module_name = "vllm_omni.diffusion.attention.backends.sage_attn"
|
||
|
|
sys.modules.pop(module_name, None)
|
||
|
|
try:
|
||
|
|
backend_module = importlib.import_module(module_name)
|
||
|
|
impl = backend_module.SageAttentionImpl(
|
||
|
|
num_heads=4,
|
||
|
|
head_size=64,
|
||
|
|
softmax_scale=1.0 / 8.0,
|
||
|
|
causal=False,
|
||
|
|
)
|
||
|
|
query = torch.randn(1, 2, 4, 64)
|
||
|
|
metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool))
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match="does not support attn_mask"):
|
||
|
|
impl.forward_cuda(query, query, query, metadata)
|
||
|
|
finally:
|
||
|
|
sys.modules.pop(module_name, None)
|