Add device/import utilities, expose vectorize/devectorize methods as public callables

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

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.
 

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)

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.

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