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


@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:
    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,
        kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4, [_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,
        )
        runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
        # vLLM 0.29 dropped attn_groups and cache_dtype here: the layout now
        # comes from the resolved CacheConfig.kv_cache_layout inside the call.
        init_kv_cache(
            runner_kv_caches,
            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=1,
            max_num_batched_tokens=32,
            max_num_blocks_per_group=[2],
            device=device,
            kernel_block_sizes=kernel_block_sizes,
        )

    block_tables.append_block_ids(0, ([3, 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=0,
            max_seq_len=19,
        ),
    )
    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)