Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 61 additions & 0 deletions tests/v1/worker/test_gpu_model_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -486,6 +486,67 @@ def test_preferred_block_size_rejects_backends_with_no_common_size():
Platform._preferred_block_size_for_backends(classes, 16, None)


@pytest.mark.parametrize("cpu", [False, True])
def test_alignment_rejected_by_a_sibling_backend_raises(monkeypatch, cpu):
# Alignment that lands on a size an exact-size sibling rejects must fail
# here, not later in select_common_block_size().
backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)]
monkeypatch.setattr(
Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends)
)

def bump(cls, vllm_config, backend_cls):
vllm_config.cache_config.block_size = 32

monkeypatch.setattr(Platform, "_align_hybrid_block_size", classmethod(bump))
cache_config = SimpleNamespace(
block_size=16, user_specified_block_size=True, kv_cache_dtype_skip_layers=None
)
vllm_config = SimpleNamespace(
cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True)
)
from vllm.platforms.cpu import CpuPlatform # overrides the method; must check too

platform = CpuPlatform if cpu else Platform
with pytest.raises(ValueError, match="MOCK_EXACT"):
platform.update_block_size_for_backend(vllm_config)


def test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises(monkeypatch):
# XPU's GDN rounding runs after super()'s check, so it must check again.
try:
from vllm.platforms.xpu import XPUPlatform
except ImportError:
pytest.skip("vllm_xpu_kernels not importable outside an XPU stack")

backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)]
monkeypatch.setattr(
Platform, "update_block_size_for_backend", classmethod(lambda cls, c: None)
)
monkeypatch.setattr(
Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends)
)
gdn_layer = SimpleNamespace(
get_attn_backend=lambda: SimpleNamespace(get_name=lambda: "GDN_ATTN")
)
monkeypatch.setattr(
"vllm.config.vllm.get_layers_from_vllm_config",
lambda vllm_config, layer_type: {"gdn": gdn_layer},
)
cache_config = SimpleNamespace(
block_size=16,
mamba_cache_mode="none",
mamba_page_size_padded=None,
user_specified_block_size=True,
kv_cache_dtype_skip_layers=None,
)
vllm_config = SimpleNamespace(
cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True)
)
with pytest.raises(ValueError, match="MOCK_EXACT"):
XPUPlatform.update_block_size_for_backend(vllm_config)


def test_set_active_mm_loras_builds_tower_and_connector_mappings():
model = Mock()
model.get_mm_lora_token_counts.side_effect = (
Expand Down
2 changes: 2 additions & 0 deletions vllm/platforms/cpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -523,7 +523,9 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
if not backend_classes:
return

pre_block_size = vllm_config.cache_config.block_size
cls._align_hybrid_block_size(vllm_config, backend_classes[0])
cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size)
Comment on lines +526 to +528

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we support xpu as well?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, done in d333dfd. XPUPlatform.update_block_size_for_backend calls super() first, so the check ran before its GDN step rounds the block size up to a multiple of 64. That rounding was unchecked, so it now runs the same check on its result. Added test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises, which skips without vllm_xpu_kernels, like the existing XPU tests. Also rebased onto current main. Still +99/-0.


@classmethod
def discover_numa_topology(cls) -> list[list[int]]:
Expand Down
31 changes: 31 additions & 0 deletions vllm/platforms/interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -698,6 +698,7 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:

# Phase 2: Align block/mamba sizes for hybrid models
# (may override user settings).
pre_block_size = cache_config.block_size
if model_config.is_hybrid:
cls._align_hybrid_block_size(vllm_config, backend_classes[0])

Expand All @@ -706,6 +707,36 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
# May override the user's --block-size.
if cache_config.kv_cache_dtype_skip_layers:
cls._align_heterogeneous_kv_block_size(vllm_config, backend_classes[0])
cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size)

@classmethod
def _check_aligned_block_size(
cls,
vllm_config: "VllmConfig",
backend_classes: "list[type[AttentionBackend]]",
pre_block_size: int,
) -> None:
"""Fail now, not at worker setup, if alignment left a rejected size.

Alignment rounds by the first backend's kernel block only, so it can
land on a size a sibling backend rejects.
"""
from vllm.config.vllm import set_current_vllm_config

block_size = vllm_config.cache_config.block_size
if block_size == pre_block_size:
return
with set_current_vllm_config(vllm_config):
rejecting = [
b.get_name()
for b in backend_classes
if not b.supports_block_size(block_size)
]
if rejecting:
raise ValueError(
f"Aligned KV cache block size {block_size} is not supported by "
f"attention backend(s): {', '.join(rejecting)}."
)

@classmethod
def _align_heterogeneous_kv_block_size(
Expand Down
5 changes: 5 additions & 0 deletions vllm/platforms/xpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -452,6 +452,7 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
if new_block_size == cache_config.block_size:
return

pre_block_size = cache_config.block_size
if cache_config.mamba_cache_mode == "align":
cache_config.mamba_block_size = new_block_size
original_mamba_page_size_padded = cache_config.mamba_page_size_padded
Expand All @@ -474,6 +475,10 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
original_mamba_page_size_padded,
cache_config.mamba_page_size_padded,
)
# This rounding runs after super()'s check, so check its result too.
cls._check_aligned_block_size(
vllm_config, cls._find_non_ssm_backends(vllm_config), pre_block_size
)

@classmethod
def support_hybrid_kv_cache(cls) -> bool:
Expand Down
Loading