fix: avoid runtime compilation in regression target encoding (ENG-1093)

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.

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,

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)

tests/test_architectures/test_regression_target_projection.py

@@ -0,0 +1,77 @@
+# Copyright (c) Prior Labs GmbH 2026.
+
+"""Regression target projections must not require a runtime compiler."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+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"])
+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, 8, device=device).T.requires_grad_()
+    kwargs = {} if version == "v3" else {"task_type": "regression"}
+    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(
+        [
+            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()