coreai_opt.palettization.config.OpKMeansPalettizerConfig

class coreai_opt.palettization.config.OpKMeansPalettizerConfig

Bases: WeightOnlyOpValidationMixin, OpCompressionConfig[PalettizationSpec]

Configuration class for palettization at the operation level.

Palettization is a weight-only compression technique that doesn’t apply to activations (inputs/outputs). Only op_state_spec is used to configure which state tensors (e.g., weights, biases) should be palettized.

op_state_spec

Palettization specifications for operation state tensors (parameters, buffers, constants). Keys are string names (e.g., “weight”, “bias”) or “*” to refer to all state inputs. Values are PalettizationSpec objects or None to disable palettization for that state tensor. Default: 4-bit palettization for “weight” and “in_proj_weight” state tensors via default_weight_palettization_spec().

Type:

dict[str, PalettizationSpec | None] | None

Example

>>> # Palettize weights with 2-bit LUT
>>> op_config = OpKMeansPalettizerConfig(
...     op_state_spec={
...         "weight": PalettizationSpec(n_bits=2)
...     }
... )
>>>
>>> # Disable palettization for a specific operation
>>> op_config = OpKMeansPalettizerConfig(
...     op_state_spec=None  # No state tensors will be palettized
... )

classmethod get_default_state_spec()

Provide default state spec for palettization.

Return type:

dict[str, PalettizationSpec | None]