ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-03-12 08:29:05 +00:00 • committed by Daniel Han
commit 4ddbb6ad80
3 changed files with 72 additions and 67 deletions
Download patch file Download diff file Expand all files Collapse all files

45
examples/kv_cache_multiturn_benchmark.py
View file

@ -13,6 +13,7 @@ The longer the conversation history, the bigger the speedup.
Run in Colab (T4/A100):
python examples/kv_cache_multiturn_benchmark.py
"""
import torch
import time
from unsloth import FastLanguageModel
@ -349,7 +350,7 @@ CONVERSATION_HISTORY = [
]
def run_benchmark(model, tokenizer, history_turns, new_question, num_runs=5):
def run_benchmark(model, tokenizer, history_turns, new_question, num_runs = 5):
"""
Run a single benchmark: compare baseline vs KV cache generation.
Returns (time_baseline, time_kv, num_history_tokens, outputs_match).
@ -359,14 +360,14 @@ def run_benchmark(model, tokenizer, history_turns, new_question, num_runs=5):
# Tokenize history and full conversation
text_history = tokenizer.apply_chat_template(
history, tokenize=False, add_generation_prompt=False
history, tokenize = False, add_generation_prompt = False
)
text_full = tokenizer.apply_chat_template(
history + new_msg, tokenize=False, add_generation_prompt=True
history + new_msg, tokenize = False, add_generation_prompt = True
)
inputs_history = tokenizer(text_history, return_tensors="pt").to("cuda")
inputs_full = tokenizer(text_full, return_tensors="pt").to("cuda")
inputs_history = tokenizer(text_history, return_tensors = "pt").to("cuda")
inputs_full = tokenizer(text_full, return_tensors = "pt").to("cuda")
len_history = inputs_history.input_ids.shape[1]
len_full = inputs_full.input_ids.shape[1]
@ -382,14 +383,14 @@ def run_benchmark(model, tokenizer, history_turns, new_question, num_runs=5):
# Pre-compute KV cache (this cost is amortized over many requests)
with torch.no_grad():
outputs_history = model(**inputs_history, use_cache=True)
outputs_history = model(**inputs_history, use_cache = True)
cached_kv = outputs_history.past_key_values
gen_kwargs = dict(max_new_tokens=50, use_cache=True, do_sample=False)
gen_kwargs = dict(max_new_tokens = 50, use_cache = True, do_sample = False)
# Warmup both paths
model.generate(**inputs_full, max_new_tokens=1)
model.generate(**inputs_full, max_new_tokens=1, past_key_values=cached_kv)
model.generate(**inputs_full, max_new_tokens = 1)
model.generate(**inputs_full, max_new_tokens = 1, past_key_values = cached_kv)
torch.cuda.synchronize()
# Benchmark baseline (no KV cache — re-processes all history tokens)
@ -409,19 +410,19 @@ def run_benchmark(model, tokenizer, history_turns, new_question, num_runs=5):
torch.cuda.synchronize()
t0 = time.perf_counter()
output_kv = model.generate(
**inputs_full, past_key_values=cached_kv, **gen_kwargs
**inputs_full, past_key_values = cached_kv, **gen_kwargs
)
torch.cuda.synchronize()
times_kv.append(time.perf_counter() - t0)
# Decode outputs
text_baseline = tokenizer.decode(
output_baseline[0][len_full:], skip_special_tokens=True
output_baseline[0][len_full:], skip_special_tokens = True
)
if output_kv.shape[1] > len_full:
text_kv = tokenizer.decode(output_kv[0][len_full:], skip_special_tokens=True)
text_kv = tokenizer.decode(output_kv[0][len_full:], skip_special_tokens = True)
else:
text_kv = tokenizer.decode(output_kv[0], skip_special_tokens=True)
text_kv = tokenizer.decode(output_kv[0], skip_special_tokens = True)
# Use median for stable timing
time_baseline = sorted(times_baseline)[len(times_baseline) // 2]
@ -446,10 +447,10 @@ def main():
print(f"Loading {model_name}...")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_name,
max_seq_length=max_seq_length,
dtype=None,
load_in_4bit=True,
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None,
load_in_4bit = True,
)
FastLanguageModel.for_inference(model)
@ -457,8 +458,8 @@ def main():
# We test with 4, 8, 12, and all 16 messages of history.
# Each step roughly doubles the cached token count.
test_cases = [
(4, "What should I look at next?"),
(8, "Can you recap what we've covered so far?"),
(4, "What should I look at next?"),
(8, "Can you recap what we've covered so far?"),
(12, "What's the single most impactful optimization?"),
(16, "Give me a 3-step action plan to go to production."),
]
@ -473,7 +474,7 @@ def main():
num_turns = num_msgs // 2 # user+assistant pairs
print(f"\n{'─' * 72}")
print(f" Conversation: {num_msgs} messages ({num_turns} turns)")
print(f" New question: \"{question}\"")
print(f' New question: "{question}"')
print(f"{'─' * 72}")
r = run_benchmark(model, tokenizer, num_msgs, question)
@ -496,7 +497,9 @@ def main():
print(f"\n{'=' * 72}")
print(" SUMMARY")
print(f"{'=' * 72}")
print(f" {'History':>8} {'New':>6} {'Baseline':>10} {'KV Cache':>10} {'Speedup':>8} {'Match':>6}")
print(
f" {'History':>8} {'New':>6} {'Baseline':>10} {'KV Cache':>10} {'Speedup':>8} {'Match':>6}"
)
print(f" {'tokens':>8} {'tokens':>6} {'(sec)':>10} {'(sec)':>10} {'':>8} {'':>6}")
print(f" {'─' * 8} {'─' * 6} {'─' * 10} {'─' * 10} {'─' * 8} {'─' * 6}")
for r in results:

