fastkmeans 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

fastkmeans/__init__.py

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

fastkmeans/kmeans.py

@@ -3,26 +3,33 @@ import time
 import torch
 import numpy as np
 
-from fastkmeans.triton_kernels import chunked_kmeans_kernel
+try:
+    from fastkmeans.triton_kernels import triton_kmeans
+
+    HAS_TRITON = True
+except ImportError:
+    triton_kmeans = None
+    HAS_TRITON = False
+
 
 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
+    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')
+    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':
+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'):
+    elif device.type == "xpu" and hasattr(torch.xpu, "is_bf16_supported"):
         return torch.xpu.is_bf16_supported()
     else:
         return False
@@ -53,8 +60,12 @@ def _kmeans_torch_double_chunked(
         Where n_samples_used can be smaller than the original if subsampling occurred.
     """
 
+    if use_triton:
+        if not HAS_TRITON:
+            raise ImportError("Triton is not available. Please install Triton and try again.")
+
     if dtype is None:
-        dtype = torch.float16 if device.type in ['cuda', 'xpu'] else torch.float32
+        dtype = torch.float16 if device.type in ["cuda", "xpu"] else torch.float32
 
     n_samples_original, n_features = data.shape
     n_samples = n_samples_original
@@ -76,12 +87,12 @@ def _kmeans_torch_double_chunked(
     centroids = data[rand_indices].clone().to(device=device, dtype=dtype)
     prev_centroids = centroids.clone()
 
-    labels = torch.empty(n_samples, dtype=torch.int64, 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)
+        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)
 
@@ -95,7 +106,7 @@ def _kmeans_torch_double_chunked(
             best_ids = torch.zeros((batch_size,), device=device, dtype=torch.int64)
 
             if use_triton:
-                chunked_kmeans_kernel(
+                triton_kmeans(
                     data_chunk=data_chunk,
                     data_chunk_norms=data_chunk_norms,
                     centroids=centroids,
@@ -103,7 +114,7 @@ def _kmeans_torch_double_chunked(
                     best_ids=best_ids,
                 )
             else:
-                best_dist = torch.full((batch_size,), float('inf'), device=device, dtype=dtype)
+                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)
@@ -116,23 +127,23 @@ def _kmeans_torch_double_chunked(
                     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])
+                    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, device=device, dtype=torch.float32))
 
-            labels[start_idx:end_idx] = best_ids.to('cpu', non_blocking=True)
+            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=dtype)
-        non_empty = (cluster_counts > 0)
+        non_empty = cluster_counts > 0
         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')
+            reinit_indices = torch.randint(0, n_samples, (len(empty_ids),), device="cpu")
             random_data = data[reinit_indices].to(device=device, dtype=dtype, non_blocking=True)
             new_centroids[empty_ids] = random_data
 
@@ -143,14 +154,16 @@ def _kmeans_torch_double_chunked(
 
         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}")
+            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})")
+                print(f"Converged after {iteration + 1} iterations (shift: {shift:.6f} < tol: {tol})")
             break
 
-    centroids_cpu = centroids.to('cpu', dtype=torch.float32)
+    centroids_cpu = centroids.to("cpu", dtype=torch.float32)
     return centroids_cpu, labels
 
 
@@ -176,9 +189,9 @@ class FastKMeans:
     max_points_per_centroid : int, optional, default=1_000_000_000
         If n_samples > k * max_points_per_centroid, the data will be subsampled to exactly
         k * max_points_per_centroid points before clustering.
-    chunk_size_data : int, default=50_000
+    chunk_size_data : int, default=10,2400
         Chunk size along the data dimension for assignment/update steps.
-    chunk_size_centroids : int, default=10_000
+    chunk_size_centroids : int, default=10,240
         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.
@@ -193,14 +206,14 @@ class FastKMeans:
         tol: float = 1e-8,
         gpu: bool = True,
         seed: int = 0,
-        max_points_per_centroid: int = 256,
-        chunk_size_data: int = 50_000,
-        chunk_size_centroids: int = 10_000,
+        max_points_per_centroid: int = 1_000_000_000,
+        chunk_size_data: int = 51_200,
+        chunk_size_centroids: int = 10_240,
         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
+        nredo: int = 1,  # for compatibility only
         use_triton: bool | None = None,
     ):
         self.d = d
@@ -216,8 +229,11 @@ class FastKMeans:
         self.dtype = dtype
         self.pin_gpu_memory = pin_gpu_memory
         self.verbose = verbose
-        if use_triton is not False:
-            use_triton = _is_bfloat16_supported(self.device) # assume triton is supported if GPU supports bfloat16
+        if use_triton is None:
+            # assume triton kernel is supported if GPU supports bfloat16
+            use_triton = HAS_TRITON and _is_bfloat16_supported(self.device)
+        if use_triton and not HAS_TRITON:
+            raise ValueError("Triton is not available. Please install Triton and try again.")
         self.use_triton = use_triton
         if nredo != 1:
             raise ValueError("nredo must be 1, redos not currently supported")
@@ -236,10 +252,10 @@ class FastKMeans:
 
         # Move data to PyTorch CPU Tensor
         data_torch = torch.from_numpy(data)
-        data_norms_torch = (data_torch ** 2).sum(dim=1)
+        data_norms_torch = (data_torch**2).sum(dim=1)
 
         device = _get_device(self.device)
-        if device == 'cuda' and self.pin_gpu_memory:
+        if device == "cuda" and self.pin_gpu_memory:
             data_torch = data_torch.pin_memory()
             data_norms_torch = data_norms_torch.pin_memory()
 
@@ -278,27 +294,29 @@ class FastKMeans:
             raise RuntimeError("Must call train() or fit() before predict().")
 
         data_torch = torch.from_numpy(data)
-        data_norms_torch = (data_torch ** 2).sum(dim=1)
+        data_norms_torch = (data_torch**2).sum(dim=1)
 
         # We'll do a chunked assignment pass, similar to the main loop, but no centroid updates
         centroids_torch = torch.from_numpy(self.centroids)
         centroids_torch = centroids_torch.to(device=self.device, dtype=torch.float32)
-        centroid_norms = (centroids_torch ** 2).sum(dim=1)
+        centroid_norms = (centroids_torch**2).sum(dim=1)
 
         n_samples = data_torch.shape[0]
-        labels = torch.empty(n_samples, dtype=torch.long, device='cpu')
+        labels = torch.empty(n_samples, dtype=torch.long, device="cpu")
 
         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=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)
+            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(
+                triton_kmeans(
                     data_chunk,
                     data_chunk_norms,
                     centroids_torch,
@@ -306,7 +324,7 @@ class FastKMeans:
                     best_ids,
                 )
             else:
-                best_dist = torch.full((batch_size,), float('inf'), device=self.device, dtype=torch.float32)
+                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:
@@ -320,10 +338,10 @@ class FastKMeans:
                     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])
+                    best_ids[improved_mask] = c_start + local_min_ids[improved_mask]
                     c_start = c_end
 
-            labels[start_idx:end_idx] = best_ids.to('cpu')
+            labels[start_idx:end_idx] = best_ids.to("cpu")
             start_idx = end_idx
 
         return labels.numpy()

fastkmeans/triton_kernels.py

@@ -1,105 +1,132 @@
+from contextlib import contextmanager
+
 import torch
 import triton
 import triton.language as tl
 
 
+@contextmanager
+def device_guard(tensor: torch.Tensor):
+    """Context manager to ensure that the Triton kernel launches on the correct device."""
+    if tensor.device.type == "cuda":  # NVIDIA or AMD/ROCm
+        with torch.cuda.device_of(tensor):
+            yield
+    elif tensor.device.type == "xpu":  # Intel GPUs
+        with torch.xpu.device_of(tensor):
+            yield
+    else:  # CPU or other back-ends
+        yield
+
+
 @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,
+        "BLOCK_M": lambda x: 128 if x["D"] <= 384 else 64,
+        "BLOCK_N": lambda x: 128 if x["D"] <= 384 else 64,
+        "BLOCK_K": lambda x: 16 if x["D"] <= 32 or x["D"] > 384 else 32,
+        "GROUP_SIZE_M": lambda x: 8 if x["D"] <= 32 else 16,
+        "num_warps": lambda x: 4,
     }
 )
 @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
+def _kmeans_kernel(
+    x_ptr,
+    x_norm_ptr,
+    c_ptr,
+    c_norm_ptr,
+    best_dist_ptr,
+    best_idx_ptr,
+    B,
+    C,
+    D: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SIZE_M: 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
+    # Map flat CTA id to (pid_m, pid_n) in “grouped” launch order
+    pid = tl.program_id(axis=0)
+
+    num_pid_m = tl.cdiv(B, BLOCK_M)  # row-tiles
+    num_pid_n = tl.cdiv(C, BLOCK_N)  # centroid-tiles
+
+    # Super-group into GROUP_SIZE_M blocks to minimize loading from global memory
+    num_pid_in_grp = GROUP_SIZE_M * num_pid_n
+    first_pid_m = (pid // num_pid_in_grp) * GROUP_SIZE_M
+    group_rows = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+
+    pid_m = first_pid_m + ((pid % num_pid_in_grp) % group_rows)  # row-tile index
+    pid_n = (pid % num_pid_in_grp) // group_rows  # centroid-tile index
+
+    row_start = pid_m * BLOCK_M
+    col_start = pid_n * BLOCK_N
+
     rows = row_start + tl.arange(0, BLOCK_M)
-    mask = rows < B
+    cols = col_start + tl.arange(0, BLOCK_N)
 
-    # 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)
+    row_mask = rows < B
+    col_mask = cols < C
 
-    # shape: [BLOCK_M]
-    x_norm = tl.load(x_norm_ptr + rows, mask=mask, other=0.0)
+    # load norms
+    x_n = tl.load(x_norm_ptr + rows, mask=row_mask, other=0.0)  # [BM]
+    c_n = tl.load(c_norm_ptr + cols, mask=col_mask, other=0.0)  # [BN]
 
-    # Prepare "best distance" + "best index"
-    best_dist = tl.full([BLOCK_M], 1e38, dtype=tl.float32)
-    best_idx = tl.zeros([BLOCK_M], dtype=tl.int64)
+    # pipelined K‑loop, will hold partial dot‑products in registers
+    dot_acc = tl.zeros([BLOCK_M, BLOCK_N], tl.float32)
 
-    # 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
+    # compute matmul tiled across SMs
+    for k0 in range(0, D, BLOCK_K):
+        k_range = k0 + tl.arange(0, BLOCK_K)
 
-        # 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 X slice
+        x_ptrs = x_ptr + rows[:, None] * D + k_range[None, :]
+        xk = tl.load(x_ptrs, mask=row_mask[:, None]).to(tl.float16)
 
-        # Load centroid norms: shape [BLOCK_N]
-        c_sqnorm = tl.load(centroids_sqnorm_ptr + cids, mask=c_mask, other=0.0).to(x.dtype)
+        # load C slice
+        c_ptrs = c_ptr + cols[:, None] * D + k_range[None, :]
+        ck = tl.load(c_ptrs, mask=col_mask[:, None]).to(tl.float16)
 
-        # 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, :])
+        # accumulate
+        dot_acc += tl.dot(xk, tl.trans(ck), out_dtype=tl.float32)
 
-        # Find the argmin along the BLOCK_N dimension
-        local_min_vals, local_min_idx = tl.min(dist_chunk, axis=1, return_indices=True)
+    # finish distance formula
+    dist = tl.fma(dot_acc, -2.0, x_n[:, None] + c_n[None, :])  # [BM, BN]
 
-        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)
+    # local arg‑min (inside this tile)
+    tile_min, tile_idx = tl.min(dist, axis=1, return_indices=True)
 
-    # 4) Write out the best centroid indices
-    tl.store(best_ids_ptr + rows, best_idx, mask=mask)
+    # compete with global best using atomics
+    prev = tl.atomic_min(best_dist_ptr + rows, tile_min, mask=row_mask)
+    improved = tile_min < prev
 
+    # update best_ids
+    tl.store(best_idx_ptr + rows, tl.where(improved, col_start + tile_idx, tl.load(best_idx_ptr + rows)), mask=row_mask)
 
-def chunked_kmeans_kernel(
+
+def triton_kmeans(
     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]
+    best_dist = torch.full((B,), 1e38, device=data_chunk.device, dtype=torch.float32)
 
     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
-    )
+        return (triton.cdiv(B, meta["BLOCK_M"]) * triton.cdiv(C, meta["BLOCK_N"]),)  # 1D grid
+
+    # Without this Triton always tries to launch from device:0 and we get
+    # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
+    with device_guard(data_chunk):
+        _kmeans_kernel[grid](
+            data_chunk,
+            data_chunk_norms,
+            centroids,
+            centroids_sqnorm,
+            best_dist,
+            best_ids,
+            B,
+            C,
+            D,
+        )

{fastkmeans-0.2.0.dist-info → fastkmeans-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fastkmeans
-Version: 0.2.0
+Version: 0.4.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>
@@ -22,7 +22,7 @@ _A fast and efficient k-means implementation for PyTorch, with support for GPU a
 
 ---
 
-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`.
+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 any PyTorch-compatible CPU or GPU, matching or outperforming `faiss` by ~4-5× on a single GPU, and is without install woes, relying on just two dependencies you already have installed: `torch` and `numpy`.
 
 ###  Get started
 
@@ -43,18 +43,22 @@ There's very, very little to this library. It provides a single interface, `Fast
 
 #### 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.
+Whenever possible, the library attempts 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/main/fastkmeans/kmeans.py#L170) to see what the arguments are, as they are 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.
 
+> Note: See the Triton section below for triton specific chunking details.
+
 ###  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).
+- `fast-pytorch-kmeans` is a great library which provides lightning fast kmeans implementation in PyTorch. 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.
@@ -65,8 +69,30 @@ There are some libraries (such as NVidia's own implementations), but they again
 - 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.
 
+### Triton Kernel
+
+`fastkmeans`'s Triton kmeans kernel is ~4-5 times faster than single-GPU `faiss` or `fastkmeans`'s PyTorch backend. On a modern GPU (Ampere or newer), the Triton backend is enabled by default.
+
+While the Triton kernel uses significantly less memory than the PyTorch implementation, increasing the chunk size above 512K can result in slower performance.
+
 ### 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).
+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, and 131072 clusters, each with w/ cluster_size*100 data points).
 
 ![fastkmeans benchmark](./benchmark_plots/4090_benchmark.png)
+
+#### Benchmarking
+
+To benchmark `fastkmeans` against `faiss` on your own machine, install faiss and PyTorch 2.5 via the `bench_env.yaml` Conda environment:
+
+```bash
+conda env create -f bench_env.yaml
+conda activate fastkmeans
+pip install fastkmeans
+```
+
+Then, run the benchmark script:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python speedbench.py --do-faiss --do-fastkmeans --do-fastkmeans-triton --do-evals
+```

fastkmeans-0.4.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+fastkmeans/__init__.py,sha256=NcMnvnLDRqvTrAXfvVJB0M845Fe8imRzI90dAOlc3MY,79
+fastkmeans/kmeans.py,sha256=l8W79LZszDiDSS8H2jFke8Wp_blFQdslCkr4Oo2CKik,14099
+fastkmeans/triton_kernels.py,sha256=vwKUApwKsGWAhN9gMh9rWzandWoEt6rV187xLAW2t3U,4119
+fastkmeans-0.4.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+fastkmeans-0.4.0.dist-info/METADATA,sha256=wKNRR_pxOUxTM4cAH2jCLnt7YgUZYYAc1kRG-4YZ8eE,7552
+fastkmeans-0.4.0.dist-info/WHEEL,sha256=GHB6lJx2juba1wDgXDNlMTyM13ckjBMKf-OnwgKOCtA,91
+fastkmeans-0.4.0.dist-info/top_level.txt,sha256=B3Zd2-kEAH_hN0hFUWgo5lO-TH7ppVol_WQ5ZT1H0js,11
+fastkmeans-0.4.0.dist-info/RECORD,,

{fastkmeans-0.2.0.dist-info → fastkmeans-0.4.0.dist-info}/WHEEL

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

fastkmeans-0.2.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-fastkmeans/__init__.py,sha256=mPLGfhksBBfB1dTkPBPrPg5E1qixixl_Qritc6A10AI,79
-fastkmeans/kmeans.py,sha256=RLDeaCIbKpCJBzV2XO3JPgLwjrhJ6vtA997Gyzm_GyA,13651
-fastkmeans/triton_kernels.py,sha256=iN8khhoaQGJ08LQy5iz4VGEXRCtvFKDfCgyGzwVqjgw,3698
-fastkmeans-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-fastkmeans-0.2.0.dist-info/METADATA,sha256=2hsTZr0t2_CNhCyrECIBk2S4ZB7ytF0YaDqgcgHDvDc,6791
-fastkmeans-0.2.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-fastkmeans-0.2.0.dist-info/top_level.txt,sha256=B3Zd2-kEAH_hN0hFUWgo5lO-TH7ppVol_WQ5ZT1H0js,11
-fastkmeans-0.2.0.dist-info/RECORD,,