84bbd57e4d
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com> Signed-off-by: dph97 <nenbaying@163.com> Signed-off-by: zhuhh97 <zhh@unknown> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: dengyunyang <584797741@qq.com>
58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend
|
|
from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_sdpa_expands_kv_when_native_gqa_kernel_is_unavailable(monkeypatch):
|
|
calls = []
|
|
|
|
def fake_sdpa(query, key, value, **kwargs):
|
|
calls.append((query.shape, key.shape, value.shape, kwargs))
|
|
return query
|
|
|
|
monkeypatch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", lambda *args: False)
|
|
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
|
|
|
|
impl = SDPAImpl(num_heads=4, num_kv_heads=2, head_size=8, softmax_scale=0.5)
|
|
output = impl.forward_cuda(
|
|
torch.randn(1, 3, 4, 8),
|
|
torch.randn(1, 3, 2, 8),
|
|
torch.randn(1, 3, 2, 8),
|
|
)
|
|
|
|
query_shape, key_shape, value_shape, kwargs = calls[0]
|
|
assert query_shape == key_shape == value_shape == (1, 4, 3, 8)
|
|
assert kwargs["enable_gqa"] is False
|
|
assert output.shape == (1, 3, 4, 8)
|
|
|
|
|
|
def test_sdpa_keeps_compressed_kv_when_native_gqa_kernel_is_available(monkeypatch):
|
|
calls = []
|
|
|
|
def fake_sdpa(query, key, value, **kwargs):
|
|
calls.append((query.shape, key.shape, value.shape, kwargs))
|
|
return query
|
|
|
|
monkeypatch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", lambda *args: True)
|
|
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
|
|
|
|
impl = SDPAImpl(num_heads=4, num_kv_heads=2, head_size=8, softmax_scale=0.5)
|
|
output = impl.forward_cuda(
|
|
torch.randn(1, 3, 4, 8),
|
|
torch.randn(1, 3, 2, 8),
|
|
torch.randn(1, 3, 2, 8),
|
|
)
|
|
|
|
query_shape, key_shape, value_shape, kwargs = calls[0]
|
|
assert query_shape == (1, 4, 3, 8)
|
|
assert key_shape == value_shape == (1, 2, 3, 8)
|
|
assert kwargs["enable_gqa"] is True
|
|
assert output.shape == (1, 3, 4, 8)
|