diff --git a/tests/test_past_kv_models.py b/tests/test_past_kv_models.py
new file mode 100644
index 0000000000..6f3663b972
--- /dev/null
+++ b/tests/test_past_kv_models.py
@@ -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()
diff --git a/tests/test_past_kv_utils.py b/tests/test_past_kv_utils.py
new file mode 100644
index 0000000000..75feefdf4f
--- /dev/null
+++ b/tests/test_past_kv_utils.py
@@ -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()
diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py
index b94ff56dec..582932c81e 100644
--- a/unsloth/kernels/flex_attention.py
+++ b/unsloth/kernels/flex_attention.py
@@ -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