From 7530a4c5d4420adf151bb607c9ea56be6ff98661 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:13:36 -0700 Subject: [PATCH] Add device/import utilities, expose vectorize/devectorize methods as public callables --- src/coreai_opt/_utils/device_utils.py | 41 ++++++ .../kmeans/kmeans_fake_palettize.py | 134 +++++++++++++++--- .../test_kmeans_fake_palettize.py | 37 +++++ tests/palettization/test_kmeans_parallel.py | 2 +- tests/test_utils/test_device_utils.py | 58 ++++++++ 5 files changed, 253 insertions(+), 19 deletions(-) create mode 100644 src/coreai_opt/_utils/device_utils.py create mode 100644 tests/test_utils/test_device_utils.py diff --git a/src/coreai_opt/_utils/device_utils.py b/src/coreai_opt/_utils/device_utils.py new file mode 100644 index 00000000..9aaaa418 --- /dev/null +++ b/src/coreai_opt/_utils/device_utils.py @@ -0,0 +1,41 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Device and build-toolchain availability probes.""" + +import os +import shutil + +import torch +from torch.utils.cpp_extension import CUDA_HOME + + +def cuda_available() -> bool: + """Return True if a CUDA device is visible to torch.""" + return bool(torch.cuda.is_available()) + + +def nvcc_available() -> bool: + """Return True if ``nvcc`` is locatable the way ``cpp_extension`` builds. + + Looks on ``PATH`` first, then under ``CUDA_HOME`` (its env vars / default + root). Independent of whether a CUDA device is present. + """ + if shutil.which("nvcc") is not None: + return True + return CUDA_HOME is not None and os.path.isfile(os.path.join(CUDA_HOME, "bin", "nvcc")) + + +def triton_available() -> bool: + """Return True if triton can be imported. + + Only ``ImportError`` counts as unavailable; other import-time errors (e.g. a + broken native install) propagate. + """ + try: + import triton # noqa: F401, PLC0415 + except ImportError: + return False + return True diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py index 566c140a..9838906c 100644 --- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py +++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py @@ -5,6 +5,7 @@ import logging from collections.abc import Callable +from dataclasses import dataclass import numpy as np import torch @@ -38,6 +39,23 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class _WeightVectorization: + """Context to invert ``vectorize``. + + Attributes: + axis (int): Axis blocks were split along. + block_shape (torch.Size): Shared ``(rows, cols)`` shape of every block. + weight_shape (torch.Size): Original (pre-scale, pre-reshape) weight shape. + weight_dtype (torch.dtype): Original weight dtype. + """ + + axis: int + block_shape: torch.Size + weight_shape: torch.Size + weight_dtype: torch.dtype + + @_FakePalettizeImplBase.register("default") class _KMeansFakePalettize(_FakePalettizeImplBase): """K-means based palettization implementation for neural network weights. @@ -303,6 +321,11 @@ def hard_assign(self, weight: torch.Tensor) -> torch.Tensor: self._maybe_refresh_indices(weight) return self._palettize(self.lut, self.indices, weight) + @property + def _resolved_axis(self) -> int: + """Palettization axis, defaulting to 0 for per-tensor granularity (``axis`` is None).""" + return self.granularity.axis if self.granularity.axis else 0 + def _blocks_to_cluster(self, weight_2d: torch.Tensor, axis: int) -> list[torch.Tensor]: """Validate cluster_dim divisibility and split a 2D weight/sensitivity tensor into per-partition blocks. @@ -319,6 +342,30 @@ def _blocks_to_cluster(self, weight_2d: torch.Tensor, axis: int) -> list[torch.T ) return self.granularity.get_blocks_to_cluster(weight_2d) + def _scale_reshape_and_block(self, weight: torch.Tensor) -> tuple[list[torch.Tensor], int]: + """Scale (if enabled), reshape to 2D, and split into per-block tensors. + + Shared prefix of the clustering, index-assignment, and vectorization paths. + + Args: + weight (torch.Tensor): Weight tensor in its original shape. + + Returns: + tuple[list[torch.Tensor], int]: The per-block 2D tensors and the + resolved palettization axis. + """ + if self.enable_per_channel_scale: + weight = self._scale_by_per_channel_scale(weight) + axis = self._resolved_axis + + # Produce a 2d tensor with output channel axis remaining as is, and all other axes flattened + # into the input channel axis. + weight_2d = self.reshape_strategy.reshape_for_kmeans(weight, axis) + + # Split along the palettization axis into per-partition blocks: a single block + # for per-tensor, or group_size-sized blocks for per-grouped-channel. + return self._blocks_to_cluster(weight_2d, axis), axis + @torch.no_grad() def _cluster_to_centroids( self, original_weights: torch.Tensor, sensitivities: torch.Tensor | None = None @@ -330,12 +377,7 @@ def _cluster_to_centroids( cluster assignment produced directly by the clustering algorithm. """ weight = original_weights.cpu() - if self.enable_per_channel_scale: - weight = self._scale_by_per_channel_scale(weight) - - axis = self.granularity.axis if self.granularity.axis else 0 - weight = self.reshape_strategy.reshape_for_kmeans(weight, axis) - block_weights_to_cluster = self._blocks_to_cluster(weight, axis) + block_weights_to_cluster, axis = self._scale_reshape_and_block(weight) if sensitivities is not None: sensitivities = sensitivities.cpu() @@ -378,18 +420,12 @@ def _assign_indices( """Nearest-centroid hard assignment of ``original_weights`` against a given ``centroids`` (P, K, D) tensor. """ - weight = original_weights.detach().cpu() - if self.enable_per_channel_scale: - weight = self._scale_by_per_channel_scale(weight) - - axis = self.granularity.axis if self.granularity.axis else 0 - weight_2d = self.reshape_strategy.reshape_for_kmeans(weight, axis) - blocks = self._blocks_to_cluster(weight_2d, axis) + blocks, axis = self._scale_reshape_and_block(original_weights.detach().cpu()) centroids_cpu = centroids.detach().cpu().float() block_indices = [] for block_idx, block_weight in enumerate(blocks): - vec = self._vectorize(block_weight) + 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)) @@ -438,7 +474,7 @@ def _palettize( along the grouped axis, so all blocks have the same size. """ clustered_weight = None - axis = self.granularity.axis if self.granularity.axis else 0 + axis = self._resolved_axis lut = lut.to(indices.device) @@ -646,13 +682,13 @@ def _cluster_weights_2d( num_clusters = 2**self.n_bits # Vectorize: reshape block_weight to (N, cluster_dim) along axis 0 - vectorized = self._vectorize(block_weight) + vectorized = self._vectorize_block(block_weight) num_clusters = min(len(vectorized), num_clusters) # Prepare sample weights from sensitivities sample_weight = None if block_sensitivity is not None: - sens_vectorized = self._vectorize(block_sensitivity) + sens_vectorized = self._vectorize_block(block_sensitivity) # Sum sensitivities along cluster_dim for per-vector importance sample_weight = sens_vectorized.sum(dim=-1, keepdim=True) @@ -674,7 +710,7 @@ def _cluster_weights_2d( return centroids, labels - def _vectorize(self, tensor: torch.Tensor) -> torch.Tensor: + def _vectorize_block(self, tensor: torch.Tensor) -> torch.Tensor: """Reshape a 2D tensor into (N, cluster_dim) vectors for k-means. Vectors are always formed along axis 0 (output channel axis). This transposes @@ -685,6 +721,68 @@ def _vectorize(self, tensor: torch.Tensor) -> torch.Tensor: return tensor.reshape(-1, 1) return tensor.transpose(0, 1).reshape(-1, self.cluster_dim) + def _vectorize(self, tensor: torch.Tensor) -> torch.Tensor: + """Alias of _vectorize_block, to be removed.""" + return self._vectorize_block(tensor) + + def _devectorize_block(self, vec: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Reconstruct a ``(rows, cols)`` block from its ``(N, cluster_dim)`` + vectors — the inverse of ``_vectorize_block``. + """ + if self.cluster_dim == 1: + return vec.reshape(rows, cols) + return vec.reshape(cols, rows).transpose(0, 1) + + def vectorize(self, weight: torch.Tensor) -> tuple[torch.Tensor, _WeightVectorization]: + """Scale (if enabled), reshape, block-split, and vectorize a weight into + stacked ``(num_blocks, vectors_per_block, cluster_dim)`` k-means vectors. + + Applies per-channel scaling when ``enable_per_channel_scale`` is set, + matching the clustering path, then produces the vector layout k-means + operates on. Device-preserving. Invert via ``devectorize``. + + Args: + weight (torch.Tensor): Weight tensor in its original shape. + + Returns: + tuple[torch.Tensor, _WeightVectorization]: The stacked ``(P, N, D)`` + vectors and the context needed to reconstruct the weight. + """ + weight_shape, weight_dtype = weight.shape, weight.dtype + blocks, axis = self._scale_reshape_and_block(weight) + block_shape = blocks[0].shape # all blocks share one shape + vectors = torch.stack([self._vectorize_block(block) for block in blocks]) + + context = _WeightVectorization(axis, block_shape, weight_shape, weight_dtype) + return vectors, context + + def devectorize(self, vectors: torch.Tensor, context: _WeightVectorization) -> torch.Tensor: + """Reconstruct a weight from stacked ``(P, N, D)`` vectors — the inverse + of ``vectorize``. + + Undoes the vectorization and block-split, restores the original shape, + then unscales when ``enable_per_channel_scale`` is set. + + Args: + vectors (torch.Tensor): Stacked ``(num_blocks, vectors_per_block, + cluster_dim)`` vectors. + context (_WeightVectorization): Context from ``vectorize``. + + Returns: + torch.Tensor: Weight in the original shape and dtype. + """ + blocks = [ + self._devectorize_block(vectors[p], *context.block_shape) + for p in range(vectors.shape[0]) + ] + clustered = torch.cat(blocks, dim=context.axis) + clustered = self.reshape_strategy.reshape_to_original( + clustered, context.axis, context.weight_shape + ) + if self.enable_per_channel_scale: + clustered = self._unscale_by_per_channel_scale(clustered) + return clustered.to(context.weight_dtype) + def _lookup_result_to_block(self, looked_up: torch.Tensor) -> torch.Tensor: """Reshape a vector LUT lookup result back to 2D weight shape. diff --git a/tests/palettization/test_kmeans_fake_palettize.py b/tests/palettization/test_kmeans_fake_palettize.py index 96351689..6b9d3629 100644 --- a/tests/palettization/test_kmeans_fake_palettize.py +++ b/tests/palettization/test_kmeans_fake_palettize.py @@ -2000,3 +2000,40 @@ def test_device_placement_on_accelerator(accelerator_device): out = palettizer.hard_assign(weight) assert out.device.type == device assert out.shape == weight.shape + + +_ROUNDTRIP_SPECS = [ + pytest.param(PerTensorGranularity(), 1, id="pt-cd1"), + pytest.param(PerTensorGranularity(), 2, id="pt-cd2"), + pytest.param(PerTensorGranularity(), 4, id="pt-cd4"), + pytest.param(PerGroupedChannelGranularity(axis=0, group_size=8), 1, id="pgc-ax0-gs8-cd1"), + pytest.param(PerGroupedChannelGranularity(axis=0, group_size=16), 4, id="pgc-ax0-gs16-cd4"), + pytest.param(PerGroupedChannelGranularity(axis=1, group_size=8), 1, id="pgc-ax1-gs8-cd1"), +] + + +@pytest.mark.parametrize("enable_per_channel_scale", [False, True], ids=["no-pcs", "pcs"]) +@pytest.mark.parametrize("granularity, cluster_dim", _ROUNDTRIP_SPECS) +def test_vectorize_devectorize_round_trip(granularity, cluster_dim, enable_per_channel_scale): + """``devectorize`` inverts ``vectorize``. + + Covers scaling on/off across granularity and cluster_dim, independent of + clustering — the round trip must reconstruct the original weight. + """ + torch.manual_seed(0) + spec = PalettizationSpec( + n_bits=4, + granularity=granularity, + cluster_dim=cluster_dim, + enable_per_channel_scale=enable_per_channel_scale, + lut_qspec=None, + ) + palettizer = _KMeansFakePalettize(**spec.__dict__) + weight = torch.randn(16, 32) + + vectors, context = palettizer.vectorize(weight) + reconstructed = palettizer.devectorize(vectors, context) + + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == weight.dtype + assert torch.allclose(reconstructed, weight, atol=1e-5) diff --git a/tests/palettization/test_kmeans_parallel.py b/tests/palettization/test_kmeans_parallel.py index d84abf44..b1a5d353 100644 --- a/tests/palettization/test_kmeans_parallel.py +++ b/tests/palettization/test_kmeans_parallel.py @@ -166,7 +166,7 @@ def test_parallel_vector_palettization_matches_sequential(self): unseeded ``torch.randint`` / ``np.random.choice`` — so spawned workers and the main process see different RNG states. To make the comparison deterministic anyway, the linear weight is hand-built so that, after - ``_vectorize`` (transpose + reshape into 2D pairs), the vectors form 4 + ``_vectorize_block`` (transpose + reshape into 2D pairs), the vectors form 4 well-separated clusters. K-means converges to those 4 centers from any reasonable initialization, so both paths produce the same reconstructed weights (up to a cluster-ID permutation), and the model output matches. diff --git a/tests/test_utils/test_device_utils.py b/tests/test_utils/test_device_utils.py new file mode 100644 index 00000000..83cea22f --- /dev/null +++ b/tests/test_utils/test_device_utils.py @@ -0,0 +1,58 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for device and toolchain availability probes.""" + +import builtins +import sys +import types + +import torch + +from coreai_opt._utils import device_utils + + +def test_cuda_available_matches_torch(): + assert device_utils.cuda_available() == torch.cuda.is_available() + + +def test_nvcc_available_on_path(monkeypatch): + monkeypatch.setattr(device_utils.shutil, "which", lambda _name: "/usr/bin/nvcc") + assert device_utils.nvcc_available() is True + + +def test_nvcc_available_via_cuda_home(monkeypatch, tmp_path): + monkeypatch.setattr(device_utils.shutil, "which", lambda _name: None) + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + (bin_dir / "nvcc").write_text("") + monkeypatch.setattr(device_utils, "CUDA_HOME", str(tmp_path)) + assert device_utils.nvcc_available() is True + + +def test_nvcc_unavailable(monkeypatch): + monkeypatch.setattr(device_utils.shutil, "which", lambda _name: None) + monkeypatch.setattr(device_utils, "CUDA_HOME", None) + assert device_utils.nvcc_available() is False + + +def test_triton_available_true(monkeypatch): + # A stub module in sys.modules makes ``import triton`` succeed regardless of + # whether triton is actually installed. + monkeypatch.setitem(sys.modules, "triton", types.ModuleType("triton")) + assert device_utils.triton_available() is True + + +def test_triton_unavailable_on_import_error(monkeypatch): + monkeypatch.delitem(sys.modules, "triton", raising=False) + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "triton": + raise ImportError("no triton") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + assert device_utils.triton_available() is False