Keep palettizer _raw_lut computation on the same device as the weight

Author: crowbat

src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py

@@ -192,7 +192,10 @@ def _initialize(self, weight: torch.Tensor) -> None:
         """
         self.centroids, indices = self._cluster_to_centroids(weight, self._sensitivities)
         self.indices = indices.detach()
+
+        # Run self.quantize_lut to seed lut quantizer qparams
         self.quantize_lut(self._raw_lut(self.centroids))
+
         self._centroids_initialized = True
         self._indices_stale = False
 
@@ -391,9 +394,9 @@ def _combine_block_indices(
     @torch.no_grad()
     def _raw_lut(self, centroids: torch.Tensor) -> torch.Tensor:
         """Reshape ``centroids`` (P, K, D) to the pre-quantization LUT shape
-        ``(P, K[, D])`` on CPU, detached.
+        ``(P, K[, D])``, detached.
         """
-        centroids = centroids.detach().cpu()
+        centroids = centroids.detach()
         return centroids if self.cluster_dim > 1 else centroids.squeeze(-1)
 
     def _palettize(
@@ -417,6 +420,8 @@ def _palettize(
         clustered_weight = None
         axis = self.granularity.axis if self.granularity.axis else 0
 
+        lut = lut.to(indices.device)
+
         # Reshape indices back to 2D for block processing (reverse of
         # reshape_to_original in _assign_indices)
         indices = self.reshape_strategy.reshape_for_kmeans(indices, axis)

tests/palettization/test_kmeans_fake_palettize.py

@@ -474,9 +474,8 @@ def test_disabled_flag_behavior(self):
 
     @pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
     def test_mps_device_handling(self):
-        """
-        Test that palettization preserves MPS device for weights while
-        using CPU for computation.
+        """Palettization keeps the LUT on the weight's device while the weight-sized
+        indices stay CPU-resident, and reconstruction returns on the original device.
         """
         # Create weights on MPS device
         weight = torch.randn(4, 8, dtype=torch.float32, device="mps")
@@ -494,14 +493,15 @@ def test_mps_device_handling(self):
             enable_per_channel_scale=spec.enable_per_channel_scale,
         )
 
-        # Calculate centroids - this should move computation to CPU internally
+        # Clustering runs on CPU internally, but centroids are placed on the weight's device.
         lut, indices = _initialize_and_get_lut_indices(palettizer, weight)
 
-        # Verify LUT and indices are on CPU (expected behavior)
-        assert lut.device.type == "cpu"
+        # The LUT follows the weight's device (kept on-device so training-time LUT
+        # quantization stays on-device); the weight-sized indices stay CPU-resident.
+        assert lut.device.type == "mps"
         assert indices.device.type == "cpu"
 
-        # Palettize the weights - this should return result on original device (MPS)
+        # Reconstruction gathers against the CPU indices and returns on the original device (MPS).
         palettized_weight = palettizer._palettize(lut, indices, weight)
 
         # Verify the palettized weights are back on MPS
@@ -1826,3 +1826,58 @@ def test_blocks_to_cluster_raises_on_indivisible_cluster_dim(self):
         weight_2d = torch.randn(4, 16)  # axis-0 size 4 not divisible by cluster_dim 3
         with pytest.raises(_IncompatibleClusterDimError):
             palettizer._blocks_to_cluster(weight_2d, axis=0)
+
+
+def _accelerator_device() -> str | None:
+    """Return an available accelerator device type ("cuda" or "mps"), else None."""
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return None
+
+
+@pytest.mark.skipif(_accelerator_device() is None, reason="requires a CUDA or MPS accelerator")
+def test_device_placement_on_accelerator():
+    """LUT quantization and reconstruction device behavior on an accelerator.
+
+    Checks in one pass that:
+      - centroids follow the model device; ``indices`` stays on CPU (memory),
+      - ``_raw_lut`` and the ``lut`` property follow the device (not forced to CPU),
+      - ``quantize_lut`` (the training-side observe path) stays on the input device
+        with no forced CPU round-trip,
+      - ``hard_assign`` (eval) gathers against the CPU ``indices`` and returns on the
+        weight's device -- no accelerator/cpu device mismatch.
+    """
+    device = _accelerator_device()
+
+    spec = PalettizationSpec(
+        n_bits=2,
+        granularity=PerTensorGranularity(),
+        cluster_dim=1,
+        lut_qspec=_make_lut_qspec(torch.int8),
+    )
+    palettizer = _KMeansFakePalettize(**spec.__dict__)
+    weight = torch.randn(8, 8, device=device)
+
+    # Clusters on CPU, places centroids on the weight's device, and seeds the LUT
+    # quantizer on-device via quantize_lut(_raw_lut(centroids)).
+    palettizer._initialize(weight)
+
+    # centroids follow the model device; the weight-sized indices stay CPU-resident.
+    assert palettizer.centroids.device.type == device
+    assert palettizer.indices.device.type == "cpu"
+
+    # _raw_lut and the lut property follow the centroids' device (not forced CPU).
+    raw_lut = palettizer._raw_lut(palettizer.centroids)
+    assert raw_lut.device.type == device
+    assert palettizer.lut.device.type == device
+
+    # Quantizing an on-device LUT stays on-device.
+    assert palettizer.quantize_lut(raw_lut).device.type == device
+
+    # Eval path: reconstruction gathers against the CPU indices and returns on the
+    # weight's device.
+    out = palettizer.hard_assign(weight)
+    assert out.device.type == device
+    assert out.shape == weight.shape