fastkmeans 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

fastkmeans/__init__.py

@@ -1,4 +1,4 @@
 from .kmeans import FastKMeans
 
 __all__ = ["FastKMeans"]
-__version__ = "0.1.0"
+__version__ = "0.3.0"

fastkmeans/kmeans.py

@@ -3,25 +3,43 @@ import time
 import torch
 import numpy as np
 
-def _get_device(preset: str = None):
-    if preset: return preset
-    if torch.cuda.is_available(): return 'cuda'
-    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): return 'mps'
-    return 'cpu'
+def _get_device(preset: str | int | torch.device | None = None):
+    if isinstance(preset, torch.device):
+        return preset
+    if isinstance(preset, str):
+        return torch.device(preset)
+    if torch.cuda.is_available(): # cuda currently handles both AMD and NVIDIA GPUs
+        return torch.device(f"cuda:{preset if isinstance(preset, int) and preset < torch.cuda.device_count() else 0}")
+    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        return torch.device('mps')
+    if hasattr(torch, 'xpu') and torch.xpu.is_available():
+        return torch.device(f'xpu:{preset if isinstance(preset, int) and preset < torch.xpu.device_count() else 0}')
+    return torch.device('cpu')
+
+
+def _is_bfloat16_supported(device:torch.device):
+    if device.type == 'cuda':
+        return torch.cuda.is_bf16_supported()
+    elif device.type == 'xpu' and hasattr(torch.xpu, 'is_bf16_supported'):
+        return torch.xpu.is_bf16_supported()
+    else:
+        return False
+
 
 @torch.inference_mode()
 def _kmeans_torch_double_chunked(
     data: torch.Tensor,
     data_norms: torch.Tensor,
     k: int,
+    device: torch.device,
+    dtype: torch.dtype | None = None,
     max_iters: int = 25,
     tol: float = 1e-8,
-    device: str = None,
-    dtype: torch.dtype = None,
     chunk_size_data: int = 50_000,
     chunk_size_centroids: int = 10_000,
     max_points_per_centroid: int = 256,
     verbose: bool = False,
+    use_triton: bool | None = None,
 ):
     """
     An efficient kmeans implementation that minimises OOM risks on modern hardware by using conversative double chunking.
@@ -32,6 +50,13 @@ def _kmeans_torch_double_chunked(
     labels_cpu    : torch.Tensor, shape (n_samples_used,), long
         Where n_samples_used can be smaller than the original if subsampling occurred.
     """
+
+    if use_triton:
+        from fastkmeans.triton_kernels import chunked_kmeans_kernel
+
+    if dtype is None:
+        dtype = torch.float16 if device.type in ['cuda', 'xpu'] else torch.float32
+
     n_samples_original, n_features = data.shape
     n_samples = n_samples_original
 
