wehub-resource-sync
84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.attention.layer import Attention
|
|
from vllm_omni.diffusion.forward_context import set_forward_context
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
def test_attention_uses_compile_boundary_for_hsdp(monkeypatch):
|
|
attention = object.__new__(Attention)
|
|
calls = []
|
|
|
|
def _boundary(query, key, value, attn_metadata=None):
|
|
calls.append("boundary")
|
|
return query
|
|
|
|
def _impl(query, key, value, attn_metadata=None):
|
|
calls.append("impl")
|
|
return query
|
|
|
|
attention._forward_hsdp_compile_boundary = _boundary
|
|
attention._forward_impl = _impl
|
|
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
|
|
|
|
config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=True))
|
|
query = torch.empty(1)
|
|
with set_forward_context(omni_diffusion_config=config):
|
|
assert Attention.forward(attention, query, query, query) is query
|
|
|
|
assert calls == ["boundary"]
|
|
|
|
|
|
def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch):
|
|
attention = object.__new__(Attention)
|
|
calls = []
|
|
|
|
def _boundary(query, key, value, attn_metadata=None):
|
|
calls.append("boundary")
|
|
return query
|
|
|
|
def _impl(query, key, value, attn_metadata=None):
|
|
calls.append("impl")
|
|
return query
|
|
|
|
attention._forward_hsdp_compile_boundary = _boundary
|
|
attention._forward_impl = _impl
|
|
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
|
|
|
|
config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=False))
|
|
query = torch.empty(1)
|
|
with set_forward_context(omni_diffusion_config=config):
|
|
assert Attention.forward(attention, query, query, query) is query
|
|
|
|
assert calls == ["impl"]
|
|
|
|
|
|
def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch):
|
|
attention = object.__new__(Attention)
|
|
calls = []
|
|
|
|
def _boundary(query, key, value, attn_metadata=None):
|
|
calls.append("boundary")
|
|
return query
|
|
|
|
def _impl(query, key, value, attn_metadata=None):
|
|
calls.append("impl")
|
|
return query
|
|
|
|
attention._forward_hsdp_compile_boundary = _boundary
|
|
attention._forward_impl = _impl
|
|
monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
|
|
|
|
query = torch.empty(1)
|
|
with set_forward_context(vllm_config=None):
|
|
assert Attention.forward(attention, query, query, query) is query
|
|
|
|
assert calls == ["impl"]
|