From 612b9d931885043aba0d57183731da96b05874f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20M=2E=20K=C3=BCbler?= Date: Thu, 1 Oct 2026 13:53:52 +0200 Subject: [PATCH 1/4] fix: avoid eager Triton compilation in regression target encoding --- src/tabpfn/architectures/tabpfn_v3.py | 8 ++- src/tabpfn/architectures/tabpfn_v3_5.py | 8 ++- .../test_regression_target_projection.py | 57 +++++++++++++++++++ 3 files changed, 69 insertions(+), 4 deletions(-) create mode 100644 tests/test_architectures/test_regression_target_projection.py diff --git a/src/tabpfn/architectures/tabpfn_v3.py b/src/tabpfn/architectures/tabpfn_v3.py index 9a0c93f6d..0ce35f6a3 100644 --- a/src/tabpfn/architectures/tabpfn_v3.py +++ b/src/tabpfn/architectures/tabpfn_v3.py @@ -2144,13 +2144,17 @@ def _embed_col_y(self, y_BN: torch.Tensor) -> torch.Tensor: """Embed y_train for the col stage → (B, T, E).""" if self.task_type == "multiclass": return self.col_y_encoder(y_BN) - return self.col_y_encoder(y_BN.unsqueeze(-1)) + # A 2D projection avoids eager Triton BMM and its C compiler requirement. + y_emb = self.col_y_encoder(y_BN.reshape(-1, 1)) + return y_emb.reshape(*y_BN.shape, y_emb.shape[-1]) def _embed_icl_y(self, y_BN: torch.Tensor) -> torch.Tensor: """Embed y_train for the ICL stage → (B, T, D).""" if self.task_type == "multiclass": return self.icl_y_encoder(y_BN) - return self.icl_y_encoder(y_BN.unsqueeze(-1)) + # A 2D projection avoids eager Triton BMM and its C compiler requirement. + y_emb = self.icl_y_encoder(y_BN.reshape(-1, 1)) + return y_emb.reshape(*y_BN.shape, y_emb.shape[-1]) def _preprocess_raw( self, diff --git a/src/tabpfn/architectures/tabpfn_v3_5.py b/src/tabpfn/architectures/tabpfn_v3_5.py index f5422849f..21009b784 100644 --- a/src/tabpfn/architectures/tabpfn_v3_5.py +++ b/src/tabpfn/architectures/tabpfn_v3_5.py @@ -2577,7 +2577,9 @@ def _embed_col_y(self, y_BN: torch.Tensor, *, task_type: TaskType) -> torch.Tens if task_type == "multiclass": y_emb = self.col_y_encoder["multiclass"](y_BN) elif task_type == "regression": - y_emb = self.col_y_encoder["regression"](y_BN.unsqueeze(-1)) + # A 2D projection avoids eager Triton BMM and its C compiler requirement. + y_emb = self.col_y_encoder["regression"](y_BN.reshape(-1, 1)) + y_emb = y_emb.reshape(*y_BN.shape, y_emb.shape[-1]) else: raise ValueError(f"Unsupported task type: {task_type}") return self.col_y_layernorm(y_emb) @@ -2587,7 +2589,9 @@ def _embed_icl_y(self, y_BN: torch.Tensor, *, task_type: TaskType) -> torch.Tens if task_type == "multiclass": y_emb = self.icl_y_encoder["multiclass"](y_BN) elif task_type == "regression": - y_emb = self.icl_y_encoder["regression"](y_BN.unsqueeze(-1)) + # A 2D projection avoids eager Triton BMM and its C compiler requirement. + y_emb = self.icl_y_encoder["regression"](y_BN.reshape(-1, 1)) + y_emb = y_emb.reshape(*y_BN.shape, y_emb.shape[-1]) else: raise ValueError(f"Unsupported task type: {task_type}") return self.icl_y_layernorm(y_emb) diff --git a/tests/test_architectures/test_regression_target_projection.py b/tests/test_architectures/test_regression_target_projection.py new file mode 100644 index 000000000..e94b6e34a --- /dev/null +++ b/tests/test_architectures/test_regression_target_projection.py @@ -0,0 +1,57 @@ +# Copyright (c) Prior Labs GmbH 2026. + +"""Regression target projections must not require a runtime compiler.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, Literal + +import pytest +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from tests.test_architectures.test_tabpfn_v3 import _get_regression_model +from tests.test_architectures.test_tabpfn_v3_5 import _get_model + + +class _RejectBmm(TorchDispatchMode): + def __torch_dispatch__( + self, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Any: + assert func is not torch.ops.aten.bmm.default, "Target projection reached BMM" + return func(*args, **(kwargs or {})) + + +@pytest.mark.parametrize("version", ["v3", "v3.5"]) +@pytest.mark.parametrize("stage", ["col", "icl"]) +@pytest.mark.parametrize("batch", [1, 8]) +@pytest.mark.parametrize("contiguous", [False, True]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_regression_target_projection( + version: str, + stage: Literal["col", "icl"], + batch: int, + contiguous: bool, + device: str, +) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + model = (_get_regression_model() if version == "v3" else _get_model()).to(device) + targets = torch.randn(1002, batch, device=device).T + if contiguous: + targets = targets.contiguous() + targets.requires_grad_() + project = getattr(model, f"_embed_{stage}_y") + kwargs = {} if version == "v3" else {"task_type": "regression"} + expected = torch.cat([project(row[None], **kwargs) for row in targets]) + expected_grad = torch.autograd.grad(expected.square().sum(), targets)[0] + with _RejectBmm(): + actual = project(targets, **kwargs) + actual_grad = torch.autograd.grad(actual.square().sum(), targets)[0] + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual_grad, expected_grad) From fd12fe76cbf60847290ac4f1ff3fe70e4156bd43 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20M=2E=20K=C3=BCbler?= Date: Thu, 1 Oct 2026 13:56:32 +0200 Subject: [PATCH 2/4] test: reject target encoder compilation in a fresh process --- .../test_regression_target_projection.py | 65 ++++++++++++++----- 1 file changed, 47 insertions(+), 18 deletions(-) diff --git a/tests/test_architectures/test_regression_target_projection.py b/tests/test_architectures/test_regression_target_projection.py index e94b6e34a..0271acfae 100644 --- a/tests/test_architectures/test_regression_target_projection.py +++ b/tests/test_architectures/test_regression_target_projection.py @@ -4,29 +4,19 @@ from __future__ import annotations -from collections.abc import Callable -from typing import Any, Literal +import os +import subprocess +import sys +from pathlib import Path +from typing import Literal import pytest import torch -from torch.utils._python_dispatch import TorchDispatchMode from tests.test_architectures.test_tabpfn_v3 import _get_regression_model from tests.test_architectures.test_tabpfn_v3_5 import _get_model -class _RejectBmm(TorchDispatchMode): - def __torch_dispatch__( - self, - func: Callable, - types: tuple[type, ...], - args: tuple[Any, ...] = (), - kwargs: dict[str, Any] | None = None, - ) -> Any: - assert func is not torch.ops.aten.bmm.default, "Target projection reached BMM" - return func(*args, **(kwargs or {})) - - @pytest.mark.parametrize("version", ["v3", "v3.5"]) @pytest.mark.parametrize("stage", ["col", "icl"]) @pytest.mark.parametrize("batch", [1, 8]) @@ -50,8 +40,47 @@ def test_regression_target_projection( kwargs = {} if version == "v3" else {"task_type": "regression"} expected = torch.cat([project(row[None], **kwargs) for row in targets]) expected_grad = torch.autograd.grad(expected.square().sum(), targets)[0] - with _RejectBmm(): - actual = project(targets, **kwargs) - actual_grad = torch.autograd.grad(actual.square().sum(), targets)[0] + actual = project(targets, **kwargs) + actual_grad = torch.autograd.grad(actual.square().sum(), targets)[0] torch.testing.assert_close(actual, expected) torch.testing.assert_close(actual_grad, expected_grad) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +def test_regression_target_projection_without_compiler(tmp_path: Path) -> None: + result = subprocess.run( + [ + sys.executable, + "-m", + "tests.test_architectures.test_regression_target_projection", + ], + cwd=Path(__file__).resolve().parents[2], + env={**os.environ, "TRITON_CACHE_DIR": str(tmp_path)}, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def _check_without_compiler() -> None: + from triton.runtime import build # noqa: PLC0415 + + def reject_compilation(*_args: object, **_kwargs: object) -> None: + raise AssertionError("Regression target encoding requested a C compiler") + + build._build = reject_compilation + targets = torch.randn(1002, 8, device="cuda").T + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + for version in ("v3", "v3.5"): + model = ( + _get_regression_model() if version == "v3" else _get_model() + ).cuda() + kwargs = {} if version == "v3" else {"task_type": "regression"} + for stage in ("col", "icl"): + getattr(model, f"_embed_{stage}_y")(targets, **kwargs) + torch.cuda.synchronize() + + +if __name__ == "__main__": + _check_without_compiler() From 193283416abf87f93d1d38f740b83c1d4e641c38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20M=2E=20K=C3=BCbler?= Date: Thu, 1 Oct 2026 13:57:52 +0200 Subject: [PATCH 3/4] docs: add regression target encoding changelog --- changelog/1338.fixed.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog/1338.fixed.md diff --git a/changelog/1338.fixed.md b/changelog/1338.fixed.md new file mode 100644 index 000000000..131e9af6a --- /dev/null +++ b/changelog/1338.fixed.md @@ -0,0 +1 @@ +Avoid a runtime C compiler requirement in batched v3 and v3.5 regression target encoding on recent PyTorch versions. From b66bb3bae03c4b3d989c7f227ae08d848e3a93a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20M=2E=20K=C3=BCbler?= Date: Thu, 1 Oct 2026 14:21:44 +0200 Subject: [PATCH 4/4] test: trim target projection coverage and gate eager Triton check --- .../test_regression_target_projection.py | 41 ++++++++----------- 1 file changed, 16 insertions(+), 25 deletions(-) diff --git a/tests/test_architectures/test_regression_target_projection.py b/tests/test_architectures/test_regression_target_projection.py index 0271acfae..a1974283f 100644 --- a/tests/test_architectures/test_regression_target_projection.py +++ b/tests/test_architectures/test_regression_target_projection.py @@ -8,45 +8,36 @@ import subprocess import sys from pathlib import Path -from typing import Literal import pytest import torch +from torch.torch_version import TorchVersion from tests.test_architectures.test_tabpfn_v3 import _get_regression_model from tests.test_architectures.test_tabpfn_v3_5 import _get_model @pytest.mark.parametrize("version", ["v3", "v3.5"]) -@pytest.mark.parametrize("stage", ["col", "icl"]) -@pytest.mark.parametrize("batch", [1, 8]) -@pytest.mark.parametrize("contiguous", [False, True]) -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_regression_target_projection( - version: str, - stage: Literal["col", "icl"], - batch: int, - contiguous: bool, - device: str, -) -> None: - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA unavailable") +def test_regression_target_projection(version: str) -> None: + device = "cuda" if torch.cuda.is_available() else "cpu" model = (_get_regression_model() if version == "v3" else _get_model()).to(device) - targets = torch.randn(1002, batch, device=device).T - if contiguous: - targets = targets.contiguous() - targets.requires_grad_() - project = getattr(model, f"_embed_{stage}_y") + targets = torch.randn(1002, 8, device=device).T.requires_grad_() kwargs = {} if version == "v3" else {"task_type": "regression"} - expected = torch.cat([project(row[None], **kwargs) for row in targets]) - expected_grad = torch.autograd.grad(expected.square().sum(), targets)[0] - actual = project(targets, **kwargs) - actual_grad = torch.autograd.grad(actual.square().sum(), targets)[0] - torch.testing.assert_close(actual, expected) - torch.testing.assert_close(actual_grad, expected_grad) + for stage in ("col", "icl"): + project = getattr(model, f"_embed_{stage}_y") + expected = torch.cat([project(row[None], **kwargs) for row in targets]) + expected_grad = torch.autograd.grad(expected.square().sum(), targets)[0] + actual = project(targets, **kwargs) + actual_grad = torch.autograd.grad(actual.square().sum(), targets)[0] + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual_grad, expected_grad) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +@pytest.mark.skipif( + TorchVersion(torch.__version__) < TorchVersion("2.13"), + reason="Eager Triton outer-product dispatch requires Torch 2.13+", +) def test_regression_target_projection_without_compiler(tmp_path: Path) -> None: result = subprocess.run( [