diff --git a/tests/diffusion/diffusion_kv/test_block_tables.py b/tests/diffusion/diffusion_kv/test_block_tables.py index f2002f1c6..aa8e07307 100644 --- a/tests/diffusion/diffusion_kv/test_block_tables.py +++ b/tests/diffusion/diffusion_kv/test_block_tables.py @@ -279,17 +279,17 @@ def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple] group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None)) return [[group], [group]], "cg-support", [4, 8] + # vLLM 0.29 dropped attn_groups and cache_dtype from init_kv_cache: the + # layout now rides on the resolved CacheConfig.kv_cache_layout instead. def init_cache( runner_kv_caches, forward_context, kv_cache_config, - attn_groups, device, - cache_dtype, kernel_block_sizes, vllm_config, ): - events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes)) + events.append(("cache", kv_cache_config, kernel_block_sizes)) runner_kv_caches.extend(["cache-0", "cache-1"]) return {"layer-0": "tensor-0", "layer-1": "tensor-1"} diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py index 861f9564d..b5f914b11 100644 --- a/tests/diffusion/diffusion_kv/test_initialization.py +++ b/tests/diffusion/diffusion_kv/test_initialization.py @@ -299,8 +299,10 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None ) assert vllm_config.model_config.original_max_model_len == -1 - # Sixteen physical pages provide ``16 * block_size`` tokens. The native - # backend chooses the block geometry, so keep this assertion backend + # Sixteen physical pages provide ``(16 - 1) * block_size`` tokens: vLLM 0.29 + # subtracts the null block that BlockPool permanently holds back before + # sizing (check_enough_kv_cache_memory), so only 15 pages are usable. The + # native backend chooses the block geometry, so keep this assertion backend # independent instead of assuming the generic 16-token block size. - expected_max_model_len = min(model_limit, spec.block_size * 16) + expected_max_model_len = min(model_limit, spec.block_size * (16 - 1)) assert vllm_config.model_config.max_model_len == expected_max_model_len diff --git a/tests/diffusion/diffusion_kv/test_layout.py b/tests/diffusion/diffusion_kv/test_layout.py new file mode 100644 index 000000000..2f410f74c --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_layout.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Physical KV cache layout selection and propagation for the Diffusion cache.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from vllm.config import CacheConfig +from vllm.v1.kv_cache_interface import FullAttentionSpec +from vllm.v1.kv_cache_layout import KVCacheLayout + +from vllm_omni.diffusion.diffusion_kv.layout import ( + DEFAULT_DIFFUSION_KV_CACHE_LAYOUT, + adopt_kv_cache_layout, + assert_backend_layout_supported, + build_kv_cache_tensor, + resolve_diffusion_kv_cache_layout, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _config(): + return SimpleNamespace(cache_config=CacheConfig()) + + +class _DenseBackend: + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return False + + +class _BlockStrideBackend: + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return True + + +def test_default_layout_is_not_block_outermost() -> None: + """The default reproduces the pre-0.29 indexes_kv_by_block_stride=False.""" + vllm_config = _config() + layout = resolve_diffusion_kv_cache_layout(vllm_config) + + assert layout is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + assert layout.is_block_outermost is False + assert vllm_config.cache_config.get_resolved_kv_cache_layout() is layout + + +def test_block_stride_backend_selects_block_outermost_layout() -> None: + layout = resolve_diffusion_kv_cache_layout(_config(), indexes_kv_by_block_stride=True) + + assert layout.is_block_outermost is True + + +def test_resolution_is_idempotent() -> None: + vllm_config = _config() + + assert resolve_diffusion_kv_cache_layout(vllm_config) is resolve_diffusion_kv_cache_layout(vllm_config) + + +def test_contradicting_requirement_raises_instead_of_misaddressing() -> None: + vllm_config = _config() + resolve_diffusion_kv_cache_layout(vllm_config) + + with pytest.raises(ValueError, match="contradicts the attention"): + resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True) + + +def test_preset_layout_is_honoured() -> None: + vllm_config = _config() + vllm_config.cache_config.kv_cache_layout = KVCacheLayout.BLNHC.name + + layout = resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True) + + assert layout is KVCacheLayout.BLNHC + + +def test_backend_assertion_is_quiet_until_the_layout_is_resolved() -> None: + """Workers collect specs before the control plane hands them the layout.""" + vllm_config = _config() + + assert_backend_layout_supported(vllm_config, _BlockStrideBackend) # unresolved: nothing to check + + resolve_diffusion_kv_cache_layout(vllm_config) + assert_backend_layout_supported(vllm_config, _DenseBackend) + assert_backend_layout_supported(vllm_config, None) + with pytest.raises(ValueError, match="contradicts the attention"): + assert_backend_layout_supported(vllm_config, _BlockStrideBackend) + + +def test_worker_adopts_the_layout_carried_on_the_kv_cache_config() -> None: + """Diffusion workers receive only a KVCacheConfig, so the name rides on it.""" + worker = _config() + shipped = SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name) + + assert adopt_kv_cache_layout(worker, shipped) is KVCacheLayout.BLNHC + assert worker.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.BLNHC + + +def test_worker_without_a_carried_layout_falls_back_to_the_same_default() -> None: + worker = _config() + + assert adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=None)) is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + + +def test_worker_disagreeing_with_the_control_plane_raises() -> None: + worker = _config() + resolve_diffusion_kv_cache_layout(worker) + + with pytest.raises(ValueError, match="disagrees with the layout"): + adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name)) + + +def test_block_stride_requirement_propagates_control_plane_to_worker() -> None: + """A block-stride backend must reach the worker as a block-outermost layout. + + Covers the whole seam: the control plane resolves with the requirement, the + name rides on each rank-local KVCacheConfig, the worker adopts it, and the + backend contract check then passes instead of raising. + """ + control_plane = _config() + layout = resolve_diffusion_kv_cache_layout(control_plane, indexes_kv_by_block_stride=True) + assert layout.is_block_outermost is True + + shipped = SimpleNamespace(kv_cache_layout=layout.name) + + worker = _config() + assert adopt_kv_cache_layout(worker, shipped) is layout + # The worker's own backend requirement now agrees with what it was handed. + assert_backend_layout_supported(worker, _BlockStrideBackend) + with pytest.raises(ValueError, match="contradicts the attention"): + assert_backend_layout_supported(worker, _DenseBackend) + + +def test_tensor_strides_come_from_the_upstream_derivation() -> None: + spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16) + + tensor = build_kv_cache_tensor(spec, 4, ["layer0", "layer1"]) + + assert tensor.layers == ["layer0", "layer1"] + assert tensor.size == spec.page_size_bytes * 2 * 4 + # Layer-outermost: each layer owns a contiguous run of its own blocks. + assert tensor.layer_stride == spec.page_size_bytes * 4 + assert tensor.block_stride == spec.page_size_bytes diff --git a/tests/diffusion/diffusion_kv/test_manager.py b/tests/diffusion/diffusion_kv/test_manager.py index 5152c2a67..8abc296f1 100644 --- a/tests/diffusion/diffusion_kv/test_manager.py +++ b/tests/diffusion/diffusion_kv/test_manager.py @@ -4,10 +4,11 @@ import pytest import torch import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest +from tests.helpers.kv_layout import build_kv_cache_tensor pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -24,7 +25,7 @@ def _config(num_blocks: int) -> KVCacheConfig: ) return KVCacheConfig( num_blocks=num_blocks, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], ) diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py index 1f242420d..a13e09051 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py @@ -18,6 +18,7 @@ from vllm.v1.kv_cache_interface import FullAttentionSpec from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module +from tests.helpers.kv_layout import layout_for_backend from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( DiffusionPagedAttentionAdapter, DiffusionPagedAttentionMetadata, @@ -1239,7 +1240,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads assert native_layer.num_heads == 4 assert native_layer.num_kv_heads == 2 assert native_layer.spec.num_kv_heads == 2 - assert native_layer.spec.indexes_kv_by_block_stride is True + # 0.29: the flag no longer rides on the spec, it selects the layout. + assert layout_for_backend(native_layer.attn_backend).is_block_outermost is True assert native_layer._q_scale_float == 1.0 assert native_layer._k_scale_float == 1.0 assert native_layer._v_scale_float == 1.0 diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py index ef5d75f4d..192e60bea 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py @@ -10,12 +10,13 @@ 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, KVCacheTensor +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 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, @@ -23,6 +24,7 @@ from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( DiffusionPagedAttentionRowBinding, ) from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig +from tests.helpers.kv_layout import build_kv_cache_tensor pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] @@ -89,10 +91,7 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N kv_cache_config = KVCacheConfig( num_blocks=4, kv_cache_tensors=[ - KVCacheTensor( - size=canonical_spec.page_size_bytes * 4, - shared_by=[_LAYER_NAME], - ) + build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME]) ], kv_cache_groups=[ KVCacheGroupSpec( @@ -101,6 +100,9 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N ) ], ) + # 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, @@ -108,13 +110,13 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N 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, - attn_groups, device, - vllm_config.cache_config.cache_dtype, kernel_block_sizes, vllm_config, ) diff --git a/tests/diffusion/diffusion_kv/test_request.py b/tests/diffusion/diffusion_kv/test_request.py index ae5609b2a..4ca6703d8 100644 --- a/tests/diffusion/diffusion_kv/test_request.py +++ b/tests/diffusion/diffusion_kv/test_request.py @@ -10,11 +10,11 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheTensor, ) from vllm.v1.request import RequestStatus from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest +from tests.helpers.kv_layout import build_kv_cache_tensor pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -119,7 +119,7 @@ def _manager(*, enable_caching: bool) -> KVCacheManager: group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec) config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["layer0"])], kv_cache_groups=[group], ) return KVCacheManager( diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py index 9485ed046..d07c41355 100644 --- a/tests/diffusion/diffusion_kv/test_spec_discovery.py +++ b/tests/diffusion/diffusion_kv/test_spec_discovery.py @@ -8,13 +8,14 @@ from unittest.mock import Mock import pytest import torch from torch import nn -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker +from tests.helpers.kv_layout import build_kv_cache_tensor, layout_for_backend pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -79,7 +80,9 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None: assert spec.num_kv_heads == 2 assert spec.head_size == 8 assert spec.dtype is torch.bfloat16 - assert spec.indexes_kv_by_block_stride is True + # vLLM 0.29 moved the block-stride decision off the spec onto the resolved + # physical layout; the backend's declaration must still select one. + assert layout_for_backend(_Backend).is_block_outermost is True assert spec.non_causal is True @@ -112,7 +115,7 @@ def test_runner_retains_matching_rank_local_config() -> None: spec = runner.get_kv_cache_spec()["image_attention"] config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["image_attention"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)], ) @@ -126,7 +129,7 @@ def test_runner_rejects_rank_local_config_for_different_layers() -> None: spec = runner.get_kv_cache_spec()["image_attention"] config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["other_attention"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)], ) diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index d6d17c18f..6bc0bdaba 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -10,7 +10,7 @@ import pytest import torch import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers from pytest_mock import MockerFixture -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode @@ -27,6 +27,7 @@ from vllm_omni.diffusion.sched import ( from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from tests.helpers.kv_layout import build_kv_cache_tensor pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -104,7 +105,7 @@ def _initialize_paged_scheduler( ) config = KVCacheConfig( num_blocks=num_blocks, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], ) scheduler.initialize( diff --git a/tests/helpers/kv_layout.py b/tests/helpers/kv_layout.py new file mode 100644 index 000000000..0306cee42 --- /dev/null +++ b/tests/helpers/kv_layout.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared access to the production ``KVCacheTensor`` layout derivation. + +Tests deliberately go through the same code production uses +(``vllm_omni.diffusion.diffusion_kv.layout``) rather than recomputing strides, +so a regression in that derivation fails the suite instead of being masked by a +parallel test-only copy. +""" + +from types import SimpleNamespace + +from vllm.config import CacheConfig +from vllm.v1.kv_cache_layout import KVCacheLayout + +from vllm_omni.diffusion.diffusion_kv.layout import ( + build_kv_cache_tensor, + resolve_diffusion_kv_cache_layout, +) + +__all__ = ["build_kv_cache_tensor", "layout_for_backend"] + + +def layout_for_backend(attn_backend) -> KVCacheLayout: + """The physical layout ``attn_backend``'s block-stride declaration selects. + + Before 0.29 a backend's ``indexes_kv_by_block_stride()`` was copied onto + ``AttentionSpec``; asserting on the resolved layout is the surviving way to + check that declaration still reaches the cache. + """ + + vllm_config = SimpleNamespace(cache_config=CacheConfig()) + return resolve_diffusion_kv_cache_layout( + vllm_config, + indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(), + ) diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 5b5b8a328..e903960d6 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -121,6 +121,13 @@ class AttentionBackend(ABC): when it unifies cache layouts across layers. Dense diffusion backends conservatively keep the default ``False``; a paged backend should override this only when its kernel actually follows that layout. + + Since vLLM 0.29 this no longer rides on ``AttentionSpec``. It selects + the block-outermost physical layout pinned on + ``CacheConfig.kv_cache_layout`` -- see + ``vllm_omni.diffusion.diffusion_kv.layout`` -- whose + ``is_block_outermost`` is what upstream reads when it sizes + ``KVCacheTensor`` regions. """ return False diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index ac541808b..4b8ca8f49 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -24,6 +24,7 @@ from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role from vllm_omni.diffusion.config import get_current_diffusion_config_or_none from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.layout import assert_backend_layout_supported from vllm_omni.diffusion.distributed.parallel_state import get_sp_group from vllm_omni.diffusion.forward_context import ( get_forward_context, @@ -261,15 +262,19 @@ class Attention(nn.Module): return None dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype # Keep backend layout discovery under the same config context used by - # upstream vLLM's attention-spec collector. + # upstream vLLM's attention-spec collector. vLLM 0.29 moved the + # block-stride decision off the spec and onto the single physical + # ``CacheConfig.kv_cache_layout`` resolved before the cache is built, + # so the backend's preference is enforced there (see + # ``vllm_omni.diffusion.diffusion_kv.initialization``) rather than + # carried per layer. with set_current_vllm_config(vllm_config): - indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride() + assert_backend_layout_supported(vllm_config, self.attn_backend) return FullAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=self.num_kv_heads, head_size=self.head_size, dtype=dtype, - indexes_kv_by_block_stride=indexes_kv_by_block_stride, non_causal=not self.causal, ) diff --git a/vllm_omni/diffusion/diffusion_kv/initialization.py b/vllm_omni/diffusion/diffusion_kv/initialization.py index 31aaaa46e..ffb0c2176 100644 --- a/vllm_omni/diffusion/diffusion_kv/initialization.py +++ b/vllm_omni/diffusion/diffusion_kv/initialization.py @@ -16,6 +16,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config from vllm_omni.platforms import current_omni_platform @@ -28,14 +29,35 @@ def build_native_kv_cache_configs( vllm_config: VllmConfig, worker_specs: list[dict[str, KVCacheSpec]], available_memory: list[int], + *, + indexes_kv_by_block_stride: bool = False, ) -> tuple[list[KVCacheConfig], KVCacheConfig]: - """Build rank-local and Scheduler-native configs using vLLM utilities.""" + """Build rank-local and Scheduler-native configs using vLLM utilities. + + ``indexes_kv_by_block_stride`` carries the native backend requirement that + 0.29 removed from ``AttentionSpec``. Every diffusion backend in tree + declares ``False``, so the default reproduces current behaviour; a caller + that learns otherwise from its workers passes ``True`` here and the + block-outermost layout propagates to every rank-local config. + """ from vllm.v1.core import single_type_kv_cache_manager + # 0.29 sizes every KVCacheTensor region from the resolved physical layout, + # so it must be pinned before get_kv_cache_configs. Idempotent, so callers + # that already resolved it (the control plane) are unaffected. + layout = resolve_diffusion_kv_cache_layout( + vllm_config, + indexes_kv_by_block_stride=indexes_kv_by_block_stride, + ) single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config) worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory) - return worker_configs, generate_scheduler_kv_cache_config(worker_configs) + scheduler_config = generate_scheduler_kv_cache_config(worker_configs) + # Workers only receive a KVCacheConfig, so the resolved layout travels on it + # and is adopted into each worker's own CacheConfig before allocation. + for config in (*worker_configs, scheduler_config): + config.kv_cache_layout = layout.name + return worker_configs, scheduler_config def initialize_diffusion_kv_control_plane( diff --git a/vllm_omni/diffusion/diffusion_kv/layout.py b/vllm_omni/diffusion/diffusion_kv/layout.py new file mode 100644 index 000000000..1ba0cd326 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/layout.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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 mis-addresses K/V + 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 " + f"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 + mis-address K/V 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(), + ) diff --git a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py index c5ab3fb16..79a0ded62 100644 --- a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py +++ b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py @@ -18,6 +18,10 @@ from vllm.v1.worker.gpu.block_table import BlockTables from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode +from vllm_omni.diffusion.diffusion_kv.layout import ( + adopt_kv_cache_layout, + assert_backend_layout_supported, +) from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( DiffusionPagedAttentionAdapter, @@ -285,13 +289,18 @@ class DiffusionKVModelRunnerBackend: layers=self._kv_cache_layer_adapters, resolve_row=self._resolve_paged_attention_row, ) + # vLLM 0.29 reads the resolved physical layout when allocating, + # and dropped attn_groups/cache_dtype from the signature. + adopt_kv_cache_layout(self.vllm_config, kv_cache_config) + for layer_adapter in self._kv_cache_layer_adapters.values(): + assert_backend_layout_supported( + self.vllm_config, getattr(layer_adapter, "attn_backend", None) + ) init_kv_cache( kv_caches, self.vllm_config.compilation_config.static_forward_context, kv_cache_config, - attn_groups, self.device, - self.vllm_config.cache_config.cache_dtype, kernel_block_sizes, self.vllm_config, ) diff --git a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py index 6a7862669..a76a5434c 100644 --- a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py +++ b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py @@ -105,10 +105,12 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase): attn_backend, ulysses_degree=ulysses_degree, ) + # vLLM 0.29 carries the block-stride decision on the resolved physical + # ``CacheConfig.kv_cache_layout`` instead of the spec; the backend's + # requirement is enforced against that layout rather than copied here. canonical_spec = replace( spec, num_kv_heads=num_kv_heads, - indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(), ) self.layer_name = layer_name self.spec = canonical_spec diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index f18a00ada..24f09210d 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -22,10 +22,10 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, - KVCacheTensor, ) from vllm.v1.request import RequestStatus +from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ( @@ -134,7 +134,7 @@ def build_kv_manager( """ layer_names = list(layer_names) group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec) - tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)] + tensors = [build_kv_cache_tensor(spec, num_blocks, layer_names)] config = KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=tensors,