| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The draft model must honour ``moe_backend`` from --speculative-config. |
| 4 | |
| 5 | The draft is loaded with the target's VllmConfig, so without an explicit |
| 6 | override it inherits the target's --moe-backend. An MTP head on a quantized |
| 7 | target is typically unquantized, and quantized-only backends reject it, so the |
| 8 | server fails to start rather than falling back. |
| 9 | """ |
| 10 | |
| 11 | from dataclasses import dataclass, field |
| 12 | from types import SimpleNamespace |
| 13 | from unittest.mock import patch |
| 14 | |
| 15 | import pytest |
| 16 | |
| 17 | from vllm.config import LoadConfig, SpeculativeConfig |
| 18 | from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model |
| 19 | |
| 20 | |
| 21 | @dataclass |
| 22 | class _KernelConfig: |
| 23 | moe_backend: str | None = None |
| 24 | |
| 25 | |
| 26 | @dataclass |
| 27 | class _CacheConfig: |
| 28 | cache_dtype: str = "auto" |
| 29 | |
| 30 | |
| 31 | @dataclass |
| 32 | class _SpeculativeConfig: |
| 33 | attention_backend: str | None = None |
| 34 | moe_backend: str | None = None |
| 35 | kv_cache_dtype: str | None = None |
| 36 | draft_model_config: object = None |
| 37 | draft_load_config: object = None |
| 38 | |
| 39 | apply_draft_overrides = SpeculativeConfig.apply_draft_overrides |
| 40 | |
| 41 | |
| 42 | @dataclass |
| 43 | class _VllmConfig: |
| 44 | kernel_config: _KernelConfig |
| 45 | cache_config: _CacheConfig |
| 46 | speculative_config: _SpeculativeConfig |
| 47 | load_config: LoadConfig = field(default_factory=LoadConfig) |
| 48 | |
| 49 | |
| 50 | def _config(target_moe: str, draft_moe: str | None) -> _VllmConfig: |
| 51 | return _VllmConfig( |
| 52 | kernel_config=_KernelConfig(moe_backend=target_moe), |
| 53 | cache_config=_CacheConfig(), |
| 54 | speculative_config=_SpeculativeConfig(moe_backend=draft_moe), |
| 55 | ) |
| 56 | |
| 57 | |
| 58 | class _Captured(Exception): |
| 59 | """Stops load_eagle_model once we hold the config it would have used.""" |
| 60 | |
| 61 | def __init__(self, vllm_config): |
| 62 | self.vllm_config = vllm_config |
| 63 | |
| 64 | |
| 65 | def _capture_draft_config(cfg): |
| 66 | def _fake_get_model(*, vllm_config, model_config): |
| 67 | raise _Captured(vllm_config) |
| 68 | |
| 69 | with ( |
| 70 | patch("vllm.v1.worker.gpu.spec_decode.eagle.utils.get_model", _fake_get_model), |
| 71 | patch( |
| 72 | "vllm.v1.worker.gpu.spec_decode.utils.get_pp_group", |
| 73 | return_value=SimpleNamespace(world_size=1), |
| 74 | ), |
| 75 | pytest.raises(_Captured) as exc, |
| 76 | ): |
| 77 | load_eagle_model(object(), cfg) |
| 78 | return exc.value.vllm_config |
| 79 | |
| 80 | |
| 81 | def test_draft_moe_backend_overrides_the_target(): |
| 82 | """A draft moe_backend must reach the draft's kernel config.""" |
| 83 | used = _capture_draft_config(_config("flashinfer_b12x", "flashinfer_cutlass")) |
| 84 | assert used.kernel_config.moe_backend == "flashinfer_cutlass" |
| 85 | |
| 86 | |
| 87 | def test_draft_inherits_target_when_no_override(): |
| 88 | """No override means the draft still inherits the target, as before.""" |
| 89 | used = _capture_draft_config(_config("flashinfer_b12x", None)) |
| 90 | assert used.kernel_config.moe_backend == "flashinfer_b12x" |
| 91 | |
| 92 | |
| 93 | def test_override_does_not_mutate_the_target_config(): |
| 94 | """The target must keep its own backend after the draft is built.""" |
| 95 | cfg = _config("flashinfer_b12x", "flashinfer_cutlass") |
| 96 | _capture_draft_config(cfg) |
| 97 | assert cfg.kernel_config.moe_backend == "flashinfer_b12x" |