| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import numpy as np |
| 4 | |
| 5 | from vllm.config import ( |
| 6 | ModelConfig, |
| 7 | SpeculativeConfig, |
| 8 | VllmConfig, |
| 9 | ) |
| 10 | from vllm.v1.spec_decode.ngram_proposer import ( |
| 11 | NgramProposer, |
| 12 | _find_longest_matched_ngram_and_propose_tokens, |
| 13 | ) |
| 14 | |
| 15 | |
| 16 | def test_find_longest_matched_ngram_and_propose_tokens(): |
| 17 | tokens = np.array([1, 2, 3, 4, 1, 2, 3, 5, 6]) |
| 18 | result = _find_longest_matched_ngram_and_propose_tokens( |
| 19 | origin_tokens=tokens, min_ngram=2, max_ngram=2, max_model_len=1024, k=2 |
| 20 | ) |
| 21 | assert len(result) == 0 |
| 22 | |
| 23 | tokens = np.array([1, 2, 3, 4, 1, 2, 3]) |
| 24 | np.testing.assert_array_equal( |
| 25 | _find_longest_matched_ngram_and_propose_tokens( |
| 26 | origin_tokens=tokens, min_ngram=2, max_ngram=2, max_model_len=1024, k=3 |
| 27 | ), |
| 28 | np.array([4, 1, 2]), |
| 29 | ) |
| 30 | np.testing.assert_array_equal( |
| 31 | _find_longest_matched_ngram_and_propose_tokens( |
| 32 | origin_tokens=tokens, min_ngram=2, max_ngram=2, max_model_len=1024, k=2 |
| 33 | ), |
| 34 | np.array([4, 1]), |
| 35 | ) |
| 36 | np.testing.assert_array_equal( |
| 37 | _find_longest_matched_ngram_and_propose_tokens( |
| 38 | origin_tokens=tokens, min_ngram=1, max_ngram=1, max_model_len=1024, k=3 |
| 39 | ), |
| 40 | np.array([4, 1, 2]), |
| 41 | ) |
| 42 | np.testing.assert_array_equal( |
| 43 | _find_longest_matched_ngram_and_propose_tokens( |
| 44 | origin_tokens=tokens, min_ngram=1, max_ngram=1, max_model_len=1024, k=2 |
| 45 | ), |
| 46 | np.array([4, 1]), |
| 47 | ) |
| 48 | |
| 49 | tokens = np.array([1, 3, 6, 2, 3, 4, 1, 2, 3]) |
| 50 | np.testing.assert_array_equal( |
| 51 | _find_longest_matched_ngram_and_propose_tokens( |
| 52 | origin_tokens=tokens, min_ngram=2, max_ngram=2, max_model_len=1024, k=3 |
| 53 | ), |
| 54 | np.array([4, 1, 2]), |
| 55 | ) |
| 56 | # Return on the first match |
| 57 | np.testing.assert_array_equal( |
| 58 | _find_longest_matched_ngram_and_propose_tokens( |
| 59 | origin_tokens=tokens, min_ngram=1, max_ngram=1, max_model_len=1024, k=2 |
| 60 | ), |
| 61 | np.array([6, 2]), |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | def test_ngram_proposer(): |
| 66 | def get_ngram_proposer(min_n: int, max_n: int, k: int) -> NgramProposer: |
| 67 | # Dummy model config. Just to set max_model_len. |
| 68 | model_config = ModelConfig(model="facebook/opt-125m") |
| 69 | return NgramProposer( |
| 70 | vllm_config=VllmConfig( |
| 71 | model_config=model_config, |
| 72 | speculative_config=SpeculativeConfig( |
| 73 | prompt_lookup_min=min_n, |
| 74 | prompt_lookup_max=max_n, |
| 75 | num_speculative_tokens=k, |
| 76 | method="ngram", |
| 77 | ), |
| 78 | ) |
| 79 | ) |
| 80 | |
| 81 | # No match. |
| 82 | token_ids_cpu = np.array([[1, 2, 3, 4, 5]]) |
| 83 | result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose( |
| 84 | num_speculative_tokens=2, |
| 85 | sampled_token_ids=[[0]], |
| 86 | num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), |
| 87 | token_ids_cpu=token_ids_cpu, |
| 88 | ) |
| 89 | assert len(result[0]) == 0 |
| 90 | |
| 91 | # No match for 4-gram. |
| 92 | token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]]) |
| 93 | result = get_ngram_proposer(min_n=4, max_n=4, k=2).propose( |
| 94 | num_speculative_tokens=2, |
| 95 | sampled_token_ids=[[0]], |
| 96 | num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), |
| 97 | token_ids_cpu=token_ids_cpu, |
| 98 | ) |
| 99 | assert len(result[0]) == 0 |
| 100 | |
| 101 | # No match for 4-gram but match for 3-gram. |
| 102 | token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]]) |
| 103 | result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose( |
| 104 | num_speculative_tokens=2, |
| 105 | sampled_token_ids=[[0]], |
| 106 | num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), |
| 107 | token_ids_cpu=token_ids_cpu, |
| 108 | ) |
| 109 | assert np.array_equal(result, np.array([[4, 1]])) |
| 110 | |
| 111 | # Match for both 4-gram and 3-gram. |
| 112 | # In this case, the proposer should return the 4-gram match. |
| 113 | token_ids_cpu = np.array([[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]]) |
| 114 | result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose( |
| 115 | num_speculative_tokens=2, |
| 116 | sampled_token_ids=[[0]], |
| 117 | num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), |
| 118 | token_ids_cpu=token_ids_cpu, |
| 119 | ) |
| 120 | assert np.array_equal(result, np.array([[1, 2]])) # Not [5, 1]] |
| 121 | |
| 122 | # Match for 2-gram and 3-gram, but not 4-gram. |
| 123 | token_ids_cpu = np.array([[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]]) |
| 124 | result = get_ngram_proposer(min_n=2, max_n=4, k=2).propose( |
| 125 | num_speculative_tokens=2, |
| 126 | sampled_token_ids=[[0]], |
| 127 | num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), |
| 128 | token_ids_cpu=token_ids_cpu, |
| 129 | ) |
| 130 | assert np.array_equal(result, np.array([[1, 2]])) # Not [5, 2]] |
| 131 | |
| 132 | # Multiple 3-gram matched, but always pick the first one. |
| 133 | token_ids_cpu = np.array([[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]]) |
| 134 | result = get_ngram_proposer(min_n=3, max_n=3, k=2).propose( |
| 135 | num_speculative_tokens=2, |
| 136 | sampled_token_ids=[[0]], |
| 137 | num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), |
| 138 | token_ids_cpu=token_ids_cpu, |
| 139 | ) |
| 140 | assert np.array_equal(result, np.array([[100, 1]])) |
| 141 | |
| 142 | # check empty input |
| 143 | token_ids_cpu = np.array([[]]) |
| 144 | result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose( |
| 145 | num_speculative_tokens=2, |
| 146 | sampled_token_ids=[[0]], |
| 147 | num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), |
| 148 | token_ids_cpu=token_ids_cpu, |
| 149 | ) |
| 150 | assert len(result[0]) == 0 |
| 151 | |
| 152 | # check multibatch input |
| 153 | # first request has 5 tokens and a match |
| 154 | # second request has 3 tokens and no match. Padded with -1 for max len 5 |
| 155 | token_ids_cpu = np.array([[1, 2, 3, 1, 2], [4, 5, 6, -1, -1]]) |
| 156 | result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose( |
| 157 | num_speculative_tokens=2, |
| 158 | sampled_token_ids=[[0], [1]], |
| 159 | num_tokens_no_spec=np.array([5, 3]), |
| 160 | token_ids_cpu=token_ids_cpu, |
| 161 | ) |
| 162 | assert len(result[0]) == 2 |
| 163 | assert np.array_equal(result[0], np.array([3, 1])) |
| 164 | assert np.array_equal(result[1], np.array([])) |
| 165 | |
| 166 | # Test non-contiguous indices: requests 0 and 2 need proposals, |
| 167 | # request 1 is in prefill |
| 168 | proposer = get_ngram_proposer(min_n=2, max_n=2, k=2) |
| 169 | max_model_len = 20 |
| 170 | token_ids_cpu = np.zeros((3, max_model_len), dtype=np.int32) |
| 171 | token_ids_cpu[0, :5] = [1, 2, 3, 1, 2] |
| 172 | token_ids_cpu[1, :3] = [4, 5, 6] |
| 173 | token_ids_cpu[2, :5] = [7, 8, 9, 7, 8] |
| 174 | num_tokens_no_spec = np.array([5, 3, 5], dtype=np.int32) |
| 175 | sampled_token_ids = [[2], [], [8]] # Empty list for request 1 simulates prefill |
| 176 | result = proposer.propose( |
| 177 | num_speculative_tokens=2, |
| 178 | sampled_token_ids=sampled_token_ids, |
| 179 | num_tokens_no_spec=num_tokens_no_spec, |
| 180 | token_ids_cpu=token_ids_cpu, |
| 181 | ) |
| 182 | assert len(result) == 3 |
| 183 | assert np.array_equal(result[0], [3, 1]) |
| 184 | assert len(result[1]) == 0 |
| 185 | assert np.array_equal(result[2], [9, 7]) |
| 186 | # Verify internal arrays written to correct indices |
| 187 | assert proposer.valid_ngram_num_drafts[0] == 2 |
| 188 | assert proposer.valid_ngram_num_drafts[1] == 0 |
| 189 | assert proposer.valid_ngram_num_drafts[2] == 2 |
| 190 | assert np.array_equal(proposer.valid_ngram_draft[0, :2], [3, 1]) |
| 191 | assert np.array_equal(proposer.valid_ngram_draft[2, :2], [9, 7]) |
| 192 | |
| 193 | # test if 0 threads available: can happen if TP size > CPU count |
| 194 | ngram_proposer = get_ngram_proposer(min_n=2, max_n=2, k=2) |
| 195 | ngram_proposer.num_numba_thread_available = 0 |
| 196 | # set max_model_len to 2 * threshold to ensure multithread is used |
| 197 | num_tokens_threshold = ngram_proposer.num_tokens_threshold |
| 198 | ngram_proposer.max_model_len = 2 * num_tokens_threshold |
| 199 | # using multibatch test |
| 200 | middle_integer = num_tokens_threshold // 2 |
| 201 | input_1 = [_ for _ in range(num_tokens_threshold)] |
| 202 | input_1 += [middle_integer, middle_integer + 1] |
| 203 | input_2 = [-1] * len(input_1) |
| 204 | input_2[:3] = [4, 5, 6] |
| 205 | token_ids_cpu = np.array([input_1, input_2]) |
| 206 | result = ngram_proposer.propose( |
| 207 | num_speculative_tokens=2, |
| 208 | sampled_token_ids=[[0], [1]], |
| 209 | num_tokens_no_spec=np.array([len(input_1), 3]), |
| 210 | token_ids_cpu=token_ids_cpu, |
| 211 | ) |
| 212 | assert len(result[0]) == 2 |
| 213 | assert np.array_equal(result[0], np.array([middle_integer + 2, middle_integer + 3])) |
| 214 | assert np.array_equal(result[1], np.array([])) |