fix: preserve per-member caches with memory saving

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.

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(

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(

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