refactor: remove coremltools dependency and use torch.utils.cpp_extension to load kmeans1d (#31)

* feat: add vendored C++ kmeans1d compiled via torch cpp_extension

Vendor coremltools' kmeans1d C++ core (_core.cpp, byte-for-byte from apple/coremltools; upstream kmeans1d 0.3.1, MIT) and JIT-compile it at runtime with torch.utils.cpp_extension, invoking its extern "C" entry points via ctypes. Provides cluster()/Clustered as a drop-in replacement for coremltools._deps._kmeans1d, with fail-loud input validation.

* refactor: switch palettization kmeans to vendored kmeans1d

Point _efficient_kmeans and kmeans_fake_palettize at coreai_opt._utils._kmeans1d instead of coremltools._deps._kmeans1d. The cluster()/Clustered call contract is unchanged.

* refactor: make coreai model palettization use vendored kmeans1d

Drop the optional third-party kmeans1d import and _HAS_KMEANS1D flag in palettize_utils; import coreai_opt._utils._kmeans1d unconditionally. The sklearn fallback for vector k-means is unchanged.

* test: add equivalence and source-hash tests for vendored kmeans1d

Compare cluster() output against the coremltools oracle (unweighted, weighted, large-n, collapse, fp16 ties), assert fixed known-good cases and input validation, and add a canary that hard-fails if the vendored _core.cpp drifts from apple/coremltools main.

* build: move coremltools to optional coreml dependency

coremltools is no longer needed at runtime; move it (and its numpy<2.4 cap) into a coreml optional-dependency + self-referencing dependency-group, kept in default-groups so the test oracle is available. Ship the vendored _core.cpp and LICENSE as package data.

* docs: add changelog fragment for coremltools optional dependency

* perf: compile vendored kmeans1d with -O2 -DNDEBUG to match coremltools

The vendored kmeans1d C++ extension was 10-15x slower per call than
coremltools' precompiled version (measured directly: ~0.9s vs ~0.06s
on a realistic-size weighted cluster() call), even with the actual
JIT-compile lock/cache fully warmed. This was the real cause of the
1.2x-2.6x palettization wall-clock regression seen in an end-to-end
LLM eval, not JIT-compile lock contention across parallel workers
(a prior "warmup()" fix targeting that, since reverted, made no
measurable difference).

Root cause: torch.utils.cpp_extension.load() does not inherit
CPython's sysconfig OPT/CFLAGS the way distutils/setuptools does, so
it defaults to -O0. coremltools' setup.py sets no explicit -O flag
either, but distutils prepends sysconfig's "-O2 -DNDEBUG" (the
standard python.org/conda default) ahead of its extra_compile_args,
so its shipped .so is effectively an -O2 build. This tight-inner-loop,
template-heavy DP algorithm is extremely sensitive to that gap.

Pass the same flags explicitly so torch's JIT build matches
coremltools' effective compile line. Numeric output is unaffected
(verified against the existing coremltools-oracle equivalence tests).

* fix: add setuptools as runtime dependency for vendored kmeans1d

torch.utils.cpp_extension.load() imports setuptools internally to JIT-compile
the vendored kmeans1d core. This was previously pulled in transitively via
coremltools, which is now an optional dependency, so it must be declared
explicitly.

* refactor: move vendored kmeans1d into src/coreai_opt/deps/_kmeans1d

Isolates the vendored third-party kmeans1d source (and its MIT LICENSE)
into a dedicated deps/ folder, separate from coreai-opt's own internal
_utils modules, per Apple legal's guidance on incorporating vendored
dependencies (mirrors how coremltools itself vendored the same code).

* refactor: update import paths for moved kmeans1d deps folder

Completes the previous move: repoints pyproject.toml package-data and
the three call sites at coreai_opt.deps._kmeans1d instead of the old
coreai_opt._utils._kmeans1d location. Also anchors the .gitignore
deps/ rule to the repo root (/deps/) so it stops shadowing the new
src/coreai_opt/deps/ package.

* docs: apply Apple copyright line to vendored kmeans1d headers

Bring the per-file headers in line with Apple legal's guidance for
incorporating vendored third-party code (mirrors coremltools' own
kmeans1d headers): every file in the vendored package now carries the
upstream MIT license text followed by a trailing Apple copyright line,
with no coreai-opt repo-wide header layered on top since this isn't
coreai-opt's own code. Bumps the trailing line from "Copyright (c) 2023
Apple Inc." (when coremltools first vendored it) to "Copyright (c) 2026
Apple Inc." (when coreai-opt vendored it from coremltools).

Also excludes src/coreai_opt/deps/ from the add-license-header
pre-commit hook, which otherwise unconditionally stamps coreai-opt's
own BSD-3 header onto any .py file lacking one.

* docs: add NOTICE.txt for kmeans1d attribution

Attributes the vendored kmeans1d source under src/coreai_opt/deps/ to
its upstream MIT license, per Apple legal's guidance for incorporating
vendored dependencies (mirrors the equivalent NOTICE.txt commit in
coremltools, which vendored the same code).

