Reduce repeated table metadata preparation during inference

sdm/processing/categorical/align.py

@@ -15,6 +15,7 @@
     Stype,
     TableTensor,
 )
+from sdm._memory import split_size
 from sdm.nn._buffer import BufferList
 from sdm.processing import EnsembleProcessor
 from sdm.relational.join import join_index
@@ -396,9 +397,8 @@ def _category_lookup(
         fitted_categories: Tensor,
         codes: Tensor,
     ) -> Tensor:
-        lookup = codes.new_full((input_categories.numel(),), -1)
         if fitted_categories.numel() == 0:
-            return lookup
+            return codes.new_full(input_categories.size(), -1)
 
         comparable_categories = input_categories
         sorted_categories = fitted_categories
@@ -414,12 +414,9 @@ def _category_lookup(
             sorted_categories,
             comparable_categories,
         )
-        position = position.clamp(max=sorted_categories.numel() - 1)
-        match = sorted_categories[position] == comparable_categories
-        left_index = match.nonzero().view(-1)
-        right_index = perm[position[left_index]]
-        lookup[left_index] = right_index.to(codes.dtype)
-        return lookup
+        position = position.clamp(max=sorted_categories.size(-1) - 1)
+        match = sorted_categories.gather(-1, position) == comparable_categories
+        return torch.where(match, perm.gather(-1, position), -1).to(codes)
 
     def _align_to_categories(
         self,
@@ -435,12 +432,22 @@ def _align_to_categories(
             {} for _ in table.categorical.categories
         ]
         string_requests: list[tuple[int, int, StringTensor, Tensor]] = []
+        numeric_requests: dict[
+            tuple[
+                int, int, torch.dtype, torch.dtype, torch.device, torch.device
+            ],
+            list[tuple[int, int, Tensor, Tensor]],
+        ] = {}
+        columns_by_width: dict[int, list[int]] = {}
         for column_index, input_categories in enumerate(
             table.categorical.categories
         ):
             if input_categories.numel() == 0:
                 continue
 
+            columns_by_width.setdefault(input_categories.numel(), []).append(
+                column_index
+            )
             lookup_by_categories = lookups_by_column[column_index]
             pending_categories: set[int] = set()
             for batch_categories in fitted_categories:
@@ -473,11 +480,50 @@ def _align_to_categories(
                         )
                     )
                     continue
