From 6f71c6210469da817551bfe2d5d7160ff6fc7465 Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Thu, 22 Oct 2026 16:52:22 -0700 Subject: [PATCH 1/2] Simplify uncached tabular benchmark inference --- benchmark/tabular/beyondarena/main.py | 7 -- benchmark/tabular/model.py | 101 +++++++------------------- benchmark/tabular/tabarena/main.py | 7 -- 3 files changed, 27 insertions(+), 88 deletions(-) diff --git a/benchmark/tabular/beyondarena/main.py b/benchmark/tabular/beyondarena/main.py index 124e8e851b..173f0044bb 100644 --- a/benchmark/tabular/beyondarena/main.py +++ b/benchmark/tabular/beyondarena/main.py @@ -50,12 +50,6 @@ help="Prediction batch size.", ) add_finetune_args(parser) -parser.add_argument( - "--enable-kv-cache", - action="store_true", - help="Cache context key/value projections at fit and reuse them " - "across prediction batches.", -) args = parser.parse_args() model_config = MODEL_CONFIGS[args.model] @@ -69,7 +63,6 @@ config = { "max_context_size": args.max_context_size, "max_columns": args.max_columns, - "kv_cache": args.enable_kv_cache, } if args.batch_size is not None: config["ag.max_batch_size"] = args.batch_size diff --git a/benchmark/tabular/model.py b/benchmark/tabular/model.py index 118e3fb42d..e74b832dea 100644 --- a/benchmark/tabular/model.py +++ b/benchmark/tabular/model.py @@ -34,7 +34,6 @@ full_finetune, kumo_small_binary_epochs, ) -from sdm.processing.execution import RecipeExecution Task = Literal["classification", "regression"] KumoTabularSize = Literal["small", "medium", "large"] @@ -225,7 +224,8 @@ def _fit( if isinstance(processor, sp.SelectColumns): processor.max_columns = params["max_columns"] - self._recipe_execution: RecipeExecution | None = None + self._recipe = recipe + self._kv_cache = params["kv_cache"] if params["kv_cache"]: with torch.amp.autocast( self._device.type, @@ -242,26 +242,9 @@ def _fit( ) return - self._recipe_execution = RecipeExecution(recipe) - with ( - torch.inference_mode(), - torch.amp.autocast(self._device.type, enabled=False), - ): - contexts = self._recipe_execution.fit_transform( - x=x_context, - y=y_context, - related_tables=None, - num_members=num_estimators, - generator=generator, - ) - self._contexts = tuple( - context._replace( - x=cast(sdm.TableTensor, context.x.cpu()), - y=cast(sdm.TableTensor, context.y.cpu()), - ) - for context in contexts - ) - # Replay the same model-side randomness as the cached fit path. + self._x_context = x_context.cpu() + self._y_context = y_context.cpu() + # Replay preprocessing and model randomness for every query batch. if generator is not None: self._rng_state = generator.get_state() elif self._device.type == "cuda": @@ -286,67 +269,41 @@ def _predict_proba( *x_query.size(), ) - if self._recipe_execution is None: - with torch.amp.autocast( + with ( + torch.inference_mode(), + torch.amp.autocast( self._device.type, self.autocast_dtype, enabled=x_query.is_cuda, - ): + ), + ): + if self._kv_cache: out = self.model.predict(x_query) - else: - with ( - torch.inference_mode(), - torch.amp.autocast(self._device.type, enabled=False), - ): - queries = self._recipe_execution.transform( - x=x_query, - related_tables=None, - ) + else: generator = torch.Generator(self._device).set_state( self._rng_state ) - outputs: list[sdm.TableTensor] = [] - size = self._estimator_batch_size(x_query) - size = len(self._contexts) if size is None else size - for start in range(0, len(self._contexts), size): - batch = [ - context._replace( - x=cast( - sdm.TableTensor, context.x.to(self._device) - ), - y=cast( - sdm.TableTensor, context.y.to(self._device) - ), - ) - for context in self._contexts[start : start + size] - ] - with torch.amp.autocast( - self._device.type, - self.autocast_dtype, - enabled=x_query.is_cuda, - ): - outputs += self.model._forward_members( - contexts=batch, - queries=queries[start : start + size], - estimator_batch_size=None, - generator=generator, - ) - - if self.problem_type == REGRESSION: - outputs = list( - self._recipe_execution.inverse_transform_target( - outputs - ) - ) - out = self._recipe_execution.transform_output(outputs) + out = self.model( + x_context=self._x_context.to(self._device), + y_context=self._y_context.to(self._device), + x_query=x_query, + recipe=self._recipe, + num_estimators=( + None if self._expand_query else self._num_estimators + ), + estimator_batch_size=self._estimator_batch_size(x_query), + generator=generator, + ) if self.problem_type == REGRESSION: return out.numerical.float().mean(dim=-1).cpu().numpy() assert self.num_classes is not None columns = out.numerical_columns - indices = [columns.index(str(i)) for i in range(self.num_classes)] - probabilities = out.numerical[..., indices].float().cpu().numpy() + probabilities = np.zeros((len(X), self.num_classes), dtype=np.float32) + probabilities[:, [int(column) for column in columns]] = ( + out.numerical.float().cpu().numpy() + ) return self._convert_proba_to_unified_form(probabilities) def _estimator_batch_size(self, x: torch.Tensor) -> int | None: @@ -376,10 +333,6 @@ def get_device(self) -> str: def _set_device(self, device: str) -> None: self.model.to(device) - if self._recipe_execution is not None: - recipe = self._recipe_execution.recipe - for processor in (recipe.features, recipe.target, recipe.output): - processor.to(device) self._device = torch.device(device) def _more_tags(self) -> dict[str, bool]: diff --git a/benchmark/tabular/tabarena/main.py b/benchmark/tabular/tabarena/main.py index 20b0248b19..798d2cc6a1 100644 --- a/benchmark/tabular/tabarena/main.py +++ b/benchmark/tabular/tabarena/main.py @@ -51,12 +51,6 @@ help="Prediction batch size.", ) add_finetune_args(parser) -parser.add_argument( - "--enable-kv-cache", - action="store_true", - help="Cache context key/value projections at fit and reuse them " - "across prediction batches.", -) args = parser.parse_args() os.environ["SDM_ENABLE_CATEGORY_CHECKS"] = "0" @@ -72,7 +66,6 @@ config = { "max_context_size": args.max_context_size, "max_columns": args.max_columns, - "kv_cache": args.enable_kv_cache, } if args.batch_size is not None: config["ag.max_batch_size"] = args.batch_size From ef108f7b58e5322df2076dd35fab33873a6e914f Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:13:22 -0700 Subject: [PATCH 2/2] Enable KV caching by default in tabular benchmarks --- benchmark/tabular/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmark/tabular/model.py b/benchmark/tabular/model.py index e74b832dea..3ba69eac4d 100644 --- a/benchmark/tabular/model.py +++ b/benchmark/tabular/model.py @@ -96,7 +96,7 @@ def _set_default_params(self) -> None: ) self._set_default_param_value("max_context_size", None) self._set_default_param_value("max_columns", None) - self._set_default_param_value("kv_cache", False) + self._set_default_param_value("kv_cache", True) self._set_default_param_value("estimator_batch_size", "auto") self._set_default_param_value("finetune", False) self._set_default_param_value("finetune_epochs", None)