potok/opencv
1
0
Fork 0
mirror of https://github.com/opencv/opencv.git synced 2026-07-30 15:53:03 +04:00
Code Packages Releases Wiki Activity

Merge pull request #29127 from Prasadayus:KV-cache

Browse Source 
Add KV cache with paged attention and prefetch #29127

Closes: https://github.com/opencv/opencv/issues/27159

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [x] There is a reference to the original bug report and related work
- [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [x] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
Prasad Ayush Kumar
2026-05-26 13:28:47 +05:30
committed by GitHub
parent 9370343d50
commit be2c53c2b7
8 changed files with 239 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samples/dnn/gemma3_inference.py
+56 -21
View File
@@ -1,7 +1,7 @@
# This file is part of OpenCV project.
# It is subject to the license terms in the LICENSE file found in the top-level directory
# of this distribution and at http://opencv.org/license.html.
# Copyright (C) 2025, BigVision LLC, all rights reserved.
# Copyright (C) 2026, BigVision LLC, all rights reserved.
# Third party copyrights are property of their respective owners.
'''
@@ -15,11 +15,17 @@ Exporting Gemma3 model to ONNX:
1. Install the required dependencies:
pip install optimum[exporters] torch transformers
pip install optimum[exporters] optimum-onnx[onnxruntime] torch transformers
2. Export the model to ONNX:
optimum-cli export onnx --model google/gemma-3-1b-it --task causal-lm gemma3_instruct_onnx/
Without KV-cache:
optimum-cli export onnx --model google/gemma-3-1b-it --task causal-lm gemma3_instruct_onnx/
With KV-cache (recommended, faster autoregressive inference):
optimum-cli export onnx --model google/gemma-3-1b-it --task causal-lm-with-past gemma3_instruct_onnx_with_past/
Run the script:
@@ -29,9 +35,18 @@ Run the script:
2. Run the script:
python gemma3_inference.py --model=<path-to-onnx-model> \
--tokenizer_path=<path-to-opencv-tokenizer-config.json> \
--prompt="What is OpenCV?"
Without KV-cache (causal-lm export):
python gemma3_inference.py --model=<path-to-onnx-model> \
--tokenizer_path=<path-to-opencv-tokenizer-config.json> \
--prompt="What is OpenCV?"
With KV-cache (causal-lm-with-past export):
python gemma3_inference.py --model=<path-to-onnx-model> \
--tokenizer_path=<path-to-opencv-tokenizer-config.json> \
--prompt="What is OpenCV?" \
--use_kv_cache
The tokenizer_path should point to an OpenCV-format config.json (e.g., from
opencv_extra/testdata/dnn/llm/gemma3/config.json), NOT the HuggingFace tokenizer_config.json.
@@ -48,6 +63,7 @@ def parse_args():
parser.add_argument('--tokenizer_path', type=str, required=True, help='Path to Gemma3 tokenizer config.json.')
parser.add_argument('--prompt', type=str, default='What is OpenCV?', help='User prompt.')
parser.add_argument('--max_new_tokens', type=int, default=64, help='Maximum number of new tokens to generate.')
parser.add_argument('--use_kv_cache', action='store_true', default=False, help='Enable KV-cache for faster inference (requires causal-lm-with-past export).')
parser.add_argument('--seed', type=int, default=0, help='Random seed.')
return parser.parse_args()
@@ -55,36 +71,55 @@ def build_gemma3_prompt(user_prompt):
'''Wrap user prompt in Gemma3 chat format.'''
return '<start_of_turn>user\n' + user_prompt + '<end_of_turn>\n<start_of_turn>model\n'
def gemma3_inference(net, prompt, max_new_tokens, tokenizer):
def gemma3_inference(net, prompt, max_new_tokens, tokenizer, use_kv_cache=True):
print("Inferencing Gemma3 model...")
tokens = tokenizer.encode(prompt)
# Prepend BOS token (id=2) as required by Gemma3
tokens = [2] + list(tokens)
tokens = np.array(tokens, dtype=np.int64).reshape(1, -1)
input_ids = np.array(tokens, dtype=np.int64).reshape(1, -1)
# Gemma3 special token IDs
eos_id = 1 # <eos>
eot_id = 106 # <end_of_turn>
stop_ids = (eos_id, eot_id)
for _ in range(max_new_tokens):
seq_len = tokens.shape[1]
attention_mask = np.ones((1, seq_len), dtype=np.int64)
generated = []
net.setInput(tokens, 'input_ids')
net.setInput(attention_mask, 'attention_mask')
logits = net.forward() # (1, seq_len, vocab_size)
logits = logits[:, -1, :] # take last token logits
if use_kv_cache:
net.enableKVCache()
prompt_len = input_ids.shape[1]
new_id = int(np.argmax(logits.reshape(-1)))
tokens = np.concatenate((tokens, np.array([[new_id]], dtype=np.int64)), axis=1)
# Prefill: process full prompt once to populate KV-cache
net.setInput(input_ids, 'input_ids')
net.setInput(np.ones((1, prompt_len), dtype=np.int64), 'attention_mask')
logits = net.forward()
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
generated = [new_id]
if new_id in stop_ids:
break
# Generate: feed one new token per step; OpenCV routes present.* -> past_key_values.*
for _ in range(max_new_tokens - 1):
if new_id in stop_ids:
break
net.setInput(np.array([[new_id]], dtype=np.int64), 'input_ids')
net.setInput(np.ones((1, prompt_len + len(generated)), dtype=np.int64), 'attention_mask')
logits = net.forward()
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
generated.append(new_id)
else:
# Without KV-cache: feed full growing sequence each step
for _ in range(max_new_tokens):
net.setInput(input_ids, 'input_ids')
net.setInput(np.ones((1, input_ids.shape[1]), dtype=np.int64), 'attention_mask')
logits = net.forward()
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
if new_id in stop_ids:
break
generated.append(new_id)
input_ids = np.concatenate([input_ids, [[new_id]]], axis=1)
return tokens
return np.array([tokens + generated], dtype=np.int64)
if __name__ == '__main__':
@@ -100,6 +135,6 @@ if __name__ == '__main__':
print(f"Prompt:\n{gemma3_prompt}")
prompt_len = len(tokenizer.encode(gemma3_prompt)) + 1 # +1 for BOS token
tokens = gemma3_inference(net, gemma3_prompt, args.max_new_tokens, tokenizer)
tokens = gemma3_inference(net, gemma3_prompt, args.max_new_tokens, tokenizer, args.use_kv_cache)
response = tokenizer.decode(tokens[0][prompt_len:].tolist())
print(f"Response:\n{response}")
samples/dnn/qwen_inference.py
+63 -28
View File
@@ -9,11 +9,17 @@ Exporting Qwen2.5 model to ONNX:
1. Install the required dependencies:
pip install optimum[exporters] torch transformers
pip install optimum[exporters] optimum-onnx[onnxruntime] torch transformers
2. Export the model to ONNX:
optimum-cli export onnx --model Qwen/Qwen2.5-0.5B-Instruct --task causal-lm qwen2.5_instruct_onnx/
Without KV-cache:
optimum-cli export onnx --model Qwen/Qwen2.5-0.5B-Instruct --task causal-lm qwen2.5_instruct_onnx/
With KV-cache (recommended, faster autoregressive inference):
optimum-cli export onnx --model Qwen/Qwen2.5-0.5B-Instruct --task causal-lm-with-past qwen2.5_instruct_onnx_with_past/
Run the script:
@@ -23,9 +29,18 @@ Run the script:
2. Run the script:
python qwen_inference.py --model=<path-to-onnx-model> \
--tokenizer_path=<path-to-qwen2.5-config.json> \
--prompt="What is OpenCV?"
Without KV-cache (causal-lm export):
python qwen_inference.py --model=<path-to-onnx-model> \
--tokenizer_path=<path-to-qwen2.5-config.json> \
--prompt="What is OpenCV?"
With KV-cache (causal-lm-with-past export):
python qwen_inference.py --model=<path-to-onnx-model> \
--tokenizer_path=<path-to-qwen2.5-config.json> \
--prompt="What is OpenCV?" \
--use_kv_cache
'''
import numpy as np
@@ -39,47 +54,66 @@ def parse_args():
parser.add_argument('--tokenizer_path', type=str, required=True, help='Path to Qwen2.5 tokenizer config.json.')
parser.add_argument('--prompt', type=str, default='What is OpenCV?', help='User prompt.')
parser.add_argument('--max_new_tokens', type=int, default=64, help='Maximum number of new tokens to generate.')
parser.add_argument('--use_kv_cache', action='store_true', default=False, help='Enable KV-cache for faster inference (requires causal-lm-with-past export).')
parser.add_argument('--seed', type=int, default=0, help='Random seed.')
return parser.parse_args()
def stable_softmax(logits):
exp_logits = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
return exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)
def build_chatml_prompt(user_prompt):
'''Wrap user prompt in Qwen2.5 ChatML format.'''
return '<|im_start|>user\n' + user_prompt + '<|im_end|>\n<|im_start|>assistant\n'
def qwen_inference(net, prompt, max_new_tokens, tokenizer):
def qwen_inference(net, prompt, max_new_tokens, tokenizer, use_kv_cache=True):
print("Inferencing Qwen2.5 model...")
tokens = tokenizer.encode(prompt)
tokens = np.array(tokens, dtype=np.int64).reshape(1, -1)
tokens = list(tokenizer.encode(prompt))
input_ids = np.array(tokens, dtype=np.int64).reshape(1, -1)
# Qwen2.5 special token IDs
im_end_id = 151645 # <|im_end|>
eos_id = 151643 # <|endoftext|>
stop_ids = (im_end_id, eos_id)
for _ in range(max_new_tokens):
seq_len = tokens.shape[1]
attention_mask = np.ones((1, seq_len), dtype=np.int64)
position_ids = np.arange(seq_len, dtype=np.int64).reshape(1, -1)
generated = []
net.setInput(tokens, 'input_ids')
net.setInput(attention_mask, 'attention_mask')
net.setInput(position_ids, 'position_ids')
logits = net.forward() # (1, seq_len, vocab_size)
logits = logits[:, -1, :] # take last token logits
if use_kv_cache:
net.enableKVCache()
prompt_len = input_ids.shape[1]
new_id = int(np.argmax(logits.reshape(-1)))
tokens = np.concatenate((tokens, np.array([[new_id]], dtype=np.int64)), axis=1)
# Prefill: process full prompt once to populate KV-cache
net.setInput(input_ids, 'input_ids')
net.setInput(np.ones((1, prompt_len), dtype=np.int64), 'attention_mask')
net.setInput(np.arange(prompt_len, dtype=np.int64).reshape(1, -1), 'position_ids')
logits = net.forward()
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
generated = [new_id]
if new_id in stop_ids:
break
# Generate: feed one new token per step; OpenCV routes present.* -> past_key_values.*
for _ in range(max_new_tokens - 1):
if new_id in stop_ids:
break
cur_len = prompt_len + len(generated)
net.setInput(np.array([[new_id]], dtype=np.int64), 'input_ids')
net.setInput(np.ones((1, cur_len), dtype=np.int64), 'attention_mask')
net.setInput(np.array([[cur_len - 1]], dtype=np.int64), 'position_ids')
logits = net.forward()
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
generated.append(new_id)
else:
# Without KV-cache: feed full growing sequence each step
for _ in range(max_new_tokens):
seq_len = input_ids.shape[1]
net.setInput(input_ids, 'input_ids')
net.setInput(np.ones((1, seq_len), dtype=np.int64), 'attention_mask')
net.setInput(np.arange(seq_len, dtype=np.int64).reshape(1, -1), 'position_ids')
logits = net.forward()
new_id = int(np.argmax(logits[:, -1, :].reshape(-1)))
if new_id in stop_ids:
break
generated.append(new_id)
input_ids = np.concatenate([input_ids, [[new_id]]], axis=1)
return tokens
return np.array([tokens + generated], dtype=np.int64)
if __name__ == '__main__':
@@ -94,6 +128,7 @@ if __name__ == '__main__':
chatml_prompt = build_chatml_prompt(args.prompt)
print(f"Prompt:\n{chatml_prompt}")
tokens = qwen_inference(net, chatml_prompt, args.max_new_tokens, tokenizer)
response = tokenizer.decode(tokens[0].tolist())
prompt_len = len(tokenizer.encode(chatml_prompt))
tokens = qwen_inference(net, chatml_prompt, args.max_new_tokens, tokenizer, args.use_kv_cache)
response = tokenizer.decode(tokens[0][prompt_len:].tolist())
print(f"Response:\n{response}")