Compute palettizer indices on original device

src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py

@@ -38,6 +38,9 @@
 
 logger = logging.getLogger(__name__)
 
+# Max size of the (rows, num_clusters) fp32 distance matrix built at once in _assign_indices.
+_CDIST_CHUNK_BYTES = 512 * 1024**2
+
 
 @dataclass(frozen=True)
 class _WeightVectorization:
@@ -576,18 +579,24 @@ def _assign_indices(
             distances = torch.cdist(vectors.float(), centroids.detach().cpu().float())
             return distances.argmin(dim=-1).flatten().to(torch.uint8)
 
-        blocks, axis = self._scale_reshape_and_block(original_weights.detach().cpu())
-        centroids_cpu = centroids.detach().cpu().float()
+        blocks, axis = self._scale_reshape_and_block(original_weights.detach())
+        centroids = centroids.detach().float()
 
         block_indices = []
         for block_idx, block_weight in enumerate(blocks):
             vec = self._vectorize_block(block_weight)
-            dist = torch.cdist(vec.float(), centroids_cpu[block_idx])
-            clusters = dist.argmin(dim=-1)
-            block_indices.append(self._build_block_indices(clusters, block_weight).to(torch.uint8))
+            centroid = centroids[block_idx]
+            chunk_rows = max(1, _CDIST_CHUNK_BYTES // (centroid.shape[0] * 4))
+            clusters = torch.cat(
+                [
+                    torch.cdist(v.float(), centroid).argmin(dim=-1).to(torch.uint8)
+                    for v in vec.split(chunk_rows)
+                ]
+            )
+            block_indices.append(self._build_block_indices(clusters, block_weight))
 
         # self.indices is always CPU-resident, matching self.lut.
-        return self._combine_block_indices(block_indices, axis, original_weights)
+        return self._combine_block_indices(block_indices, axis, original_weights).cpu()
 
     def _combine_block_indices(
         self,

tests/palettization/test_kmeans_fake_palettize.py

@@ -2069,6 +2069,79 @@ def test_indices_consistent_with_centroids(self):
         reassigned = palettizer._assign_indices(weight, centroids)
         assert torch.equal(indices, reassigned)
 
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_assign_indices_chunking_matches_unchunked(self, dtype):
+        weight = torch.randn(8, 32).to(dtype)
+        spec = PalettizationSpec(n_bits=3, granularity=PerTensorGranularity())
+        palettizer = _KMeansFakePalettize(**spec.__dict__)
+        centroids, _ = palettizer._cluster_to_centroids(weight)
+
+        expected = palettizer._assign_indices(weight, centroids)
+        with patch("coreai_opt.palettization.kmeans.kmeans_fake_palettize._CDIST_CHUNK_BYTES", 1):
+            chunked = palettizer._assign_indices(weight, centroids)
+        assert torch.equal(expected, chunked)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_assign_indices_bounds_cdist_size(self, dtype):
+        weight = torch.randn(64, 32).to(dtype)
+        spec = PalettizationSpec(n_bits=3, granularity=PerTensorGranularity())
+        palettizer = _KMeansFakePalettize(**spec.__dict__)
+        centroids, _ = palettizer._cluster_to_centroids(weight)
+        chunk_rows = 100
+        budget = chunk_rows * centroids.shape[1] * 4
+        cdist_calls = []
+        real_cdist = torch.cdist
+
+        def spy_cdist(x, y):
+            dist = real_cdist(x, y)
+            cdist_calls.append((x.dtype, dist.numel() * 4))
+            return dist
+
+        with (
+            patch(
+                "coreai_opt.palettization.kmeans.kmeans_fake_palettize._CDIST_CHUNK_BYTES", budget
+            ),
+            patch("torch.cdist", spy_cdist),
+        ):
+            palettizer._assign_indices(weight, centroids)
+
+        assert len(cdist_calls) == -(-weight.numel() // chunk_rows)
+        assert all(dt == torch.float32 and size <= budget for dt, size in cdist_calls)
+
+    @pytest.mark.cuda
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"])
+    def test_assign_indices_peak_memory_bounded_by_chunk_budget(self, dtype):
+        """Peak GPU memory above the resident tensors stays within a few chunk budgets,
+        far below the unchunked ``(N, K)`` distance matrix.
+        """
+        side, num_clusters, budget = 2048, 16, 4 * 1024**2
+        num_weights = side * side
+        spec = PalettizationSpec(n_bits=4, granularity=PerTensorGranularity())
+        palettizer = _KMeansFakePalettize(**spec.__dict__)
+        weight = torch.randn(side, side, device="cuda", dtype=dtype)
+        centroids = torch.randn(1, num_clusters, 1, device="cuda")
+
+        with patch(
+            "coreai_opt.palettization.kmeans.kmeans_fake_palettize._CDIST_CHUNK_BYTES", budget
+        ):
+            palettizer._assign_indices(weight, centroids)  # warm-up
+            torch.cuda.synchronize()
+            baseline = torch.cuda.memory_allocated()
+            torch.cuda.reset_peak_memory_stats()
+            palettizer._assign_indices(weight, centroids)
+            torch.cuda.synchronize()
+        extra = torch.cuda.max_memory_allocated() - baseline
+
+        # One chunk's distance matrix plus cdist's temporaries (1.5x budget), and the uint8
+        # indices (1 byte per weight).
+        slack = 2 * 1024**2
+        limit = 1.5 * budget + num_weights + slack
+        unchunked = num_weights * num_clusters * 4
+        assert extra <= limit, (
+            f"peak {extra / 1024**2:.1f} MiB above resident exceeds {limit / 1024**2:.1f} MiB "
+            f"(unchunked distance matrix would be {unchunked / 1024**2:.0f} MiB)"
+        )
+
     def test_blocks_to_cluster_raises_on_indivisible_cluster_dim(self):
         spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity(), cluster_dim=3)
         palettizer = _KMeansFakePalettize(**spec.__dict__)