| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for SpecDecodeBaseProposer.initialize_attn_backend. |
| 4 | |
| 5 | Block tables are stored at kernel-block granularity, so the proposer's |
| 6 | ``block_size`` (used for slot-mapping math) must be the kernel block size, |
| 7 | not the KV cache manager's block size — the two differ when manager blocks |
| 8 | are split for the attention kernel. The value must also be deterministic: |
| 9 | ``_draft_attn_layer_names`` is a set, whose iteration order varies across |
| 10 | processes, so anything derived from iteration order must not leak into |
| 11 | ``block_size``. |
| 12 | """ |
| 13 | |
| 14 | from types import SimpleNamespace |
| 15 | |
| 16 | import pytest |
| 17 | |
| 18 | import vllm.v1.spec_decode.llm_base_proposer as llm_base_proposer |
| 19 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 20 | from vllm.v1.spec_decode.eagle import EagleProposer |
| 21 | |
| 22 | SCHEDULER_BLOCK_SIZE = 256 |
| 23 | KERNEL_BLOCK_SIZE = 64 |
| 24 | |
| 25 | |
| 26 | class _FakeAttentionGroup: |
| 27 | def __init__(self, backend, layer_names, kv_cache_spec, kv_cache_group_id): |
| 28 | self.backend = backend |
| 29 | self.layer_names = list(layer_names) |
| 30 | self.kv_cache_spec = kv_cache_spec |
| 31 | self.kv_cache_group_id = kv_cache_group_id |
| 32 | self.kernel_block_size = None |
| 33 | |
| 34 | def create_metadata_builders(self, vllm_config, device, kernel_block_size=None): |
| 35 | self.kernel_block_size = kernel_block_size |
| 36 | |
| 37 | def get_metadata_builder(self): |
| 38 | return SimpleNamespace(kv_cache_spec=self.kv_cache_spec) |
| 39 | |
| 40 | |
| 41 | def _make_proposer( |
| 42 | monkeypatch: pytest.MonkeyPatch, layer_names: set[str] |
| 43 | ) -> EagleProposer: |
| 44 | fake_layers = {} |
| 45 | for name in layer_names: |
| 46 | backend = SimpleNamespace(full_cls_name=lambda: "FakeBackend") |
| 47 | fake_layers[name] = SimpleNamespace( |
| 48 | get_attn_backend=lambda backend=backend: backend |
| 49 | ) |
| 50 | monkeypatch.setattr( |
| 51 | llm_base_proposer, "get_layers_from_vllm_config", lambda *a, **k: fake_layers |
| 52 | ) |
| 53 | monkeypatch.setattr(llm_base_proposer, "AttentionGroup", _FakeAttentionGroup) |
| 54 | |
| 55 | proposer = EagleProposer.__new__(EagleProposer) |
| 56 | proposer.vllm_config = None |
| 57 | proposer.device = None |
| 58 | proposer._draft_attn_layer_names = set(layer_names) |
| 59 | proposer.kv_cache_gid = -1 |
| 60 | proposer.draft_attn_groups = [] |
| 61 | proposer.block_size = -1 |
| 62 | return proposer |
| 63 | |
| 64 | |
| 65 | def _make_kv_cache_config(layer_names: set[str]) -> KVCacheConfig: |
| 66 | spec = SimpleNamespace(block_size=SCHEDULER_BLOCK_SIZE) |
| 67 | group = SimpleNamespace(layer_names=list(layer_names), kv_cache_spec=spec) |
| 68 | return SimpleNamespace(kv_cache_groups=[group]) # type: ignore[return-value] |
| 69 | |
| 70 | |
| 71 | def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch): |
| 72 | """The proposer's slot-mapping math runs against the kernel-granularity |
| 73 | block table, so block_size must come from kernel_block_sizes.""" |
| 74 | layer_names = {"draft.0.self_attn.attn"} |
| 75 | proposer = _make_proposer(monkeypatch, layer_names) |
| 76 | |
| 77 | proposer.initialize_attn_backend( |
| 78 | _make_kv_cache_config(layer_names), |
| 79 | kernel_block_sizes=[KERNEL_BLOCK_SIZE], |
| 80 | ) |
| 81 | |
| 82 | assert proposer.block_size == KERNEL_BLOCK_SIZE |
| 83 | assert proposer.block_size != SCHEDULER_BLOCK_SIZE |
| 84 | # The metadata builder keeps receiving the kernel block size as well. |
| 85 | group = proposer.draft_attn_groups[0] |
| 86 | assert isinstance(group, _FakeAttentionGroup) |
| 87 | assert group.kernel_block_size == KERNEL_BLOCK_SIZE |
| 88 | |
| 89 | |
| 90 | def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch): |
| 91 | layer_names = {"draft.0.self_attn.attn"} |
| 92 | proposer = _make_proposer(monkeypatch, layer_names) |
| 93 | |
| 94 | proposer.initialize_attn_backend( |
| 95 | _make_kv_cache_config(layer_names), kernel_block_sizes=None |
| 96 | ) |
| 97 | |
| 98 | assert proposer.block_size == SCHEDULER_BLOCK_SIZE |
| 99 | |
| 100 | |
| 101 | def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch): |
| 102 | """_draft_attn_layer_names is a set; the attention groups built from it |
| 103 | must not depend on its (process-random) iteration order.""" |
| 104 | layer_names = {"draft.c.attn", "draft.a.attn", "draft.b.attn"} |
| 105 | expected_order = sorted(layer_names) |
| 106 | |
| 107 | for insertion_order in (expected_order, expected_order[::-1]): |
| 108 | proposer = _make_proposer(monkeypatch, set(insertion_order)) |
| 109 | proposer.initialize_attn_backend( |
| 110 | _make_kv_cache_config(set(insertion_order)), |
| 111 | kernel_block_sizes=[KERNEL_BLOCK_SIZE], |
| 112 | ) |
| 113 | assert len(proposer.draft_attn_groups) == 1 |
| 114 | assert proposer.draft_attn_groups[0].layer_names == expected_order |
| 115 | assert proposer.block_size == KERNEL_BLOCK_SIZE |