vllm-omni/tests/diffusion/attention/test_cudnn_attn.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from contextlib import contextmanager, nullcontext
from typing import Any

import pytest
import torch
from torch.nn.attention import SDPBackend

import vllm_omni.diffusion.attention.backends.cudnn_attn as cudnn_backend
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend, CuDNNAttentionImpl

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]


@pytest.mark.cpu
def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
    """Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V."""
    selected_backends = []

    @contextmanager
    def fake_sdpa_kernel(backends):
        selected_backends.append(tuple(backends))
        yield

    def fake_sdpa(query, key, value, **kwargs):
        return query

    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel)
    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)

    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
    query = torch.randn(1, 2, 2, 8)
    singleton_kv = torch.randn(1, 1, 2, 8)

    output = impl.forward_cuda(query, singleton_kv, singleton_kv)

    assert output.shape == query.shape
    assert selected_backends == [(SDPBackend.MATH,)]


@pytest.mark.cpu
def test_explicit_cudnn_rejects_kv_seq_len_one():
    impl = CuDNNAttentionImpl(
        num_heads=2,
        head_size=8,
        softmax_scale=0.5,
        backend_explicit=True,
    )
    query = torch.randn(1, 2, 2, 8)
    singleton_kv = torch.randn(1, 1, 2, 8)

    with pytest.raises(ValueError, match="explicitly selected.*sequence length 1"):
        impl.forward_cuda(query, singleton_kv, singleton_kv)


@pytest.mark.cpu
def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
    selected_backends = []

    @contextmanager
    def fake_sdpa_kernel(backends):
        selected_backends.append(tuple(backends))
        yield

    def reject_shape(*args, **kwargs):
        raise RuntimeError("No available kernel. Aborting execution.")

    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel)
    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", reject_shape)

    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
    tensors = torch.randn(1, 2, 2, 8)

    with pytest.raises(RuntimeError, match="No available kernel"):
        impl.forward_cuda(tensors, tensors, tensors)

    assert selected_backends == [(SDPBackend.CUDNN_ATTENTION,)]


@pytest.mark.cpu
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
    observed: dict[str, Any] = {}

    def fake_sdpa(query, key, value, **kwargs):
        observed.update(query=query, key=key, value=value, kwargs=kwargs)
        return query

    monkeypatch.setattr(
        "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
        lambda _backends: nullcontext(),
    )
    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
    impl = CuDNNAttentionImpl(
        num_heads=2,
        head_size=4,
        softmax_scale=0.5,
    )
    query = torch.randn(1, 8, 2, 4)
    key = torch.randn_like(query)
    value = torch.randn_like(query)

    output = impl.forward_cuda(
        query,
        key,
        value,
        AttentionMetadata(extra={"valid_kv_length": 5}),
    )

    assert output.shape == query.shape
    assert observed["query"].shape == (1, 2, 8, 4)
    assert observed["key"].shape == (1, 2, 5, 4)
    assert observed["value"].shape == (1, 2, 5, 4)
    assert observed["kwargs"]["attn_mask"] is None


@pytest.mark.cpu
def test_cudnn_rejects_invalid_valid_kv_length():
    impl = CuDNNAttentionImpl(
        num_heads=2,
        head_size=4,
        softmax_scale=0.5,
    )
    query = torch.randn(1, 8, 2, 4)

    with pytest.raises(ValueError, match="valid_kv_length"):
        impl.forward_cuda(
            query,
            query,
            query,
            AttentionMetadata(extra={"valid_kv_length": 9}),
        )


@pytest.mark.parametrize("head_size", [8, 64, 128, 256])
@pytest.mark.cpu
def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size):
    assert CuDNNAttentionBackend.supports_head_size(head_size)
    assert head_size in CuDNNAttentionBackend.get_supported_head_sizes()


@pytest.mark.parametrize("head_size", [0, 7, 12, 320])
@pytest.mark.cpu
def test_cudnn_backend_rejects_incompatible_head_sizes(head_size):
    assert not CuDNNAttentionBackend.supports_head_size(head_size)
    assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes()