73
tests/test_past_kv_models.py
View file

@ -10,6 +10,7 @@ Or run individual model tests:
python -m pytest tests/test_past_kv_models.py -v -s -k "Gemma2"
python -m pytest tests/test_past_kv_models.py -v -s -k "Llama"
"""
import unittest
import torch
@ -19,7 +20,7 @@ def _skip_if_no_cuda():
raise unittest.SkipTest("CUDA not available")
def _run_past_kv_test(test_case, model_name, load_in_4bit=True):
def _run_past_kv_test(test_case, model_name, load_in_4bit = True):
"""
Shared test logic: generate with baseline vs past_key_values and verify
outputs match (or at minimum, that no errors are raised).
@ -27,10 +28,10 @@ def _run_past_kv_test(test_case, model_name, load_in_4bit=True):
from unsloth import FastLanguageModel
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_name,
max_seq_length=2048,
dtype=None,
load_in_4bit=load_in_4bit,
model_name = model_name,
max_seq_length = 2048,
dtype = None,
load_in_4bit = load_in_4bit,
)
FastLanguageModel.for_inference(model)
@ -45,15 +46,15 @@ def _run_past_kv_test(test_case, model_name, load_in_4bit=True):
# Tokenize history alone
text_history = tokenizer.apply_chat_template(
messages_history, tokenize=False, add_generation_prompt=False
messages_history, tokenize = False, add_generation_prompt = False
)
inputs_history = tokenizer(text_history, return_tensors="pt").to("cuda")
inputs_history = tokenizer(text_history, return_tensors = "pt").to("cuda")
# Tokenize full conversation
text_full = tokenizer.apply_chat_template(
messages_history + messages_new, tokenize=False, add_generation_prompt=True
messages_history + messages_new, tokenize = False, add_generation_prompt = True
)
inputs_full = tokenizer(text_full, return_tensors="pt").to("cuda")
inputs_full = tokenizer(text_full, return_tensors = "pt").to("cuda")
len_history = inputs_history.input_ids.shape[1]
len_full = inputs_full.input_ids.shape[1]
@ -61,35 +62,33 @@ def _run_past_kv_test(test_case, model_name, load_in_4bit=True):
# Pre-compute KV cache for history
with torch.no_grad():
outputs_history = model(**inputs_history, use_cache=True)
outputs_history = model(**inputs_history, use_cache = True)
past_kv = outputs_history.past_key_values
# Baseline generation (no custom KV)
output_baseline = model.generate(
**inputs_full,
max_new_tokens=30,
use_cache=True,
do_sample=False,
max_new_tokens = 30,
use_cache = True,
do_sample = False,
)
text_baseline = tokenizer.decode(
output_baseline[0][len_full:], skip_special_tokens=True
output_baseline[0][len_full:], skip_special_tokens = True
)
print(f" Baseline: {text_baseline.strip()}")
# KV cache generation
output_kv = model.generate(
**inputs_full,
max_new_tokens=30,
past_key_values=past_kv,
use_cache=True,
do_sample=False,
max_new_tokens = 30,
past_key_values = past_kv,
use_cache = True,
do_sample = False,
)
if output_kv.shape[1] > len_full:
text_kv = tokenizer.decode(
output_kv[0][len_full:], skip_special_tokens=True
)
text_kv = tokenizer.decode(output_kv[0][len_full:], skip_special_tokens = True)
else:
text_kv = tokenizer.decode(output_kv[0], skip_special_tokens=True)
text_kv = tokenizer.decode(output_kv[0], skip_special_tokens = True)
print(f" KV Cache: {text_kv.strip()}")
# Both should produce coherent output (not crash)
@ -107,7 +106,7 @@ def _run_past_kv_test(test_case, model_name, load_in_4bit=True):
torch.cuda.empty_cache()
def _run_tuple_kv_test(test_case, model_name, load_in_4bit=True):
def _run_tuple_kv_test(test_case, model_name, load_in_4bit = True):
"""
Test that passing tuple past_key_values (not DynamicCache) works.
This validates the _ensure_cache_is_dynamic v5 compat path.
@ -115,19 +114,19 @@ def _run_tuple_kv_test(test_case, model_name, load_in_4bit=True):
from unsloth import FastLanguageModel
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_name,
max_seq_length=2048,
dtype=None,
load_in_4bit=load_in_4bit,
model_name = model_name,
max_seq_length = 2048,
dtype = None,
load_in_4bit = load_in_4bit,
)
FastLanguageModel.for_inference(model)
prompt = "The capital of France is"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
inputs = tokenizer(prompt, return_tensors = "pt").to("cuda")
# Get KV cache from forward pass
with torch.no_grad():
outputs = model(**inputs, use_cache=True)
outputs = model(**inputs, use_cache = True)
past_kv = outputs.past_key_values
# Convert DynamicCache to tuple format (simulating user-provided tuple KV)
@ -137,16 +136,16 @@ def _run_tuple_kv_test(test_case, model_name, load_in_4bit=True):
tuple_kv = past_kv # Already tuple
# This should NOT raise ValueError even on transformers v5
next_token = tokenizer(" Paris", return_tensors="pt").to("cuda")
full_input = torch.cat([inputs.input_ids, next_token.input_ids], dim=1)
next_token = tokenizer(" Paris", return_tensors = "pt").to("cuda")
full_input = torch.cat([inputs.input_ids, next_token.input_ids], dim = 1)
output = model.generate(
input_ids=full_input,
max_new_tokens=10,
past_key_values=tuple_kv,
use_cache=True,
do_sample=False,
input_ids = full_input,
max_new_tokens = 10,
past_key_values = tuple_kv,
use_cache = True,
do_sample = False,
)
text = tokenizer.decode(output[0], skip_special_tokens=True)
text = tokenizer.decode(output[0], skip_special_tokens = True)
print(f"\n Tuple KV output: {text.strip()}")
test_case.assertGreater(len(text.strip()), 0)

21
tests/test_past_kv_utils.py
View file

@ -5,6 +5,7 @@ Self-contained — does NOT import unsloth, so runs without a GPU.
Run with:
python -m pytest tests/test_past_kv_utils.py -v
"""
import unittest
import torch
from transformers.cache_utils import DynamicCache, Cache
@ -14,6 +15,7 @@ from transformers.cache_utils import DynamicCache, Cache
# These match the implementations in unsloth/models/llama.py exactly.
# Kept inline so the test suite can run on any machine (no GPU needed).
def _ensure_cache_is_dynamic(past_key_values):
"""Convert list/tuple of (K, V) pairs to DynamicCache for transformers v5 compat."""
if past_key_values is None:
@ -34,15 +36,16 @@ def _slice_position_ids(position_ids, input_ids):
return None
if position_ids.dim() == 2:
if position_ids.shape[1] > input_ids.shape[1]:
position_ids = position_ids[:, -input_ids.shape[1]:]
position_ids = position_ids[:, -input_ids.shape[1] :]
elif position_ids.dim() == 1:
if position_ids.shape[0] > input_ids.shape[1]:
position_ids = position_ids[-input_ids.shape[1]:]
position_ids = position_ids[-input_ids.shape[1] :]
return position_ids
# ── Tests ──────────────────────────────────────────────────────────────
class TestEnsureCacheIsDynamic(unittest.TestCase):
"""Tests for _ensure_cache_is_dynamic conversion utility."""
@ -102,18 +105,18 @@ class TestSlicePositionIds(unittest.TestCase):
"""Tests for _slice_position_ids utility."""
def test_none_passthrough(self):
input_ids = torch.zeros(1, 5, dtype=torch.long)
input_ids = torch.zeros(1, 5, dtype = torch.long)
self.assertIsNone(_slice_position_ids(None, input_ids))
def test_2d_no_slice_needed(self):
input_ids = torch.zeros(1, 10, dtype=torch.long)
input_ids = torch.zeros(1, 10, dtype = torch.long)
position_ids = torch.arange(10).unsqueeze(0)
result = _slice_position_ids(position_ids, input_ids)
self.assertTrue(torch.equal(result, position_ids))
def test_2d_slice_needed(self):
"""position_ids longer than input_ids — should take last N."""
input_ids = torch.zeros(1, 3, dtype=torch.long)
input_ids = torch.zeros(1, 3, dtype = torch.long)
position_ids = torch.arange(10).unsqueeze(0) # shape (1, 10)
result = _slice_position_ids(position_ids, input_ids)
self.assertEqual(result.shape, (1, 3))
@ -121,13 +124,13 @@ class TestSlicePositionIds(unittest.TestCase):
self.assertTrue(torch.equal(result, expected))
def test_1d_no_slice_needed(self):
input_ids = torch.zeros(1, 5, dtype=torch.long)
input_ids = torch.zeros(1, 5, dtype = torch.long)
position_ids = torch.arange(5)
result = _slice_position_ids(position_ids, input_ids)
self.assertTrue(torch.equal(result, position_ids))
def test_1d_slice_needed(self):
input_ids = torch.zeros(1, 3, dtype=torch.long)
input_ids = torch.zeros(1, 3, dtype = torch.long)
position_ids = torch.arange(10) # shape (10,)
result = _slice_position_ids(position_ids, input_ids)
self.assertEqual(result.shape, (3,))
@ -136,14 +139,14 @@ class TestSlicePositionIds(unittest.TestCase):
def test_shorter_position_ids_passthrough(self):
"""position_ids shorter than input_ids — should pass through unchanged."""
input_ids = torch.zeros(1, 10, dtype=torch.long)
input_ids = torch.zeros(1, 10, dtype = torch.long)
position_ids = torch.arange(5).unsqueeze(0)
result = _slice_position_ids(position_ids, input_ids)
self.assertTrue(torch.equal(result, position_ids))
def test_exact_match(self):
"""Exact same length — no slicing."""
input_ids = torch.zeros(2, 7, dtype=torch.long)
input_ids = torch.zeros(2, 7, dtype = torch.long)
position_ids = torch.arange(7).unsqueeze(0).expand(2, -1)
result = _slice_position_ids(position_ids, input_ids)
self.assertEqual(result.shape, (2, 7))