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

Add transformers v5 compat, Qwen3/Gemma2 past_key_values support, and shared utilities

Browse source 
- Add _ensure_cache_is_dynamic to convert tuple/list KV caches to DynamicCache
  before transformers v5's _get_cache rejects them
- Wrap generate() via fix_prepare_inputs_for_generation for all model types
- Fix Qwen3 RoPE: index cos/sin by position_ids (both branches were identical)
- Fix Gemma2 softcapping attention: handle 4D masks and Q_len != K_len during
  prefill with past_key_values
- Add _slice_position_ids shared utility, replace inline duplication in
  PeftModel_fast_forward, MistralForCausalLM_fast_forward, and
  CausalLM_fast_forward (covers Llama/Qwen3/Gemma2)
- Remove redundant seq_len assignment in _fast_prepare_inputs_for_generation
- Add unit tests (test_past_kv_utils.py) and GPU integration tests
  (test_past_kv_models.py) for Llama, Qwen3, and Gemma2
This commit is contained in:
vivekkalyanarangan30 2026-02-25 10:37:59 +05:30 • committed by Daniel Han
commit cea52f18ac
3 changed files with 378 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

201
tests/test_past_kv_models.py Normal file
View file

@ -0,0 +1,201 @@
"""
Integration tests for past_key_values support across model architectures.
Requires a CUDA GPU. Best run in Colab or a GPU-equipped machine.
Run with:
python -m pytest tests/test_past_kv_models.py -v -s
Or run individual model tests:
python -m pytest tests/test_past_kv_models.py -v -s -k "Qwen3"
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
def _skip_if_no_cuda():
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA not available")
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).
"""
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,
)
FastLanguageModel.for_inference(model)
# Build a conversation with history
messages_history = [
{"role": "user", "content": "Remember: the secret code is ALPHA-7."},
{"role": "assistant", "content": "Got it, the secret code is ALPHA-7."},
]
messages_new = [
{"role": "user", "content": "What is the secret code?"},
]
# Tokenize history alone
text_history = tokenizer.apply_chat_template(
messages_history, tokenize=False, add_generation_prompt=False
)
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
)
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]
print(f"\n History tokens: {len_history}, Full tokens: {len_full}")
# Pre-compute KV cache for history
with torch.no_grad():
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,
)
text_baseline = tokenizer.decode(
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,
)
if output_kv.shape[1] > len_full:
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)
print(f" KV Cache: {text_kv.strip()}")
# Both should produce coherent output (not crash)
test_case.assertGreater(len(text_kv.strip()), 0, "KV cache output is empty")
# Outputs should match
test_case.assertEqual(
text_baseline.strip(),
text_kv.strip(),
"Baseline and KV cache outputs differ",
)
# Cleanup
del model, tokenizer
torch.cuda.empty_cache()
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.
"""
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,
)
FastLanguageModel.for_inference(model)
prompt = "The capital of France is"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# Get KV cache from forward pass
with torch.no_grad():
outputs = model(**inputs, use_cache=True)
past_kv = outputs.past_key_values
# Convert DynamicCache to tuple format (simulating user-provided tuple KV)
if hasattr(past_kv, "get_seq_length"):
tuple_kv = tuple(past_kv[i] for i in range(len(past_kv)))
else:
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)
output = model.generate(
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)
print(f"\n Tuple KV output: {text.strip()}")
test_case.assertGreater(len(text.strip()), 0)
del model, tokenizer
torch.cuda.empty_cache()
class TestPastKVLlama(unittest.TestCase):
def setUp(self):
_skip_if_no_cuda()
def test_past_kv_generation(self):
"""Test past_key_values with Llama model."""
try:
_run_past_kv_test(self, "unsloth/Llama-3.2-1B-Instruct")
except Exception as e:
self.skipTest(f"Model loading failed: {e}")
def test_tuple_kv_v5_compat(self):
"""Test tuple KV cache conversion (v5 compat) with Llama."""
try:
_run_tuple_kv_test(self, "unsloth/Llama-3.2-1B-Instruct")
except Exception as e:
self.skipTest(f"Model loading failed: {e}")
class TestPastKVQwen3(unittest.TestCase):
def setUp(self):
_skip_if_no_cuda()
def test_past_kv_generation(self):
"""Test past_key_values with Qwen3 model (validates RoPE position_ids fix)."""
try:
_run_past_kv_test(self, "unsloth/Qwen3-0.6B")
except Exception as e:
self.skipTest(f"Model loading failed: {e}")
class TestPastKVGemma2(unittest.TestCase):
def setUp(self):
_skip_if_no_cuda()
def test_past_kv_generation(self):
"""Test past_key_values with Gemma2 model (validates 4D mask fix)."""
try:
_run_past_kv_test(self, "unsloth/gemma-2-2b-it")
except Exception as e:
self.skipTest(f"Model loading failed: {e}")
if __name__ == "__main__":
unittest.main()

153
tests/test_past_kv_utils.py Normal file
View file

@ -0,0 +1,153 @@
"""
Unit tests for past_key_values utilities.
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
# ── Inline copies of the functions under test ──────────────────────────
# 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:
return None
if isinstance(past_key_values, Cache):
return past_key_values
if isinstance(past_key_values, (tuple, list)) and len(past_key_values) > 0:
cache = DynamicCache()
for layer_idx, layer_kv in enumerate(past_key_values):
cache.update(layer_kv[0], layer_kv[1], layer_idx)
return cache
return past_key_values
def _slice_position_ids(position_ids, input_ids):
"""Slice position_ids to match input_ids length if needed."""
if position_ids is None:
return None
if position_ids.dim() == 2:
if position_ids.shape[1] > 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]:]
return position_ids
# ── Tests ──────────────────────────────────────────────────────────────
class TestEnsureCacheIsDynamic(unittest.TestCase):
"""Tests for _ensure_cache_is_dynamic conversion utility."""
def test_none_passthrough(self):
self.assertIsNone(_ensure_cache_is_dynamic(None))
def test_dynamic_cache_passthrough(self):
cache = DynamicCache()
k = torch.randn(1, 4, 8, 16)
v = torch.randn(1, 4, 8, 16)
cache.update(k, v, 0)
result = _ensure_cache_is_dynamic(cache)
self.assertIs(result, cache)
def test_tuple_conversion(self):
"""Tuple of (K, V) pairs should be converted to DynamicCache."""
n_layers = 3
layers = []
for _ in range(n_layers):
k = torch.randn(1, 4, 8, 16)
v = torch.randn(1, 4, 8, 16)
layers.append((k, v))
past_kv = tuple(layers)
result = _ensure_cache_is_dynamic(past_kv)
self.assertIsInstance(result, DynamicCache)
for i in range(n_layers):
cached_k, cached_v = result[i]
self.assertTrue(torch.equal(cached_k, layers[i][0]))
self.assertTrue(torch.equal(cached_v, layers[i][1]))
def test_list_conversion(self):
"""List of (K, V) pairs should be converted to DynamicCache."""
layers = [(torch.randn(1, 4, 8, 16), torch.randn(1, 4, 8, 16))]
result = _ensure_cache_is_dynamic(layers)
self.assertIsInstance(result, DynamicCache)
cached_k, cached_v = result[0]
self.assertTrue(torch.equal(cached_k, layers[0][0]))
def test_empty_tuple_passthrough(self):
result = _ensure_cache_is_dynamic(())
self.assertEqual(result, ())
def test_empty_list_passthrough(self):
result = _ensure_cache_is_dynamic([])
self.assertEqual(result, [])
def test_seq_length_preserved(self):
"""Verify DynamicCache reports correct sequence length after conversion."""
seq_len = 42
layers = [(torch.randn(1, 4, seq_len, 16), torch.randn(1, 4, seq_len, 16))]
result = _ensure_cache_is_dynamic(tuple(layers))
self.assertEqual(result.get_seq_length(), seq_len)
class TestSlicePositionIds(unittest.TestCase):
"""Tests for _slice_position_ids utility."""
def test_none_passthrough(self):
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)
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)
position_ids = torch.arange(10).unsqueeze(0) # shape (1, 10)
result = _slice_position_ids(position_ids, input_ids)
self.assertEqual(result.shape, (1, 3))
expected = torch.tensor([[7, 8, 9]])
self.assertTrue(torch.equal(result, expected))
def test_1d_no_slice_needed(self):
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)
position_ids = torch.arange(10) # shape (10,)
result = _slice_position_ids(position_ids, input_ids)
self.assertEqual(result.shape, (3,))
expected = torch.tensor([7, 8, 9])
self.assertTrue(torch.equal(result, expected))
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)
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)
position_ids = torch.arange(7).unsqueeze(0).expand(2, -1)
result = _slice_position_ids(position_ids, input_ids)
self.assertEqual(result.shape, (2, 7))
if __name__ == "__main__":
unittest.main()

