From 2f41437f00dae6367944475b1c0be57bf481c9ac Mon Sep 17 00:00:00 2001 From: Manish Mathai Date: Tue, 20 Oct 2026 00:35:49 -0700 Subject: [PATCH] fix: count pooled blocks from the restored keys, not the donor table A conversation restored from the SSD tier below its end kept the donor's whole pooled indexer table, and the cache counted its valid blocks from that table's length. The indexer then computed no blocks for the new tokens and scored attention against another conversation's keys, so the model lost the request (#635). --- CHANGELOG.md | 6 +++++ mtplx/models/qwen4_exp.py | 11 +++++++- tests/test_qsa_state_restore_frontier.py | 34 ++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 tests/test_qsa_state_restore_frontier.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5039769..750351679 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 + +- **A conversation restored from the SSD tier below its end no longer attends through another conversation's block keys.** On the Flash-Next pair, a sub-prefix restore decoded the keys and values up to the restore point but kept the donor's whole pooled indexer table, and the cache took its count of valid pooled blocks from that table's length. The indexer then found every block of the new tokens already pooled, computed none, and scored attention against the donor's keys, so the answer lost the request and cited facts that are not in the prompt (#635). The count now follows the restored keys, as it already did after a RAM restore and trim. + ## [2.12.2] - 2026-10-03 ### Fixed diff --git a/mtplx/models/qwen4_exp.py b/mtplx/models/qwen4_exp.py index 1d80eded7..3dfbc1331 100644 --- a/mtplx/models/qwen4_exp.py +++ b/mtplx/models/qwen4_exp.py @@ -3268,7 +3268,16 @@ def state(self, v): self.kv.state = (keys, values) self.raw_keys = raw self.pooled = pooled - self.pooled_len = 0 if pooled is None else pooled.shape[1] + # The valid pooled blocks are those the restored KV covers. A snapshot + # decoded for a sub-prefix restore (cold tier, ``decode_payload_prefix``) + # cuts keys and values to the restore point but hands over the donor's + # whole pooled table: trusting its length made the indexer skip the + # blocks of every later token and score against the donor's keys. + self.pooled_len = ( + 0 + if pooled is None + else min(pooled.shape[1], self.kv.offset // self.ratio) + ) # Derived mirror: rebuilt lazily on the first pooled_f32_view read. # Restored snapshots stay 4-leaf — the state contract is unchanged. self.pooled_f32_t = None diff --git a/tests/test_qsa_state_restore_frontier.py b/tests/test_qsa_state_restore_frontier.py new file mode 100644 index 000000000..ba41e4719 --- /dev/null +++ b/tests/test_qsa_state_restore_frontier.py @@ -0,0 +1,34 @@ +"""A QSACache restored from a prefix-decoded snapshot must not trust the donor's pooled table length.""" +import mlx.core as mx +from mtplx.models.qwen4_exp import QSACache + + +def _leaves(kv_rows, donor_rows, ratio=4, dim=8): + keys = mx.zeros((1, 2, kv_rows, 16), dtype=mx.bfloat16) + values = mx.zeros((1, 2, kv_rows, 16), dtype=mx.bfloat16) + raw = mx.zeros((1, donor_rows, dim), dtype=mx.bfloat16) + pooled = mx.zeros((1, donor_rows // ratio, dim), dtype=mx.bfloat16) + return keys, values, raw, pooled + + +def test_prefix_restore_clamps_pooled_len_to_the_restored_kv(): + cache = QSACache(compress_ratio=4) + cache.state = _leaves(kv_rows=2048, donor_rows=58073) + assert cache.offset == 2048 + assert cache.pooled_len == 2048 // 4 + + +def test_full_state_round_trip_keeps_its_pooled_blocks(): + cache = QSACache(compress_ratio=4) + cache.state = _leaves(kv_rows=2048, donor_rows=2048) + assert cache.pooled_len == 512 + again = QSACache(compress_ratio=4) + again.state = cache.state + assert again.pooled_len == 512 + + +def test_snapshot_without_pooled_has_no_valid_blocks(): + cache = QSACache(compress_ratio=4) + keys, values, raw, _ = _leaves(kv_rows=256, donor_rows=256) + cache.state = (keys, values, raw, None) + assert cache.pooled_len == 0