@@ -47,19 +72,16 @@ def _kmeans_torch_double_chunked(
     if n_samples < k:
         raise ValueError(f"Number of training points ({n_samples}) is less than k ({k}).")
 
-    if dtype is None: dtype = torch.float16 if device == 'cuda' else torch.float32
-
     # centroid init -- random is the only supported init
     rand_indices = torch.randperm(n_samples)[:k]
     centroids = data[rand_indices].clone().to(device=device, dtype=dtype)
     prev_centroids = centroids.clone()
 
-    labels = torch.empty(n_samples, dtype=torch.long, device='cpu')  # Keep labels on CPU
-
+    labels = torch.empty(n_samples, dtype=torch.int64, device='cpu')  # Keep labels on CPU
 
     for iteration in range(max_iters):
         iteration_start_time = time.time()
-        
+
         centroid_norms = (centroids ** 2).sum(dim=1)
         cluster_sums = torch.zeros((k, n_features), device=device, dtype=torch.float32)
         cluster_counts = torch.zeros((k,), device=device, dtype=torch.float32)
@@ -71,58 +93,62 @@ def _kmeans_torch_double_chunked(
             data_chunk = data[start_idx:end_idx].to(device=device, dtype=dtype, non_blocking=True)
             data_chunk_norms = data_norms[start_idx:end_idx].to(device=device, dtype=dtype, non_blocking=True)
             batch_size = data_chunk.size(0)
-
-            best_dist = torch.full((batch_size,), float('inf'), device=device, dtype=dtype)
-            best_ids = torch.zeros((batch_size,), device=device, dtype=torch.long)
-
-            c_start = 0
-            while c_start < k:
-                c_end = min(c_start + chunk_size_centroids, k)
-                centroid_chunk = centroids[c_start:c_end]
-                centroid_chunk_norms = centroid_norms[c_start:c_end]
-
-                dist_chunk = data_chunk_norms.unsqueeze(1) + centroid_chunk_norms.unsqueeze(0)
-                dist_chunk = dist_chunk.addmm_(
-                    data_chunk, centroid_chunk.t(), alpha=-2.0, beta=1.0
+            best_ids = torch.zeros((batch_size,), device=device, dtype=torch.int64)
+
+            if use_triton:
+                chunked_kmeans_kernel(
+                    data_chunk=data_chunk,
+                    data_chunk_norms=data_chunk_norms,
+                    centroids=centroids,
+                    centroids_sqnorm=centroid_norms,
+                    best_ids=best_ids,
                 )
+            else:
+                best_dist = torch.full((batch_size,), float('inf'), device=device, dtype=dtype)
+                c_start = 0
+                while c_start < k:
+                    c_end = min(c_start + chunk_size_centroids, k)
+                    centroid_chunk = centroids[c_start:c_end]
+                    centroid_chunk_norms = centroid_norms[c_start:c_end]
 
-                local_min_vals, local_min_ids = torch.min(dist_chunk, dim=1)
-                improved_mask = local_min_vals < best_dist
-                best_dist[improved_mask] = local_min_vals[improved_mask]
-                best_ids[improved_mask] = (c_start + local_min_ids[improved_mask])
+                    dist_chunk = data_chunk_norms.unsqueeze(1) + centroid_chunk_norms.unsqueeze(0)
+                    dist_chunk = dist_chunk.addmm_(data_chunk, centroid_chunk.t(), alpha=-2.0, beta=1.0)
 
-                c_start = c_end
+                    local_min_vals, local_min_ids = torch.min(dist_chunk, dim=1)
+                    improved_mask = local_min_vals < best_dist
+                    best_dist[improved_mask] = local_min_vals[improved_mask]
+                    best_ids[improved_mask] = (c_start + local_min_ids[improved_mask])
+
+                    c_start = c_end
 
             cluster_sums.index_add_(0, best_ids, data_chunk.float())
-            cluster_counts.index_add_(0, best_ids, torch.ones_like(best_ids, dtype=torch.float32))
+            cluster_counts.index_add_(0, best_ids, torch.ones_like(best_ids, device=device, dtype=torch.float32))
 
             labels[start_idx:end_idx] = best_ids.to('cpu', non_blocking=True)
             start_idx = end_idx
 
-        new_centroids = torch.zeros_like(centroids, device=device, dtype=torch.float32)
+        new_centroids = torch.zeros_like(centroids, device=device, dtype=dtype)
         non_empty = (cluster_counts > 0)
-        new_centroids[non_empty] = (
-            cluster_sums[non_empty] / cluster_counts[non_empty].unsqueeze(1)
-        )
+        new_centroids[non_empty] = (cluster_sums[non_empty] / cluster_counts[non_empty].unsqueeze(1)).to(dtype=dtype)
 
         empty_ids = (~non_empty).nonzero(as_tuple=True)[0]
         if len(empty_ids) > 0:
             reinit_indices = torch.randint(0, n_samples, (len(empty_ids),), device='cpu')
-            random_data = data[reinit_indices].to(device=device, dtype=torch.float32)
+            random_data = data[reinit_indices].to(device=device, dtype=dtype, non_blocking=True)
             new_centroids[empty_ids] = random_data
 
-        new_centroids = new_centroids.to(dtype=dtype)
-
         shift = torch.norm(new_centroids - prev_centroids.to(new_centroids.device), dim=1).sum().item()
         centroids = new_centroids
 
         prev_centroids = centroids.clone()
-        
+
         iteration_time = time.time() - iteration_start_time
-        if verbose: print(f"Iteration {iteration+1}/{max_iters} took {iteration_time:.4f}s, total time: {time.time() - iteration_start_time + iteration_time:.4f}s, shift: {shift:.6f}")
-        
+        if verbose:
+            print(f"Iteration {iteration+1}/{max_iters} took {iteration_time:.4f}s, total time: {time.time() - iteration_start_time + iteration_time:.4f}s, shift: {shift:.6f}")
+
         if shift < tol:
-            if verbose: print(f"Converged after {iteration+1} iterations (shift: {shift:.6f} < tol: {tol})")
+            if verbose:
+                print(f"Converged after {iteration+1} iterations (shift: {shift:.6f} < tol: {tol})")
             break
 
     centroids_cpu = centroids.to('cpu', dtype=torch.float32)
@@ -155,6 +181,9 @@ class FastKMeans:
         Chunk size along the data dimension for assignment/update steps.
     chunk_size_centroids : int, default=10_000
         Chunk size along the centroid dimension for assignment/update steps.
+    use_triton : bool | None, default=None
+       Use the fast Triton backend for the assignment/update steps.
+       If None, the Triton backend will be enabled for modern GPUs.
     """
 
     def __init__(
@@ -168,33 +197,36 @@ class FastKMeans:
         max_points_per_centroid: int = 256,
         chunk_size_data: int = 50_000,
         chunk_size_centroids: int = 10_000,
-        device: str = None,
+        device: str | int | torch.device | None = None,
         dtype: torch.dtype = None,
         pin_gpu_memory: bool = True,
         verbose: bool = False,
         nredo: int = 1, # for compatibility only
+        use_triton: bool | None = None,
     ):
         self.d = d
         self.k = k
         self.niter = niter
         self.tol = tol
-        self.gpu = gpu
         self.seed = seed
         self.max_points_per_centroid = max_points_per_centroid
         self.chunk_size_data = chunk_size_data
         self.chunk_size_centroids = chunk_size_centroids
+        self.device = _get_device("cpu" if gpu is False else device)
         self.centroids = None
-        if device not in [None, 'cuda'] and self.gpu: print("Warning: device is set to 'cuda' but gpu is True, ignoring 'device' argument and setting it to 'cuda'!")
-        self.device = 'cuda' if self.gpu else device
         self.dtype = dtype
         self.pin_gpu_memory = pin_gpu_memory
-        if nredo != 1: raise ValueError("nredo must be 1, redos not currently supported")
         self.verbose = verbose
+        if use_triton is not False:
+            use_triton = _is_bfloat16_supported(self.device) # assume triton is supported if GPU supports bfloat16
+        self.use_triton = use_triton
+        if nredo != 1:
+            raise ValueError("nredo must be 1, redos not currently supported")
 
     def train(self, data: np.ndarray):
         """
         Trains (fits) the KMeans model on the given data and sets `self.centroids`. Designed to mimic faiss's `train()` method.
-        
+
         Parameters
         ----------
         data : np.ndarray of shape (n_samples, d), float32
@@ -224,6 +256,7 @@ class FastKMeans:
             chunk_size_centroids=self.chunk_size_centroids,
             max_points_per_centroid=self.max_points_per_centroid,
             verbose=self.verbose,
+            use_triton=self.use_triton,
         )
         self.centroids = centroids.numpy()
 
@@ -242,6 +275,8 @@ class FastKMeans:
         -------
         labels : np.ndarray of shape (n_samples,), int64
         """
+        if self.use_triton:
+            from fastkmeans.triton_kernels import chunked_kmeans_kernel
         if self.centroids is None:
             raise RuntimeError("Must call train() or fit() before predict().")
 
@@ -250,11 +285,7 @@ class FastKMeans:
 
         # We'll do a chunked assignment pass, similar to the main loop, but no centroid updates
         centroids_torch = torch.from_numpy(self.centroids)
-        device = centroids_torch.device.type
-        if device == 'cpu' and self.gpu and torch.cuda.is_available():
-            device = 'cuda'  # If user asked for GPU, put centroids there
-
-        centroids_torch = centroids_torch.to(device=device, dtype=torch.float32)
+        centroids_torch = centroids_torch.to(device=self.device, dtype=torch.float32)
         centroid_norms = (centroids_torch ** 2).sum(dim=1)
 
         n_samples = data_torch.shape[0]
@@ -263,30 +294,37 @@ class FastKMeans:
         start_idx = 0
         while start_idx < n_samples:
             end_idx = min(start_idx + self.chunk_size_data, n_samples)
-            data_chunk = data_torch[start_idx:end_idx].to(device=device, dtype=torch.float32, non_blocking=True)
-            data_chunk_norms = data_norms_torch[start_idx:end_idx].to(device=device, dtype=torch.float32, non_blocking=True)
-            batch_size = data_chunk.size(0)
 
-            best_dist = torch.full((batch_size,), float('inf'), device=device, dtype=torch.float32)
-            best_ids = torch.zeros((batch_size,), device=device, dtype=torch.long)
-
-            c_start = 0
-            k = centroids_torch.shape[0]
-            while c_start < k:
-                c_end = min(c_start + self.chunk_size_centroids, k)
-                centroid_chunk = centroids_torch[c_start:c_end]
-                centroid_chunk_norms = centroid_norms[c_start:c_end]
-                
-                dist_chunk = data_chunk_norms.unsqueeze(1) + centroid_chunk_norms.unsqueeze(0)
-                dist_chunk = dist_chunk.addmm_(
-                    data_chunk, centroid_chunk.t(), alpha=-2.0, beta=1.0
+            data_chunk = data_torch[start_idx:end_idx].to(device=self.device, dtype=torch.float32, non_blocking=True)
+            data_chunk_norms = data_norms_torch[start_idx:end_idx].to(device=self.device, dtype=torch.float32, non_blocking=True)
+            batch_size = data_chunk.size(0)
+            best_ids = torch.zeros((batch_size,), device=self.device, dtype=torch.long)
+
+            if self.use_triton:
+                chunked_kmeans_kernel(
+                    data_chunk,
+                    data_chunk_norms,
+                    centroids_torch,
+                    centroid_norms,
+                    best_ids,
                 )
-
-                local_min_vals, local_min_ids = torch.min(dist_chunk, dim=1)
-                improved_mask = local_min_vals < best_dist
-                best_dist[improved_mask] = local_min_vals[improved_mask]
-                best_ids[improved_mask] = (c_start + local_min_ids[improved_mask])
-                c_start = c_end
+            else:
+                best_dist = torch.full((batch_size,), float('inf'), device=self.device, dtype=torch.float32)
+                c_start = 0
+                k = centroids_torch.shape[0]
+                while c_start < k:
+                    c_end = min(c_start + self.chunk_size_centroids, k)
+                    centroid_chunk = centroids_torch[c_start:c_end]
+                    centroid_chunk_norms = centroid_norms[c_start:c_end]
+
+                    dist_chunk = data_chunk_norms.unsqueeze(1) + centroid_chunk_norms.unsqueeze(0)
+                    dist_chunk = dist_chunk.addmm_(data_chunk, centroid_chunk.t(), alpha=-2.0, beta=1.0)
+
+                    local_min_vals, local_min_ids = torch.min(dist_chunk, dim=1)
+                    improved_mask = local_min_vals < best_dist
+                    best_dist[improved_mask] = local_min_vals[improved_mask]
+                    best_ids[improved_mask] = (c_start + local_min_ids[improved_mask])
+                    c_start = c_end
 
             labels[start_idx:end_idx] = best_ids.to('cpu')
             start_idx = end_idx

fastkmeans/triton_kernels.py

@@ -0,0 +1,105 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.heuristics(
+    {
+        "BLOCK_M": lambda a: 128 if a["D"] <= 128 else (64 if a["D"] <= 512 else 32),
+        "BLOCK_N": lambda a: 64 if a["D"] <= 128 else 32,
+        "num_warps": lambda a: 4 if a["D"] <= 64 else 8,
+        "num_stages": lambda a: 2 if a["D"] < 64 else 1,
+    }
+)
+@triton.jit
+def _chunked_kmeans_kernel(
+    data_ptr,  # [B, D], row-major
+    x_norm_ptr,  # [B], precomputed L2 norms of data
+    centroids_ptr,  # [C, D], row-major
+    centroids_sqnorm_ptr,  # [C], precomputed L2 norms of centroids
+    best_ids_ptr,  # [B], int32 (to store best centroid indices)
+    B,  # number of data points
+    C,  # number of centroids
+    D: tl.constexpr,  # dimension, or number of features
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """
+    Each Triton block processes BLOCK_M rows of data. The kernel:
+      1) loads those rows (and their precomputed norms from x_norm_ptr),
+      2) loops over all centroids in chunks of BLOCK_N,
+      3) computes distances, finds the best centroid,
+      4) writes out the best centroid index for each data point.
+    """
+    # 1) Identify which data rows this block handles
+    block_id = tl.program_id(axis=0)
+    row_start = block_id * BLOCK_M
+    rows = row_start + tl.arange(0, BLOCK_M)
+    mask = rows < B
+
+    # 2) Load data rows and precomputed x_norm: shape [BLOCK_M, D]
+    row_offsets = rows[:, None] * D + tl.arange(0, D)
+    x = tl.load(data_ptr + row_offsets, mask=mask[:, None], other=0.0)
+
+    # shape: [BLOCK_M]
+    x_norm = tl.load(x_norm_ptr + rows, mask=mask, other=0.0)
+
+    # Prepare "best distance" + "best index"
+    best_dist = tl.full([BLOCK_M], 1e38, dtype=tl.float32)
+    best_idx = tl.zeros([BLOCK_M], dtype=tl.int64)
+
+    # 3) Iterate over the centroids in chunks of BLOCK_N
+    for chunk in range(0, C, BLOCK_N):
+        cids = chunk + tl.arange(0, BLOCK_N)
+        c_mask = cids < C
+
+        # Load sub-block of centroids: shape [BLOCK_N, D]
+        c_offsets = cids[:, None] * D + tl.arange(0, D)
+        cvals = tl.load(centroids_ptr + c_offsets, mask=c_mask[:, None], other=0.0).to(x.dtype)
+
+        # Load centroid norms: shape [BLOCK_N]
+        c_sqnorm = tl.load(centroids_sqnorm_ptr + cids, mask=c_mask, other=0.0).to(x.dtype)
+
+        # Compute distance = x_norm + c_sqnorm - 2 * dot(x, c)
+        dots = tl.dot(x, tl.trans(cvals))  # shape [BLOCK_M, BLOCK_N]
+        dist_chunk = tl.fma(dots, -2.0, x_norm[:, None] + c_sqnorm[None, :])
+
+        # Find the argmin along the BLOCK_N dimension
+        local_min_vals, local_min_idx = tl.min(dist_chunk, axis=1, return_indices=True)
+
+        improved = local_min_vals < best_dist
+        best_dist = tl.where(improved, local_min_vals, best_dist)
+        best_idx = tl.where(improved, chunk + local_min_idx, best_idx)
+
+    # 4) Write out the best centroid indices
+    tl.store(best_ids_ptr + rows, best_idx, mask=mask)
+
+
+def chunked_kmeans_kernel(
+    data_chunk: torch.Tensor,
+    data_chunk_norms: torch.Tensor,
+    centroids: torch.Tensor,
+    centroids_sqnorm: torch.Tensor,
+    best_ids: torch.Tensor,
+):
+    """
+    Launches the Triton kernel to assign each point to its nearest centroid in one pass.
+
+    best_ids: pre-allocated [B] (int32) to store the best centroid ID of each point.
+    """
+    B, D = data_chunk.shape
+    C = centroids.shape[0]
+
+    def grid(meta):
+        return (triton.cdiv(B, meta["BLOCK_M"]),)
+
+    _chunked_kmeans_kernel[grid](
+        data_chunk,
+        data_chunk_norms,
+        centroids,
+        centroids_sqnorm,
+        best_ids,
+        B,  # num_points
+        C,  # num_centroids
+        D,  # dimension, or number of features
+    )

fastkmeans-0.3.0.dist-info/METADATA

@@ -0,0 +1,72 @@
+Metadata-Version: 2.4
+Name: fastkmeans
+Version: 0.3.0
+Summary: Add your description here
+Author-email: Ben Clavié <bc@answer.ai>, Benjamin Warner <bw@answer.ai>
+Maintainer-email: Ben Clavié <bc@answer.ai>, Benjamin Warner <bw@answer.ai>
+License: Apache-2.0
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: numpy
+Dynamic: license-file
+
+# fastkmeans
+
+![Python Versions](https://img.shields.io/badge/Python-3.8_3.9_3.10_3.11_3.12_3.13-blue)
+[![Twitter Follow](https://img.shields.io/twitter/follow/bclavie?style=social)](https://twitter.com/bclavie)
+<!-- [![Downloads](https://static.pepy.tech/badge/fastkmeans/month)](https://pepy.tech/project/fastkmeans) -->
+
+_A fast and efficient k-means implementation for PyTorch, with support for GPU and CPU._
+
+---
+
+Welcome to `fastkmeans`! This is an extremely tiny library, meant to be slotted-in anywhere you need "fast-enough" pytorch native k-means clustering. It's compatible with both CPU and GPU, matching or outspeeding `faiss` (except in multi-GPU settings), and it comes without any install woes, relying on just two dependencies you already have installed anyway: `torch` and `numpy`.
+
+###  Get started
+
+```sh
+[uv] pip install fastkmeans
+```
+
+... and that's all you need to do! `FastKMeans` is now ready to use.
+
+###  So what does this do?
+
+There's very, very little to this library. It provides a single interface, `FastKMeans`, which you can use by importing it from `fastkmeans`. This interface has been designed to be a slot-in replacement for existing FAISS implementations, while being *mostly* sklearn-compliant as well. Effectively, this means that three methods are exposed:
+
+- train(): mimics the FAISS API, training the model on a dataset.
+- fit(): mimics the sklearn API, training the model on a dataset.
+- predict(): mimics the sklearn API, use the trained clusters to predict where new points belong.
+- fit_predict(): mimics the sklearn API, chaining the two calls above.
+
+#### Behaviour
+
+Whenever possible, the library will attempt to mimic the FAISS API, albeit with a bit more flexibility. We encourage you to check out the [API docstring](https://github.com/AnswerDotAI/fastkmeans/blob/17c2a1b4cabc84c7bb0cb392fd2d2e3ec4d1b825/fastkmeans/kmeans.py#L132) to see what the arguments are, as they're very straightforward. The default behaviour of the library mostly follows faiss's, including downsampling data to a maximum of 256 points per centroid to speed up calculations, which can be freely modified and/or disabled.  The only major difference is that, by default, the library does adopt an early stopping mechanism based on a `tol` parameter, which stops the algorithm when the centroids don't move more than `tol` between iterations. This is unlike faiss', whose default behaviour is to run for a fixed number of iterations no matter what -- you can restore this behaviour by setting `tol` to -1.
+
+#### Chunking
+
+The algorith is implemented with a double-chunking logics, where both the data points and the centroids are split into moderately-sized chunks, avoiding the risks of OOMs. The defaults allow you to cluster 26_214_400 128-dimensional points into 262_144 clusters with ~11GB memory usage (including storing the data in fp32). You can check out the available arguments [here](https://github.com/AnswerDotAI/fastkmeans/blob/17c2a1b4cabc84c7bb0cb392fd2d2e3ec4d1b825/fastkmeans/kmeans.py#L132) to see how to tweak these. As a rule of thumb, increasing chunk sizes will speed up computations, at the cost of memory usage, and decreasing it will have the reverse effect.
+
+###  Why `fastkmeans`?
+
+The main motivation behind `fastkmeans` is having a considerably easier way to package late-interaction models, such as ColBERT, in both its Stanford implementation, its PyLate implentation, and for the RAGatouille high-level API. The use of clustering for ColBERT is perhaps somewhat peculiar, as it relies on **large numbers of clusters** for relatively few data points (~100 per cluster centre). This has been a major problem in getting ColBERT to be more usable, as the existing alternatives, while great in their own merit, have flaws for this particular use:
+
+- `faiss` is highly-optimized and is the original library used by the ColBERT authors and most implementations nowadays. However, it is rather difficult to install as there are no "perfect" PyPi wheels, with many segfault issues reported, as the official install is only supported via conda or from source. It can also be finnicky, causing issues with PyTorch if not installed via conda too. Finally, and this is a major problem for maintaining libraries such as ragatouille: it requires different packages and different install methods for its `cpu` and `gpu` variant, meaning additional friction for users and the inability to provide a nice default.
+- `fast-pytorch-kmeans` is a great library which provides lightning fast kmeans implementation in PyTorch: in fact, it's faster than this library! However, it relies on highly vectorized operations which are exceedingly memory hungry, and consistently OOMs on consumer hardware when trying to index even a moderate number of colbert documents (or produces suboptimal clusters with minibatching).
+- `scikit-learn`, while being the ML giant whose shoulders we all stand on, only supports CPU. This becomes unusably slow when indexing larger volumes of documents, especially as there'll (almost) always be a GPU available in these situations.
+
+There are some libraries (such as NVidia's own implementations), but they again require more dependencies than we'd like for nimble packaging, and/or are less flexible in terms of hardware.
+
+###  Limitations
+
+- On a few toy datasets & MNIST, `fastkmeans` reaches roughly the same NMI as `faiss` and `scikit-learn`, indicating that it creates at least somewhat coherent clusters. However, it's not extensively tested, especially in non-ColBERT uses, so your mileage may vary.
+- The "chunking" defaults to avoid OOMs is rather simplistic, and you might need to tweak the numbers depending on your hardware and dataset size.
+- The library currently assumes you'll be using it either on a CPU or a single GPU. Multiple GPUs don't currently provide a major speed-up, this might change in the future, though we expect users wanting to index 10M+ documents to likely have the more robust `faiss` available on their machine anyway.
+
+### Speed
+
+Below is `fastkmeans` benchmarked against `faiss` on a single RTX 4090 GPU, with 128-dimensional data points at various data scales that will commonly be used in ColBERT-style indexing (8192, 16384, 32768, 65536, 131072 and 262144 clusters, each with w/ cluster_size*100 data points).
+
+![fastkmeans benchmark](./benchmark_plots/4090_benchmark.png)

fastkmeans-0.3.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+fastkmeans/__init__.py,sha256=EKvVznV3rTwe5Yz_tqlF_bJMZsh-5vLSjVvMtCx3Xhk,79
+fastkmeans/kmeans.py,sha256=_FbU_avJPxIEfikhPMPgBibj9ckKej21iWJCm-YNEUM,13778
+fastkmeans/triton_kernels.py,sha256=iN8khhoaQGJ08LQy5iz4VGEXRCtvFKDfCgyGzwVqjgw,3698
+fastkmeans-0.3.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+fastkmeans-0.3.0.dist-info/METADATA,sha256=KOyw3KsNPCdF_6spJl2SByB8mVtZ3I6GE1H8N2KMUpM,6791
+fastkmeans-0.3.0.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+fastkmeans-0.3.0.dist-info/top_level.txt,sha256=B3Zd2-kEAH_hN0hFUWgo5lO-TH7ppVol_WQ5ZT1H0js,11
+fastkmeans-0.3.0.dist-info/RECORD,,

{fastkmeans-0.1.0.dist-info → fastkmeans-0.3.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (78.1.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

fastkmeans-0.1.0.dist-info/METADATA

@@ -1,13 +0,0 @@
-Metadata-Version: 2.4
-Name: fastkmeans
-Version: 0.1.0
-Summary: Add your description here
-Author-email: Ben Clavié <bc@answer.ai>
-Maintainer-email: Ben Clavié <bc@answer.ai>
-License: Apache-2.0
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: torch
-Requires-Dist: numpy
-Dynamic: license-file

fastkmeans-0.1.0.dist-info/RECORD

@@ -1,7 +0,0 @@
-fastkmeans/__init__.py,sha256=eIDvwiAarlwyo4n6AvHsV3cPqkbLuILUE1gFAjpzItY,79
-fastkmeans/kmeans.py,sha256=HO44oEOyI8esRDpUJYIq6tHXXVgysSRazoMSIJ2HkX8,12016
-fastkmeans-0.1.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-fastkmeans-0.1.0.dist-info/METADATA,sha256=Fy5bypSuzqwdPs-Qw2fTJ8RZn2b9AKj9POyvv4U7gkk,344
-fastkmeans-0.1.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-fastkmeans-0.1.0.dist-info/top_level.txt,sha256=B3Zd2-kEAH_hN0hFUWgo5lO-TH7ppVol_WQ5ZT1H0js,11
-fastkmeans-0.1.0.dist-info/RECORD,,