tzhouamandClaude Opus 5
9a399adba8
Scoped pre-commit over the 27 files this rebase touches: ruff format reflow, the vLLM-Omni SPDX copyright line on the new modules, and wording the typos hook rejected. The Python 3.10 mypy hook still reports 24 errors across five of these files. All 24 sit on pre-existing lines; none fall on a line this rebase changed, which the repository documents as the CI-skipped condition. 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-Omni 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(),
|
|
)
|