diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5039769..aca58d1f5 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/mtplx/generation.py b/mtplx/generation.py index ebbd8eca7..0b93442c8 100644 --- a/mtplx/generation.py +++ b/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}" diff --git a/mtplx/models/qwen4_exp.py b/mtplx/models/qwen4_exp.py index 1d80eded7..a23895a34 100644 --- a/mtplx/models/qwen4_exp.py +++ b/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( diff --git a/mtplx/one_copy.py b/mtplx/one_copy.py index 2c8635279..c3b564223 100644 --- a/mtplx/one_copy.py +++ b/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,8 +166,9 @@ 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) @@ -171,7 +176,13 @@ def resize_bill(cache: Any, rows: int) -> int | 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 diff --git a/tests/test_qsa_cache_resize_frontier.py b/tests/test_qsa_cache_resize_frontier.py new file mode 100644 index 000000000..f9661057c --- /dev/null +++ b/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)