vllm-omni/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

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

from __future__ import annotations

from types import SimpleNamespace

import pytest
import torch
from torch import nn
from torch.nn import functional as F
from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
from vllm.v1.worker.gpu.block_table import BlockTables

from tests.helpers.kv_layout import build_kv_cache_tensor
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout
from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
    DiffusionPagedAttentionAdapter,
    DiffusionPagedAttentionLayerAdapter,
    DiffusionPagedAttentionRow,
    DiffusionPagedAttentionRowBinding,
)
from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig

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

_LAYER_NAME = "model.layers.0.attn"
_NUM_HEADS = 2
_HEAD_SIZE = 64
_BLOCK_SIZE = 16


class _SmokeDiffusionAttention(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.num_heads = _NUM_HEADS
        self.num_kv_heads = _NUM_HEADS
        self.head_size = _HEAD_SIZE
        self.softmax_scale = _HEAD_SIZE**-0.5


def _native_adapter(num_rows: int = 1):
    device = torch.device("cuda", torch.accelerator.current_device_index())
    vllm_config = VllmConfig(
        cache_config=CacheConfig(
            block_size=_BLOCK_SIZE,
            cache_dtype="float16",
            enable_prefix_caching=False,
        ),
        attention_config=AttentionConfig(use_non_causal=True),
    )
    model_config = _DiffusionVllmModelConfig(
        model=None,
        dtype=torch.float16,
        max_model_len=32,
        original_max_model_len=32,
        hf_config=SimpleNamespace(model_type="diffusion_paged_attention_smoke"),
    )
    vllm_config.model_config = model_config
    model_config.set_attention_geometry(
        num_heads=_NUM_HEADS,
        num_kv_heads=_NUM_HEADS,
        head_size=_HEAD_SIZE,
    )
    vllm_config.compilation_config.static_forward_context.clear()

    diffusion_layer = _SmokeDiffusionAttention().to(device)
    spec = FullAttentionSpec(
        block_size=_BLOCK_SIZE,
        num_kv_heads=_NUM_HEADS,
        head_size=_HEAD_SIZE,
        dtype=torch.float16,
        non_causal=True,
    )
    with set_current_vllm_config(vllm_config):
        native_layer = DiffusionPagedAttentionLayerAdapter(
            layer_name=_LAYER_NAME,
            layer=diffusion_layer,
            spec=spec,
            vllm_config=vllm_config,
            device=device,
        )
    vllm_config.compilation_config.static_forward_context[_LAYER_NAME] = native_layer

    canonical_spec = native_layer.spec
    kv_cache_config = KVCacheConfig(
        num_blocks=4 * num_rows,
        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4 * num_rows, [_LAYER_NAME])],
        kv_cache_groups=[
            KVCacheGroupSpec(
                layer_names=[_LAYER_NAME],
                kv_cache_spec=canonical_spec,
            )
        ],
    )
    # 0.29 reads the resolved physical layout when allocating; the diffusion
    # control plane pins it, so this standalone harness must too.
    resolve_diffusion_kv_cache_layout(vllm_config)
    with set_current_vllm_config(vllm_config):
        attn_groups, _, kernel_block_sizes = init_attn_backend(
            kv_cache_config,
            vllm_config,
            device,
        )
        # f2aad6aa70 (#56888): init_kv_cache no longer mutates a runner list;
        # first arg is forward_context and the allocated caches are returned.
        init_kv_cache(
            vllm_config.compilation_config.static_forward_context,
            kv_cache_config,
            device,
            kernel_block_sizes,
            vllm_config,
        )
        block_tables = BlockTables(
            block_sizes=[_BLOCK_SIZE],
            max_num_reqs=num_rows,
            max_num_batched_tokens=32 * num_rows,
            max_num_blocks_per_group=[2],
            device=device,
            kernel_block_sizes=kernel_block_sizes,
        )

    for row in range(num_rows):
        block_tables.append_block_ids(row, ([4 * row + 3, 4 * row + 1],), overwrite=True)
    block_tables.apply_staged_writes()
    adapter = DiffusionPagedAttentionAdapter(
        vllm_config=vllm_config,
        device=device,
        kv_cache_config=kv_cache_config,
        block_tables=block_tables,
        attn_groups=attn_groups,
        layers={_LAYER_NAME: native_layer},
        resolve_row=lambda _request_id, sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
            row_index=sequence_id,
            max_seq_len=19,
        ),
    )
    return adapter, diffusion_layer, device


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
    adapter, diffusion_layer, device = _native_adapter()
    prefix_batch = adapter.prepare_batch(
        [
            DiffusionPagedAttentionRow(
                request_id="req-0",
                sequence_id=0,
                query_len=17,
                seq_len=17,
            )
        ]
    )

    torch.manual_seed(1)
    query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
    packed_qkv = torch.randn(
        19,
        3 * _NUM_HEADS * _HEAD_SIZE,
        dtype=torch.float16,
        device=device,
    )
    _, key_flat, value_flat = packed_qkv.split(_NUM_HEADS * _HEAD_SIZE, dim=-1)
    key = key_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
    value = value_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
    assert not key.is_contiguous()
    assert not value.is_contiguous()
    omni_backend = FlashAttentionImpl(
        num_heads=_NUM_HEADS,
        head_size=_HEAD_SIZE,
        softmax_scale=diffusion_layer.softmax_scale,
        num_kv_heads=_NUM_HEADS,
    )
    with adapter.activate(prefix_batch):
        prefix_context = adapter.prepare_layer_context(_LAYER_NAME, query[:17], key[:17], value[:17])
        prefix_output = omni_backend.forward_paged(prefix_context)
    prefix_slot_mappings = prefix_batch.slot_mappings.clone()

    suffix_batch = adapter.prepare_batch(
        [
            DiffusionPagedAttentionRow(
                request_id="req-0",
                sequence_id=0,
                query_len=2,
                seq_len=19,
                kv_start_pos=17,
            )
        ]
    )
    with adapter.activate(suffix_batch):
        suffix_context = adapter.prepare_layer_context(_LAYER_NAME, query[17:], key[17:], value[17:])
        suffix_output = omni_backend.forward_paged(suffix_context)
    suffix_slot_mappings = suffix_batch.slot_mappings.clone()

    piecewise_batch = adapter.prepare_batch(
        [
            DiffusionPagedAttentionRow(
                request_id="req-0",
                sequence_id=0,
                query_len=19,
                seq_len=19,
            )
        ]
    )
    mixed_mask = torch.ones(19, 19, dtype=torch.bool, device=device).tril()
    mixed_mask[5:10, 5:10] = True
    mixed_metadata = AttentionMetadata(
        attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
        full_attn_spans=[[(5, 10)]],
    )
    with adapter.activate(piecewise_batch):
        piecewise_context = adapter.prepare_layer_context(
            _LAYER_NAME,
            query,
            key,
            value,
            omni_attn_metadata=mixed_metadata,
        )
        piecewise_output = omni_backend.forward_paged(piecewise_context)

    prefix_reference = (
        F.scaled_dot_product_attention(
            query[:17].transpose(0, 1).unsqueeze(0),
            key[:17].transpose(0, 1).unsqueeze(0),
            value[:17].transpose(0, 1).unsqueeze(0),
            is_causal=False,
            scale=diffusion_layer.softmax_scale,
        )
        .squeeze(0)
        .transpose(0, 1)
    )
    suffix_reference = (
        F.scaled_dot_product_attention(
            query[17:].transpose(0, 1).unsqueeze(0),
            key.transpose(0, 1).unsqueeze(0),
            value.transpose(0, 1).unsqueeze(0),
            is_causal=False,
            scale=diffusion_layer.softmax_scale,
        )
        .squeeze(0)
        .transpose(0, 1)
    )
    piecewise_reference = (
        F.scaled_dot_product_attention(
            query.transpose(0, 1).unsqueeze(0),
            key.transpose(0, 1).unsqueeze(0),
            value.transpose(0, 1).unsqueeze(0),
            attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
            scale=diffusion_layer.softmax_scale,
        )
        .squeeze(0)
        .transpose(0, 1)
    )
    torch.accelerator.synchronize(device)

    assert prefix_slot_mappings.cpu().tolist() == [list(range(3 * _BLOCK_SIZE, 4 * _BLOCK_SIZE)) + [_BLOCK_SIZE]]
    assert suffix_slot_mappings.cpu().tolist() == [[_BLOCK_SIZE + 1, _BLOCK_SIZE + 2]]
    torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2)
    torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2)
    torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
