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

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

import pytest
import torch
import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
from vllm.v1.core.kv_cache_manager import KVCacheManager
from vllm.v1.core.kv_cache_utils import BlockHash
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
)
from vllm.v1.request import RequestStatus

from tests.helpers.kv_layout import build_kv_cache_tensor
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest

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

BLOCK_SIZE = 4


def _request(**overrides) -> DiffusionKVRequest:
    values = dict(
        request_id="req/diffusion-kv/0",
        sequence_id=0,
        prefix_len=4,
        target_len=4,
        seq_len=8,
    )
    values.update(overrides)
    return DiffusionKVRequest(**values)


def test_request_exposes_native_request_and_diffusion_semantics() -> None:
    request = _request()

    assert request.sequence_id == 0
    assert request.prefix_len == 4
    assert request.target_len == 4
    assert request.num_tokens == request.seq_len == 8
    assert request.num_prompt_tokens == 4
    assert request.num_computed_tokens == 0
    assert request.block_hashes == []
    assert request.kv_contexts == ()
    assert request.skip_reading_prefix_cache is True
    assert request.shared_prefix_boundary == 0
    assert request.status is RequestStatus.WAITING


@pytest.mark.parametrize(
    ("overrides", "message"),
    [
        ({"sequence_id": -1}, "sequence_id"),
        ({"prefix_len": -1}, "prefix_len"),
        ({"target_len": 0}, "target_len"),
        ({"seq_len": 0}, "seq_len"),
        ({"prefix_len": 5, "target_len": 4, "seq_len": 8}, "must not exceed"),
    ],
)
def test_request_rejects_invalid_lengths(overrides, message: str) -> None:
    with pytest.raises(ValueError, match=message):
        _request(**overrides)


@pytest.mark.parametrize(
    ("overrides", "message"),
    [
        ({"context_id": ""}, "context_id"),
        ({"cache_role": ""}, "cache_role"),
        ({"num_tokens": 0}, "num_tokens"),
    ],
)
def test_context_rejects_invalid_identity_or_length(overrides, message: str) -> None:
    values = dict(context_id="text", cache_role="cross.text", num_tokens=8)
    values.update(overrides)

    with pytest.raises(ValueError, match=message):
        DiffusionKVContext(**values)


def test_request_owns_independent_kv_contexts() -> None:
    text = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
    image = DiffusionKVContext(context_id="image", cache_role="cross.image", num_tokens=4)

    request = _request(kv_contexts=(text, image))

    assert request.kv_contexts == (text, image)


def test_request_rejects_duplicate_context_ids() -> None:
    contexts = (
        DiffusionKVContext(context_id="condition", cache_role="cross.text", num_tokens=8),
        DiffusionKVContext(context_id="condition", cache_role="cross.image", num_tokens=4),
    )

    with pytest.raises(ValueError, match="unique context_id"):
        _request(kv_contexts=contexts)


def test_request_rejects_non_context_values() -> None:
    with pytest.raises(TypeError, match="DiffusionKVContext"):
        _request(kv_contexts=(object(),))


def _manager(*, enable_caching: bool) -> KVCacheManager:
    # vLLM 0.26 lazily registers its built-in KV specs. Another test may have
    # already registered an out-of-tree spec, making the registry non-empty
    # before FullAttentionSpec is registered. Initialize the native mappings.
    native_kv_managers.register_all_kvcache_specs(None)

    spec = FullAttentionSpec(
        block_size=BLOCK_SIZE,
        num_kv_heads=2,
        head_size=8,
        dtype=torch.bfloat16,
    )
    group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)
    config = KVCacheConfig(
        num_blocks=8,
        kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["layer0"])],
        kv_cache_groups=[group],
    )
    return KVCacheManager(
        config,
        max_model_len=64,
        scheduler_block_size=BLOCK_SIZE,
        hash_block_size=BLOCK_SIZE,
        enable_caching=enable_caching,
    )


def test_empty_hash_request_conforms_to_native_request_local_allocation() -> None:
    manager = _manager(enable_caching=True)
    free_before = manager.block_pool.get_num_free_blocks()
    request = _request()

    # vLLM 0.26 adds a shared-prefix-boundary return value.
    computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        request,
        num_new_tokens=request.num_tokens,
        new_computed_blocks=computed_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )

    assert blocks is not None
    assert len(manager.get_block_ids(request.request_id)[0]) == 2
    # Empty hashes intentionally mean request-local residency only. Calling
    # native cache_blocks() would require a canonical hash for the full prefix
    # block and is therefore outside this request's contract.
    manager.free(request)
    assert manager.block_pool.get_num_free_blocks() == free_before


def test_hashed_request_conforms_to_native_delayed_prefix_commit() -> None:
    manager = _manager(enable_caching=True)
    prefix_hash = BlockHash(b"prefix".ljust(32, b"\0"))
    target_hash = BlockHash(b"target".ljust(32, b"\0"))
    request = _request(block_hashes=[prefix_hash, target_hash])

    computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        request,
        num_new_tokens=request.num_tokens,
        new_computed_blocks=computed_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )

    assert blocks is not None
    manager.cache_blocks(request, request.prefix_len)
    manager.free(request)

    cached_request = _request(
        request_id="req/diffusion-kv/cached",
        block_hashes=[prefix_hash, target_hash],
    )
    cached_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(cached_request)

    # Only the stable prefix was published. The dynamic target remains a miss.
    assert num_computed_tokens == cached_request.prefix_len
    assert len(cached_blocks.blocks[0]) == 1
    cached_allocation = manager.allocate_slots(
        cached_request,
        num_new_tokens=cached_request.num_tokens - num_computed_tokens,
        num_new_computed_tokens=num_computed_tokens,
        new_computed_blocks=cached_blocks,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )
    assert cached_allocation is not None
    manager.free(cached_request)