| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the EAGLE speculator's draft attention metadata builder. |
| 4 | |
| 5 | These tests guard the regression where ``_build_uniform_attn_metadata`` did |
| 6 | not populate ``seq_lens_cpu_upper_bound`` on the per-step |
| 7 | ``CommonAttentionMetadata``. Several downstream attention backends and |
| 8 | helpers (``split_decodes_prefills_and_extends``, the MLA indexer, |
| 9 | flex-attention, cross-attention) assert this field is non-None, so |
| 10 | omitting it caused crashes at the start of draft decode for certain |
| 11 | backends (e.g. ``ROCM_AITER_FA`` with eagle/eagle3 spec decode): |
| 12 | |
| 13 | AssertionError: assert common_attn_metadata.seq_lens_cpu_upper_bound is not None |
| 14 | """ |
| 15 | |
| 16 | from types import MethodType, SimpleNamespace |
| 17 | from unittest.mock import patch |
| 18 | |
| 19 | import numpy as np |
| 20 | import pytest |
| 21 | import torch |
| 22 | |
| 23 | from vllm.config.compilation import CUDAGraphMode |
| 24 | from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor |
| 25 | from vllm.v1.worker.gpu.spec_decode import speculator as base_speculator |
| 26 | from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator |
| 27 | |
| 28 | |
| 29 | def _make_fake_speculator( |
| 30 | *, |
| 31 | max_num_reqs: int = 8, |
| 32 | max_num_tokens: int = 16, |
| 33 | max_model_len: int = 1024, |
| 34 | draft_max_seq_len: int = 1024, |
| 35 | ) -> SimpleNamespace: |
| 36 | """Build a fake EagleSpeculator with just the attributes used by |
| 37 | ``_build_uniform_attn_metadata``. We deliberately avoid constructing a |
| 38 | real ``EagleSpeculator`` because that requires a full ``VllmConfig`` |
| 39 | and a draft model. |
| 40 | """ |
| 41 | fake_input_buffers = SimpleNamespace( |
| 42 | query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32), |
| 43 | seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), |
| 44 | dcp_local_seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), |
| 45 | ) |
| 46 | fake_block_tables = SimpleNamespace( |
| 47 | input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)], |
| 48 | slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64), |
| 49 | cp_size=1, |
| 50 | cp_rank=0, |
| 51 | cp_interleave=1, |
| 52 | ) |
| 53 | fake = SimpleNamespace( |
| 54 | arange_np=np.arange(max_num_reqs + 1, dtype=np.int32), |
| 55 | draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool), |
| 56 | block_tables=fake_block_tables, |
| 57 | dcp_size=1, |
| 58 | input_buffers=fake_input_buffers, |
| 59 | attn_groups=[], |
| 60 | kv_cache_config=SimpleNamespace(kv_cache_groups=[]), |
| 61 | max_model_len=max_model_len, |
| 62 | draft_max_seq_len=draft_max_seq_len, |
| 63 | ) |
| 64 | # The uniform wrapper delegates through self; bind the real implementation. |
| 65 | fake._build_attn_metadata = MethodType(EagleSpeculator._build_attn_metadata, fake) |
| 66 | return fake |
| 67 | |
| 68 | |
| 69 | def _run_build( |
| 70 | fake, |
| 71 | *, |
| 72 | num_reqs, |
| 73 | num_reqs_padded, |
| 74 | base, |
| 75 | step, |
| 76 | num_query_per_req=1, |
| 77 | cg_mode=CUDAGraphMode.FULL, |
| 78 | ): |
| 79 | captured: dict[str, object] = {} |
| 80 | |
| 81 | def fake_build_attn_metadata(**kwargs): |
| 82 | captured.update(kwargs) |
| 83 | return {} |
| 84 | |
| 85 | # Request padding only occurs under FULL graphs, where the captured batch |
| 86 | # is a uniform decode of num_reqs_padded * num_query_per_req tokens. |
| 87 | batch_desc = BatchExecutionDescriptor( |
| 88 | cg_mode=cg_mode, |
| 89 | num_tokens=num_reqs_padded * num_query_per_req, |
| 90 | num_reqs=num_reqs_padded, |
| 91 | uniform_token_count=num_query_per_req, |
| 92 | ) |
| 93 | |
| 94 | with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata): |
| 95 | EagleSpeculator._build_uniform_attn_metadata( |
| 96 | fake, |
| 97 | batch_desc=batch_desc, |
| 98 | num_reqs=num_reqs, |
| 99 | num_query_per_req=num_query_per_req, |
| 100 | seq_lens_cpu_upper_bound=base, |
| 101 | step=step, |
| 102 | ) |
| 103 | return captured |
| 104 | |
| 105 | |
| 106 | def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): |
| 107 | """The fix: every per-step ``CommonAttentionMetadata`` carries a non-None |
| 108 | ``seq_lens_cpu_upper_bound`` derived from the target-side upper bound plus |
| 109 | the current draft-step offset. Padded entries are zeroed (matching the |
| 110 | main model runner's convention).""" |
| 111 | fake = _make_fake_speculator() |
| 112 | base = torch.tensor([100, 200, 300, 0], dtype=torch.int32) |
| 113 | |
| 114 | captured = _run_build(fake, num_reqs=3, num_reqs_padded=4, base=base, step=2) |
| 115 | |
| 116 | bound = captured["seq_lens_cpu_upper_bound"] |
| 117 | assert isinstance(bound, torch.Tensor), ( |
| 118 | "seq_lens_cpu_upper_bound must be a tensor, not None" |
| 119 | ) |
| 120 | assert bound.shape == (4,), ( |
| 121 | f"expected shape (num_reqs_padded=4,), got {bound.shape}" |
| 122 | ) |
| 123 | assert bound.device.type == "cpu" |
| 124 | assert bound.dtype == torch.int32 |
| 125 | # base[:num_reqs] + step, padded tail zeroed. |
| 126 | assert torch.equal(bound, torch.tensor([102, 202, 302, 0], dtype=torch.int32)) |
| 127 | |
| 128 | |
| 129 | def test_build_draft_attn_metadata_zeroes_padded_upper_bound_tail(): |
| 130 | """The padded tail of the upper-bound tensor is zeroed, so it stays a |
| 131 | valid tensor of length ``num_reqs_padded`` regardless of padding.""" |
| 132 | fake = _make_fake_speculator() |
| 133 | base = torch.zeros(2, dtype=torch.int32) |
| 134 | |
| 135 | captured = _run_build(fake, num_reqs=1, num_reqs_padded=2, base=base, step=1) |
| 136 | |
| 137 | bound = captured["seq_lens_cpu_upper_bound"] |
| 138 | assert isinstance(bound, torch.Tensor) |
| 139 | assert bound.shape == (2,) |
| 140 | # base is all zeros, so the real entry is 0 + step and the pad entry is 0. |
| 141 | assert torch.equal(bound, torch.tensor([1, 0], dtype=torch.int32)) |
| 142 | |
| 143 | |
| 144 | def test_build_draft_attn_metadata_clamps_to_max_model_len(): |
| 145 | """The per-request upper bound (target bound + step) is clamped to the |
| 146 | model length so it never exceeds the allocated KV range.""" |
| 147 | fake = _make_fake_speculator(max_model_len=1024) |
| 148 | base = torch.tensor([1023, 500], dtype=torch.int32) |
| 149 | |
| 150 | captured = _run_build(fake, num_reqs=2, num_reqs_padded=2, base=base, step=3) |
| 151 | |
| 152 | bound = captured["seq_lens_cpu_upper_bound"] |
| 153 | # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected. |
| 154 | assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32)) |
| 155 | |
| 156 | |
| 157 | @pytest.mark.parametrize("draft_dcp_size", [1, 2]) |
| 158 | def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens(draft_dcp_size): |
| 159 | fake = _make_fake_speculator() |
| 160 | fake.dcp_size = draft_dcp_size |
| 161 | fake.block_tables.cp_size = 2 |
| 162 | fake.block_tables.cp_rank = 1 |
| 163 | fake.block_tables.cp_interleave = 4 |
| 164 | fake.input_buffers.seq_lens[:3] = torch.tensor([5, 9, 16]) |
| 165 | |
| 166 | def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave): |
| 167 | assert draft_dcp_size > 1 |
| 168 | assert seq_lens is fake.input_buffers.seq_lens |
| 169 | assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (3, 2, 1, 4) |
| 170 | out[:num_reqs].copy_(torch.tensor([1, 4, 8], dtype=torch.int32)) |
| 171 | out[num_reqs:].zero_() |
| 172 | return out |
| 173 | |
| 174 | with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare): |
| 175 | captured = _run_build( |
| 176 | fake, |
| 177 | num_reqs=3, |
| 178 | num_reqs_padded=4, |
| 179 | base=torch.tensor([5, 9, 16]), |
| 180 | step=0, |
| 181 | ) |
| 182 | |
| 183 | local = captured["dcp_local_seq_lens"] |
| 184 | if draft_dcp_size == 1: |
| 185 | assert local is None |
| 186 | return |
| 187 | assert isinstance(local, torch.Tensor) |
| 188 | assert local.data_ptr() == fake.input_buffers.dcp_local_seq_lens.data_ptr() |
| 189 | assert local.tolist() == [1, 4, 8, 0] |