42
unsloth/kernels/flex_attention.py
View file

@ -43,17 +43,18 @@ except:
if not HAS_FLEX_ATTENTION:
# Logit softcapping
@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options)
def slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len):
def slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, kv_len):
n_heads = self.config.num_attention_heads
head_dim = self.head_dim
n_kv_heads = self.config.num_key_value_heads
n_groups = self.num_key_value_groups
actual_q_len = Q.shape[-2]
# Grouped query attention
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
K = K.reshape(bsz, n_heads, q_len, head_dim)
V = V.reshape(bsz, n_heads, q_len, head_dim)
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_len, head_dim)
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_len, head_dim)
K = K.reshape(bsz, n_heads, kv_len, head_dim)
V = V.reshape(bsz, n_heads, kv_len, head_dim)
# See https://github.com/google/gemma_pytorch/commit/03e657582d17cb5a8617ebf333c1c16f3694670e
# Gemma 9b should use 256 and not 224 (hs / nah). 27b uses the below
@ -65,13 +66,15 @@ if not HAS_FLEX_ATTENTION:
Q = Q * torch.tensor(s**-0.5, dtype = Q.dtype) # Follow Keras exactly
A = torch.matmul(Q, K.transpose(2, 3))
A = t * torch.tanh(A / t) # Logit softcapping
A += causal_mask[:q_len, :q_len]
# Much slower in torch compile!
# A.masked_fill_(causal_mask[:q_len, :q_len], -float("inf"))
# Handle both 2D static masks and 4D dynamic masks
if causal_mask.dim() >= 3:
A += causal_mask
else:
A += causal_mask[:actual_q_len, :kv_len]
A = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32).to(Q.dtype)
A = torch.matmul(A, V)
A = A.transpose(1, 2).contiguous()
A = A.reshape(bsz, q_len, n_heads * head_dim)
A = A.reshape(bsz, actual_q_len, n_heads * head_dim)
return A
create_flex_attention_causal_mask = None
@ -151,17 +154,18 @@ torch_tanh = torch.tanh
torch_nn_functional_softmax = torch.nn.functional.softmax
def slow_inference_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len):
def slow_inference_attention_softcapping(Q, K, V, causal_mask, self, bsz, kv_len):
n_heads = self.config.num_attention_heads
head_dim = self.head_dim
n_kv_heads = self.config.num_key_value_heads
n_groups = self.num_key_value_groups
actual_q_len = Q.shape[-2]
# Grouped query attention
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
K = K.reshape(bsz, n_heads, q_len, head_dim)
V = V.reshape(bsz, n_heads, q_len, head_dim)
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_len, head_dim)
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_len, head_dim)
K = K.reshape(bsz, n_heads, kv_len, head_dim)
V = V.reshape(bsz, n_heads, kv_len, head_dim)
# See https://github.com/google/gemma_pytorch/commit/03e657582d17cb5a8617ebf333c1c16f3694670e
# Gemma 9b should use 256 and not 224 (hs / nah). 27b uses the below
@ -177,11 +181,13 @@ def slow_inference_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len)
A /= t
torch_tanh(A, out = A)
A *= t
A += causal_mask[:q_len, :q_len]
# Much slower in torch compile!
# A.masked_fill_(causal_mask[:q_len, :q_len], -float("inf"))
# Handle both 2D static masks and 4D dynamic masks
if causal_mask.dim() >= 3:
A += causal_mask
else:
A += causal_mask[:actual_q_len, :kv_len]
A = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32).to(Q.dtype)
A = torch_matmul(A, V)
A = A.transpose(1, 2).contiguous()
A = A.reshape(bsz, q_len, n_heads * head_dim)
A = A.reshape(bsz, actual_q_len, n_heads * head_dim)
return A