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. 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