| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The draft must honour ``attention_backend`` from --speculative-config. |
| 4 | |
| 5 | ``init_attn_backend`` reads the backend off each constructed layer, so these |
| 6 | assert on the config ``load_eagle_model`` hands to ``get_model``. |
| 7 | """ |
| 8 | |
| 9 | from dataclasses import dataclass, field |
| 10 | from types import SimpleNamespace |
| 11 | from unittest.mock import patch |
| 12 | |
| 13 | import pytest |
| 14 | |
| 15 | from vllm.config import LoadConfig, SpeculativeConfig |
| 16 | from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model |
| 17 | |
| 18 | |
| 19 | @dataclass |
| 20 | class _AttentionConfig: |
| 21 | backend: str | None = None |
| 22 | |
| 23 | |
| 24 | @dataclass |
| 25 | class _KernelConfig: |
| 26 | moe_backend: str | None = None |
| 27 | |
| 28 | |
| 29 | @dataclass |
| 30 | class _CacheConfig: |
| 31 | cache_dtype: str = "auto" |
| 32 | |
| 33 | |
| 34 | @dataclass |
| 35 | class _SpeculativeConfig: |
| 36 | attention_backend: str | None = None |
| 37 | moe_backend: str | None = None |
| 38 | kv_cache_dtype: str | None = None |
| 39 | draft_model_config: object = None |
| 40 | draft_load_config: object = None |
| 41 | |
| 42 | apply_draft_overrides = SpeculativeConfig.apply_draft_overrides |
| 43 | |
| 44 | |
| 45 | @dataclass |
| 46 | class _VllmConfig: |
| 47 | attention_config: _AttentionConfig |
| 48 | kernel_config: _KernelConfig |
| 49 | cache_config: _CacheConfig |
| 50 | speculative_config: _SpeculativeConfig |
| 51 | load_config: LoadConfig = field(default_factory=LoadConfig) |
| 52 | |
| 53 | |
| 54 | def _config(target_backend: str, draft_backend: str | None) -> _VllmConfig: |
| 55 | return _VllmConfig( |
| 56 | attention_config=_AttentionConfig(backend=target_backend), |
| 57 | kernel_config=_KernelConfig(), |
| 58 | cache_config=_CacheConfig(), |
| 59 | speculative_config=_SpeculativeConfig(attention_backend=draft_backend), |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | class _Captured(Exception): |
| 64 | def __init__(self, vllm_config): |
| 65 | self.vllm_config = vllm_config |
| 66 | |
| 67 | |
| 68 | def _capture_draft_config(cfg): |
| 69 | def _fake_get_model(*, vllm_config, model_config): |
| 70 | raise _Captured(vllm_config) |
| 71 | |
| 72 | with ( |
| 73 | patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model), |
| 74 | patch( |
| 75 | "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group", |
| 76 | return_value=SimpleNamespace(world_size=1), |
| 77 | ), |
| 78 | pytest.raises(_Captured) as exc, |
| 79 | ): |
| 80 | load_eagle_model(object(), cfg) |
| 81 | return exc.value.vllm_config |
| 82 | |
| 83 | |
| 84 | def test_draft_attention_backend_overrides_the_target(): |
| 85 | used = _capture_draft_config(_config("FLASHINFER", "TRITON_ATTN")) |
| 86 | assert used.attention_config.backend == "TRITON_ATTN" |
| 87 | |
| 88 | |
| 89 | def test_unset_leaves_the_target_backend_in_place(): |
| 90 | """Clearing it lets the draft autoselect a KV layout the target lacks.""" |
| 91 | cfg = _config("FLEX_ATTENTION", None) |
| 92 | used = _capture_draft_config(cfg) |
| 93 | assert used is cfg |
| 94 | assert used.attention_config.backend == "FLEX_ATTENTION" |
| 95 | |
| 96 | |
| 97 | def test_override_does_not_mutate_the_target_config(): |
| 98 | cfg = _config("FLASHINFER", "TRITON_ATTN") |
| 99 | _capture_draft_config(cfg) |
| 100 | assert cfg.attention_config.backend == "FLASHINFER" |