fix: guard QSA cache frontiers during resize and extension

CHANGELOG.md

@@ -4,6 +4,12 @@ All notable user-facing changes to MTPLX. The format is based on
 [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow
 [Semantic Versioning](https://semver.org/).
 
+## [Unreleased]
+
+### Fixed
+
+- Flash-Next cache rewinds no longer retain pooled-block counts beyond the KV frontier. Capacity planning can decline an unsafe in-place resize, and a resize rejected during verification falls back to eager verification while preserving live rows (#592, #605).
+
 ## [2.12.2] - 2026-10-03
 
 ### Fixed

mtplx/generation.py

@@ -107,6 +107,7 @@
 )
 from .native_mlp import set_native_mlp_context
 from .one_copy import (
+    QSAResizeWouldCut,
     hand_back_on_raise,
     held_qsa_rows,
     layers_to_resize,
@@ -936,6 +937,18 @@ def _qwen4_fixed_m4_resize_held(
         # leaves that layer as it was: the request answers on the eager
         # verifier, as a refused layer does. Anything else is not the
         # resize's to absorb.
+        if isinstance(exc, QSAResizeWouldCut):
+            # If the live frontier changed after the estimate, keep its
+            # buffers intact and use eager verification.
+            receipt.update(
+                engaged=False, reason="cache_resize_guard",
+                resize_error=f"{type(exc).__name__}: {exc}",
+            )
+            _announce_qwen4_fixed_m4_skip(
+                f"QSA resize to {rows} rows would cut live cache; "
+                "keeping eager verification"
+            )
+            return False
         if not is_allocation_failure(exc):
             raise
         failure = f"{type(exc).__name__}: {exc}"

mtplx/models/qwen4_exp.py

@@ -3380,7 +3380,11 @@ def _extend_pooled(self, cache: QSACache, total: int) -> Optional[mx.array]:
         if getattr(cache, "fixed_capacity", False):
             return self._extend_pooled_fixed(cache, total)
         nb_total = total // self.ratio
-        nb_old = min(cache.pooled_len, nb_total)
+        # After a KV rewind, recompute newly completed blocks instead of
+        # selecting stale pooled keys.
+        # The indexer runs before KV advances through this forward's rows.
+        nb_old = min(cache.pooled_len, cache.kv.offset // self.ratio, nb_total)
+        cache.pooled_len = nb_old
         if nb_total > nb_old:
             fresh = cache.raw_keys[:, nb_old * self.ratio : nb_total * self.ratio, :]
             if self._prepare_kernel_supported(

mtplx/one_copy.py

@@ -34,6 +34,10 @@
 _DISABLE_VALUES = {"0", "false", "no", "off"}
 
 
+class QSAResizeWouldCut(ValueError):
+    """The requested capacity would truncate a cache's live logical rows."""
+
+
 def one_copy_enabled() -> bool:
     """The one-copy conversation store (MTPLX_ONE_COPY, on unless set off)."""
 
@@ -162,16 +166,23 @@ def resize_bill(cache: Any, rows: int) -> int | None:
     buffers written beside its old ones. On the 2026-10-01 turn (12 layers
     of 147,456 rows, a 148,480-row bank) that is 0.38 GB, where the padded
     copy of the same bank was priced at 4.22 GB. None when a QSA layer of
-    ``cache`` cannot be resized (a promoted bank, a missing buffer): the
-    promotion then copies, as it always did.
+    ``cache`` cannot be resized (a promoted bank, a missing buffer, or a
+    target below the live frontier): the promotion then copies, as it
+    always did.
     """
 
     entries = _resizable(cache)
     if entries is None:
         return None
     gained = peak = 0
     for entry in entries:
-        new, old = _layer_bytes(entry, rows)
+        try:
+            new, old = _layer_bytes(entry, rows)
+        except QSAResizeWouldCut:
+            # Capacity planning may decline an in-place resize. Let the
+            # caller price its copy/eager path; an actual resize must still
+            # reject cutting live rows.
+            return None
         peak = max(peak, gained + new)
         gained += new - old
     return peak
@@ -204,18 +215,26 @@ def _resized_buffers(entry: Any, rows: int) -> list[tuple[str, Any, int, int]]:
 
     rows = int(rows)
     offset = int(entry.kv.offset)
+    ratio = max(1, int(entry.ratio))
+    # A rewound lease can retain its previous pooled count. Blocks beyond
+    # the KV frontier are invalid, just as in QSACache.trim/state and
+    # TensorOffsetQSACache.demote (#592).
+    valid_pooled = min(int(entry.pooled_len), offset // ratio)
+    entry.pooled_len = valid_pooled
     wanted = (
         ("keys", entry.kv.keys, 2, rows, offset),
         ("values", entry.kv.values, 2, rows, offset),
         ("raw_keys", entry.raw_keys, 1, rows, offset),
-        ("pooled", entry.pooled, 1, rows // max(1, int(entry.ratio)), int(entry.pooled_len)),
+        ("pooled", entry.pooled, 1, rows // ratio, valid_pooled),
     )
     resized = []
     for name, buffer, axis, target, held in wanted:
         if int(buffer.shape[axis]) == target:
             continue
         if target < held:
-            raise ValueError(f"QSA {name} hold {held} rows; resizing to {target} would cut them")
+            raise QSAResizeWouldCut(
+                f"QSA {name} hold {held} rows; resizing to {target} would cut them"
+            )
         resized.append((name, buffer, axis, target))
     return resized
 

tests/test_qsa_cache_resize_frontier.py

@@ -0,0 +1,148 @@
+"""QSA cache rewinds and declined resizes must preserve the logical frontier."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import mlx.core as mx
+
+import mtplx.models.qwen4_exp as qwen
+import mtplx.one_copy as oc
+from mtplx import generation
+
+
+def cache_at(offset, pooled_len, capacity=80000, ratio=4):
+    cache = qwen.QSACache(ratio)
+    cache.kv.keys = mx.arange(capacity * 2).reshape(1, 1, capacity, 2)
+    cache.kv.values = -cache.kv.keys
+    cache.kv.offset = offset
+    cache.raw_keys = mx.arange(capacity * 8, dtype=mx.float32).reshape(1, capacity, 8)
+    cache.pooled = mx.full((1, capacity // ratio, 8), -123.0)
+    cache.pooled_len = pooled_len
+    mx.eval(cache.kv.keys, cache.kv.values, cache.raw_keys, cache.pooled)
+    return cache
+
+
+class CacheRegressionTests(unittest.TestCase):
+    def test_exact_reported_stale_length_resize_preserves_active_prefix(self):
+        cache = cache_at(27957, 19843)
+        expected = cache.state
+        self.assertIsNotNone(oc.resize_bill([cache], 32768))
+        self.assertEqual(oc.resize_qsa_buffers([cache], 32768), 0)
+        self.assertEqual(cache.pooled_len, 27957 // 4)
+        self.assertEqual(cache.pooled.shape[1], 8192)
+        for actual, before in zip(cache.state, expected):
+            self.assertTrue(mx.array_equal(actual, before).item())
+
+    def test_normal_resize_does_not_change_logical_state(self):
+        cache = cache_at(17, 4, capacity=32)
+        expected = cache.state
+        oc.resize_qsa_buffers([cache], 64)
+        self.assertEqual(cache.pooled_len, 4)
+        for actual, before in zip(cache.state, expected):
+            self.assertTrue(mx.array_equal(actual, before).item())
+
+    def test_already_matching_buffers_still_normalize_stale_count(self):
+        cache = cache_at(17, 7, capacity=32)
+        self.assertEqual(oc.resize_qsa_buffers([cache], 32), 0)
+        self.assertEqual(cache.pooled_len, 4)
+
+    def test_other_compression_ratio(self):
+        cache = cache_at(17, 4, capacity=64, ratio=8)
+        oc.resize_qsa_buffers([cache], 32)
+        self.assertEqual(cache.pooled_len, 2)
+        self.assertEqual(cache.pooled.shape[1], 4)
+
+    def test_real_live_rows_cannot_be_truncated(self):
+        cache = cache_at(48, 12, capacity=64)
+        before = cache.state
+        self.assertIsNone(oc.resize_bill([cache], 32))
+        with self.assertRaises(oc.QSAResizeWouldCut):
+            oc.resize_qsa_buffers([cache], 32)
+        self.assertEqual(cache.kv.offset, 48)
+        for actual, old in zip(cache.state, before):
+            self.assertTrue(mx.array_equal(actual, old).item())
+
+    def test_promotion_estimate_falls_back_to_copy(self):
+        cache = cache_at(48, 12, capacity=64)
+        with (
+            patch.object(
+                generation, "_qwen4_fixed_m4_adoption_rows", return_value=None
+            ),
+            patch.object(qwen, "_qsa_gather_enabled", return_value=True),
+            patch.object(qwen, "_qsa_gather_min_context", return_value=0),
+        ):
+            result = generation._qwen4_fixed_m4_promotion_bill(
+                None, 28, SimpleNamespace(), 32, [cache], 64, 100
+            )
+        self.assertEqual(result, ("copied", 3200, 32))
+
+    def test_execution_guard_keeps_eager_path_on_unsafe_resize(self):
+        cache = cache_at(48, 12, capacity=64)
+        receipt = {}
+        with patch.object(generation, "_announce_qwen4_fixed_m4_skip"):
+            result = generation._qwen4_fixed_m4_resize_held(
+                None,
+                [cache],
+                SimpleNamespace(resize_rows=32),
+                receipt,
+                session_bank=None,
+                protect_ids=[],
+            )
+        self.assertFalse(result)
+        self.assertEqual(receipt["reason"], "cache_resize_guard")
+        self.assertEqual(cache.kv.offset, 48)
+
+    def indexer(self):
+        return qwen.QSAIndexer(
+            qwen.TextArgs(
+                hidden_size=8,
+                head_dim=8,
+                num_hidden_layers=1,
+                indexer_n_heads=1,
+                indexer_kv_heads=1,
+                indexer_head_dim=8,
+                indexer_budget=8,
+                partial_rotary_factor=1.0,
+            )
+        )
+
+    def test_rewound_extension_recomputes_new_blocks_and_mirror(self):
+        cache = cache_at(8, 7, capacity=32)
+        indexer = self.indexer()
+        # Preserve two valid blocks, overwrite stale blocks after offset 8.
+        old_prefix = cache.pooled[:, :2, :]
+        cache.pooled_f32_view(7)
+        expected = indexer._pool_keys_eager(cache.raw_keys[:, 8:16, :], 2, 4)
+        with patch.object(
+            qwen.QSAIndexer, "_prepare_kernel_supported", return_value=False
+        ):
+            result = indexer._extend_pooled(cache, 16)
+        self.assertTrue(mx.array_equal(result[:, :2, :], old_prefix).item())
+        self.assertTrue(mx.array_equal(result[:, 2:4, :], expected).item())
+        mirror = mx.swapaxes(cache.pooled_f32_view(4)[:, 0], 1, 2)
+        self.assertTrue(mx.array_equal(mirror, result.astype(mx.float32)).item())
+        self.assertEqual(cache.pooled_len, 4)
+
+    def test_rewound_extension_without_complete_block_clears_length(self):
+        cache = cache_at(0, 7, capacity=32)
+        result = self.indexer()._extend_pooled(cache, 3)
+        self.assertIsNone(result)
+        self.assertEqual(cache.pooled_len, 0)
+
+    def test_regular_extension_matches_full_recomputation(self):
+        cache = cache_at(8, 2, capacity=32)
+        indexer = self.indexer()
+        cache.write_pooled(
+            indexer._pool_keys_eager(cache.raw_keys[:, :8, :], 0, 2), 0, 2
+        )
+        expected = indexer._pool_keys_eager(cache.raw_keys[:, :16, :], 0, 4)
+        with patch.object(
+            qwen.QSAIndexer, "_prepare_kernel_supported", return_value=False
+        ):
+            result = indexer._extend_pooled(cache, 16)
+        self.assertTrue(mx.array_equal(result, expected).item())
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)