A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Core] [N-gram SD Optimization][1/n] Propose tokens with a single KMP (#22437)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Aug 13, 2025 at 14:44 UTC
31a500c86fb38417a2696d516a8ea1a642a5df06
6 files changed
+389
-207
benchmarks/benchmark_block_pool.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import gc
4
+
5
+from tabulate import tabulate
6
+
7
+from benchmark_utils import TimeCollector
8
+from vllm.utils import FlexibleArgumentParser
9
+from vllm.v1.core.block_pool import BlockPool
10
+
11
+
12
+def main(args):
13
+ rows = []
14
+ for allocate_block in args.allocate_blocks:
15
+ # Enforce a GC collect ahead to minimize the impact among runs
16
+ gc.collect()
17
+ block_pool = BlockPool(num_gpu_blocks=args.num_gpu_blocks, enable_caching=True)
18
+
19
+ get_blocks_times = TimeCollector(TimeCollector.US)
20
+ free_blocks_times = TimeCollector(TimeCollector.US)
21
+ for _ in range(args.num_iteration):
22
+ with get_blocks_times:
23
+ blocks = block_pool.get_new_blocks(allocate_block)
24
+ with free_blocks_times:
25
+ block_pool.free_blocks(blocks)
26
+
27
+ rows.append(
28
+ [get_blocks_times.cnt, args.num_gpu_blocks, allocate_block]
29
+ + get_blocks_times.dump_avg_max()
30
+ + free_blocks_times.dump_avg_max()
31
+ )
32
+
33
+ print(
34
+ tabulate(
35
+ rows,
36
+ headers=[
37
+ "Iterations",
38
+ "Total\nBlocks",
39
+ "Allocated\nBlocks",
40
+ "Get Blocks\nAvg (us)",
41
+ "Get Blocks\nMax (us)",
42
+ "Free Blocks\nAvg (us)",
43
+ "Free Blocks\nMax (us)",
44
+ ],
45
+ tablefmt="grid",
46
+ floatfmt=".3f",
47
+ )
48
+ )
49
+
50
+
51
+def invoke_main() -> None:
52
+ parser = FlexibleArgumentParser(
53
+ description="Benchmark the performance of BlockPool for KV Cache."
54
+ )
55
+ parser.add_argument("--num-gpu-blocks", type=int, default=100000)
56
+ parser.add_argument(
57
+ "--num-iteration",
58
+ type=int,
59
+ default=1000,
60
+ help="Number of iterations to run to stablize final data readings",
61
+ )
62
+ parser.add_argument(
63
+ "--allocate-blocks",
64
+ type=int,
65
+ nargs="*",
66
+ default=[10, 50, 100, 500, 1000],
67
+ help="Number of blocks to allocate",
68
+ )
69
+ args = parser.parse_args()
70
+ main(args)
71
+
72
+
73
+if __name__ == "__main__":
74
+ invoke_main() # pragma: no cover
benchmarks/benchmark_ngram_proposer.py
new
+112
@@ -0,0 +1,112 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import gc
4
+
5
+import numpy as np
6
+from tabulate import tabulate
7
+
8
+from benchmark_utils import TimeCollector
9
+from vllm.config import ModelConfig, SpeculativeConfig, VllmConfig
10
+from vllm.utils import FlexibleArgumentParser
11
+from vllm.v1.spec_decode.ngram_proposer import NgramProposer
12
+
13
+
14
+def main(args):
15
+ rows = []
16
+ for max_ngram in args.max_ngram:
17
+ collector = TimeCollector(TimeCollector.US)
18
+
19
+ model_config = ModelConfig(
20
+ model="facebook/opt-125m",
21
+ task="generate",
22
+ max_model_len=args.num_token + args.num_spec_token,
23
+ tokenizer="facebook/opt-125m",
24
+ tokenizer_mode="auto",
25
+ dtype="auto",
26
+ seed=None,
27
+ trust_remote_code=False,
28
+ )
29
+ proposer = NgramProposer(
30
+ vllm_config=VllmConfig(
31
+ model_config=model_config,
32
+ speculative_config=SpeculativeConfig(
33
+ prompt_lookup_min=args.min_ngram,
34
+ prompt_lookup_max=max_ngram,
35
+ num_speculative_tokens=args.num_spec_token,
36
+ method="ngram",
37
+ ),
38
+ )
39
+ )
40
+
41
+ # Warm up
42
+ proposer.propose(np.random.randint(0, 20, (args.num_token,)))
43
+
44
+ gc.collect()
45
+ for _ in range(args.num_iteration):
46
+ tokens = np.random.randint(0, 20, (args.num_req, args.num_token))
47
+ with collector:
48
+ for i in range(args.num_req):
49
+ proposer.propose(tokens[i, :])
50
+ rows.append(
51
+ [args.num_req, args.num_token, args.min_ngram, max_ngram]
52
+ + collector.dump_avg_max()
53
+ )
54
+
55
+ print(
56
+ tabulate(
57
+ rows,
58
+ headers=[
59
+ "# Request",
60
+ "# Token",
61
+ "Min Ngram",
62
+ "Max Ngram",
63
+ "Avg (us)",
64
+ "Max (us)",
65
+ ],
66
+ tablefmt="grid",
67
+ floatfmt=".3f",
68
+ )
69
+ )
70
+
71
+
72
+def invoke_main() -> None:
73
+ parser = FlexibleArgumentParser(
74
+ description="Benchmark the performance of N-gram speculative decode drafting"
75
+ )
76
+ parser.add_argument(
77
+ "--num-iteration",
78
+ type=int,
79
+ default=100,
80
+ help="Number of iterations to run to stablize final data readings",
81
+ )
82
+ parser.add_argument(
83
+ "--num-req", type=int, default=128, help="Number of requests in the batch"
84
+ )
85
+ parser.add_argument(
86
+ "--num-token", type=int, default=1500, help="Number of tokens for each request"
87
+ )
88
+ parser.add_argument(
89
+ "--min-ngram",
90
+ type=int,
91
+ default=3,
92
+ help="Minimum n-gram to match",
93
+ )
94
+ parser.add_argument(
95
+ "--max-ngram",
96
+ type=int,
97
+ nargs="*",
98
+ default=[5, 7, 10, 15, 20],
99
+ help="Maximum n-gram to match",
100
+ )
101
+ parser.add_argument(
102
+ "--num-spec-token",
103
+ type=int,
104
+ default=3,
105
+ help="Number of speculative tokens to generate",
106
+ )
107
+ args = parser.parse_args()
108
+ main(args)
109
+
110
+
111
+if __name__ == "__main__":
112
+ invoke_main() # pragma: no cover
benchmarks/benchmark_utils.py
+53
-2
@@ -1,11 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
import argparse
4
import json
5
import math
6
import os
8
-from typing import Any
7
+import time
8
+from types import TracebackType
9
+from typing import Any, Optional, Union
10
11
12
def convert_to_pytorch_benchmark_format(
@@ -72,3 +73,53 @@ def write_to_json(filename: str, records: list) -> None:
73
cls=InfEncoder,
74
default=lambda o: f"<{type(o).__name__} object is not JSON serializable>",
75
)
76
+
77
+
78
+# Collect time and generate time metrics
79
+#
80
+# Example Usage:
81
+# collector = TimeCollector(TimeCollector.US)
82
+# for _ in range(total_iteration):
83
+# with collector:
84
+# ...
85
+# collector.dump_avg_max()
86
+class TimeCollector:
87
+ NS: int = 1
88
+ US: int = NS * 1000
89
+ MS: int = US * 1000
90
+ S: int = MS * 1000
91
+
92
+ def __init__(self, scale: int) -> None:
93
+ self.cnt: int = 0
94
+ self._sum: int = 0
95
+ self._max: Optional[int] = None
96
+ self.scale = scale
97
+ self.start_time: int = time.monotonic_ns()
98
+
99
+ def collect(self, v: int) -> None:
100
+ self.cnt += 1
101
+ self._sum += v
102
+ if self._max is None:
103
+ self._max = v
104
+ else:
105
+ self._max = max(self._max, v)
106
+
107
+ def avg(self) -> Union[float, str]:
108
+ return self._sum * 1.0 / self.cnt / self.scale if self.cnt > 0 else "N/A"
109
+
110
+ def max(self) -> Union[float, str]:
111
+ return self._max / self.scale if self._max else "N/A"
112
+
113
+ def dump_avg_max(self) -> list[Union[float, str]]:
114
+ return [self.avg(), self.max()]
115
+
116
+ def __enter__(self) -> None:
117
+ self.start_time = time.monotonic_ns()
118
+
119
+ def __exit__(
120
+ self,
121
+ exc_type: Optional[type[BaseException]],
122
+ exc_value: Optional[BaseException],
123
+ exc_traceback: Optional[TracebackType],
124
+ ) -> None:
125
+ self.collect(time.monotonic_ns() - self.start_time)
benchmarks/kv_cache/benchmark_block_pool.py
deleted
-108
@@ -1,108 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import gc
4
-import time
5
-from typing import Optional
6
-
7
-from tabulate import tabulate
8
-
9
-from vllm.utils import FlexibleArgumentParser
10
-from vllm.v1.core.block_pool import BlockPool
11
-
12
-
13
-class Metric:
14
- def __init__(self) -> None:
15
- self.cnt: int = 0
16
- self.sum_v: int = 0
17
- self.max_v: Optional[int] = None
18
-
19
- def update(self, v: int) -> None:
20
- self.cnt += 1
21
- self.sum_v += v
22
- if self.max_v is None:
23
- self.max_v = v
24
- else:
25
- self.max_v = max(self.max_v, v)
26
-
27
- def avg_v(self) -> float:
28
- return self.sum_v * 1.0 / self.cnt
29
-
30
-
31
-def main(args):
32
- rows = []
33
- for allocate_block in args.allocate_blocks:
34
- # Enforce a GC collect ahead to minimize the impact among runs
35
- gc.collect()
36
- block_pool = BlockPool(num_gpu_blocks=args.num_gpu_blocks, enable_caching=True)
37
-
38
- get_blocks_metric: Metric = Metric()
39
- free_blocks_metric: Metric = Metric()
40
- for _ in range(args.num_iteration):
41
- t1 = time.monotonic_ns()
42
- blocks = block_pool.get_new_blocks(allocate_block)
43
- t2 = time.monotonic_ns()
44
- block_pool.free_blocks(blocks)
45
- t3 = time.monotonic_ns()
46
- get_blocks_metric.update(t2 - t1)
47
- free_blocks_metric.update(t3 - t2)
48
-
49
- if get_blocks_metric.max_v is not None and free_blocks_metric.max_v is not None:
50
- rows.append(
51
- [
52
- get_blocks_metric.cnt,
53
- args.num_gpu_blocks,
54
- allocate_block,
55
- get_blocks_metric.avg_v() / 1000000,
56
- get_blocks_metric.max_v / 1000000.0,
57
- free_blocks_metric.avg_v() / 1000000,
58
- free_blocks_metric.max_v / 1000000.0,
59
- ]
60
- )
61
- else:
62
- print(
63
- "No valid metrics found."
64
- f" {get_blocks_metric.max_v=} {free_blocks_metric.max_v=}"
65
- )
66
-
67
- print(
68
- tabulate(
69
- rows,
70
- headers=[
71
- "Iterations",
72
- "Total\nBlocks",
73
- "Allocated\nBlocks",
74
- "Get Blocks\nAvg (ms)",
75
- "Get Blocks\nMax (ms)",
76
- "Free Blocks\nAvg (ms)",
77
- "Free Blocks\nMax (ms)",
78
- ],
79
- tablefmt="grid",
80
- floatfmt=".6f",
81
- )
82
- )
83
-
84
-
85
-def invoke_main() -> None:
86
- parser = FlexibleArgumentParser(
87
- description="Benchmark the performance of BlockPool for KV Cache."
88
- )
89
- parser.add_argument("--num-gpu-blocks", type=int, default=100000)
90
- parser.add_argument(
91
- "--num-iteration",
92
- type=int,
93
- default=1000,
94
- help="Number of iterations to run to stablize final data readings",
95
- )
96
- parser.add_argument(
97
- "--allocate-blocks",
98
- type=int,
99
- nargs="*",
100
- default=[10, 50, 100, 500, 1000],
101
- help="Number of blocks to allocate",
102
- )
103
- args = parser.parse_args()
104
- main(args)
105
-
106
-
107
-if __name__ == "__main__":
108
- invoke_main() # pragma: no cover
tests/v1/spec_decode/test_ngram.py
+65
-37
@@ -1,43 +1,63 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
import numpy as np
4
5
from vllm.config import ModelConfig, SpeculativeConfig, VllmConfig
7
-from vllm.v1.spec_decode.ngram_proposer import (NgramProposer,
8
- _find_subarray_kmp,
9
- _kmp_lps_array)
6
+from vllm.v1.spec_decode.ngram_proposer import (
7
+ NgramProposer, _find_longest_matched_ngram_and_propose_tokens)
8
9
12
-def test_kmp_lps_array():
13
- np.testing.assert_array_equal(_kmp_lps_array(np.array([])), np.array([]))
14
- np.testing.assert_array_equal(_kmp_lps_array(np.array([1])), np.array([0]))
15
- np.testing.assert_array_equal(_kmp_lps_array(np.array([1, 1, 1])),
16
- np.array([0, 1, 2]))
17
- np.testing.assert_array_equal(_kmp_lps_array(np.array([1, 2, 3, 4])),
18
- np.array([0, 0, 0, 0]))
19
- np.testing.assert_array_equal(_kmp_lps_array(np.array([1, 2, 1, 2, 3])),
20
- np.array([0, 0, 1, 2, 0]))
10
+def test_find_longest_matched_ngram_and_propose_tokens():
11
+ tokens = np.array([1, 2, 3, 4, 1, 2, 3, 5, 6])
12
+ assert _find_longest_matched_ngram_and_propose_tokens(origin_tokens=tokens,
13
+ min_ngram=2,
14
+ max_ngram=2,
15
+ max_model_len=1024,
16
+ k=2) is None
17
18
+ tokens = np.array([1, 2, 3, 4, 1, 2, 3])
19
+ np.testing.assert_array_equal(
20
+ _find_longest_matched_ngram_and_propose_tokens(origin_tokens=tokens,
21
+ min_ngram=2,
22
+ max_ngram=2,
23
+ max_model_len=1024,
24
+ k=3),
25
+ np.array([4, 1, 2]))
26
+ np.testing.assert_array_equal(
27
+ _find_longest_matched_ngram_and_propose_tokens(origin_tokens=tokens,
28
+ min_ngram=2,
29
+ max_ngram=2,
30
+ max_model_len=1024,
31
+ k=2), np.array([4, 1]))
32
+ np.testing.assert_array_equal(
33
+ _find_longest_matched_ngram_and_propose_tokens(origin_tokens=tokens,
34
+ min_ngram=1,
35
+ max_ngram=1,
36
+ max_model_len=1024,
37
+ k=3),
38
+ np.array([4, 1, 2]))
39
+ np.testing.assert_array_equal(
40
+ _find_longest_matched_ngram_and_propose_tokens(origin_tokens=tokens,
41
+ min_ngram=1,
42
+ max_ngram=1,
43
+ max_model_len=1024,
44
+ k=2), np.array([4, 1]))
45
23
-def test_find_subarray_kmp():
24
- X = np.array([1, 2, 3, 4, 1, 2, 3, 5, 6])
25
- assert _find_subarray_kmp(X, 2, 2) is None
26
- X = np.array([1, 2, 3, 4, 1, 2, 3])
27
- np.testing.assert_array_equal(_find_subarray_kmp(X, 2, 3),
28
- np.array([4, 1, 2]))
29
- np.testing.assert_array_equal(_find_subarray_kmp(X, 2, 2), np.array([4,
30
- 1]))
31
- np.testing.assert_array_equal(_find_subarray_kmp(X, 1, 3),
32
- np.array([4, 1, 2]))
33
- np.testing.assert_array_equal(_find_subarray_kmp(X, 1, 2), np.array([4,
34
- 1]))
35
- X = np.array([1, 3, 6, 2, 3, 4, 1, 2, 3])
36
- np.testing.assert_array_equal(_find_subarray_kmp(X, 2, 3),
37
- np.array([4, 1, 2]))
46
+ tokens = np.array([1, 3, 6, 2, 3, 4, 1, 2, 3])
47
+ np.testing.assert_array_equal(
48
+ _find_longest_matched_ngram_and_propose_tokens(origin_tokens=tokens,
49
+ min_ngram=2,
50
+ max_ngram=2,
51
+ max_model_len=1024,
52
+ k=3),
53
+ np.array([4, 1, 2]))
54
# Return on the first match
39
- np.testing.assert_array_equal(_find_subarray_kmp(X, 1, 3),
40
- np.array([6, 2, 3]))
55
+ np.testing.assert_array_equal(
56
+ _find_longest_matched_ngram_and_propose_tokens(origin_tokens=tokens,
57
+ min_ngram=1,
58
+ max_ngram=1,
59
+ max_model_len=1024,
60
+ k=2), np.array([6, 2]))
61
62
63
def test_ngram_proposer():
@@ -56,27 +76,35 @@ def test_ngram_proposer():
76
77
# No match.
78
result = ngram_proposer(
59
- 2, 2, 2).propose(context_token_ids=np.array([1, 2, 3, 4, 5]))
79
+ min_n=2, max_n=2,
80
+ k=2).propose(context_token_ids=np.array([1, 2, 3, 4, 5]))
81
assert result is None
82
83
# No match for 4-gram.
84
result = ngram_proposer(
64
- 4, 4, 2).propose(context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]))
85
+ min_n=4, max_n=4,
86
+ k=2).propose(context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]))
87
assert result is None
88
89
# No match for 4-gram but match for 3-gram.
90
result = ngram_proposer(
69
- 3, 4, 2).propose(context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]))
91
+ min_n=3, max_n=4,
92
+ k=2).propose(context_token_ids=np.array([1, 2, 3, 4, 1, 2, 3]))
93
assert np.array_equal(result, np.array([4, 1]))
94
95
# Match for both 4-gram and 3-gram.
96
# In this case, the proposer should return the 4-gram match.
74
- result = ngram_proposer(3, 4, 2).propose(
97
+ result = ngram_proposer(min_n=3, max_n=4, k=2).propose(
98
context_token_ids=np.array([2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]))
99
assert np.array_equal(result, np.array([1, 2])) # Not [5, 1]
100
101
# Match for 2-gram and 3-gram, but not 4-gram.
79
- result = ngram_proposer(
80
- 2, 4,
81
- 2).propose(context_token_ids=np.array([3, 4, 5, 2, 3, 4, 1, 2, 3, 4]))
102
+ result = ngram_proposer(min_n=2, max_n=4, k=2).propose(
103
+ context_token_ids=np.array([3, 4, 5, 2, 3, 4, 1, 2, 3, 4]))
104
assert np.array_equal(result, np.array([1, 2])) # Not [5, 2]
105
+
106
+ # Multiple 3-gram matched, but always pick the first one.
107
+ result = ngram_proposer(
108
+ min_n=3, max_n=3, k=2).propose(context_token_ids=np.array(
109
+ [1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]))
110
+ assert np.array_equal(result, np.array([100, 1]))
vllm/v1/spec_decode/ngram_proposer.py
+85
-60
@@ -11,6 +11,10 @@ from vllm.config import VllmConfig
11
class NgramProposer:
12
13
def __init__(self, vllm_config: VllmConfig):
14
+ assert vllm_config.speculative_config is not None
15
+ assert vllm_config.speculative_config.prompt_lookup_min is not None
16
+ assert vllm_config.speculative_config.prompt_lookup_max is not None
17
+
18
# Minimum length of the n-gram to match.
19
self.min_n = vllm_config.speculative_config.prompt_lookup_min
20
# Maximum length of the n-gram to match.
@@ -54,17 +58,13 @@ class NgramProposer:
58
followed that pattern. Here we will return [4,2,3] because
59
we only have three tokens after the match.
60
"""
57
- # Do not generate draft tokens beyond the max model length.
58
- k = min(self.k, self.max_model_len - context_token_ids.shape[0])
59
- if k <= 0:
60
- return None
61
-
61
# TODO(woosuk): Optimize this.
63
- for n in range(self.max_n, self.min_n - 1, -1):
64
- result = _find_subarray_kmp(context_token_ids, n, k)
65
- if result is not None:
66
- return result
67
- return None
62
+ return _find_longest_matched_ngram_and_propose_tokens(
63
+ origin_tokens=context_token_ids,
64
+ min_ngram=self.min_n,
65
+ max_ngram=self.max_n,
66
+ max_model_len=self.max_model_len,
67
+ k=self.k)
68
69
def load_model(self, *args, **kwargs):
70
# No model to load.
@@ -72,61 +72,86 @@ class NgramProposer:
72
73
74
@jit(nopython=True)
75
-def _kmp_lps_array(pattern: np.ndarray) -> np.ndarray:
75
+def _find_longest_matched_ngram_and_propose_tokens(
76
+ origin_tokens: np.ndarray, min_ngram: int, max_ngram: int,
77
+ max_model_len: int, k: int) -> Optional[np.ndarray]:
78
"""
77
- Build the lps (longest proper prefix which is also suffix)
78
- array for the pattern.
79
+ Find the longest n-gram which matches the suffix of the given tokens
80
+ whose length is within [min_ngram, max_ngram] (inclusive).
81
+
82
+ If found, we will extract k right after the matched ngram.
83
"""
80
- lps = np.zeros(len(pattern), dtype=np.int32)
81
- prev_lps = 0 # length of the previous longest prefix suffix
82
- i = 1
84
+ # Do not generate draft tokens is context is shorter than minimum n-gram
85
+ total_token = origin_tokens.shape[0]
86
+ if total_token < min_ngram:
87
+ return None
88
+
89
+ # Do not generate draft tokens beyond the max model length.
90
+ k = min(k, max_model_len - total_token)
91
+ if k <= 0:
92
+ return None
93
+
94
+ # Flip tokens, and the goal become to find longest ngram
95
+ # on the rightmost position which matches the prefix with
96
+ # length [min_n, max_n] (inclusive).
97
+ tokens = origin_tokens[::-1]
98
84
- while i < len(pattern):
85
- if pattern[i] == pattern[prev_lps]:
99
+ # Longest prefix (not including itself) which is a suffix of
100
+ # the current position.
101
+ # lps[i] = max{v, where tokens[0:v] == tokens[i+1-v:i+1]}
102
+ #
103
+ # As ngram is capped by max_ngram to save memory, we only need to
104
+ # store lps for the first max_ngram prefix.
105
+ lps = np.zeros(max_ngram, dtype=np.int32)
106
+
107
+ longest_ngram = 0
108
+ position = 0
109
+
110
+ # lps[0] always equal to 0, we starts with index 1
111
+ prev_lps = 0
112
+ i = 1
113
+ while i < total_token:
114
+ # tokens[:prev_lps] is the longest prefix as a suffix of tokens[:i]
115
+ if tokens[prev_lps] == tokens[i]:
116
+ # Token match: tokens[:prev_lps+1] is the longest prefix as
117
+ # a suffix of tokens[:i+1]
118
prev_lps += 1
87
- lps[i] = prev_lps
119
+ # Check if we found a longer valid ngram.
120
+ #
121
+ # Update position when longest_ngram matched prev_lps,
122
+ # as we want to get the target n-gram of the earliest position
123
+ # in the original tokens (i.e.
124
+ # latest position in the reversed tokens)
125
+ if prev_lps >= longest_ngram:
126
+ longest_ngram = prev_lps
127
+ position = i
128
+ if i < max_ngram:
129
+ # Store LPS for the first max_ngram prefix
130
+ lps[i] = prev_lps
131
+ if prev_lps == max_ngram:
132
+ # When prev_lps reached max_ngram, update prev_lps
133
+ # to lps[max_ngram-1] to avoid matching ngram
134
+ # longer than max_ngram
135
+ prev_lps = lps[max_ngram - 1]
136
i += 1
137
+ elif prev_lps != 0:
138
+ # Token mismatch: try the second longest prefix
139
+ # among all suffix of tokens[:i],
140
+ # which is the longest prefix of tokens[:prev_lps]
141
+ prev_lps = lps[prev_lps - 1]
142
else:
90
- if prev_lps != 0:
91
- prev_lps = lps[prev_lps - 1]
92
- else:
93
- lps[i] = 0
94
- i += 1
95
- return lps
96
-
97
-
98
-@jit(nopython=True)
99
-def _find_subarray_kmp(
100
- context_token_ids: np.ndarray,
101
- n: int,
102
- k: int,
103
-) -> Optional[np.ndarray]:
104
- context_len = context_token_ids.shape[0]
105
- assert n > 0
106
-
107
- pattern = context_token_ids[-n:]
108
- # Precompute lps array for Y
109
- lps = _kmp_lps_array(pattern)
110
-
111
- i = 0
112
- j = 0
113
- # -n because the last n tokens are used as pattern
114
- while i < context_len - n:
115
- if context_token_ids[i] == pattern[j]:
143
+ # Token mismatch, and no more prefix (except empty string)
144
+ # as a suffix of tokens[:i]
145
i += 1
117
- j += 1
146
119
- # If we have matched the entire Y
120
- if j == n:
121
- # Found pattern in context, gather the next K elements
122
- return context_token_ids[i:i + k]
123
- else:
124
- # Mismatch
125
- if j != 0:
126
- # Use the lps array to avoid re-checking elements
127
- j = lps[j - 1]
128
- else:
129
- i += 1
130
-
131
- # Y not found
132
- return None
147
+ if longest_ngram < min_ngram:
148
+ # No valid ngram is found
149
+ return None
150
+
151
+ # Flip the position back, so in origin_tokens,
152
+ # origin_tokens[total_token-1-position:total_token-1-position+longest_ngram]
153
+ # is the matched ngram, so we should start drafting tokens from
154
+ # total_token-1-position+longest_ngram
155
+ start_position = total_token - 1 - position + longest_ngram
156
+ k = min(k, total_token - start_position)
157
+ return origin_tokens[start_position:start_position + k]