Add initial support for joint sparsity pipeline

src/coreai_opt/palettization/kmeans/_prepare_for_export.py

@@ -6,6 +6,7 @@
 from dataclasses import dataclass
 from os import PathLike
 from pathlib import Path
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -44,6 +45,34 @@ class PalettizationInfo:
     lut_quantization: LUTQuantizationInfo | None = None
 
 
+class _SparsePalettizeReconstruction(nn.Module):
+    """Parametrization module inserted to reconstruct a sparse-palettized weight.
+
+    Traces ``coreai.lut_to_dense`` and ``coreai.sparse_to_dense`` in that order.
+    """
+
+    def __init__(
+        self,
+        nonzero_indices: torch.Tensor,
+        lut: torch.Tensor,
+        mask: torch.Tensor,
+        vector_axis: int | None,
+    ) -> None:
+        super().__init__()
+        self.register_buffer("nonzero_indices", nonzero_indices)
+        # nonzero_indices is rank 1 (flattened by masking), so lut must be
+        # reshaped to rank 3 (lut_to_dense requires lut.rank == indices.rank + 2).
+        self.register_buffer("lut", lut.reshape(1, lut.shape[-2], lut.shape[-1]))
+        self.register_buffer("mask", mask)
+        self.vector_axis = 0 if vector_axis is None else vector_axis
+
+    def forward(self, _: Any) -> torch.Tensor:
+        nonzero_values = torch.ops.coreai.lut_to_dense(
+            self.nonzero_indices, self.lut, self.vector_axis
+        )
+        return torch.ops.coreai.sparse_to_dense(nonzero_values, self.mask)
+
+
 def _expand_rank(
     tensor: torch.Tensor,
     target_rank: int,
@@ -208,6 +237,7 @@ def _insert_mlir_custom_op(
     module_name: str,
     param_name: str,
     palett_info: PalettizationInfo,
+    fake_palett_mod: _FakePalettizeImplBase,
     fake_palett_idx: int,
     mmap_dir: str | PathLike[str] | None,
 ) -> None:
@@ -225,6 +255,10 @@ def _insert_mlir_custom_op(
     4. Both: lut_to_dense(int LUT) + constexpr_blockwise_shift_scale(fused_scale)
        where fused_scale = lut_scale * per_channel_scale
 
+    When ``fake_palett_mod.sparsity`` is set, the LUT lookup runs on the
+    nonzero-only indices and the result is packed via ``coreai::sparse_to_dense``
+    instead of installing a plain Palettize/ScaledPalettize parametrization.
+
     When ``mmap_dir`` is provided, the new MLIR module is serialized to a
     safetensors file under that directory and reloaded via mmap before being
     swapped in.
@@ -259,7 +293,18 @@ def _import_coreai_torch_modules():
 
     vector_axis = _DEFAULT_VECTOR_AXIS if palett_info.cluster_dim > 1 else None
 
-    if needs_scale:
+    if fake_palett_mod.sparsity is not None:
+        # Reuses the mask from prepare()'s forward pass. needs_scale is always False here:
+        # PalettizationSpec rejects lut_qspec/enable_per_channel_scale combined with sparsity.
+        mask = fake_palett_mod._sparsity_mask.to(torch.bool)
+        nonzero_indices = palett_info.indices[mask]
+        mlir_palett_mod = _SparsePalettizeReconstruction(
+            nonzero_indices=nonzero_indices,
+            lut=palett_info.lut,
+            mask=mask,
+            vector_axis=vector_axis,
+        )
+    elif needs_scale:
         lut, scale, zero_point = _resolve_mlir_lut_and_scale(palett_info)
         mlir_palett_mod = ScaledPalettizeParametrization(
             indices=palett_info.indices,
@@ -334,7 +379,7 @@ def _process_palettized_parameter(
         _register_mil_compression_metadata(module, param_name, palett_info)
     elif backend == ExportBackend.CoreAI:
         _insert_mlir_custom_op(
-            module, module_name, param_name, palett_info, fake_palett_idx, mmap_dir
+            module, module_name, param_name, palett_info, fake_palett_mod, fake_palett_idx, mmap_dir
         )
 
 

src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py

@@ -82,13 +82,15 @@ def __init__(
         rounding_precision: int = 4,
         op_to_optimize: Callable | None = None,
         training_strategy_spec: TrainingStrategySpec | None = None,
+        sparsity: float | None = None,
     ):
         super().__init__(
             n_bits=n_bits,
             lut_qspec=lut_qspec,
             granularity=granularity,
             cluster_dim=cluster_dim,
             enable_per_channel_scale=enable_per_channel_scale,
+            sparsity=sparsity,
         )
 
         self.enable_fast_kmeans_mode = enable_fast_kmeans_mode

src/coreai_opt/palettization/kmeans/palettizer.py

@@ -493,6 +493,7 @@ def _spec_to_partial(
         # Serialize the spec, then layer in the owning module's compressor-specific
         # settings (e.g. enable_fast_kmeans_mode, rounding_precision).
         args = spec.model_dump_preserve_objects()
+        args["sparsity"] = spec._sparsity
         args.update(module_config._get_fake_module_kwargs())
         return _KMeansFakePalettize.with_args(**args)
 

src/coreai_opt/palettization/spec/fake_palettize.py

@@ -18,6 +18,7 @@
 from coreai_opt.palettization.spec import (
     PalettizationGranularity,
 )
+from coreai_opt.pruning.spec import PruneImplBase, Unstructured
 from coreai_opt.quantization.spec import QuantizationSpec
 
 
@@ -38,6 +39,7 @@ def __init__(
         granularity: PalettizationGranularity,
         cluster_dim: int,
         enable_per_channel_scale: bool,
+        sparsity: float | None = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -46,6 +48,7 @@ def __init__(
         self.granularity = granularity
         self.cluster_dim = cluster_dim
         self.enable_per_channel_scale = enable_per_channel_scale
+        self.sparsity = sparsity
 
         self.register_buffer("fake_palett_enabled", torch.tensor([1], dtype=torch.uint8))
         # Non-persistent (kept out of new checkpoints); when set to 1 (at runtime or
@@ -54,6 +57,7 @@ def __init__(
             "observer_enabled", torch.tensor([0], dtype=torch.uint8), persistent=False
         )
         self._disabled = False
+        self.register_buffer("_sparsity_mask", None, persistent=False)
 
         self.register_buffer("indices", None)
         self.register_buffer("per_channel_scale", None)
@@ -71,6 +75,12 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         if self._disabled:
             return tensor
 
+        if self.sparsity is not None:
+            self._sparsity_mask = PruneImplBase.resolve("default").compute_mask(
+                tensor, self.sparsity, Unstructured()
+            )
+            tensor = tensor * self._sparsity_mask
+
         self.ensure_initialized(tensor)
 
         # Check for self._disabled again in case ensure_initialized disabled the palettizer.

src/coreai_opt/palettization/spec/spec.py

@@ -108,6 +108,27 @@ class PalettizationSpec(CompressionSpec):
     # Private attribute for compression type
     _compression_type: CompressionType = PrivateAttr(default=CompressionType.PALETTIZATION)
 
+    # Sparsity level, in [0, 1]. Set via the `_sparsity` constructor/dict key.
+    _sparsity: float | None = PrivateAttr(default=None)
+
+    def __init__(self, **data: Any) -> None:
+        sparsity = data.pop("_sparsity", None)
+        super().__init__(**data)
+        if sparsity is not None:
+            self._validate_sparsity(sparsity)
+            self._sparsity = sparsity
+
+    def _validate_sparsity(self, sparsity: float) -> None:
+        """Reject sparsity combined with a position-dependent LUT/scale mapping."""
+        if not (0.0 <= sparsity <= 1.0):
+            raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}")
+        if self.lut_qspec is not None:
+            raise ValueError("lut_qspec not supported for joint sparsity.")
+        if not isinstance(self.granularity, PerTensorGranularity):
+            raise ValueError(f"granularity={self.granularity} not supported for joint sparsity.")
+        if self.enable_per_channel_scale:
+            raise ValueError("enable_per_channel_scale not supported for joint sparsity.")
+
     @model_validator(mode="after")
     def validate_lut_qspec(self) -> "PalettizationSpec":
         """Validate that lut_qspec only uses supported configurations."""

src/coreai_opt/quantization/_graph/_prepare_for_export.py

@@ -267,8 +267,13 @@ def _import_coreai_custom_ops():
     if minval is not None:
         minval = minval.to(dtype=_compute_dtype_for_export)
 
-    # Construct quantized weights
+    # Construct quantized weights, reusing the mask computed during
+    # prepare()'s forward pass if sparsity is set.
     dense_weight = resolve_attr(model, input_node.target).data
+    mask: torch.Tensor | None = None
+    if fake_quant_mod.sparsity is not None:
+        mask = fake_quant_mod._sparsity_mask.to(torch.bool)
+        dense_weight = dense_weight * mask
     quantized_data = fake_quant_mod.quantize(dense_weight, scale, zero_point, minval)
 
     # Drop one of the offsets so that the export
@@ -284,13 +289,28 @@ def _import_coreai_custom_ops():
 
     # Register buffers and get buffer names
     param_name = str(input_node.target).replace(".", "_")
-    buffer_names = _register_quantization_buffers(
-        model, param_name, scale, zero_point, quantized_data, minval
-    )
+    if mask is not None:
+        nonzero_data = quantized_data[mask]
+        buffer_names = _register_quantization_buffers(
+            model, param_name, scale, zero_point, minval=minval
+        )
+        model.register_buffer(f"{param_name}_nonzero", nonzero_data)
+        model.register_buffer(f"{param_name}_mask", mask)
+    else:
+        buffer_names = _register_quantization_buffers(
+            model, param_name, scale, zero_point, quantized_data, minval
+        )
 
     # Create graph nodes and replace fake quantization
     with model.graph.inserting_before(node):
-        quantized_data_node = model.graph.get_attr(buffer_names["quantized_data"])
+        if mask is not None:
+            nonzero_node = model.graph.get_attr(f"{param_name}_nonzero")
+            mask_node = model.graph.get_attr(f"{param_name}_mask")
+            quantized_data_node = model.graph.call_function(
+                coreai.sparse_to_dense, (nonzero_node, mask_node)
+            )
+        else:
+            quantized_data_node = model.graph.get_attr(buffer_names["quantized_data"])
         scale_node = model.graph.get_attr(buffer_names["scale"])
 
         if zero_point is not None:

src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py

@@ -64,6 +64,7 @@
     FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls",
     FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls",
     FieldName.SCALE_DTYPE: "scale_dtype",
+    FieldName.SPARSITY: "_sparsity",
 }
 
 

src/coreai_opt/quantization/_graph/_qspec_constraints.py

@@ -320,6 +320,7 @@ def _policy_float_range_union(proposals: Sequence[FieldValue]) -> FieldValue:
     FieldName.QPARAM_CALCULATOR_CLS: _policy_priority_wins,
     FieldName.RANGE_CALCULATOR_CLS: _policy_priority_wins,
     FieldName.SCALE_DTYPE: _policy_priority_wins,
+    FieldName.SPARSITY: _policy_priority_wins,
     # Covering every member's values is a correctness constraint, not a
     # preference, so this one unions instead of deferring to priority.
     FieldName.FLOAT_RANGE: _policy_float_range_union,

src/coreai_opt/quantization/_graph/_qspec_resolution.py

@@ -178,6 +178,7 @@ def _shared_spec_pointing_at(anchor: NodeSlot) -> _SharedQuantizationSpec:
     FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls",
     FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls",
     FieldName.SCALE_DTYPE: "scale_dtype",
+    FieldName.SPARSITY: "_sparsity",
 }
 
 

src/coreai_opt/quantization/_graph/_qspec_types.py

@@ -61,6 +61,7 @@ class FieldName(enum.Enum):
     QPARAM_CALCULATOR_CLS = enum.auto()
     RANGE_CALCULATOR_CLS = enum.auto()
     SCALE_DTYPE = enum.auto()
+    SPARSITY = enum.auto()  # QuantizationSpec's private, settable `_sparsity` input
     # Weight or activation, set by which config dict the spec came from. Not a
     # QuantizationSpec attribute, but an input to construct_partial alongside
     # them, so resolution needs it to rebuild the observer.

src/coreai_opt/quantization/spec/factory.py

@@ -209,6 +209,7 @@ def create_fake_quantizer(
             "quant_max": spec.quant_max,
             "qparams_calculator": qparams_calculator,
             "n_bits": spec.n_bits,
+            "sparsity": spec._sparsity,
         }
 
         # Automatically detect and include any extra arguments
@@ -289,6 +290,7 @@ def create_fake_quantizer_partial(
             "quant_min": spec.quant_min,
             "quant_max": spec.quant_max,
             "n_bits": spec.n_bits,
+            "sparsity": spec._sparsity,
         }
 
         # Automatically detect and include any extra arguments

src/coreai_opt/quantization/spec/fake_quantize.py

@@ -26,6 +26,7 @@
     is_float_quant_dtype as _is_float_quant_dtype,
 )
 from coreai_opt.config.spec import CompressionSimulatorBase, CompressionTargetTensor
+from coreai_opt.pruning.spec import PruneImplBase, Unstructured
 from coreai_opt.quantization._utils import get_quantization_shapes as _get_quantization_shapes
 from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError
 from coreai_opt.quantization.spec.qscheme import QuantizationScheme
@@ -54,6 +55,7 @@ def __init__(
         quant_max: int | float,
         qparams_calculator: QParamsCalculatorBase,
         n_bits: int | None = None,
+        sparsity: float | None = None,
         **kwargs,
     ):
         super().__init__()
@@ -64,7 +66,9 @@ def __init__(
         self.quant_min = quant_min
         self.quant_max = quant_max
         self.qparams_calculator = qparams_calculator
+        self.sparsity = sparsity
         self.register_buffer("_disabled", torch.tensor(False))
+        self.register_buffer("_sparsity_mask", None, persistent=False)
 
         # Infer n_bits from dtype if not provided
         if n_bits is None:
@@ -146,6 +150,12 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         if self._disabled.item():
             return tensor
 
+        if self.sparsity is not None:
+            self._sparsity_mask = PruneImplBase.resolve("default").compute_mask(
+                tensor, self.sparsity, Unstructured()
+            )
+            tensor = tensor * self._sparsity_mask
+
         if self.observer_enabled[0] == 1:
             # Call the forward function of the qparams_calculator
             # to collect observer statistics when the observer is

src/coreai_opt/quantization/spec/spec.py

@@ -373,6 +373,16 @@ class type: MinMaxRangeCalculator or custom registered class type
     # Private attribute for compression type
     _compression_type: CompressionType = PrivateAttr(default=CompressionType.QUANTIZATION)
 
+    # Sparsity level, in [0, 1]. Set via the `_sparsity` constructor/dict key.
+    _sparsity: float | None = PrivateAttr(default=None)
+
+    def __init__(self, **data: Any) -> None:
+        sparsity = data.pop("_sparsity", None)
+        super().__init__(**data)
+        if sparsity is not None:
+            self._validate_sparsity_zero_preserving(sparsity)
+            self._sparsity = sparsity
+
     # Supported dtypes for quantization (class attribute for testing extensibility)
     SUPPORTED_DTYPES: ClassVar[set[torch.dtype]] = {
         # Signed integer types
@@ -556,6 +566,22 @@ def validate_scale_dtype(self) -> QuantizationSpec:
 
         return self
 
+    def _validate_sparsity_zero_preserving(self, sparsity: float) -> None:
+        """Reject sparsity unless a raw 0 dequantizes to exactly 0.0."""
+        if not (0.0 <= sparsity <= 1.0):
+            raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}")
+        if _is_float4_dtype(self.dtype):
+            raise ValueError("FP4 dtype not supported for joint sparsity.")
+        if self.dtype.is_floating_point:
+            return
+
+        if self.qformulation != QuantizationFormulation.ZP:
+            raise ValueError(f"qformulation={self.qformulation} not supported for joint sparsity.")
+        if self.qscheme == QuantizationScheme.ASYMMETRIC:
+            raise ValueError(f"qscheme={self.qscheme} not supported for joint sparsity.")
+        if not self.dtype.is_signed:
+            raise ValueError(f"unsigned dtype={self.dtype} not supported for joint sparsity.")
+
     def get_extra_args(self) -> dict[str, Any]:
         """
         Automatically detect and return fields beyond base QuantizationSpec.