Accelerating Local LLM Conversations with KV Cache Reuse on AMD Ryzen AI

		DEFAULT_PROMPTS = [
    " My name is Alex, I work as a data scientist at Orbital AI.",
    " I live in Berlin and enjoy playing the piano in my free time.",
    " Where do I live and what do I enjoy doing?",
    " What is my profession?"
]
	

		```python
generator = og.Generator(model, params)

for i, prompt in enumerate(prompts):
    # Append only the new prompt tokens — the generator retains all prior context
    generator.append_tokens(tokenizer.encode(prompt))

    loop = 0
    while not generator.is_done() and loop < 120:
        if loop == 0:
            in_pos = generator.get_sequence(0)  # mark start of this response
        loop += 1
        generator.generate_next_token()

    # Extract only the tokens generated for this prompt
    output_tokens = generator.get_sequence(0)[len(in_pos):]
    output_text = tokenizer.decode(output_tokens)

    print(f"\nPrompt #{i+1}: {prompt}")
    print("Output:\n", output_text)

    results.append({
        "prompt": prompt,
        "response": output_text,
        "tokens": len(output_tokens)
    })
    total_tokens += len(output_tokens)
```
	

		```bash
conda activate ryzen-ai-1.7.1
cd "C:\Program Files\RyzenAI\1.7.1\LLM\example"
python run_kvreuse.py -m D:\model\qwen2.5-3B
```
	

		```
Normal flow:
  Q1 → A1 → Q2 → A2 → Q3 → A3
  KV cache: [Q1][A1][Q2][A2][Q3][A3]

After rewind_to(len(Q1 + A1)):
  KV cache: [Q1][A1]          ← everything after A1 is discarded

Re-ask Q2:
  KV cache: [Q1][A1][Q2][A2'] ← model is in the same state as after the original A1

If the model is deterministic: A2' == A2
```
	

		```python
# Phase 1: Run Q1, Q2, Q3 in sequence to build up the full KV cache
for i in range(3):
    generator.append_tokens(tokenizer.encode(prompts[i]))
    loop = 0
    while not generator.is_done() and loop < 120:
        loop += 1
        generator.generate_next_token()

    if i == 0:
        # Save the sequence length after Q1+A1 — this is our rewind target
        output_tokens0 = generator.get_sequence(0)
    if i == 1:
        # Save A2's text so we can compare it after the rewind
        output_tokens1 = generator.get_sequence(0)
        output_text1 = tokenizer.decode(output_tokens1)

# Phase 2: Rewind to the position after Q1+A1
generator.rewind_to(len(output_tokens0))

# Re-ask Q2 from this rewound state
generator.append_tokens(tokenizer.encode(prompts[1]))
loop = 0
while not generator.is_done() and loop < 120:
    loop += 1
    generator.generate_next_token()

output_tokens1_rewind = generator.get_sequence(0)
output_text1_rewind = tokenizer.decode(output_tokens1_rewind)

# Phase 3: Verify the outputs match
if output_text1 == output_text1_rewind:
    print("rewind test OK")
return
```
	

		```bash
python run_rewind.py -m D:\model\qwen2.5-3B
```
	

		```
rewind test OK
```
	

		```python
def load_prompts(prompt_input):
    """Loads prompts from a .txt file, a direct string, or falls back to default prompts."""
    if prompt_input:
        if os.path.exists(prompt_input):
            with open(prompt_input, "r", encoding="utf-8") as f:
                prompts = [line.strip() for line in f.readlines() if line.strip()]
            if prompts:
                return prompts
        else:
            # Treat as a direct string input
            return [prompt_input]
    print("Warning: Invalid or missing prompt input. Using default prompts.")
    return DEFAULT_PROMPTS
 
def load_model_and_tokenizer(model_path, verbose=False):
    """Loads the ONNX model and tokenizer, determining the model type from the config."""
    config_path = os.path.join(model_path, 'genai_config.json')
 
    # Read the model type from the configuration file
    with open(config_path, 'r') as config_file:
        config = json.load(config_file)
        model_type = config['model']['type']
 
    if verbose:
        print(f"Loading {model_type} model from {model_path}...")
 
    model = og.Model(model_path)
    tokenizer = get_tokenizer(model_path, model_type, model)
    return model, tokenizer, model_type
 
def _tokens_to_list(tokens):
    """Normalize tokenizer output to a flat Python list for append_tokens."""
    if isinstance(tokens, np.ndarray):
        return tokens.flatten().tolist()
    return list(tokens)
 
def _generate_up_to(generator, max_new_tokens):
    """Generate up to max_new_tokens; stop early if is_done. Returns count generated."""
    generated = 0
    while generated < max_new_tokens and not generator.is_done():
        generator.generate_next_token()
        generated += 1
    return generated
 
def _new_tokens_since(generator, start_len):
    return _tokens_to_list(generator.get_sequence(0))[start_len:]
 
def generate_text(model, tokenizer, prompts, model_type, args):
  """KV benchmark: slice text1/text2 from benchmark_test.txt tokens, then time 10 runs."""
  TEXT1_LEN = 2048
  TEXT2_LEN = 256
  OUT1_MAX = 256          # max new tokens after text1
  OUT2_MAX = 128
  BENCH_LOOPS = 10
  NEED_TOKENS = TEXT1_LEN + TEXT2_LEN
 
  script_dir = os.path.dirname(os.path.abspath(__file__))
  benchmark_path = os.path.join(script_dir, "benchmark_test.txt")
  if not os.path.exists(benchmark_path):
    raise FileNotFoundError(f"benchmark_test.txt not found: {benchmark_path}")
 
  with open(benchmark_path, "r", encoding="utf-8") as f:
    benchmark_text = f.read()
 
  # --- Phase 1: tokenize file only; take fixed slices (no model input of full file) ---
  all_tokens = _tokens_to_list(tokenizer.encode(benchmark_text))
  print(
    f"Slicing from {benchmark_path}: encoded {len(all_tokens)} tokens, "
    f"using [{0}:{TEXT1_LEN}) + [{TEXT1_LEN}:{NEED_TOKENS}), ignoring the rest."
  )
  if len(all_tokens) < NEED_TOKENS:
    raise RuntimeError(
      f"benchmark_test.txt too short after encode: got {len(all_tokens)} tokens, "
      f"need at least {NEED_TOKENS} ({TEXT1_LEN} + {TEXT2_LEN})."
    )
 
  text1 = all_tokens[0:TEXT1_LEN]
  text2 = all_tokens[TEXT1_LEN:NEED_TOKENS]
 
  print(f"text1 ready: {len(text1)} tokens | text2 ready: {len(text2)} tokens")
 
  # --- Phase 2: one generator — text1 -> out1, then text2 -> out2 (KV retained) ---
  bench_max_len = TEXT1_LEN + TEXT2_LEN + OUT1_MAX + OUT2_MAX + 256
 
  bench_params = og.GeneratorParams(model)
  bench_search = {
    "do_sample": args.do_random_sampling,
    "max_length": bench_max_len,
    "min_length": args.min_length,
    "top_p": args.top_p,
    "top_k": args.top_k,
    "temperature": args.temperature,
    "repetition_penalty": args.repetition_penalty,
  }
  bench_params.set_search_options(**{k: v for k, v in bench_search.items() if v is not None})
  bench_params.try_graph_capture_with_max_batch_size(1)
 
  total_time_text1 = 0.0
  total_time_text2 = 0.0
  out1 = None
  out2 = None
 
  print(f"\n[Phase 2] KV-retained benchmark ({BENCH_LOOPS} iterations)...\n")
 
  for i in range(BENCH_LOOPS):
    generator = og.Generator(model, bench_params)
 
    # text1 -> out1 (up to OUT1_MAX new tokens, stop if is_done early)
    t0 = time.time()
    generator.append_tokens(text1)
    in_pos = len(_tokens_to_list(generator.get_sequence(0)))
    _generate_up_to(generator, OUT1_MAX)
    out1 = _new_tokens_since(generator, in_pos)
    total_time_text1 += time.time() - t0
 
    # text2 -> out2 (same generator, KV from text1 retained)
    t0 = time.time()
    generator.append_tokens(text2)
    in_pos2 = len(_tokens_to_list(generator.get_sequence(0)))
    _generate_up_to(generator, OUT2_MAX)
    out2 = _new_tokens_since(generator, in_pos2)
    total_time_text2 += time.time() - t0
 
  print("\n--- Phase 2 results (single generator, text1 then text2) ---")
  print(f"avg time text1 -> out1: {total_time_text1 / BENCH_LOOPS:.4f}s")
  print(f"avg time text2 -> out2: {total_time_text2 / BENCH_LOOPS:.4f}s")
 
  # --- Phase 3: text1 -> 256 out1; new generator -> text1+out1+text2 -> out2 ---
  total_time_p3_text1 = 0.0
  total_time_p3_combo = 0.0
  p3_out1 = None
  p3_out2 = None
 
  print(f"\n[Phase 3] split-generator benchmark ({BENCH_LOOPS} iterations)...\n")
 
  for i in range(BENCH_LOOPS):
    # Step A: text1 -> exactly OUT1_MAX (256) new tokens
    gen_a = og.Generator(model, bench_params)
    t0 = time.time()
    gen_a.append_tokens(text1)
    pos_a = len(_tokens_to_list(gen_a.get_sequence(0)))
    _generate_up_to(gen_a, OUT1_MAX)
    p3_out1 = _new_tokens_since(gen_a, pos_a)
    total_time_p3_text1 += time.time() - t0
 
    # Step B: fresh generator, text1 + out1 + text2 -> out2
    gen_b = og.Generator(model, bench_params)
    t0 = time.time()
    gen_b.append_tokens(text1 + p3_out1 + text2)
    pos_b = len(_tokens_to_list(gen_b.get_sequence(0)))
    _generate_up_to(gen_b, OUT2_MAX)
    p3_out2 = _new_tokens_since(gen_b, pos_b)
    total_time_p3_combo += time.time() - t0
 
  print("\n--- Phase 3 results (text1->512; new gen; text1+out1+text2->out2) ---")
  print(f"avg time step A text1 -> out1: {total_time_p3_text1 / BENCH_LOOPS:.4f}s")
  print(f"avg time step B text1+out1+text2 -> out2: {total_time_p3_combo / BENCH_LOOPS:.4f}s")
```