From c9f406e33ce3211106638797b3dde70945c2bf7d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonas=20M=2E=20K=C3=BCbler?= <jonas@priorlabs.ai>
Date: Thu, 8 Oct 2026 08:08:49 +0200
Subject: [PATCH 1/2] fix: keep cached attention aligned with ensemble members

---
 src/tabpfn/architectures/tabpfn_v3.py     | 25 ++++----------
 src/tabpfn/architectures/tabpfn_v3_5.py   | 25 ++++----------
 tests/test_architectures/test_kv_cache.py | 42 +++++++++++++++++++++++
 3 files changed, 54 insertions(+), 38 deletions(-)

diff --git a/src/tabpfn/architectures/tabpfn_v3.py b/src/tabpfn/architectures/tabpfn_v3.py
index 0ce35f6a3..10ba61e59 100644
--- a/src/tabpfn/architectures/tabpfn_v3.py
+++ b/src/tabpfn/architectures/tabpfn_v3.py
@@ -1354,27 +1354,14 @@ def forward(
             )
             x_BRE = x_BRE + attn_out
         elif cached_kv is not None:
-            # Use cached KV -- chunking over test batch is fine
-            # TODO: Performance test this as it might not be needed.
-            def _attn_fn_cached(
-                x: torch.Tensor,
-                single_eval_pos: int | None = None,
-            ) -> torch.Tensor:
-                out, _ = self.icl_attention(
-                    self.layernorm(x),
-                    single_eval_pos=single_eval_pos,
-                    cached_kv=cached_kv,
-                )
-                return out
-
-            x_BRE = chunked_evaluate_maybe_inplace(
-                _attn_fn_cached,
-                x_BRE,
-                save_peak_memory_factor=save_peak_memory_factor,
-                residual=True,
-                batch_dims=1,
+            # Query and cache batch indices must stay aligned; the generic
+            # chunking helper slices only the queries, not their cached K/V.
+            attn_out, _ = self.icl_attention(
+                self.layernorm(x_BRE),
                 single_eval_pos=single_eval_pos,
+                cached_kv=cached_kv,
             )
+            x_BRE = x_BRE + attn_out
         else:
             # Default path -- no cache
             def _attn_fn(
diff --git a/src/tabpfn/architectures/tabpfn_v3_5.py b/src/tabpfn/architectures/tabpfn_v3_5.py
index 234d6dc00..a420437f2 100644
--- a/src/tabpfn/architectures/tabpfn_v3_5.py
+++ b/src/tabpfn/architectures/tabpfn_v3_5.py
@@ -1620,27 +1620,14 @@ def forward(
             )
             x_BRE = x_BRE + attn_out
         elif cached_kv is not None:
-            # Use cached KV -- chunking over test batch is fine
-            # TODO: Performance test this as it might not be needed.
-            def _attn_fn_cached(
-                x: torch.Tensor,
-                single_eval_pos: int | None = None,
-            ) -> torch.Tensor:
-                out, _ = self.icl_attention(
-                    self.layernorm(x),
-                    single_eval_pos=single_eval_pos,
-                    cached_kv=cached_kv,
-                )
-                return out
-
-            x_BRE = chunked_evaluate_maybe_inplace(
-                _attn_fn_cached,
-                x_BRE,
-                save_peak_memory_factor=save_peak_memory_factor,
-                residual=True,
-                batch_dims=1,
+            # Query and cache batch indices must stay aligned; the generic
+            # chunking helper slices only the queries, not their cached K/V.
+            attn_out, _ = self.icl_attention(
+                self.layernorm(x_BRE),
                 single_eval_pos=single_eval_pos,
+                cached_kv=cached_kv,
             )
+            x_BRE = x_BRE + attn_out
         else:
             # Default path -- no cache
             def _attn_fn(
diff --git a/tests/test_architectures/test_kv_cache.py b/tests/test_architectures/test_kv_cache.py
index cccf76c26..8fafae7b8 100644
--- a/tests/test_architectures/test_kv_cache.py
+++ b/tests/test_architectures/test_kv_cache.py
@@ -24,6 +24,48 @@
 from tabpfn.inference import _resolve_kv_cache_precision
 
 
+@pytest.mark.parametrize(
+    "block_type",
+    [tabpfn_v3.ICLTransformerBlock, tabpfn_v3_5.ICLTransformerBlock],
+    ids=["v3", "v3.5"],
+)
+@pytest.mark.parametrize("quantized", [False, True])
+@torch.no_grad()
+def test__cached_attention__memory_saving_preserves_each_members_cache(
+    block_type: type[tabpfn_v3.ICLTransformerBlock]
+    | type[tabpfn_v3_5.ICLTransformerBlock],
+    quantized: bool,
+) -> None:
+    block = block_type(
+        emsize=8, nhead=1, dim_feedforward=16, norm_factory=torch.nn.LayerNorm
+    ).eval()
+    # The default zero projection would hide incorrect attention output.
+    block.icl_attention.out_projection.weight.copy_(torch.eye(8))
+    x = torch.randn(2, 1, 8)
+    caches = [
+        KVCacheEntry(key=torch.randn(1, 5, 1, 8), value=torch.full((1, 5, 1, 8), value))
+        for value in (0.0, 1.0)
+    ]
+    batched_cache = KVCacheEntry.concatenate(caches)
+    expected = torch.cat(
+        [
+            block(
+                x[i : i + 1].clone(),
+                single_eval_pos=0,
+                cached_kv=cache.quantize() if quantized else cache,
+            )[0]
+            for i, cache in enumerate(caches)
+        ]
+    )
+    actual, _ = block(
+        x.clone(),
+        single_eval_pos=0,
+        save_peak_memory_factor=8,
+        cached_kv=batched_cache.quantize() if quantized else batched_cache,
+    )
+    torch.testing.assert_close(actual, expected)
+
+
 def _kv_tensor(seed: int = 0) -> torch.Tensor:
     torch.manual_seed(seed)
     # (B, N, num_kv_heads, head_dim), with per-head magnitude variation so

From f7801821a7eaa5bbc75a2b76254a5c5f74c40441 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonas=20M=2E=20K=C3=BCbler?= <jonas@priorlabs.ai>
Date: Thu, 8 Oct 2026 08:11:02 +0200
Subject: [PATCH 2/2] docs: add cached attention fix changelog

---
 changelog/1365.fixed.md | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog/1365.fixed.md

diff --git a/changelog/1365.fixed.md b/changelog/1365.fixed.md
new file mode 100644
index 000000000..bd55a9e69
--- /dev/null
+++ b/changelog/1365.fixed.md
@@ -0,0 +1 @@
+Fix incorrect predictions and repeated cache dequantization for batched TabPFN-3 and TabPFN-3.5 cached inference with memory saving enabled.