From e4f887bc7e0005a4cedf05b50c34cf61f26f81d5 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Tue, 20 Oct 2026 09:54:25 -0700 Subject: [PATCH 1/4] Compute palettizer indices on original device --- .../kmeans/kmeans_fake_palettize.py | 23 +++++++++++++------ .../test_kmeans_fake_palettize.py | 11 +++++++++ 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 7d31ab6c..2b31ef25 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/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: @@ -531,18 +534,24 @@ def _assign_indices( """Nearest-centroid hard assignment of ``original_weights`` against a given ``centroids`` (P, K, D) tensor. """ - 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)) + vec = self._vectorize_block(block_weight).float() + centroid = centroids[block_idx] + chunk_rows = max(1, _CDIST_CHUNK_BYTES // (centroid.shape[0] * 4)) + clusters = torch.cat( + [ + torch.cdist(v, 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, diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index c746c631..da7c4f50 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -2052,6 +2052,17 @@ def test_indices_consistent_with_centroids(self): reassigned = palettizer._assign_indices(weight, centroids) assert torch.equal(indices, reassigned) + def test_assign_indices_chunking_matches_unchunked(self): + weight = torch.randn(8, 32) + 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) + 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__) From 26073419a7cae8dc1adf38d9164f9859b9b71124 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:51:50 -0700 Subject: [PATCH 2/4] Address review comments --- src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py | 4 ++-- tests/palettization/test_kmeans_fake_palettize.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 2b31ef25..d92cd749 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -539,12 +539,12 @@ def _assign_indices( block_indices = [] for block_idx, block_weight in enumerate(blocks): - vec = self._vectorize_block(block_weight).float() + vec = self._vectorize_block(block_weight) centroid = centroids[block_idx] chunk_rows = max(1, _CDIST_CHUNK_BYTES // (centroid.shape[0] * 4)) clusters = torch.cat( [ - torch.cdist(v, centroid).argmin(dim=-1).to(torch.uint8) + torch.cdist(v.float(), centroid).argmin(dim=-1).to(torch.uint8) for v in vec.split(chunk_rows) ] ) diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index da7c4f50..d34a37e5 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -2052,8 +2052,9 @@ def test_indices_consistent_with_centroids(self): reassigned = palettizer._assign_indices(weight, centroids) assert torch.equal(indices, reassigned) - def test_assign_indices_chunking_matches_unchunked(self): - weight = torch.randn(8, 32) + @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) From 7aec5725b88c84453b0fe05c41f8591d39cf5213 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Tue, 27 Oct 2026 15:02:39 -0700 Subject: [PATCH 3/4] Add additional unit tests --- .../test_kmeans_fake_palettize.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index d34a37e5..f5aac6fb 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -2064,6 +2064,67 @@ def test_assign_indices_chunking_matches_unchunked(self, dtype): 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 (~2.5x budget), and the uint8 + # indices (1 byte per weight) held twice while concatenating. + slack = 2 * 1024**2 + limit = 2.5 * budget + 2 * 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__) From 9c08d4cdcd7bc2702fabb67d4a05df0b7db6184d Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Tue, 27 Oct 2026 15:19:29 -0700 Subject: [PATCH 4/4] Tighten unit test bounds --- tests/palettization/test_kmeans_fake_palettize.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index f5aac6fb..d0186fd9 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -2115,10 +2115,10 @@ def test_assign_indices_peak_memory_bounded_by_chunk_budget(self, dtype): torch.cuda.synchronize() extra = torch.cuda.max_memory_allocated() - baseline - # One chunk's distance matrix plus cdist's temporaries (~2.5x budget), and the uint8 - # indices (1 byte per weight) held twice while concatenating. + # 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 = 2.5 * budget + 2 * num_weights + slack + 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 "