From 35cf17878331f2455c1535195f0d81507a57ae60 Mon Sep 17 00:00:00 2001
From: Can <cansolakoglu130@gmail.com>
Date: Sat, 17 Jan 2026 08:15:32 +0300
Subject: [PATCH] Add packed sequence fallback for seq_kv_cache streaming

Detects `packed_seq_lengths` in forward_inputs and bypasses seq_kv_cache chunking to avoid KV cache corruption with packed sequences. Falls back to batch microbatching (if configured) or full reference forward pass. Adds test verifying fallback triggers `_compute_kl_batch_micro` with microbatch_size=1 when packed sequences present.
---
 tests/test_asft.py     | 60 ++++++++++++++++++++++++++++++++++++++++++
 unsloth/losses/asft.py | 41 +++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+)

diff --git a/tests/test_asft.py b/tests/test_asft.py
index 6428604971..ebc9ed7794 100644
--- a/tests/test_asft.py
+++ b/tests/test_asft.py
@@ -965,6 +965,66 @@ class TestSeqKVCacheStreaming:
 
         assert kl.shape == (batch_size, seq_len)
 
+    def test_seq_kv_cache_falls_back_with_packed_sequences(self):
+        """Test that packed sequences bypass seq_kv_cache chunking."""
+        batch_size, seq_len, vocab_size = 2, 4, 3
+        cur_logits = torch.randn(batch_size, seq_len, vocab_size)
+        shift_labels = torch.zeros(batch_size, seq_len, dtype = torch.long)
+        valid_mask = shift_labels != -100
+        input_ids = torch.arange(seq_len).repeat(batch_size, 1)
+        packed_seq_lengths = torch.tensor([2, 2], dtype = torch.int32)
+
+        class DummyModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = SimpleNamespace(
+                    use_cache = True,
+                    final_logit_softcapping = 0,
+                    logit_scale = 0,
+                )
+
+        model = DummyModel()
+        ref_forward = MagicMock()
+        forward_inputs = {
+            "input_ids": input_ids,
+            "packed_seq_lengths": packed_seq_lengths,
+        }
+
+        def batch_side_effect(
+            model,
+            cur_logits,
+            shift_labels,
+            valid_mask,
+            ref_forward,
+            forward_inputs,
+            microbatch_size,
+            logit_softcapping = 0,
+            logit_scaling = 0,
+            force_fp32 = True,
+            kl_direction = "forward",
+        ):
+            batch, seq_len = shift_labels.shape
+            return torch.zeros(batch, seq_len, device = shift_labels.device)
+
+        with patch(
+            "unsloth.losses.asft._compute_kl_batch_micro",
+            side_effect = batch_side_effect,
+        ) as batch_mock:
+            kl = _compute_kl_seq_kv_cache(
+                model,
+                cur_logits,
+                shift_labels,
+                valid_mask,
+                ref_forward,
+                forward_inputs,
+                seq_chunk_size = 2,
+            )
+
+        assert batch_mock.called
+        assert batch_mock.call_args[0][6] == 1
+        assert not ref_forward.called
+        assert kl.shape == (batch_size, seq_len)
+
     def test_config_immutability_when_none_values(self, simple_model):
         """Test that streaming_config is not mutated when values are None."""
         config = ASFTStreamingConfig(
diff --git a/unsloth/losses/asft.py b/unsloth/losses/asft.py
index 4b4b47bda8..bf23fc8dbb 100644
--- a/unsloth/losses/asft.py
+++ b/unsloth/losses/asft.py
@@ -603,6 +603,47 @@ def _compute_kl_seq_kv_cache(
     batch_size, seq_len, vocab_size = cur_logits.shape
     device = cur_logits.device
 
+    packed_seq_lengths = forward_inputs.get("packed_seq_lengths", None)
+    if packed_seq_lengths is not None:
+        # Avoid seq_kv_cache with packed sequences; fall back to batch/full reference.
+        fallback_microbatch = None
+        if microbatch_size is not None and microbatch_size < batch_size:
+            fallback_microbatch = microbatch_size
+        elif allow_auto_microbatch_fallback:
+            fallback_microbatch = max(
+                1, batch_size // _DEFAULT_REF_MICROBATCH_DIVISOR
+            )
+            if fallback_microbatch >= batch_size:
+                fallback_microbatch = None
+        if fallback_microbatch is not None:
+            return _compute_kl_batch_micro(
+                model,
+                cur_logits,
+                shift_labels,
+                valid_mask,
+                ref_forward,
+                forward_inputs,
+                fallback_microbatch,
+                logit_softcapping,
+                logit_scaling,
+                force_fp32,
+                kl_direction,
+            )
+        ref_outputs = ref_forward(**forward_inputs)
+        ref_logits, _ = _unwrap_reference_outputs(ref_outputs)
+        kl_full = _compute_kl_divergence(
+            cur_logits,
+            ref_logits,
+            model,
+            logit_softcapping,
+            logit_scaling,
+            force_fp32,
+            kl_direction,
+        )
+        if kl_full.dim() == 1:
+            kl_full = kl_full.view(batch_size, seq_len)
+        return kl_full
+
     if microbatch_size is not None:
         microbatch_size = max(1, microbatch_size)
     if microbatch_size is not None and microbatch_size < batch_size: