tzhouamandClaude Opus 5
c306445ff2
vLLM 0.29 removed AttentionSpec.indexes_kv_by_block_stride and reshaped KVCacheTensor from (size, shared_by) to an explicit layout: layers in L order plus layer_stride and block_stride. The block-stride decision now lives on one physical CacheConfig.kv_cache_layout that the engine core resolves once, and layout.is_block_outermost is the direct successor of the old per-spec flag. Diffusion backends subclass the Omni AttentionBackend ABC, so upstream get_supported_kv_cache_layouts never sees them and the engine core never resolves a layout for them. Add diffusion_kv/layout.py to pin it before get_kv_cache_configs runs, defaulting to LBNHC, which reproduces the indexes_kv_by_block_stride=False that every dense diffusion backend declared. The backend hook is kept as an enforced contract rather than deleted: a backend needing block-stride indexing selects a block-outermost layout, and a contradiction raises instead of mis-addressing K/V. Workers receive only a KVCacheConfig, so the resolved name travels on it (KVCacheConfig.kv_cache_layout) and is adopted into each worker CacheConfig before allocation. Without that, init_kv_cache raises in the worker because get_resolved_kv_cache_layout has nothing to return. init_kv_cache also lost attn_groups and cache_dtype in 0.29; the production caller in model_runner_backend and the test stubs move to the six-argument form. Tensor construction goes through one derivation built on the upstream compute_layout_strides helper, shared by production and tests, so a wrong stride cannot be masked by a parallel test-only copy. Also updates the sizing expectation: 0.29 subtracts the null block that BlockPool permanently holds back before estimating, so sixteen physical pages now yield fifteen usable ones. One review point is addressed only in part, deliberately. build_native_kv_cache_configs now takes indexes_kv_by_block_stride and propagates a block-outermost layout end to end when it is set, with a test covering control plane to worker adoption. Automatic discovery of that requirement from the workers is not added: 0.29 removed the field that used to carry it, every diffusion backend in tree declares False, and inventing an executor RPC for a capability nothing requests would ship untestable plumbing. Until a backend needs it, the seam is explicit and the mismatch raises at worker init rather than corrupting addressing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0146Xu5Sgj4vXxLxvLjV5YPj Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
# 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(),
|
|
)
|