* test: normalize Apple copyright year in kmeans1d drift canary

The vendored-source-matches-upstream canary hashes our _core.cpp
against coremltools' live upstream copy. Now that coreai-opt's copy
carries its own vendoring-year copyright line ("Copyright (c) 2026
Apple Inc.") rather than coremltools' ("Copyright (c) 2023 Apple
Inc."), the raw byte hashes never match. Strip that one known,
intentional line from both sides before hashing so the canary still
catches real drift in the kmeans1d logic without permanently failing
on the copyright year.

* chore: remove unused deps/ gitignore rule

The unanchored `deps/` rule (from the initial commit, alongside other
personal-scratch entries like scratch/ and /local_stash) silently
shadowed the newly-added, tracked src/coreai_opt/deps/ package: search
tools like ripgrep skip gitignored paths by default, and any future
new file added under deps/ would need an explicit force-add. Nothing
in the repo (Makefile, CI, scripts, docs) reads or writes a deps/
folder, so remove the rule outright rather than just anchoring it.

* docs: note C++ compiler requirement in changelog fragment

The vendored kmeans1d is JIT-compiled via torch.utils.cpp_extension,
which needs a C++ toolchain on the host at runtime; previously this
requirement rode in transitively via the coremltools wheel, so it's
worth calling out now that coremltools is optional.

* docs: trim changelog wording per review feedback

Drop "(copied byte-for-byte from coremltools)" -- attribution already
lives in the vendored files' headers and NOTICE.txt, and the phrase
will go stale once the core is further optimized (per aseemw's PR #31
review).

* docs: drop coremltools-diff comment from vendored core.py

aseemw's review: comments explaining implementation choices by diffing
against what coremltools' build does are unnecessary now that this is
a fresh adaptation, not a diff against coremltools. Drop the comment
entirely rather than rephrasing it to be self-contained; the flags
themselves are unchanged.

* build: drop unneeded coreml group from docs dependencies

Re-verified via explain-coreai-opt (aseemw asked why docs needs the
coreml group): Sphinx's autodoc/autosummary machinery does genuinely
import every public coreai_opt module at build time, but nothing under
src/coreai_opt/ imports coremltools anywhere -- the only remaining
kmeans1d/coremltools attribution now lives in comments and the vendored
package's own LICENSE/header, none of which autodoc touches. No CI job
builds docs either, so this has no build-time effect to verify against.

* test: remove coremltools as a test dependency for kmeans1d

aseemw's PR #31 review: kmeans1d correctness tests importing
coremltools as a live oracle defeats the purpose of removing it as a
required dependency, and the "must stay byte-identical to upstream"
canary asserts an invariant that isn't actually required (the vendored
code can be modified/adapted going forward).

Collapses TestEquivalence/TestKmeans1DBehavior/TestsKmeans1dFromCoremltools
into one TestKmeans1D class and deletes TestVendoredSourceMatchesUpstream
entirely. Small discrete-input cases (k>n clamping, fewer-than-k-cluster
collapse, explicit duplicates) now assert hardcoded literal values
directly, captured by running the vendored cluster() once (already
validated against coremltools in CI, so a known-good golden value).

