From 4ddbb6ad8028a43b822a71d9e6982a8d8e8bf456 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 12 Mar 2026 08:29:05 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- examples/kv_cache_multiturn_benchmark.py | 45 ++++++++------- tests/test_past_kv_models.py | 73 ++++++++++++------------ tests/test_past_kv_utils.py | 21 ++++--- 3 files changed, 72 insertions(+), 67 deletions(-) diff --git a/examples/kv_cache_multiturn_benchmark.py b/examples/kv_cache_multiturn_benchmark.py index 5cad92389b..7da7d2aaa8 100644 --- a/examples/kv_cache_multiturn_benchmark.py +++ b/examples/kv_cache_multiturn_benchmark.py @@ -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: diff --git a/tests/test_past_kv_models.py b/tests/test_past_kv_models.py index 6f3663b972..1af3016eb6 100644 --- a/tests/test_past_kv_models.py +++ b/tests/test_past_kv_models.py @@ -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) diff --git a/tests/test_past_kv_utils.py b/tests/test_past_kv_utils.py index 75feefdf4f..9817073909 100644 --- a/tests/test_past_kv_utils.py +++ b/tests/test_past_kv_utils.py @@ -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))