| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """End-to-end checks that ``backend_per_kind`` selects the requested attention |
| 4 | backend for each KV-cache group at runtime. |
| 5 | |
| 6 | Uses ``google/gemma-3-1b-it``, which interleaves full-attention and |
| 7 | sliding-window layers, so the model produces separate ``full_attention`` and |
| 8 | ``sliding_window`` KV-cache groups. |
| 9 | """ |
| 10 | |
| 11 | import pytest |
| 12 | |
| 13 | from vllm import LLM |
| 14 | from vllm.config.attention import AttentionConfig |
| 15 | from vllm.platforms import current_platform |
| 16 | |
| 17 | MODEL = "google/gemma-3-1b-it" |
| 18 | |
| 19 | |
| 20 | def _collect_group_backends(worker) -> list[tuple[str, str]]: |
| 21 | """Runs on the worker: returns (spec_kind, backend_name) per attn group.""" |
| 22 | from vllm.v1.kv_cache_interface import get_kv_cache_spec_kind |
| 23 | |
| 24 | out: list[tuple[str, str]] = [] |
| 25 | for kv_group in worker.model_runner.attn_groups: |
| 26 | for attn_group in kv_group: |
| 27 | kind = get_kv_cache_spec_kind(attn_group.kv_cache_spec) |
| 28 | out.append((kind.value, attn_group.backend.get_name())) |
| 29 | return out |
| 30 | |
| 31 | |
| 32 | @pytest.mark.skipif( |
| 33 | not current_platform.is_cuda(), reason="backend names are CUDA-specific" |
| 34 | ) |
| 35 | @pytest.mark.parametrize( |
| 36 | "backend_per_kind", |
| 37 | [ |
| 38 | {"full_attention": "FLASH_ATTN", "sliding_window": "TRITON_ATTN"}, |
| 39 | # Swapped, to prove the mapping is causal rather than the default. |
| 40 | {"full_attention": "TRITON_ATTN", "sliding_window": "FLASH_ATTN"}, |
| 41 | ], |
| 42 | ) |
| 43 | def test_backend_per_kind_splits_groups(backend_per_kind, monkeypatch): |
| 44 | # collective_rpc ships the callable to the EngineCore subprocess; the |
| 45 | # secure msgpack encoder can't serialize functions, so opt into the |
| 46 | # pickle fallback (same pattern as test_pooling_chunked_prefill). |
| 47 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 48 | llm = LLM( |
| 49 | model=MODEL, |
| 50 | attention_config=AttentionConfig(backend_per_kind=backend_per_kind), |
| 51 | enforce_eager=True, |
| 52 | max_model_len=2048, |
| 53 | gpu_memory_utilization=0.4, |
| 54 | ) |
| 55 | |
| 56 | group_backends = llm.llm_engine.collective_rpc(_collect_group_backends)[0] |
| 57 | kinds = {kind for kind, _ in group_backends} |
| 58 | |
| 59 | # gemma3 must actually split into both kinds for this test to be meaningful. |
| 60 | assert "full_attention" in kinds |
| 61 | assert "sliding_window" in kinds |
| 62 | |
| 63 | for kind, backend_name in group_backends: |
| 64 | if kind in backend_per_kind: |
| 65 | assert backend_name == backend_per_kind[kind], ( |
| 66 | f"{kind} group used {backend_name}, expected {backend_per_kind[kind]}" |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | @pytest.mark.skipif( |
| 71 | not current_platform.is_cuda(), reason="backend names are CUDA-specific" |
| 72 | ) |
| 73 | def test_backend_per_kind_overrides_global_backend(monkeypatch): |
| 74 | """A per-kind entry wins over the global ``backend`` for its kind; other |
| 75 | kinds fall back to the global backend.""" |
| 76 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 77 | llm = LLM( |
| 78 | model=MODEL, |
| 79 | attention_config=AttentionConfig( |
| 80 | backend="FLASH_ATTN", |
| 81 | backend_per_kind={"sliding_window": "TRITON_ATTN"}, |
| 82 | ), |
| 83 | enforce_eager=True, |
| 84 | max_model_len=2048, |
| 85 | gpu_memory_utilization=0.4, |
| 86 | ) |
| 87 | |
| 88 | group_backends = llm.llm_engine.collective_rpc(_collect_group_backends)[0] |
| 89 | |
| 90 | for kind, backend_name in group_backends: |
| 91 | if kind == "sliding_window": |
| 92 | assert backend_name == "TRITON_ATTN" |
| 93 | elif kind == "full_attention": |
| 94 | assert backend_name == "FLASH_ATTN" |