-                lookup_by_categories[identity] = self._category_lookup(
-                    input_categories,
-                    column_categories,
+                pending_categories.add(identity)
+                key = (
+                    input_categories.numel(),
+                    column_categories.numel(),
+                    input_categories.dtype,
+                    column_categories.dtype,
+                    input_categories.device,
+                    column_categories.device,
+                )
+                numeric_requests.setdefault(key, []).append(
+                    (
+                        column_index,
+                        identity,
+                        input_categories,
+                        column_categories,
+                    )
+                )
+
+        # Equal vocabulary widths need no padding and preserve row-wise sorts.
+        for requests in numeric_requests.values():
+            input_categories, column_categories = requests[0][2:]
+            # Bound stacked vocabularies, sort/search buffers and lookups.
+            size = split_size(
+                num_items=len(requests),
+                item_bytes=8
+                * (input_categories.numel() + column_categories.numel())
+                * max(
+                    input_categories.element_size(),
+                    column_categories.element_size(),
+                    8,
+                ),
+                device=input_categories.device,
+            )
+            for start in range(0, len(requests), size):
+                chunk = requests[start : start + size]
+                lookups = self._category_lookup(
+                    torch.stack([request[2] for request in chunk]),
+                    torch.stack([request[3] for request in chunk]),
                     codes,
                 )
+                for (column_index, identity, _, _), lookup in zip(
+                    chunk, lookups.unbind(), strict=True
+                ):
+                    lookups_by_column[column_index][identity] = lookup
 
         string_lookups = self._string_category_lookups(
             input_categories=tuple(
@@ -497,25 +543,42 @@ def _align_to_categories(
         ):
             lookups_by_column[column_index][identity] = lookup
 
-        for column_index, input_categories in enumerate(
-            table.categorical.categories
-        ):
-            if input_categories.numel() == 0:
-                continue
-            lookup_by_categories = lookups_by_column[column_index]
-            lookups = [
-                lookup_by_categories[id(batch_categories[column_index])]
-                for batch_categories in fitted_categories
-            ]
-            column_codes = codes[..., column_index]
-            mask = column_codes >= 0
-            lookup = torch.stack(lookups)
-            indices = column_codes.clamp_min(0).long()
-            aligned_codes[..., column_index] = torch.where(
-                mask,
-                lookup.gather(1, indices),
-                -1,
+        for width, columns in columns_by_width.items():
+            # Bound gathered codes, integer indices, masks and mapped outputs.
+            size = split_size(
+                num_items=len(columns),
+                item_bytes=8
+                * (codes[..., :1].numel() + len(fitted_categories) * width)
+                * max(codes.element_size(), 8),
+                device=codes.device,
             )
+            for start in range(0, len(columns), size):
+                column_indices = columns[start : start + size]
+                # Map [column, batch, row] codes with vocabulary lookups.
+                lookup = torch.stack(
+                    [
+                        lookups_by_column[column_index][
+                            id(batch_categories[column_index])
+                        ]
+                        for column_index in column_indices
+                        for batch_categories in fitted_categories
+                    ]
+                ).view(len(column_indices), len(fitted_categories), width)
+                # Upload known column metadata without waiting for device work.
+                column_index = torch.tensor(column_indices, device="cpu").to(
+                    device=codes.device, non_blocking=True
+                )
+                column_codes = codes.index_select(-1, column_index).movedim(
+                    -1, 0
+                )
+                mask = column_codes >= 0
+                indices = column_codes.clamp_min(0).long()
+                mapped_codes = torch.where(
+                    mask,
+                    lookup.gather(-1, indices),
+                    -1,
+                ).movedim(0, -1)
+                aligned_codes.index_copy_(-1, column_index, mapped_codes)
 
         return tuple(
             (table if single_table else table[batch_index]).replace_blocks(

sdm/tensor/table.py

@@ -125,6 +125,20 @@ def __hash__(self) -> int:
         return hash(frozenset(self.columns.items()))
 
 
+@functools.lru_cache(maxsize=128)
+def _column_locations(
+    columns: tuple[tuple[Stype, tuple[str, ...]], ...],
+) -> dict[str, tuple[Stype, int]]:
+    locations = {
+        name: (stype, index)
+        for stype, names in columns
+        for index, name in enumerate(names)
+    }
+    if sum(len(names) for _, names in columns) != len(locations):
+        raise ValueError("Expected column names to be unique")
+    return locations
+
+
 class TableTensor(Tensor):
     r"""A :class:`torch.Tensor` for tensorized, lossless table data.
 
@@ -328,18 +342,12 @@ def __new__(
                     f"{len(columns[stype])} {_columns} (got {block.size(-1)})"
                 )
 
-        column_names = list(chain.from_iterable(columns.values()))
-        column_to_loc: dict[str, tuple[Stype, int]] = {}
-        for stype_like, names in columns.items():
-            stype = Stype(stype_like)
-            for i, name in enumerate(names):
-                column_to_loc[name] = (stype, i)
-        if len(column_names) != len(column_to_loc):
-            raise ValueError("Expected column names to be unique")
+        # Only immutable schema keys are retained; each table owns its lookup.
+        column_to_loc = _column_locations(tuple(columns.items())).copy()
 
         out = Tensor._make_wrapper_subclass(
             cls,
-            size=(*size, len(column_names)),
+            size=(*size, len(column_to_loc)),
             dtype=numerical.dtype,
             device=numerical.device,
             requires_grad=False,

test/processing/categorical/test_align.py

@@ -1,15 +1,15 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Literal
+from typing import Literal, cast
 
 import pandas as pd
 import pytest
 import torch
 
 from sdm import CategoricalTensor, EnsembleTable, StringTensor, TableTensor
 from sdm.processing import AlignCategories
-from sdm.testing import withCUDA
+from sdm.testing import onlyCUDA, withCUDA
 
 
 def _table(
@@ -110,6 +110,39 @@ def test_align_categories_removes_query_only_joint_vocabulary() -> None:
     assert query.categorical.code.squeeze(-1).tolist() == [-1, 1]
 
 
+@onlyCUDA
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+def test_align_numeric_categories_without_host_synchronization(
+    dtype: torch.dtype,
+) -> None:
+    context = TableTensor(
+        categorical=CategoricalTensor(
+            code=torch.tensor([[0], [1], [0]], device="cuda", dtype=dtype),
+            categories=(torch.tensor([9, 3, 7], device="cuda"),),
+        ),
+    )
+    query = TableTensor(
+        categorical=CategoricalTensor(
+            code=torch.tensor(
+                [[0], [1], [2], [-1]], device="cuda", dtype=dtype
+            ),
+            categories=(torch.tensor([3, 7, 9], device="cuda"),),
+        ),
+    )
+    processor = AlignCategories(sort_by="value").fit(context)
+    previous = torch.cuda.get_sync_debug_mode()
+    try:
+        torch.cuda.set_sync_debug_mode("error")
+        with torch.device("cuda"):
+            output = processor.transform(query)
+    finally:
+        torch.cuda.set_sync_debug_mode(previous)
+
+    assert output.categorical.code.dtype == dtype
+    assert output.categorical.code.tolist() == [[0], [-1], [1], [-1]]
+    assert output.categorical.categories[0].tolist() == [3, 9]
+
+
 @pytest.mark.parametrize("sort_by", ["code", "frequency", "value"])
 def test_align_categories_orders_joint_vocabulary(
     sort_by: Literal["code", "frequency", "value"],
@@ -504,3 +537,76 @@ def test_align_categories_ensemble_matches_member_fits(
         expected_query = reference.transform(query)
         assert query_output[member_id].equal(expected_query)
         assert fitted_query_output[member_id].equal(expected_query)
+
+
+@withCUDA
+@pytest.mark.parametrize("sort_by", ["code", "frequency", "value"])
+@pytest.mark.parametrize("min_frequency", [1, 2])
+@pytest.mark.parametrize("memory_fraction", ["0.05", "1e-10"])
+def test_align_numeric_columns_preserves_independent_vocabularies(
+    device: torch.device,
+    sort_by: Literal["code", "frequency", "value"],
+    min_frequency: int,
+    memory_fraction: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("SDM_CHUNK_MEMORY_FRACTION", memory_fraction)
+    context_columns: list[torch.Tensor] = []
+    query_columns: list[torch.Tensor] = []
+    expected_columns: list[torch.Tensor] = []
+    for index in range(7):
+        dtype = torch.float64 if index % 2 else torch.int64
+        context_values = [5 + 20 * index, 1 + 20 * index, 9 + 20 * index]
+        query_values = [1 + 20 * index, 7 + 20 * index, 5 + 20 * index]
+        if index == 6:
+            context_values.append(13 + 20 * index)
+            query_values.append(13 + 20 * index)
+        context = TableTensor(
+            columns={"categorical": [str(index)]},
+            categorical=CategoricalTensor(
+                code=torch.tensor(
+                    [[0], [1], [2], [1], [2], [2], [-1]],
+                    dtype=torch.int32,
+                    device=device,
+                ),
+                categories=(
+                    torch.tensor(context_values, dtype=dtype, device=device),
+                ),
+            ),
+        )
+        query = TableTensor(
+            columns={"categorical": [str(index)]},
+            categorical=CategoricalTensor(
+                code=torch.tensor(
+                    [[0], [1], [2], [-1]], dtype=torch.int32, device=device
+                ),
+                categories=(
+                    torch.tensor(query_values, dtype=dtype, device=device),
+                ),
+            ),
+        )
+        expected_columns.append(
+            AlignCategories(sort_by, min_frequency=min_frequency)
+            .fit(context)
+            .transform(query)
+        )
+        context_columns.append(context)
+        query_columns.append(query)
+
+    context = cast(TableTensor, torch.cat(context_columns, dim=-1))
+    query = cast(TableTensor, torch.cat(query_columns, dim=-1))
+    processor = AlignCategories(sort_by, min_frequency=min_frequency).fit(
+        context
+    )
+    with torch.no_grad():
+        actual = processor.transform(query)
+    expected = cast(TableTensor, torch.cat(expected_columns, dim=-1))
+    assert torch.equal(actual.categorical.code, expected.categorical.code)
+    assert all(
+        torch.equal(actual_values, expected_values)
+        for actual_values, expected_values in zip(
+            actual.categorical.categories,
+            expected.categorical.categories,
+            strict=True,
+        )
+    )

test/tensor/test_table.py

@@ -103,6 +103,36 @@ def test_column_names() -> None:
         )
 
 
+def test_repeated_schema_preserves_column_ownership_and_validation() -> None:
+    names = ["left", "right"]
+    values = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    first = TableTensor(columns={"numerical": names}, numerical=values)
+    columns = {str(stype): names for stype, names in first.columns.items()}
+    second = TableTensor(columns=columns, numerical=values + 1)
+    names.reverse()
+    changed = columns.copy()
+    changed["numerical"] = tuple(names)
+    reversed_table = TableTensor(columns=changed, numerical=values)
+
+    assert first.columns[Stype.numerical] == ("left", "right")
+    assert second.columns == first.columns
+    assert reversed_table.columns[Stype.numerical] == ("right", "left")
+    torch.testing.assert_close(first[["left"]].numerical, values[:, :1])
+    torch.testing.assert_close(second[["left"]].numerical, values[:, :1] + 1)
+    torch.testing.assert_close(
+        reversed_table[["left"]].numerical, values[:, 1:]
+    )
+
+    with pytest.raises(ValueError, match="hold 2 columns"):
+        TableTensor(columns=columns, numerical=torch.ones(2, 3))
+    with pytest.raises(ValueError, match="to be unique"):
+        TableTensor(
+            columns={"numerical": ["left"], "datetime": ["left"]},
+            numerical=torch.ones(2, 1),
+            datetime=torch.ones(2, 1, dtype=torch.int64),
+        )
+
+
 def test_equal() -> None:
     numerical = torch.randn(2, 2)
     categorical = CategoricalTensor(