def test_hunyuan_homogeneous_cfg_paged_attention_graph_replay() -> None:
    """Real paged kernels, strided QKV, mixed image spans and changed replay inputs."""
    adapter, layer, device = _native_adapter(num_rows=2)
    packed = torch.randn(2, 19, 3, _NUM_HEADS, _HEAD_SIZE, device=device, dtype=torch.float16)
    query, key, value = packed.unbind(dim=2)
    backend = FlashAttentionImpl(
        num_heads=_NUM_HEADS, head_size=_HEAD_SIZE, softmax_scale=layer.softmax_scale, num_kv_heads=_NUM_HEADS
    )
    batch = adapter.prepare_batch(
        [DiffusionPagedAttentionRow(request_id="cfg", sequence_id=i, query_len=19, seq_len=19) for i in range(2)]
    )
    metadata = AttentionMetadata(full_attn_spans=[[(5, 10)], [(5, 10)]])
    with adapter.activate(batch):
        context = adapter.prepare_layer_context(_LAYER_NAME, query, key, value, omni_attn_metadata=metadata)
    assert context.piecewise_plan.homogeneous_batch_shape == (2, 19)
    warmup = torch.cuda.Stream()
    warmup.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(warmup):
        for _ in range(3):
            backend.forward_paged(context)
    torch.cuda.current_stream().wait_stream(warmup)
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        captured = backend.forward_paged(context)
    output_ptr = captured.data_ptr()
    mask = torch.ones(19, 19, device=device, dtype=torch.bool).tril()
    mask[5:10, 5:10] = True
    for _ in range(2):
        packed.copy_(torch.randn_like(packed))
        graph.replay()
        actual = captured.clone()
        eager = backend.forward_paged(context)
        reference = F.scaled_dot_product_attention(
            query.transpose(1, 2),
            key.transpose(1, 2),
            value.transpose(1, 2),
            attn_mask=mask,
            scale=layer.softmax_scale,
        ).transpose(1, 2)
        torch.testing.assert_close(actual.reshape_as(reference), reference, rtol=2e-2, atol=2e-2)
        torch.testing.assert_close(actual, eager, rtol=2e-2, atol=2e-2)
        assert captured.data_ptr() == output_ptr