@pytest.mark.parametrize("batch_size", [1, 2])
@pytest.mark.parametrize("kv_heads", [1, 2])
@pytest.mark.parametrize("query_offset", [0, 4])
@pytest.mark.cpu
def test_piecewise_matches_dense_block_causal_mask(monkeypatch, batch_size, kv_heads, query_offset):
    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
    torch.manual_seed(42)
    query = torch.randn(batch_size, 12, 2, 8)
    key = torch.randn(batch_size, 12, kv_heads, 8)
    value = torch.randn_like(key)
    spans = [(2, 5), (7, 10)]
    positions = torch.arange(12)
    mask = positions[:, None] >= positions[None, :]
    for start, end in spans:
        mask[start:end, start:end] = True
    expected = torch.nn.functional.scaled_dot_product_attention(
        query[:, query_offset:].transpose(1, 2),
        key.transpose(1, 2),
        value.transpose(1, 2),
        attn_mask=mask[None, None, query_offset:],
        scale=0.5,
        enable_gqa=kv_heads == 1,
    ).transpose(1, 2)
    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
    actual = impl.forward_cuda(
        query[:, query_offset:], key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size)
    )
    torch.testing.assert_close(actual, expected)


@pytest.mark.cpu
def test_piecewise_preserves_discontinuous_query_ranges(monkeypatch):
    from vllm_omni.diffusion.attention.backends.abstract import QueryRange

    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
    query, key, value = [torch.randn(2, 12, 2, 8) for _ in range(3)]
    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
    metadata = AttentionMetadata(full_attn_spans=[[(2, 5), (7, 12)]] * 2)
    expected = impl.forward_cuda(query, key, value, metadata)
    metadata.query_ranges = (QueryRange(0, 3, 0), QueryRange(3, 6, 8))
    actual = impl.forward_cuda(torch.cat([query[:, :3], query[:, 8:11]], dim=1), key, value, metadata)
    torch.testing.assert_close(actual, torch.cat([expected[:, :3], expected[:, 8:11]], dim=1))


@pytest.mark.cpu
def test_piecewise_rejects_padding_mask(monkeypatch):
    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
    query = torch.randn(1, 8, 2, 8)
    metadata = AttentionMetadata(full_attn_spans=[[(2, 8)]], attn_mask=torch.ones(1, 8, dtype=torch.bool))
    with pytest.raises(ValueError, match="padding mask"):
        impl.forward_cuda(query, query, query, metadata)


@pytest.mark.cpu
def test_dense_mask_takes_precedence_over_spans(monkeypatch):
    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", lambda _: nullcontext())
    query, key, value = [torch.randn(2, 8, 2, 8) for _ in range(3)]
    mask = torch.ones(2, 1, 8, 8, dtype=torch.bool).tril()
    mask[1, :, :, 3] = False
    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
    expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask))
    actual = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask, full_attn_spans=[[(2, 8)]] * 2))
    torch.testing.assert_close(actual, expected, rtol=0, atol=0)


@pytest.mark.cuda
@pytest.mark.gpu
@pytest.mark.parametrize("batch_size", [1, 2])
def test_cudnn_piecewise_matches_dense_on_device(batch_size):
    generator = torch.Generator(device="cuda").manual_seed(42)
    query, key, value = [
        torch.randn(batch_size, 96, 2, 128, generator=generator, device="cuda", dtype=torch.bfloat16) for _ in range(3)
    ]
    spans = [(8, 40), (48, 96)]
    mask = torch.ones(96, 96, device="cuda", dtype=torch.bool).tril()
    for start, end in spans:
        mask[start:end, start:end] = True
    impl = CuDNNAttentionImpl(num_heads=2, head_size=128, softmax_scale=128**-0.5)
    actual = impl.forward_cuda(query, key, value, AttentionMetadata(full_attn_spans=[spans] * batch_size))
    expected = impl.forward_cuda(query, key, value, AttentionMetadata(attn_mask=mask[None, None]))
    torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.004)