vllm-omni/vllm_omni/diffusion/diffusion_kv/layout.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Physical KV cache layout selection for the native Diffusion paged cache.

vLLM 0.29 removed ``AttentionSpec.indexes_kv_by_block_stride`` and replaced it
with a single physical layout per model, resolved once and recorded on
``CacheConfig.kv_cache_layout``.  ``KVCacheLayout.is_block_outermost`` is the
direct successor of the old per-spec flag: upstream derives
``interleaved_block_stride`` from it when it lays out ``KVCacheTensor`` regions
(``vllm/v1/core/kv_cache_utils.py``).

Diffusion never reaches vLLM's engine-core resolution.  Its backends subclass
Omni's own ``AttentionBackend`` ABC, which upstream's
``get_supported_kv_cache_layouts`` never sees, so the layout is pinned here
before ``get_kv_cache_configs`` runs.
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING

from vllm.v1.kv_cache_interface import KVCacheSpec, KVCacheTensor, compute_layout_strides
from vllm.v1.kv_cache_layout import KVCacheLayout

if TYPE_CHECKING:
    from vllm.config import VllmConfig

# Layer-outermost, matching the head of upstream's ``_DEFAULT_LAYOUT_PREFERENCE``
# and reproducing the ``indexes_kv_by_block_stride=False`` that every dense
# diffusion backend declared before 0.29.
DEFAULT_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.LBNHC
# Block-outermost counterpart, for a paged backend whose kernel really does read
# K/V pages by the runtime block stride.
BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.BLNHC


def resolve_diffusion_kv_cache_layout(
    vllm_config: VllmConfig,
    *,
    indexes_kv_by_block_stride: bool = False,
) -> KVCacheLayout:
    """Pin the physical KV layout for the Diffusion cache and return it.

    A layout already present on the config wins, mirroring upstream's
    ``resolve_kv_cache_layout``; it is validated against the requirement rather
    than silently overwritten, because an inconsistent layout reads K/V
    from the wrong offsets instead of raising.
    """

    required_block_outermost = bool(indexes_kv_by_block_stride)
    cache_config = getattr(vllm_config, "cache_config", None)
    if cache_config is None:
        # A config double with nowhere to record the choice; report what the
        # requirement implies without pretending it was pinned.
        return BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT

    current = getattr(cache_config, "kv_cache_layout", None)
    if current is not None:
        layout = KVCacheLayout[current]
        if layout.is_block_outermost != required_block_outermost:
            raise ValueError(
                "Diffusion KV cache layout "
                f"{layout.name} (is_block_outermost="
                f"{layout.is_block_outermost}) contradicts the attention "
                "backend, which requires is_block_outermost="
                f"{required_block_outermost}."
            )
        return layout

    layout = BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
    cache_config.kv_cache_layout = layout.name
    return layout


def adopt_kv_cache_layout(vllm_config: VllmConfig, kv_cache_config) -> KVCacheLayout:
    """Adopt the control plane's resolved layout inside a worker.

    The engine core resolves the layout once and every worker adopts it before
    allocating (``CacheConfig.kv_cache_layout``).  Diffusion workers receive only
    a ``KVCacheConfig``, so the name rides along on ``kv_cache_config`` and is
    copied onto the worker's own config here.  Without this, anything calling
    ``get_resolved_kv_cache_layout`` in a worker -- ``init_kv_cache`` does --
    raises.
    """

    name = getattr(kv_cache_config, "kv_cache_layout", None)
    if name is None:
        # Older configs, or a standalone harness: fall back to the same default
        # the control plane pins so both sides still agree.
        return resolve_diffusion_kv_cache_layout(vllm_config)

    cache_config = getattr(vllm_config, "cache_config", None)
    if cache_config is None:
        return KVCacheLayout[name]

    current = getattr(cache_config, "kv_cache_layout", None)
    if current is None:
        cache_config.kv_cache_layout = name
    elif current != name:
        raise ValueError(
            f"Worker KV cache layout {current!r} disagrees with the layout resolved by the control plane ({name!r})."
        )
    return KVCacheLayout[name]


def build_kv_cache_tensor(
    spec: KVCacheSpec,
    num_blocks: int,
    layers: Sequence[str],
    *,
    layout: KVCacheLayout | None = None,
    offset: int = 0,
) -> KVCacheTensor:
    """A ``KVCacheTensor`` covering ``layers`` for ``num_blocks`` blocks.

    Before 0.29 a tensor named the layers sharing it
    (``KVCacheTensor(size=..., shared_by=[...])``); 0.29 replaced that with an
    explicit layout, so layer ``l``'s block ``b`` starts at
    ``offset + l * layer_stride + b * block_stride``.  Note this lays the layers
    out at distinct offsets rather than aliasing them, which is what upstream
    now does for a single cache group.

    The strides come from upstream's own ``compute_layout_strides``, called the
    way ``vllm.v1.core.kv_cache_utils`` calls it, because hand-computed strides
    read K/V from the wrong offsets silently instead of raising.
    """

    layout = layout or DEFAULT_DIFFUSION_KV_CACHE_LAYOUT
    layer_names = list(layers)
    # Mirrors kv_cache_utils: the block stride is pinned up front only for
    # block-outermost layouts, where every block repeats the same packing.
    bytes_per_block = spec.page_size_bytes * len(layer_names)
    interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
    layer_stride, block_stride, _, _, _ = compute_layout_strides(
        spec,
        num_blocks,
        len(layer_names),
        layout,
        fixed_strides=(None, interleaved_block_stride, None, None, None),
    )
    return KVCacheTensor(
        size=bytes_per_block * num_blocks,
        layers=layer_names,
        layer_stride=layer_stride,
        block_stride=block_stride,
        offset=offset,
    )


def assert_backend_layout_supported(vllm_config: VllmConfig, attn_backend: type | None) -> None:
    """Fail loudly when a backend's block-stride need contradicts the layout.

    Skipped while the layout is still unresolved -- worker processes collect
    specs before the control plane hands them the resolved name -- so this
    enforces the contract wherever it is knowable and never crashes a rank that
    has simply not been told yet.
    """

    if attn_backend is None:
        return
    # Spec discovery runs against lightweight config doubles in tests and against
    # worker configs that have not been handed the resolved name yet; in both
    # cases there is simply nothing to check.
    cache_config = getattr(vllm_config, "cache_config", None)
    if cache_config is None or getattr(cache_config, "kv_cache_layout", None) is None:
        return
    indexes_by_block_stride = getattr(attn_backend, "indexes_kv_by_block_stride", None)
    if indexes_by_block_stride is None:
        return
    resolve_diffusion_kv_cache_layout(
        vllm_config,
        indexes_kv_by_block_stride=indexes_by_block_stride(),
    )