Larger/random-array cases (unweighted/weighted equivalence grids,
large-n up to 100k, fp16 duplicates, weighted-integer-counts) would
require thousands of literal values to hardcode inline, so their
inputs and expected coremltools-oracle output are instead checked into
tests/palettization/assets/kmeans1d/*.npz and loaded at test time
(~3.3MB total, mostly the large-n case).

Net effect: this test file has zero coremltools imports (confirmed via
AST inspection), so coreml stops being an implicit test dependency for
kmeans1d correctness.

Author: guru-desh

.gitignore

@@ -47,7 +47,6 @@ uv.lock
 *.nc
 .python-version
 /.tox
-deps/
 scratch/
 /local_stash
 /model_weights

.pre-commit-config.yaml

@@ -23,6 +23,7 @@ repos:
         language: system
         entry: scripts/pre_commit/add_license_header.py --license-file configs/BSD-3-LICENSE-HEADER-TEMPLATE --start-year 2026
         files: '\.(py|sh|js|css|html)$|(^|/)(GNUmakefile|[Mm]akefile)$'
+        exclude: ^src/coreai_opt/deps/
 
   # ----------------------------------------------------------------------------
   # 0.2 Update LICENSE year

NOTICE.txt

@@ -0,0 +1,25 @@
+Copyright 2026 Apple Inc.
+
+This project contains content adapted from kmeans1d (https://github.com/dstein64/kmeans1d), the license for which follows:
+
+MIT License
+
+Copyright (c) 2019 Daniel Steinberg
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

changelog.d/31.changed

@@ -0,0 +1 @@
+Replace the coremltools-based 1D k-means used by palettization with a vendored C++ core that is JIT-compiled at runtime via `torch.utils.cpp_extension`. `coremltools` is no longer a runtime dependency (it is now an optional dependency, installable via the `coreml` extra). This requires a C++ compiler to be available on the host at runtime.

pyproject.toml

@@ -31,11 +31,16 @@ classifiers = [
 ]
 dynamic = [ "version" ]
 dependencies = [
-    "coremltools>=8.3",
-    "numpy>=2,<2.4",              # TODO: Remove once coremltools has released a version > 9.0
+    # Required at runtime by torch.utils.cpp_extension to JIT-compile the vendored
+    # kmeans1d C++ core (a C++ toolchain must also be present on the host).
+    "ninja>=1.11",
+    "numpy>=2",
     "pydantic>=2.0.0",
+    "pyyaml>=6.0",
     "rich>=13.0.0",
     "safetensors>=0.5.3,<=0.7.0",
+    # Required at runtime by torch.utils.cpp_extension
+    "setuptools>=42",
     # PyTorch >= 2.9.0 requires torchao >= 0.15.0
     # Python's standard dependency specification (PEP 508) doesn't support conditional dependencies
     # based on other package versions. We can either 1) add a stricter check to require the newer torchao
@@ -54,6 +59,10 @@ coreai = [
     "coreai-torch==0.4.1",
     "scikit-learn>=1.7.2",
 ]
+coreml = [
+    "coremltools>=8.3",
+    "numpy>=2,<2.4",    # TODO: Remove once coremltools has released a version > 9.0
+]
 [project.urls]
 Changelog = "https://github.com/apple/coreai-optimization/blob/main/CHANGELOG.md"
 Repository = "https://github.com/apple/coreai-optimization"
@@ -95,7 +104,9 @@ docs = [
 # installs CoreAI by default (see tool.uv.default-groups). The extra is the
 # single source of truth; this group self-references it to avoid duplicating the
 # package list.
+# This applies to the `coreml` group as well.
 coreai = [ "coreai-opt[coreai]" ]
+coreml = [ "coreai-opt[coreml]" ]
 # Used in CI to force latest mimimum supported torch version
 # These torch versions must be in bounds of torch versions listed in project dependencies
 highest_tested_torch = [
@@ -145,6 +156,7 @@ tutorial = [
     "papermill>=2.7.0",
     "torchinfo>=1.8.0",
     { include-group = "coreai" },
+    { include-group = "coreml" },
     { include-group = "torchvision" },
 ]
 
@@ -153,14 +165,14 @@ package-dir = { "" = "src" }
 [tool.setuptools.dynamic]
 version = { attr = "coreai_opt._about.__version__" }
 [tool.setuptools.package-data]
-coreai_opt = [ "py.typed" ]
+coreai_opt = [ "deps/_kmeans1d/_core.cpp", "deps/_kmeans1d/LICENSE", "py.typed" ]
 [tool.setuptools.packages]
 find.exclude = [ "coreai_opt_benchmarking*", "tests*" ]
 find.where = [ "src" ]
 
 # make env group installed by default with uv sync
 [tool.uv]
-default-groups = [ "dev", "coreai" ]
+default-groups = [ "dev", "coreai", "coreml" ]
 index = [
     { explicit = true, name = "pytorch-cpu", url = "https://download.pytorch.org/whl/cpu" },
     { explicit = true, name = "pytorch-cu128", url = "https://download.pytorch.org/whl/cu128" },

src/coreai_opt/coreai_utils/_utils/palettize_utils.py

@@ -24,19 +24,12 @@
 )
 from coreai_opt.coreai_utils._utils.graph_utils import _select_input_output_channel_axis
 from coreai_opt.coreai_utils.common import CompressionGranularity as _CompressionGranularity
+from coreai_opt.deps import _kmeans1d
 
 logger = logging.getLogger(__name__)
 
 _CONV2D_OP = "coreai.conv2d"
 
-try:
-    import kmeans1d as _kmeans1d
-
-    _HAS_KMEANS1D = True
-except ImportError:
-    _kmeans1d = None  # type: ignore[assignment]
-    _HAS_KMEANS1D = False
-
 LutParams = namedtuple("LutParams", "indices lut vector_axis")
 
 _SUPPORTED_NBITS: tuple[int, ...] = (1, 2, 3, 4, 6, 8)
@@ -99,8 +92,7 @@ def _get_kmeans_lookup_table_and_weight(
         weight.shape[1] == 1 and num_weights >= 10_000 and weight.dtype == np.float16
     )
 
-    if (is_better_to_use_kmeans1d and _HAS_KMEANS1D) or force_kmeans1d:
-        assert _HAS_KMEANS1D, "Unable to import kmeans1d, please make sure it's installed."
+    if is_better_to_use_kmeans1d or force_kmeans1d:
         values, indices, counts = np.unique(weight, return_inverse=True, return_counts=True)
         indices = indices.flatten()
         n_clusters = min(len(values), lut_len)

src/coreai_opt/deps/_kmeans1d/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Daniel Steinberg
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/coreai_opt/deps/_kmeans1d/__init__.py

@@ -0,0 +1,27 @@
+# MIT License
+#
+# Copyright (c) 2019 Daniel Steinberg
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+# Copyright © 2026 Apple Inc.
+
+from coreai_opt.deps._kmeans1d.core import Clustered, cluster
+
+__all__ = ["Clustered", "cluster"]