Skip to content

one_copy: fall back to copying when a QSA resize would cut held bank rows - #619

Open
oliverlavery wants to merge 5 commits into
youssofal:mainfrom
oliverlavery:fix/qsa-hold-cut-admission
Open

oliverlavery wants to merge 5 commits into
youssofal:mainfrom
oliverlavery:fix/qsa-hold-cut-admission

Conversation

@oliverlavery

Copy link
Copy Markdown

Problem

On 2.12.2, a QSA resize that would reduce row count below rows owned by a held bank entry raises ValueError: QSA {name} hold {N} rows; resizing to {M} would cut them (one_copy.py _resized_buffers). The raise escapes through the admission path: generation.py prices the fixed-M4 bank via resize_bill() unguarded, so the ValueError surfaces as a server_error and the request dies with 0 tokens.

Repro (Flash-Next MTPLX build, 262k ctx): serve, run a ~50k-token session, then send a short prompt from a new session — fails with the hold-cut error until the engine restarts. Every differently-shaped new session fails, so one long session poisons the lane. Hit in production via OpenClaw subagent spawning (3 consecutive failures, 0 tokens each).

Fix

_layer_bytes() returns None for the hold-cut case instead of raising; resize_bill() propagates None so admission selects the copy path — the same fallback its docstring already promises for promoted banks and missing buffers. _pending_layers() skips the unresizable layer so the execution path skips it too. The fixed-M4 resize caller additionally degrades to the eager verifier if a hold-cut ValueError still surfaces (admission/execution race).

Validation

  • All 63 tests in test_one_copy_admission.py, test_one_copy_conversation.py, test_one_copy_cancelled_answer.py, test_demotions_qsa_prefill_sites.py, test_memguard_dense_qsa.py pass.
  • End-to-end: big session → immediate new-session spawn → served (previously server_error). Sustained mixed-shape traffic clean.

…rows

A QSA layer resize refuses (ValueError) when the target rows would cut
into rows a held bank entry owns — an entry from an earlier, longer
conversation still in the session bank. The admission path priced the
bank via resize_bill() unguarded, so the ValueError surfaced as a
server_error and failed every new session whose context shape differed
from the held entry's, until the engine restarted.

_layer_bytes() now returns None for the hold-cut case; resize_bill()
propagates None so admission selects the copy path (the same fallback a
promoted bank or a missing buffer already takes), and _pending_layers()
excludes the unresizable layer so the execution path skips it. The
fixed-M4 resize caller additionally degrades to the eager verifier if a
hold-cut ValueError still surfaces (admission/execution race).

Repro: serve Flash-Next, run a ~50k-token session, then send a short
prompt from a new session — pre-patch this fails with 'QSA pooled hold
N rows; resizing to M would cut them' (server_error, 0 tokens); post-
patch it serves.
@jerryjvl

Copy link
Copy Markdown

Hi @oliverlavery, I'm hitting this on 2.12.2 too (Flash-Next, M5 Ultra). I see the fix was reverted on 14 Oct (78ebeb2). Could you say why? I suspect the copy fallback: at 128k+ tokens it prices a multi-GB padded bank, which can evict idle sessions or force eager verify.

The problem seems to be a stale pooled_len, not real data. pooled_len is written before attention advances kv.offset, so it can be left ahead of it. Every other reader clamps it to offset // ratio, but _resized_buffers uses the raw value. In my failures the keys check passed (offset <= rows), yet pooled claimed 42,593 blocks against a 32,768 target, so those extra blocks can't be valid.

A narrower fix: clamp it, and store the clamp. Against 2.12.2:

--- a/mtplx/one_copy.py
+++ b/mtplx/one_copy.py
@@ -194,6 +194,13 @@ def _resizable(cache: Any) -> list[Any] | None:
     return entries

 
+def _valid_pooled_len(entry: Any) -> int:
+    """Pooled blocks ``entry`` really holds: ``pooled_len`` up to the frontier."""
+
+    return min(int(entry.pooled_len), int(entry.kv.offset) // max(1, int(entry.ratio)))
+
+
 def _resized_buffers(entry: Any, rows: int) -> list[tuple[str, Any, int, int]]:
     """(name, buffer, row axis, rows) for each buffer of ``entry`` not at the bank's rows.
 
@@ -208,7 +215,7 @@ def _resized_buffers(entry: Any, rows: int) -> list[tuple[str, Any, int, int]]:
         ("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 // max(1, int(entry.ratio)), _valid_pooled_len(entry)),
     )
     resized = []
     for name, buffer, axis, target, held in wanted:
@@ -262,6 +269,9 @@ def _resize_layer(entry: Any, rows: int) -> None:
         entry.raw_keys = resized["raw_keys"]
     if "pooled" in resized:
         entry.pooled = resized["pooled"]
+        # Blocks past the frontier went with the cut; say so, as
+        # ``QSACache.trim`` does (restore staging reads it unclamped).
+        entry.pooled_len = _valid_pooled_len(entry)
         # Derived from ``pooled``; rebuilt on the next stock read.
         entry.pooled_f32_t = None

The second change matters: without it, qsa_restore_staging (which reads pooled_len unclamped) raises later.

The in-place resize then goes ahead with no copy. This matches what the bank's own copy promotion already does: from_qsa_cache cuts pooled to capacity // ratio without reading pooled_len. In a test with the real QSACache, the stale case resizes with the valid blocks unchanged, and a genuine cut still raises. I'm running it locally now.

A QSA resize that would cut rows a layer holds means the held cache
belongs to an earlier, longer conversation than the prompt being
admitted. Adopting, resizing, or copying it would attend over (and
answer from) the wrong conversation.

- resize_bill / _resized_buffers raise QSACutError (a ValueError) on a
  cut instead of ever pricing it as a copy of the stale cache
- resize_cut_rows exposes the cut predicate to admission
- _qwen4_fixed_m4_promotion_bill returns ("refused", 0, 0) for a
  foreign held cache; _qwen4_fixed_m4_lane_fits refuses the lane and
  the request verifies eagerly (recomputing the right state)
- _qwen4_fixed_m4_resize_held catches QSACutError from the
  admission/execution race: receipt records held_cache_foreign, the
  lane skips, nothing is silently half-resized
- _pending_layers can never silently skip a cut layer: the raise
  propagates through resize_qsa_buffers/layers_to_resize
- tests: pin the new contract (loud refusal, buffers untouched, the
  bill never prices a cut as a copy) + a lane-refusal end-to-end test

Fixes the stale previous-turn answers and the under-billed
copy-induced MLX wedges seen with a9255ec (reverted in b86ceb3).

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants