From 5fc0fdea9c080fcf6375d8dad9204dcc34c6b159 Mon Sep 17 00:00:00 2001 From: zeyuyang8 Date: Sat, 25 Jul 2026 16:24:44 -0700 Subject: [PATCH 1/3] Flatten assign/update launch grids to remove 65535 batch limit The assign (euclid/cosine, non-split) and shared _centroid_update_chunk_kernel used a 2D launch grid (n_tiles, B) with the batch dimension on grid.y, which CUDA caps at 65535. Any problem with B > 65535 fails to launch with 'CUDA: invalid argument'. Flatten to a 1D grid (B * n_tiles,) and decode program_id(0) as b = flat_id // n_tiles, tile = flat_id % n_tiles. grid.x is capped at 2^31-1, so both large-B and large-N launch. The block->(b,tile) linearization is unchanged, so per-program work and atomic contention are identical. Split-D assign kernels keep their 2D launch (they decode a 2D grid). benchmarks/grid_fix/ adds: - bench_grid.py: before/after throughput on the non-split assign + update paths (run on main and this branch, then --compare). - bench_large_b.py: capability check that B > 65535 launches. --- benchmarks/grid_fix/bench_grid.py | 184 +++++++++++++++++++++++++ benchmarks/grid_fix/bench_large_b.py | 52 +++++++ flash_kmeans/assign_euclid_triton.py | 32 +++-- flash_kmeans/centroid_update_triton.py | 14 +- 4 files changed, 267 insertions(+), 15 deletions(-) create mode 100644 benchmarks/grid_fix/bench_grid.py create mode 100644 benchmarks/grid_fix/bench_large_b.py diff --git a/benchmarks/grid_fix/bench_grid.py b/benchmarks/grid_fix/bench_grid.py new file mode 100644 index 0000000..d50300b --- /dev/null +++ b/benchmarks/grid_fix/bench_grid.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python +"""Regression benchmark for the flattened 1D launch-grid change. + +The grid change only rewrites how `program_id` is decoded and how the grid is +launched for the *non-split* assign kernels (`euclid`/`cosine`) and the shared +`_centroid_update_chunk_kernel`. The math each program does is unchanged, so +this script exists to prove throughput is not regressed on the paths every +current user hits. + +Run it on `main` and on the grid-fix branch, then compare the two JSON files: + + git checkout main + python benchmarks/grid_fix/bench_grid.py --out /tmp/grid_main.json + git checkout + python benchmarks/grid_fix/bench_grid.py --out /tmp/grid_fix.json + python benchmarks/grid_fix/bench_grid.py --compare /tmp/grid_main.json /tmp/grid_fix.json + +All shapes use D <= 512 so they stay on the non-split kernels (the ones the fix +touches). The split-D kernels keep their 2D launch and are out of scope here. +""" +import argparse +import json +import subprocess + +import torch + +from flash_kmeans.assign_euclid_triton import euclid_assign_triton, cosine_assign_triton +from flash_kmeans.centroid_update_triton import ( + triton_centroid_update_sorted_euclid, + triton_centroid_update_sorted_cosine, +) + +# (B, N, K, D) — a spread of batch/point/cluster/dim sizes on the non-split path. +SHAPES = [ + (1, 1_000_000, 256, 128), # single big batch, large N + (8, 131_072, 256, 128), # medium batch + (64, 16_384, 256, 128), # many batches + (1, 4_000_000, 1024, 64), # very large N, K=1024 + (256, 8_192, 128, 128), # large B (2D grid.y would be 256 here — still legal) + (4, 262_144, 512, 256), # larger D +] + +DTYPE = torch.float16 +WARMUP = 5 +ITERS = 20 + + +def _time_ms(fn, warmup=WARMUP, iters=ITERS): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def _mpts_per_s(B, N, ms): + """Millions of points processed per second (B*N points per call).""" + return (B * N) / (ms * 1e-3) / 1e6 + + +def bench_case(B, N, K, D): + dev = "cuda" + x = torch.randn(B, N, D, device=dev, dtype=DTYPE) + centroids = torch.randn(B, K, D, device=dev, dtype=DTYPE) + x_sq = (x.float() ** 2).sum(-1) + cluster_ids = torch.randint(0, K, (B, N), device=dev, dtype=torch.int64) + + # normalized copy for cosine centroid update + x_norm = torch.nn.functional.normalize(x.float(), dim=-1).to(DTYPE) + + results = {} + + results["euclid_assign"] = _time_ms( + lambda: euclid_assign_triton(x, centroids, x_sq) + ) + results["cosine_assign"] = _time_ms( + lambda: cosine_assign_triton(x, centroids) + ) + results["euclid_update"] = _time_ms( + lambda: triton_centroid_update_sorted_euclid(x, cluster_ids, centroids) + ) + results["cosine_update"] = _time_ms( + lambda: triton_centroid_update_sorted_cosine(x_norm, cluster_ids, centroids) + ) + + del x, centroids, x_sq, cluster_ids, x_norm + torch.cuda.empty_cache() + return results + + +def run(out_path): + commit = subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"] + ).decode().strip() + branch = subprocess.check_output( + ["git", "rev-parse", "--abbrev-ref", "HEAD"] + ).decode().strip() + + gpu = torch.cuda.get_device_name(0) + print(f"GPU: {gpu} | branch: {branch} | commit: {commit}\n") + + all_results = {"gpu": gpu, "branch": branch, "commit": commit, "cases": {}} + + header = f"{'shape (B,N,K,D)':>26} | {'euclid_assign':>14} | {'cosine_assign':>14} | {'euclid_update':>14} | {'cosine_update':>14}" + print(header) + print("-" * len(header)) + for shape in SHAPES: + r = bench_case(*shape) + all_results["cases"][str(shape)] = r + B, N = shape[0], shape[1] + print( + f"{str(shape):>26} | " + f"{r['euclid_assign']:>7.3f} ms {_mpts_per_s(B,N,r['euclid_assign']):>4.0f}M/s | " + f"{r['cosine_assign']:>7.3f} ms {_mpts_per_s(B,N,r['cosine_assign']):>4.0f}M/s | " + f"{r['euclid_update']:>7.3f} ms {_mpts_per_s(B,N,r['euclid_update']):>4.0f}M/s | " + f"{r['cosine_update']:>7.3f} ms {_mpts_per_s(B,N,r['cosine_update']):>4.0f}M/s" + ) + + if out_path: + with open(out_path, "w") as f: + json.dump(all_results, f, indent=2) + print(f"\nSaved -> {out_path}") + + +def compare(before_path, after_path): + with open(before_path) as f: + before = json.load(f) + with open(after_path) as f: + after = json.load(f) + + print(f"BEFORE: {before['branch']} @ {before['commit']}") + print(f"AFTER : {after['branch']} @ {after['commit']}") + print(f"GPU : {after['gpu']}\n") + print("Numbers are AFTER/BEFORE latency ratio (<1.0 = faster, >1.0 = slower).\n") + + kernels = ["euclid_assign", "cosine_assign", "euclid_update", "cosine_update"] + header = f"{'shape (B,N,K,D)':>26} | " + " | ".join(f"{k:>14}" for k in kernels) + print(header) + print("-" * len(header)) + + worst = 0.0 + for shape, b in before["cases"].items(): + a = after["cases"].get(shape) + if a is None: + continue + ratios = [] + cells = [] + for k in kernels: + ratio = a[k] / b[k] + ratios.append(ratio) + worst = max(worst, ratio) + cells.append(f"{ratio:>13.3f}x") + print(f"{shape:>26} | " + " | ".join(cells)) + + print(f"\nWorst-case slowdown ratio across all cases/kernels: {worst:.3f}x") + if worst <= 1.05: + print("=> No meaningful regression (within 5% noise band).") + else: + print("=> Potential regression >5%; inspect the case above.") + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--out", default=None, help="write results JSON to this path") + p.add_argument( + "--compare", nargs=2, metavar=("BEFORE", "AFTER"), + help="compare two result JSON files instead of benchmarking", + ) + args = p.parse_args() + + if args.compare: + compare(*args.compare) + else: + run(args.out) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/grid_fix/bench_large_b.py b/benchmarks/grid_fix/bench_large_b.py new file mode 100644 index 0000000..6fdc353 --- /dev/null +++ b/benchmarks/grid_fix/bench_large_b.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python +"""Capability check for the flattened launch grid. + +The old 2D grid put the batch dimension on `grid.y`, which CUDA caps at 65535. +Any problem with B > 65535 therefore fails to launch on `main`. The flattened +1D grid moves everything onto `grid.x` (cap 2^31-1), so it launches. + +Run on both branches: + main -> expect a launch failure (the bug this PR fixes) + grid-fix branch -> expect PASS + + python benchmarks/grid_fix/bench_large_b.py +""" +import torch + +from flash_kmeans.assign_euclid_triton import euclid_assign_triton, cosine_assign_triton +from flash_kmeans.centroid_update_triton import triton_centroid_update_sorted_euclid + +# B > 65535 so the old grid.y cap is exceeded; keep N/D/K tiny to stay small. +B, N, K, D = 70_000, 64, 8, 32 +DTYPE = torch.float16 + + +def main(): + print(f"GPU: {torch.cuda.get_device_name(0)}") + print(f"Case: B={B} (> 65535), N={N}, K={K}, D={D}\n") + dev = "cuda" + x = torch.randn(B, N, D, device=dev, dtype=DTYPE) + centroids = torch.randn(B, K, D, device=dev, dtype=DTYPE) + x_sq = (x.float() ** 2).sum(-1) + cluster_ids = torch.randint(0, K, (B, N), device=dev, dtype=torch.int64) + + ok = True + for name, fn in [ + ("euclid_assign", lambda: euclid_assign_triton(x, centroids, x_sq)), + ("cosine_assign", lambda: cosine_assign_triton(x, centroids)), + ("euclid_update", lambda: triton_centroid_update_sorted_euclid(x, cluster_ids, centroids)), + ]: + try: + out = fn() + torch.cuda.synchronize() + print(f" [PASS] {name} launched, out shape {tuple(out.shape)}") + except Exception as e: + ok = False + msg = str(e).splitlines()[0] if str(e) else type(e).__name__ + print(f" [FAIL] {name}: {msg}") + + print(f"\n=> {'ALL LAUNCHED (grid fix works)' if ok else 'LAUNCH FAILED (expected on main)'}") + + +if __name__ == "__main__": + main() diff --git a/flash_kmeans/assign_euclid_triton.py b/flash_kmeans/assign_euclid_triton.py index 596e45e..02272cf 100644 --- a/flash_kmeans/assign_euclid_triton.py +++ b/flash_kmeans/assign_euclid_triton.py @@ -833,8 +833,12 @@ def _euclid_assign_kernel( maintains the running minimum distance as well as the corresponding index for every point in the tile. """ - pid_n = tl.program_id(0) # tile index along N dimension - pid_b = tl.program_id(1) # batch index + # FIX: 1D flattened grid — no 65535 limit on either B or N. + # grid.x limit is 2^31-1, handles both big-B and big-N cases. + flat_id = tl.program_id(0) + n_tiles = tl.cdiv(N, BLOCK_N) + pid_b = (flat_id // n_tiles) + pid_n = (flat_id % n_tiles) pid_b = pid_b.to(tl.int64) n_start = pid_n * BLOCK_N @@ -1062,8 +1066,12 @@ def _cosine_assign_kernel( maintains the running minimum distance as well as the corresponding index for every point in the tile. """ - pid_n = tl.program_id(0) # tile index along N dimension - pid_b = tl.program_id(1) # batch index + # FIX: 1D flattened grid — no 65535 limit on either B or N. + # grid.x limit is 2^31-1, handles both big-B and big-N cases. + flat_id = tl.program_id(0) + n_tiles = tl.cdiv(N, BLOCK_N) + pid_b = (flat_id // n_tiles) + pid_n = (flat_id % n_tiles) pid_b = pid_b.to(tl.int64) n_start = pid_n * BLOCK_N @@ -1289,7 +1297,11 @@ def euclid_assign_triton( stride_csq_b, stride_csq_k = c_sq.stride() stride_out_b, stride_out_n = out.stride() - grid = lambda META: (triton.cdiv(N, META["BLOCK_N"]), B) + grid = lambda META: (B * triton.cdiv(N, META["BLOCK_N"]),) + # split-D kernels decode a 2D grid (program_id(0)=N-tile, program_id(1)=batch), + # so they keep the 2D launch; the flattened-grid fix only applies to the + # non-split kernels above. + grid_split_d = lambda META: (triton.cdiv(N, META["BLOCK_N"]), B) use_split_d = _need_split_d(D, x.dtype, x.device) @@ -1328,7 +1340,7 @@ def euclid_assign_triton( if use_split_d: if selected_config is None: - _euclid_assign_kernel_split_d_autotuned[grid]( + _euclid_assign_kernel_split_d_autotuned[grid_split_d]( x, centroids, x_sq, c_sq, out, B, N, K, D, stride_x_b, stride_x_n, stride_x_d, @@ -1338,7 +1350,7 @@ def euclid_assign_triton( stride_out_b, stride_out_n, ) else: - _euclid_assign_kernel_split_d[grid]( + _euclid_assign_kernel_split_d[grid_split_d]( x, centroids, x_sq, c_sq, out, B, N, K, D, stride_x_b, stride_x_n, stride_x_d, @@ -1443,10 +1455,12 @@ def cosine_assign_triton(x: torch.Tensor, centroids: torch.Tensor, out: torch.Te stride_c_b, stride_c_k, stride_c_d = centroids.stride() stride_out_b, stride_out_n = out.stride() - grid = lambda META: (triton.cdiv(N, META["BLOCK_N"]), B) + grid = lambda META: (B * triton.cdiv(N, META["BLOCK_N"]),) + # split-D kernel decodes a 2D grid; keep its 2D launch (see euclid path). + grid_split_d = lambda META: (triton.cdiv(N, META["BLOCK_N"]), B) if _need_split_d(D, x.dtype, x.device): - _cosine_assign_kernel_split_d_autotuned[grid]( + _cosine_assign_kernel_split_d_autotuned[grid_split_d]( x, centroids, out, diff --git a/flash_kmeans/centroid_update_triton.py b/flash_kmeans/centroid_update_triton.py index 2611459..a8df1eb 100644 --- a/flash_kmeans/centroid_update_triton.py +++ b/flash_kmeans/centroid_update_triton.py @@ -258,12 +258,14 @@ def _centroid_update_chunk_kernel( next-power-of-two padding waste (the inner loop iterates over the real D and masks only the final partial tile). """ - # program indices – 2-D launch grid: (chunk_id, batch_id) - pid_chunk = tl.program_id(axis=0) - pid_b = tl.program_id(axis=1) + # FIX: 1D flattened grid — no 65535 limit on either B or N. + flat_id = tl.program_id(axis=0) + n_tiles = tl.cdiv(N, BLOCK_N) + pid_b = flat_id // n_tiles + pid_chunk = flat_id % n_tiles b = pid_b.to(tl.int64) - chunk_start = (pid_chunk * BLOCK_N).to(tl.int64) # position of the first token handled by this program + chunk_start = (pid_chunk * BLOCK_N).to(tl.int64) # Nothing to do – out of range if chunk_start >= N: @@ -338,7 +340,7 @@ def triton_centroid_update_sorted_cosine(x_norm: torch.Tensor, cluster_ids: torc centroid_sums = torch.zeros((B, K, D), device=x_norm.device, dtype=torch.float32) centroid_cnts = torch.zeros((B, K), device=x_norm.device, dtype=torch.int32) - grid = (triton.cdiv(N, BLOCK_N), B) + grid = (B * triton.cdiv(N, BLOCK_N),) _centroid_update_chunk_kernel[grid]( x_norm, sorted_idx_int, @@ -409,7 +411,7 @@ def triton_centroid_update_sorted_euclid(x: torch.Tensor, cluster_ids: torch.Ten else: assert centroid_cnts.shape == (B, K) - grid = (triton.cdiv(N, BLOCK_N), B) + grid = (B * triton.cdiv(N, BLOCK_N),) _centroid_update_chunk_kernel[grid]( x, # original features sorted_idx_int, # gather indices From ce41377780631fd982e10f7a53e3e3720584d376 Mon Sep 17 00:00:00 2001 From: Zeyu Yang Date: Sat, 25 Jul 2026 16:29:04 -0700 Subject: [PATCH 2/3] Add weighted k-means, k-means++ init, and weighted centroid-update kernel Rebased onto the flattened-launch-grid fix (now split into its own PR), so this PR contains only the opt-in feature work: - Weighted k-means: new Triton kernel triton_centroid_update_sorted_euclid_weighted + batch_kmeans_Euclid_weighted, exposed via FlashKMeans.fit(data, weights=...). Torch-native fallback supports weights too. - k-means++ init: init={random,scalable-kmeans++,standard-kmeans++}, n_init best-of-N restarts by per-batch inertia. The weighted centroid-update kernel streams D in BLOCK_D tiles with masking (matching _centroid_update_chunk_kernel), so non-power-of-two D (e.g. 80/96/192) works and large D + wide dtype stays within a bounded per-program footprint. It uses a quarter of the shared tile budget because it keeps two fp32 [BLOCK_N, BLOCK_D] tiles live at once (features and features*weights). Tests: tests/test_weighted_centroid.py (pow2/non-pow2/large-D correctness, ones==unweighted, empty-cluster fallback). benchmarks/weighted/bench_weighted.py covers correctness + throughput. --- benchmarks/weighted/bench_weighted.py | 163 +++++++++++++ flash_kmeans/__init__.py | 4 + flash_kmeans/centroid_update_triton.py | 175 ++++++++++++++ flash_kmeans/interface.py | 121 ++++++++-- flash_kmeans/kmeans_triton_impl.py | 142 ++++++++++-- flash_kmeans/torch_fallback.py | 303 ++++++++++++++++++++++++- tests/test_kmeanspp.py | 247 ++++++++++++++++++++ tests/test_weighted_centroid.py | 96 ++++++++ 8 files changed, 1213 insertions(+), 38 deletions(-) create mode 100644 benchmarks/weighted/bench_weighted.py create mode 100644 tests/test_kmeanspp.py create mode 100644 tests/test_weighted_centroid.py diff --git a/benchmarks/weighted/bench_weighted.py b/benchmarks/weighted/bench_weighted.py new file mode 100644 index 0000000..2eaeadd --- /dev/null +++ b/benchmarks/weighted/bench_weighted.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python +"""Correctness + throughput for the new weighted centroid-update kernel. + +Correctness is checked three ways against a pure-torch reference: + 1. random positive weights -> matches torch weighted mean per cluster + 2. all-ones weights -> matches the *unweighted* triton kernel + 3. an empty cluster -> falls back to old_centroids for that k + +Throughput reports the weighted kernel latency next to the unweighted kernel so +the extra cost of carrying per-point weights is visible. + + python benchmarks/weighted/bench_weighted.py +""" +import torch + +from flash_kmeans.centroid_update_triton import ( + triton_centroid_update_sorted_euclid, + triton_centroid_update_sorted_euclid_weighted, +) + +DTYPE = torch.float16 +WARMUP = 5 +ITERS = 20 + + +def torch_weighted_reference(x, cluster_ids, old_centroids, weights): + """Per-cluster weighted mean in fp32; empty clusters keep old_centroids.""" + B, N, D = x.shape + K = old_centroids.shape[1] + xf = x.float() + wf = weights.float() + out = torch.empty((B, K, D), device=x.device, dtype=torch.float32) + for b in range(B): + for k in range(K): + mask = cluster_ids[b] == k + wsum = wf[b][mask].sum() + if wsum <= 0: + out[b, k] = old_centroids[b, k].float() + else: + out[b, k] = (xf[b][mask] * wf[b][mask, None]).sum(0) / wsum + return out.to(x.dtype) + + +def _time_ms(fn): + for _ in range(WARMUP): + fn() + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(ITERS): + fn() + e.record() + torch.cuda.synchronize() + return s.elapsed_time(e) / ITERS + + +def check(name, ok, extra=""): + tag = "PASS" if ok else "FAIL" + print(f" [{tag}] {name} {extra}") + return ok + + +def _correctness_for_shape(B, N, K, D, dtype, dev): + """Run the three correctness checks for one (shape, dtype) and return ok.""" + torch.manual_seed(0) + all_ok = True + tag = f"D={D:<4} {str(dtype).split('.')[-1]:>7}" + + # 1) random positive weights vs torch reference + x = torch.randn(B, N, D, device=dev, dtype=dtype) + cluster_ids = torch.randint(0, K, (B, N), device=dev, dtype=torch.int64) + old_c = torch.randn(B, K, D, device=dev, dtype=dtype) + weights = torch.rand(B, N, device=dev) + 0.1 # positive + + got = triton_centroid_update_sorted_euclid_weighted(x, cluster_ids, old_c, weights) + ref = torch_weighted_reference(x, cluster_ids, old_c, weights) + # tolerance scaled by dtype (fp16 centroids quantize coarsely) + tol = 5e-2 if dtype == torch.float16 else 1e-4 + max_abs = (got.float() - ref.float()).abs().max().item() + all_ok &= check(f"{tag} | random weights vs torch reference", + max_abs < tol, f"(max abs diff = {max_abs:.2e})") + + # 2) all-ones weights == unweighted kernel + ones = torch.ones(B, N, device=dev) + got_w = triton_centroid_update_sorted_euclid_weighted(x, cluster_ids, old_c, ones) + got_u = triton_centroid_update_sorted_euclid(x, cluster_ids, old_c) + max_abs2 = (got_w.float() - got_u.float()).abs().max().item() + all_ok &= check(f"{tag} | all-ones weights == unweighted kernel", + max_abs2 < tol, f"(max abs diff = {max_abs2:.2e})") + + # 3) empty cluster falls back to old_centroids + cluster_ids2 = cluster_ids.clone() + cluster_ids2[cluster_ids2 == 0] = 1 # cluster 0 now empty in every batch + got_e = triton_centroid_update_sorted_euclid_weighted(x, cluster_ids2, old_c, weights) + fell_back = torch.allclose(got_e[:, 0].float(), old_c[:, 0].float(), atol=1e-3) + all_ok &= check(f"{tag} | empty cluster -> old_centroids fallback", fell_back) + return all_ok + + +def correctness(): + print("Correctness (D includes non-power-of-two + large-D/fp32):") + dev = "cuda" + # (B, N, K, D, dtype) — cover pow2, non-pow2, and large-D/fp32 (issue #19). + cases = [ + (4, 20_000, 64, 128, torch.float16), # pow2 baseline + (4, 20_000, 64, 64, torch.float16), # pow2 small + (4, 20_000, 64, 96, torch.float16), # non-pow2 (would crash pre-fix) + (4, 20_000, 64, 192, torch.float16), # non-pow2 + (2, 20_000, 64, 80, torch.float16), # non-pow2 head dim + (2, 20_000, 64, 1024, torch.float32), # large D + fp32 (untiled would spill) + ] + all_ok = True + for B, N, K, D, dtype in cases: + all_ok &= _correctness_for_shape(B, N, K, D, dtype, dev) + print(f"\n => {'ALL PASSED' if all_ok else 'FAILURES PRESENT'}\n") + return all_ok + + +def throughput(): + print("Throughput (weighted vs unweighted centroid update):") + dev = "cuda" + shapes = [ + (1, 1_000_000, 256, 128), + (8, 131_072, 256, 128), + (4, 262_144, 512, 256), + (4, 262_144, 512, 96), # non-pow2 + (2, 262_144, 512, 192), # non-pow2 + ] + header = f"{'shape (B,N,K,D)':>26} | {'unweighted':>18} | {'weighted':>18} | {'overhead':>9}" + print(header) + print("-" * len(header)) + for B, N, K, D in shapes: + x = torch.randn(B, N, D, device=dev, dtype=DTYPE) + cluster_ids = torch.randint(0, K, (B, N), device=dev, dtype=torch.int64) + old_c = torch.randn(B, K, D, device=dev, dtype=DTYPE) + weights = torch.rand(B, N, device=dev) + 0.1 + + t_u = _time_ms(lambda: triton_centroid_update_sorted_euclid(x, cluster_ids, old_c)) + t_w = _time_ms(lambda: triton_centroid_update_sorted_euclid_weighted(x, cluster_ids, old_c, weights)) + mpts_u = (B * N) / (t_u * 1e-3) / 1e6 + mpts_w = (B * N) / (t_w * 1e-3) / 1e6 + print( + f"{str((B,N,K,D)):>26} | " + f"{t_u:>7.3f} ms {mpts_u:>5.0f}M/s | " + f"{t_w:>7.3f} ms {mpts_w:>5.0f}M/s | " + f"{(t_w/t_u - 1)*100:>7.1f}%" + ) + del x, cluster_ids, old_c, weights + torch.cuda.empty_cache() + + +def main(): + print(f"GPU: {torch.cuda.get_device_name(0)}\n") + ok = correctness() + print() + throughput() + if not ok: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/flash_kmeans/__init__.py b/flash_kmeans/__init__.py index bccea1f..9a76dad 100644 --- a/flash_kmeans/__init__.py +++ b/flash_kmeans/__init__.py @@ -3,12 +3,14 @@ try: from .kmeans_triton_impl import ( batch_kmeans_Euclid, + batch_kmeans_Euclid_weighted, batch_kmeans_Cosine, batch_kmeans_Dot, ) from .centroid_update_triton import ( triton_centroid_update_euclid, triton_centroid_update_sorted_euclid, + triton_centroid_update_sorted_euclid_weighted, ) from .kmeans_large import kmeans_largeN, kmeans_largeN_assign except Exception: @@ -36,10 +38,12 @@ def no_torch_fallback(): __all__ = [ "batch_kmeans_Euclid", + "batch_kmeans_Euclid_weighted", "batch_kmeans_Cosine", "batch_kmeans_Dot", "triton_centroid_update_euclid", "triton_centroid_update_sorted_euclid", + "triton_centroid_update_sorted_euclid_weighted", "FlashKMeans", "kmeans_largeN", "kmeans_largeN_assign", diff --git a/flash_kmeans/centroid_update_triton.py b/flash_kmeans/centroid_update_triton.py index a8df1eb..16efe3c 100644 --- a/flash_kmeans/centroid_update_triton.py +++ b/flash_kmeans/centroid_update_triton.py @@ -440,6 +440,181 @@ def triton_centroid_update_sorted_euclid(x: torch.Tensor, cluster_ids: torch.Ten # ------------------------------ END new implementation ------------------------------ +@triton.jit +def _centroid_update_chunk_weighted_kernel( + x_ptr, # *f16 / *f32 [B, N, D] – ORIGINAL ORDER + sorted_idx_ptr, # *i32 [B, N] + sorted_cluster_ptr, # *i32 [B, N] + weight_ptr, # *f32 [B, N] – per-point weights in ORIGINAL order + sum_ptr, # *f32 [B, K, D] + weight_sum_ptr, # *f32 [B, K] + # strides + stride_x_b, stride_x_n, stride_x_d, + stride_idx_b, stride_idx_n, stride_cluster_b, stride_cluster_n, + stride_w_b, stride_w_n, + stride_sum_b, stride_sum_k, stride_sum_d, + stride_ws_b, stride_ws_k, + B: tl.constexpr, + N: tl.constexpr, + D: tl.constexpr, + K: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Weighted variant of _centroid_update_chunk_kernel. + + Multiplies features by per-point weight inside the kernel to avoid + materialising a full weighted-data copy on the host side. + Accumulates float32 weight sums instead of int counts. + + The feature dimension is streamed in ``BLOCK_D`` chunks (split-D, matching + ``_centroid_update_chunk_kernel``), so the per-program ``[BLOCK_N, BLOCK_D]`` + tile is bounded regardless of D or dtype, and non-power-of-two D is handled + by masking the final partial tile rather than requiring a power-of-two + ``tl.arange(0, D)``. + """ + # FIX: 1D flattened grid — no 65535 limit on either B or N. + flat_id = tl.program_id(axis=0) + n_tiles = tl.cdiv(N, BLOCK_N) + pid_b = flat_id // n_tiles + pid_chunk = flat_id % n_tiles + + b = pid_b.to(tl.int64) + chunk_start = (pid_chunk * BLOCK_N).to(tl.int64) + + if chunk_start >= N: + return + + idx_batch_base = sorted_idx_ptr + b * stride_idx_b + cid_batch_base = sorted_cluster_ptr + b * stride_cluster_b + x_batch_base = x_ptr + b * stride_x_b + w_batch_base = weight_ptr + b * stride_w_b + + offs_token = tl.arange(0, BLOCK_N).to(tl.int64) + # Hoisted feature-lane base; the (constexpr) d_start offset is added per + # D-tile below (see _centroid_update_chunk_kernel). + base_dim = tl.arange(0, BLOCK_D).to(tl.int64) + + token_idx = chunk_start + offs_token + valid_tok = token_idx < N + first_token_idx = chunk_start + last_token_idx = tl.minimum(chunk_start + BLOCK_N, N) - 1 + + first_id = tl.load(cid_batch_base + first_token_idx) + last_id = tl.load(cid_batch_base + last_token_idx) + all_ids = tl.load(cid_batch_base + token_idx * stride_cluster_n, + mask=valid_tok, other=-1) + + all_tokens_idxs = tl.load(idx_batch_base + token_idx * stride_idx_n, + mask=valid_tok, other=-1) + all_tokens_idxs = all_tokens_idxs.to(tl.int64) + + # Load per-point weights (original order) once for the whole chunk + all_weights = tl.load(w_batch_base + all_tokens_idxs * stride_w_n, + mask=valid_tok, other=0.0) + all_weights = all_weights.to(tl.float32) + + for cid in range(first_id, last_id + 1): + cluster_mask = all_ids == cid + cluster_size = tl.sum(cluster_mask.to(tl.int32)) + if cluster_size != 0: + # Per-token weights for this cluster (D-independent; computed once). + cluster_weights = tl.where(cluster_mask, all_weights, 0.0) + for d_start in range(0, D, BLOCK_D): + offs_dim = d_start + base_dim + d_mask = offs_dim < D + row_ptrs = (x_batch_base + + all_tokens_idxs[:, None] * stride_x_n + + offs_dim[None, :] * stride_x_d) + cluster_feats = tl.load( + row_ptrs, + mask=cluster_mask[:, None] & d_mask[None, :], + other=0.0, + ) # [BLOCK_N, BLOCK_D] + cluster_feats = cluster_feats.to(tl.float32) + + weighted_feats = cluster_feats * cluster_weights[:, None] + sum_feats = tl.sum(weighted_feats, axis=0) # [BLOCK_D] + + dest_ptr = (sum_ptr + b * stride_sum_b + + cid * stride_sum_k + offs_dim * stride_sum_d) + tl.atomic_add(dest_ptr, sum_feats, mask=d_mask) + + w_sum = tl.sum(cluster_weights) + tl.atomic_add(weight_sum_ptr + b * stride_ws_b + + cid * stride_ws_k, w_sum) + + +def triton_centroid_update_sorted_euclid_weighted( + x: torch.Tensor, + cluster_ids: torch.Tensor, + old_centroids: torch.Tensor, + weights: torch.Tensor, + *, + BLOCK_N: int = 256, +): + """Weighted centroid update using a dedicated Triton kernel. + + Avoids materialising a full weighted-data copy by multiplying features + by per-point weights inside the kernel. Also eliminates the + fp16->fp32->fp16->fp32 precision round-trip of the old approach. + + Parameters + ---------- + x : Tensor [B, N, D] + Input feature vectors. + cluster_ids : LongTensor [B, N] + Cluster assignment for each point. + old_centroids : Tensor [B, K, D] + Previous centroids (used to fill empty clusters). + weights : Tensor [B, N] + Per-sample weights (positive). + """ + assert x.is_cuda and cluster_ids.is_cuda and weights.is_cuda + B, N, D = x.shape + K = old_centroids.shape[1] + + sorted_cluster_ids, sorted_idx = torch.sort(cluster_ids, dim=-1) + sorted_idx_int = sorted_idx.to(torch.int32) + + centroid_sums = torch.zeros((B, K, D), device=x.device, dtype=torch.float32) + weight_sums = torch.zeros((B, K), device=x.device, dtype=torch.float32) + weights_f32 = weights.float() + + grid = (B * triton.cdiv(N, BLOCK_N),) + _centroid_update_chunk_weighted_kernel[grid]( + x, + sorted_idx_int, + sorted_cluster_ids.to(torch.int32), + weights_f32, + centroid_sums, + weight_sums, + x.stride(0), x.stride(1), x.stride(2), + sorted_idx_int.stride(0), sorted_idx_int.stride(1), + sorted_cluster_ids.stride(0), sorted_cluster_ids.stride(1), + weights_f32.stride(0), weights_f32.stride(1), + centroid_sums.stride(0), centroid_sums.stride(1), centroid_sums.stride(2), + weight_sums.stride(0), weight_sums.stride(1), + B, N, D, K, + BLOCK_N=BLOCK_N, + # The weighted kernel keeps two fp32 [BLOCK_N, BLOCK_D] tiles live at once + # (cluster_feats and cluster_feats * weights), vs one in the unweighted + # kernel, so it needs a tighter per-tile budget to avoid spilling. A + # quarter of the shared budget caps BLOCK_D at 128 for D=256 fp16, which + # matches the unweighted kernel's throughput (a full-D tile is ~2x slower). + BLOCK_D=_choose_block_d( + D, BLOCK_N, _dtype_bytes(x.dtype), + budget_bytes=_CHUNK_TILE_BUDGET_BYTES // 4, + ), + ) + + centroids = centroid_sums / weight_sums.unsqueeze(-1).clamp(min=1e-8) + empty_mask = (weight_sums == 0).unsqueeze(-1) + centroids = torch.where(empty_mask, old_centroids.float(), centroids) + + return centroids.to(x.dtype) + + def main(): torch.manual_seed(0) diff --git a/flash_kmeans/interface.py b/flash_kmeans/interface.py index cef0cd8..3df3e6d 100644 --- a/flash_kmeans/interface.py +++ b/flash_kmeans/interface.py @@ -3,6 +3,12 @@ from typing import Optional from flash_kmeans.torch_fallback import euclid_assign_torch_native_chunked, batch_kmeans_Euclid_torch_native + +try: + from flash_kmeans.kmeans_triton_impl import batch_kmeans_Euclid_weighted as _batch_kmeans_Euclid_weighted + _HAS_WEIGHTED = True +except Exception: + _HAS_WEIGHTED = False import torch try: @@ -53,6 +59,15 @@ class FlashKMeans: the chunk size of n_samples when copying data from CPU to GPU in chunks. verbose : bool, default=False Whether to print per-iteration info. + init : str, default="random" + Centroid initialization method. + - "random": uniform random selection from data points. + - "scalable-kmeans++": scalable kmeans++ / K-Means|| (Bahmani et al., 2012), matches cuml. + - "standard-kmeans++": standard greedy kmeans++ (Arthur & Vassilvitskii, 2007), matches scikit-learn. + n_init : int | str, default="auto" + Number of times k-means is run with different initializations. The + result with the lowest inertia is kept. "auto" = 10 for random init, + 1 for kmeans++ (matching sklearn). dtype : torch.dtype, optional Compute Data type for algorithm. device : torch.device | None @@ -72,6 +87,8 @@ def __init__( chunk_size_centroids: int = 1024, chunk_size_data_cpu: int = 1048576, verbose: bool = False, + init: str = "random", + n_init: int | str = "auto", dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, ): @@ -85,8 +102,14 @@ def __init__( self.chunk_size_centroids = int(chunk_size_centroids) self.chunk_size_data_cpu = int(chunk_size_data_cpu) self.verbose = bool(verbose) + self.init = init self.dtype = dtype + if n_init == "auto": + self.n_init = 10 if init == "random" else 1 + else: + self.n_init = int(n_init) + if self.use_triton: try: _require_triton_cuda() @@ -103,7 +126,7 @@ def __init__( self.device = device - def train(self, data: torch.Tensor): + def train(self, data: torch.Tensor, weights: torch.Tensor = None): """ Fit KMeans on data and store centroids. @@ -117,6 +140,10 @@ def train(self, data: torch.Tensor): if data is from GPU, it will process directly on GPU. if data is from CPU, it will copy & process data on GPU by chunk_size_data_cpu. + weights : torch.Tensor, optional + Per-sample weights for weighted k-means. + Shape: (n_samples,) or (batch_size, n_samples) + """ if data.ndim == 2: @@ -129,15 +156,79 @@ def train(self, data: torch.Tensor): else: raise ValueError("data must be of shape (n_samples, n_features) or (batch_size, n_samples, n_features)") - # Set random seed - torch.manual_seed(self.seed) - torch.cuda.manual_seed_all(self.seed) + # Normalize weights shape + if weights is not None: + if weights.ndim == 1: + weights_b = weights.unsqueeze(0) + else: + weights_b = weights + weights_b = weights_b.to(device=self.device, dtype=torch.float32, copy=False) + else: + weights_b = None + + best_inertia = None # (B_int,) per-batch inertia + best_centroids_b = None + best_cluster_ids_b = None + B_int = x_b.shape[0] - if data.device.type == "cpu" and N > self.chunk_size_data_cpu: - # handle for large N on CPU + for run_i in range(self.n_init): + torch.manual_seed(self.seed + run_i) + torch.cuda.manual_seed_all(self.seed + run_i) + + cluster_ids_b, centroids_b = self._run_kmeans_once( + x_b, N, B, data, weights_b, + ) + + if self.n_init > 1: + # Per-batch inertia so we can pick the best run independently + # for each batch element, not just the best run overall. + D = x_b.shape[-1] + x_eval = x_b.to(device=self.device, dtype=centroids_b.dtype, copy=False) + assigned = centroids_b.gather( + 1, cluster_ids_b.unsqueeze(-1).expand(-1, -1, D) + ) + inertia = ((x_eval - assigned) ** 2).sum(dim=(-1, -2)) # (B_int,) + + if best_inertia is None: + best_inertia = inertia + best_centroids_b = centroids_b + best_cluster_ids_b = cluster_ids_b + else: + improved = inertia < best_inertia # (B_int,) + best_inertia = torch.where(improved, inertia, best_inertia) + # Update centroids and assignments only for improved batches + mask_c = improved[:, None, None].expand_as(centroids_b) + mask_id = improved[:, None].expand_as(cluster_ids_b) + best_centroids_b = torch.where(mask_c, centroids_b, best_centroids_b) + best_cluster_ids_b = torch.where(mask_id, cluster_ids_b, best_cluster_ids_b) + else: + best_centroids_b = centroids_b + best_cluster_ids_b = cluster_ids_b + + self.centroids_b = best_centroids_b + self.cluster_ids_b = best_cluster_ids_b + self._batch_size = B + + def _run_kmeans_once(self, x_b, N, B, data, weights_b): + """Run a single k-means pass and return (cluster_ids_b, centroids_b).""" + if weights_b is not None: + assert _HAS_WEIGHTED, "Weighted k-means requires Triton implementation" + compute_dtype = self.dtype or x_b.dtype + x_b = x_b.to(device=self.device, dtype=compute_dtype, copy=False) + cluster_ids_b, centroids_b, _ = _batch_kmeans_Euclid_weighted( + x_b, + self.k, + weights_b, + max_iters=self.niter, + tol=self.tol, + init_centroids=None, + verbose=self.verbose, + init=self.init, + ) + elif data.device.type == "cpu" and N > self.chunk_size_data_cpu: assert B is None, "Batched data with large N on CPU is not supported yet." - assert self.use_triton, "process large N data requires triton implementation." - cluster_ids_b, centroids_b = kmeans_largeN( + assert self.use_triton, "process large N data requires triton implementation." + cluster_ids_b, centroids_b = kmeans_largeN( x_b[0], self.k, max_iters=self.niter, @@ -150,12 +241,10 @@ def train(self, data: torch.Tensor): centroids_b.unsqueeze_(0) cluster_ids_b.unsqueeze_(0) else: - # Ensure CUDA + dtype compute_dtype = self.dtype or x_b.dtype x_b = x_b.to(device=self.device, dtype=compute_dtype, copy=False) if self.use_triton: - # Run batched Triton KMeans (Euclidean) cluster_ids_b, centroids_b, _ = batch_kmeans_Euclid( x_b, self.k, @@ -163,9 +252,9 @@ def train(self, data: torch.Tensor): tol=self.tol, init_centroids=None, verbose=self.verbose, + init=self.init, ) else: - # Run batched PyTorch KMeans (Euclidean) cluster_ids_b, centroids_b, _ = batch_kmeans_Euclid_torch_native( x_b, self.k, @@ -175,15 +264,13 @@ def train(self, data: torch.Tensor): verbose=self.verbose, chunk_size_N=self.chunk_size_data, chunk_size_K=self.chunk_size_centroids, + init=self.init, ) - - self.centroids_b = centroids_b - self.cluster_ids_b = cluster_ids_b - self._batch_size = B + return cluster_ids_b, centroids_b - def fit(self, data: torch.Tensor): + def fit(self, data: torch.Tensor, weights: torch.Tensor = None): """Alias for train; returns self.""" - self.train(data) + self.train(data, weights=weights) return self def predict(self, data: torch.Tensor) -> torch.LongTensor: diff --git a/flash_kmeans/kmeans_triton_impl.py b/flash_kmeans/kmeans_triton_impl.py index 43c4046..76103f3 100644 --- a/flash_kmeans/kmeans_triton_impl.py +++ b/flash_kmeans/kmeans_triton_impl.py @@ -2,8 +2,15 @@ import torch.nn.functional as F from torch.cuda import nvtx from flash_kmeans.assign_euclid_triton import euclid_assign_triton, cosine_assign_triton -from flash_kmeans.centroid_update_triton import triton_centroid_update_cosine, triton_centroid_update_euclid, triton_centroid_update_sorted_euclid, triton_centroid_update_sorted_cosine +from flash_kmeans.centroid_update_triton import ( + triton_centroid_update_cosine, + triton_centroid_update_euclid, + triton_centroid_update_sorted_euclid, + triton_centroid_update_sorted_cosine, + triton_centroid_update_sorted_euclid_weighted, +) from tqdm import trange +from flash_kmeans.torch_fallback import scalable_kmeans_pp, standard_kmeans_pp # -------------------- Compiled single-iteration kernels -------------------- @@ -36,6 +43,13 @@ def _dot_iter(x, centroids): shift = (centroids_new - centroids).norm(dim=-1).max() return centroids_new, shift, cluster_ids +def _euclid_iter_weighted(x, x_sq, centroids, weights, use_heuristic=True): + cluster_ids = euclid_assign_triton(x, centroids, x_sq, use_heuristic=use_heuristic) + centroids_new = triton_centroid_update_sorted_euclid_weighted(x, cluster_ids, centroids, weights) + shift = (centroids_new - centroids).norm(dim=-1).max() + return centroids_new, shift, cluster_ids + + COMPILE_FLAG = False try: @@ -61,6 +75,7 @@ def batch_kmeans_Euclid( verbose=False, *, use_heuristic=True, + init="random", ): """ Batched KMeans clustering in PyTorch using Euclidean distance. @@ -76,7 +91,15 @@ def batch_kmeans_Euclid( cluster_ids: (B, N) LongTensor, cluster assignment for each point. centroids: (B, n_clusters, D) final cluster centers. """ - B, N, D = x.shape + B, N, D_orig = x.shape + + # Triton tl.dot requires inner dim >= 16; pad with zeros if needed + if D_orig < _MIN_TRITON_D: + pad = _MIN_TRITON_D - D_orig + x = F.pad(x, (0, pad)) # (B, N, D_padded) + if init_centroids is not None: + init_centroids = F.pad(init_centroids, (0, pad)) + D = x.shape[-1] # Pre-compute squared L2 norm of all points (constant during iterations). # Done in chunks to avoid materializing a full (B, N, D) `x ** 2` temp, @@ -87,13 +110,19 @@ def batch_kmeans_Euclid( x_sq[:, i:i + XSQ_CHUNK] = (x[:, i:i + XSQ_CHUNK] ** 2).sum(dim=-1) if init_centroids is None: - # Randomly select initial centers from x - indices = torch.randint(0, N, (B, n_clusters), device=x.device) - centroids = torch.gather( - x, - dim=1, - index=indices[..., None].expand(-1, -1, D) - ) # (B, n_clusters, D) + if init == "scalable-kmeans++": + centroids = scalable_kmeans_pp(x, n_clusters, x_sq) + elif init == "standard-kmeans++": + centroids = standard_kmeans_pp(x, n_clusters, x_sq) + else: + # Randomly select initial centers from x (without replacement, matching sklearn) + uniform = torch.ones(B, N, device=x.device) + indices = torch.multinomial(uniform, n_clusters, replacement=False) + centroids = torch.gather( + x, + dim=1, + index=indices[..., None].expand(-1, -1, D) + ) # (B, n_clusters, D) else: centroids = init_centroids @@ -112,6 +141,91 @@ def batch_kmeans_Euclid( break centroids = centroids_new.clone() + # Strip padding from centroids + if D_orig < _MIN_TRITON_D: + centroids = centroids[..., :D_orig] + + return cluster_ids, centroids, it + 1 + + +_MIN_TRITON_D = 16 + + +def batch_kmeans_Euclid_weighted( + x, + n_clusters, + weights, + max_iters=100, + tol=0.0, + init_centroids=None, + verbose=False, + *, + use_heuristic=True, + init="random", +): + """ + Batched weighted KMeans clustering using Euclidean distance. + + Args: + x: Tensor of shape (B, N, D), batch_size B, N points per batch, D dims. + n_clusters: Number of clusters. + weights: Tensor of shape (B, N), per-sample weights (positive). + max_iters: Max number of iterations. + tol: Relative tolerance for center movement. + init_centroids: Optional initial centroids (B, n_clusters, D). + verbose: Print loss for each iter. + use_heuristic: Use heuristic Triton config (skip autotune). + Returns: + cluster_ids: (B, N) LongTensor, cluster assignment for each point. + centroids: (B, n_clusters, D) final cluster centers. + n_iters: number of iterations performed. + """ + B, N, D_orig = x.shape + + # Triton tl.dot requires inner dim >= 16; pad with zeros if needed + if D_orig < _MIN_TRITON_D: + pad = _MIN_TRITON_D - D_orig + x = F.pad(x, (0, pad)) # (B, N, D_padded) + if init_centroids is not None: + init_centroids = F.pad(init_centroids, (0, pad)) + D = x.shape[-1] + + x_sq = (x ** 2).sum(dim=-1) # (B, N) + + if init_centroids is None: + if init == "scalable-kmeans++": + centroids = scalable_kmeans_pp(x, n_clusters, x_sq, weights=weights) + elif init == "standard-kmeans++": + centroids = standard_kmeans_pp(x, n_clusters, x_sq, weights=weights) + else: + # Weighted random initialization + probs = weights.float() / weights.float().sum(dim=-1, keepdim=True) + indices = torch.multinomial(probs, n_clusters, replacement=False) + centroids = torch.gather( + x, + dim=1, + index=indices[..., None].expand(-1, -1, D) + ) + else: + centroids = init_centroids + + centroids = centroids.view(B, n_clusters, D) + + for it in range(max_iters): + centroids_new, center_shift, cluster_ids = _euclid_iter_weighted( + x, x_sq, centroids, weights, use_heuristic + ) + + if verbose: + print(f"Iter {it}, center shift: {center_shift.item():.6f}") + if center_shift < tol: + break + centroids = centroids_new.clone() + + # Strip padding from centroids + if D_orig < _MIN_TRITON_D: + centroids = centroids[..., :D_orig] + return cluster_ids, centroids, it + 1 @@ -135,8 +249,9 @@ def batch_kmeans_Cosine(x, n_clusters, max_iters=100, tol=0.0, init_centroids=No x_norm = F.normalize(x, p=2, dim=-1) # (B, N, D) if init_centroids is None: - # Randomly select initial centers from x_norm - indices = torch.randint(0, N, (B, n_clusters), device=x.device) + # Randomly select initial centers from x_norm (without replacement, matching sklearn) + uniform = torch.ones(B, N, device=x.device) + indices = torch.multinomial(uniform, n_clusters, replacement=False) centroids = torch.gather( x_norm, dim=1, @@ -170,8 +285,9 @@ def batch_kmeans_Dot(x, n_clusters, max_iters=100, tol=0.0, init_centroids=None, B, N, D = x.shape if init_centroids is None: - # 随机初始化中心 - indices = torch.randint(0, N, (B, n_clusters), device=x.device) + # Randomly select initial centers (without replacement, matching sklearn) + uniform = torch.ones(B, N, device=x.device) + indices = torch.multinomial(uniform, n_clusters, replacement=False) centroids = torch.gather( x, dim=1, diff --git a/flash_kmeans/torch_fallback.py b/flash_kmeans/torch_fallback.py index 5cc39c3..b83e5f3 100644 --- a/flash_kmeans/torch_fallback.py +++ b/flash_kmeans/torch_fallback.py @@ -1,6 +1,287 @@ +import math + import torch import torch.nn.functional as F + +def _kmeanspp_sequential(x, n_clusters, x_sq=None, weights=None): + """ + Standard batched kmeans++ initialization (Arthur & Vassilvitskii, 2007). + + K sequential rounds, each sampling one centroid proportional to min distance². + Used internally by scalable kmeans++ for the final candidate reduction step. + + Args: + x: (B, N, D) input points. + n_clusters: number of clusters K. + x_sq: (B, N) precomputed ||x||^2, optional. + weights: (B, N) per-sample weights, optional. + + Returns: + centroids: (B, K, D) initial centroids. + """ + B, N, D = x.shape + device = x.device + if x_sq is None: + x_sq = (x ** 2).sum(dim=-1) + + centroids = torch.empty((B, n_clusters, D), device=device, dtype=x.dtype) + batch_arange = torch.arange(B, device=device) + w_f = weights.float() if weights is not None else None + + # First centroid + if w_f is not None: + w_probs = w_f / w_f.sum(dim=-1, keepdim=True).clamp_min(1e-30) + first_idx = torch.multinomial(w_probs, 1).squeeze(-1) + else: + first_idx = torch.randint(0, N, (B,), device=device) + centroids[:, 0] = x[batch_arange, first_idx] + + if n_clusters == 1: + return centroids + + c = centroids[:, 0:1, :] + c_sq = (c ** 2).sum(dim=-1) + min_dists = (x_sq - 2 * torch.bmm(x, c.transpose(1, 2)).squeeze(-1) + c_sq).float() + min_dists.clamp_min_(0) + + for k in range(1, n_clusters): + probs = min_dists * w_f if w_f is not None else min_dists + prob_sums = probs.sum(dim=-1, keepdim=True) + # Avoid CPU-GPU sync: use torch.where instead of if zero_rows.any() + probs = torch.where(prob_sums > 0, probs / prob_sums.clamp_min(1e-30), 1.0 / N) + + idx = torch.multinomial(probs, 1).squeeze(-1) + centroids[:, k] = x[batch_arange, idx] + + new_c = centroids[:, k:k+1, :] + new_c_sq = (new_c ** 2).sum(dim=-1) + new_dists = (x_sq - 2 * torch.bmm(x, new_c.transpose(1, 2)).squeeze(-1) + new_c_sq).float() + new_dists.clamp_min_(0) + torch.minimum(min_dists, new_dists, out=min_dists) + + return centroids + + +def _update_min_dists_batched(x, x_sq, new_cands, min_dists, max_bytes=512 * 1024 * 1024): + """Update min_dists in-place with distances to new_cands, chunking along + the candidate dimension to keep peak memory under max_bytes.""" + B, N, _ = x.shape + n_cands = new_cands.shape[1] + # Chunk size so (B, N, chunk) float32 tensor fits in max_bytes + chunk_l = max(1, max_bytes // (B * N * 4)) + + for l_start in range(0, n_cands, chunk_l): + l_end = min(l_start + chunk_l, n_cands) + chunk = new_cands[:, l_start:l_end, :] + chunk_sq = (chunk ** 2).sum(dim=-1) # (B, chunk) + # (B, N, chunk) distances in input dtype, cast to float32 + dists = ( + x_sq.unsqueeze(-1) + - 2 * torch.bmm(x, chunk.transpose(1, 2)) + + chunk_sq.unsqueeze(-2) + ).float() + dists.clamp_min_(0) + chunk_min = dists.min(dim=-1).values # (B, N) + torch.minimum(min_dists, chunk_min, out=min_dists) + + +def standard_kmeans_pp(x, n_clusters, x_sq=None, weights=None, n_local_trials=None): + """ + Standard kmeans++ with greedy local trials (Arthur & Vassilvitskii, 2007). + Matches scikit-learn's _kmeans_plusplus implementation. + + Each step samples n_local_trials candidates and picks the one that + minimizes the total potential (sum of weighted min distances²). + + Args: + x: (B, N, D) input points. + n_clusters: K. + x_sq: (B, N) precomputed ||x||^2, optional. + weights: (B, N) per-sample weights, optional. + n_local_trials: candidates per step. Default: 2 + int(log(K)). + + Returns: + centroids: (B, K, D) initial centroids. + """ + B, N, D = x.shape + device = x.device + if x_sq is None: + x_sq = (x ** 2).sum(dim=-1) + + if n_local_trials is None: + n_local_trials = 2 + int(math.log(n_clusters)) + + centroids = torch.empty((B, n_clusters, D), device=device, dtype=x.dtype) + batch_arange = torch.arange(B, device=device) + w_f = weights.float() if weights is not None else None + + # First centroid + if w_f is not None: + w_probs = w_f / w_f.sum(dim=-1, keepdim=True).clamp_min(1e-30) + first_idx = torch.multinomial(w_probs, 1).squeeze(-1) + else: + first_idx = torch.randint(0, N, (B,), device=device) + centroids[:, 0] = x[batch_arange, first_idx] + + if n_clusters == 1: + return centroids + + # Initial min distances: dist² from each point to first centroid + c = centroids[:, 0:1, :] # (B, 1, D) + c_sq = (c ** 2).sum(dim=-1) # (B, 1) + closest_dist_sq = ( + x_sq - 2 * torch.bmm(x, c.transpose(1, 2)).squeeze(-1) + c_sq + ).float() + closest_dist_sq.clamp_min_(0) + + for k in range(1, n_clusters): + # Sampling probabilities: dist² * weight + # Avoids allocating ones tensor for unweighted case + weighted_dists = closest_dist_sq * w_f if w_f is not None else closest_dist_sq + prob_sums = weighted_dists.sum(dim=-1, keepdim=True) + # Avoid CPU-GPU sync: use torch.where instead of if zero_rows.any() + probs = torch.where(prob_sums > 0, weighted_dists / prob_sums.clamp_min(1e-30), 1.0 / N) + + # Sample n_local_trials candidates (with replacement, matching sklearn) + n_trials = min(n_local_trials, N) + candidate_ids = torch.multinomial(probs, n_trials, replacement=True) # (B, n_trials) + candidates = torch.gather( + x, 1, candidate_ids.unsqueeze(-1).expand(-1, -1, D) + ) # (B, n_trials, D) + + # Distances from each candidate to all points: (B, n_trials, N) + cand_sq = (candidates ** 2).sum(dim=-1) # (B, n_trials) + dist_to_cands = ( + x_sq.unsqueeze(1) + - 2 * torch.bmm(candidates, x.transpose(1, 2)) + + cand_sq.unsqueeze(-1) + ).float() + dist_to_cands.clamp_min_(0) + + # For each candidate, new min distances + new_min_dists = torch.min( + closest_dist_sq.unsqueeze(1), dist_to_cands + ) # (B, n_trials, N) + + # Potential per candidate: sum of weighted new_min_dists + # Weighted: bmm with weight vector. Unweighted: simple sum (avoids bmm with ones). + if w_f is not None: + candidates_pot = torch.bmm( + new_min_dists, w_f.unsqueeze(-1) + ).squeeze(-1) # (B, n_trials) + else: + candidates_pot = new_min_dists.sum(dim=-1) # (B, n_trials) + + # Pick best candidate per batch (lowest potential) + best_trial = candidates_pot.argmin(dim=-1) # (B,) + best_id = candidate_ids[batch_arange, best_trial] + centroids[:, k] = x[batch_arange, best_id] + closest_dist_sq = new_min_dists[batch_arange, best_trial] # (B, N) + + return centroids + + +def scalable_kmeans_pp(x, n_clusters, x_sq=None, weights=None, + oversampling_factor=2.0, n_rounds=8): + """ + Scalable K-Means++ (K-Means||) initialization (Bahmani et al., 2012). + Matches cuml/cuvs implementation. + + Performs n_rounds passes, each sampling l = oversampling_factor * K + candidates in parallel, then reduces via weighted sequential kmeans++. + + Args: + x: (B, N, D) input points. + n_clusters: K. + x_sq: (B, N) precomputed ||x||^2, optional. + weights: (B, N) per-sample weights, optional. + oversampling_factor: candidates per round = oversampling_factor * K (default 2.0). + n_rounds: number of oversampling rounds (default 8, matching cuml). + + Returns: + centroids: (B, K, D) initial centroids. + """ + B, N, D = x.shape + device = x.device + l = max(1, int(oversampling_factor * n_clusters)) + + if x_sq is None: + x_sq = (x ** 2).sum(dim=-1) + + w_f = weights.float() if weights is not None else None + batch_arange = torch.arange(B, device=device) + + # --- Step 1: first center --- + if w_f is not None: + w_probs = w_f / w_f.sum(dim=-1, keepdim=True).clamp_min(1e-30) + first_idx = torch.multinomial(w_probs, 1).squeeze(-1) + else: + first_idx = torch.randint(0, N, (B,), device=device) + + # For very small K, sequential is cheaper than the oversampling overhead + if n_clusters <= 3: + centroids = torch.empty((B, 1, D), device=device, dtype=x.dtype) + centroids[:, 0] = x[batch_arange, first_idx] + if n_clusters == 1: + return centroids + return _kmeanspp_sequential(x, n_clusters, x_sq, weights) + + candidates_list = [x[batch_arange, first_idx].unsqueeze(1)] # [(B, 1, D)] + + # Initial min distances + c = candidates_list[0] + c_sq = (c ** 2).sum(dim=-1) + min_dists = (x_sq - 2 * torch.bmm(x, c.transpose(1, 2)).squeeze(-1) + c_sq).float() + min_dists.clamp_min_(0) + + # --- Step 2: oversampling rounds --- + for _ in range(n_rounds): + probs = min_dists * w_f if w_f is not None else min_dists + prob_sums = probs.sum(dim=-1, keepdim=True) + # Avoid CPU-GPU sync: use torch.where instead of if zero_rows.any() + probs = torch.where(prob_sums > 0, probs / prob_sums.clamp_min(1e-30), 1.0 / N) + + n_samples = min(l, N) + new_idx = torch.multinomial(probs, n_samples, replacement=False) # (B, l) + new_cands = torch.gather( + x, 1, new_idx.unsqueeze(-1).expand(-1, -1, D) + ) # (B, l, D) + candidates_list.append(new_cands) + + # Update min distances with all new candidates at once (chunked) + _update_min_dists_batched(x, x_sq, new_cands, min_dists) + + # --- Step 3: collect candidates --- + all_candidates = torch.cat(candidates_list, dim=1) # (B, C, D) + C = all_candidates.shape[1] + + if C <= n_clusters: + # Degenerate: pad with random points + extra_idx = torch.randint(0, N, (B, n_clusters - C), device=device) + extra = torch.gather(x, 1, extra_idx.unsqueeze(-1).expand(-1, -1, D)) + return torch.cat([all_candidates, extra], dim=1) + + # --- Step 4: weight candidates by assigned-point count --- + assignments = euclid_assign_torch_native_chunked( + x, all_candidates, x_sq + ) # (B, N) int32 + + cand_weights = torch.zeros((B, C), device=device, dtype=torch.float32) + if w_f is not None: + cand_weights.scatter_add_(1, assignments.long(), w_f) + else: + cand_weights.scatter_add_( + 1, assignments.long(), + torch.ones((B, N), device=device, dtype=torch.float32), + ) + + # --- Step 5: reduce C candidates to K centroids --- + # Weighted sequential kmeans++ on candidates (matches cuml's reduction) + cand_sq = (all_candidates.float() ** 2).sum(-1) + return _kmeanspp_sequential(all_candidates, n_clusters, cand_sq, weights=cand_weights) + + def euclid_assign_torch_native_chunked(x, centroids, x_sq, chunk_size_N=32768, chunk_size_K=1024): """ Torch naive implementation for assignment with chunking to avoid OOM. @@ -134,7 +415,7 @@ def _euclid_iter_torch_naive(x, x_sq, centroids, chunk_size_N=32768, chunk_size_ return centroids_new, shift, cluster_ids -def batch_kmeans_Euclid_torch_native(x, n_clusters, max_iters=100, tol=0.0, init_centroids=None, verbose=False, chunk_size_N=32768, chunk_size_K=1024): +def batch_kmeans_Euclid_torch_native(x, n_clusters, max_iters=100, tol=0.0, init_centroids=None, verbose=False, chunk_size_N=32768, chunk_size_K=1024, init="random"): """ Batched KMeans clustering in PyTorch using Euclidean distance. @@ -154,13 +435,19 @@ def batch_kmeans_Euclid_torch_native(x, n_clusters, max_iters=100, tol=0.0, init x_sq = (x ** 2).sum(dim=-1) # (B, N) if init_centroids is None: - # Randomly select initial centers from x - indices = torch.randint(0, N, (B, n_clusters), device=x.device) - centroids = torch.gather( - x, - dim=1, - index=indices[..., None].expand(-1, -1, D) - ) # (B, n_clusters, D) + if init == "scalable-kmeans++": + centroids = scalable_kmeans_pp(x, n_clusters, x_sq) + elif init == "standard-kmeans++": + centroids = standard_kmeans_pp(x, n_clusters, x_sq) + else: + # Randomly select initial centers from x (without replacement, matching sklearn) + uniform = torch.ones(B, N, device=x.device) + indices = torch.multinomial(uniform, n_clusters, replacement=False) + centroids = torch.gather( + x, + dim=1, + index=indices[..., None].expand(-1, -1, D) + ) # (B, n_clusters, D) else: centroids = init_centroids diff --git a/tests/test_kmeanspp.py b/tests/test_kmeanspp.py new file mode 100644 index 0000000..b988c68 --- /dev/null +++ b/tests/test_kmeanspp.py @@ -0,0 +1,247 @@ +"""Correctness tests for scalable kmeans++ initialization.""" + +import torch +import torch.testing +from flash_kmeans.torch_fallback import ( + scalable_kmeans_pp, + standard_kmeans_pp, + euclid_assign_torch_native_chunked, +) + + +def _make_clustered_data(B, N_per_cluster, K, D, spread=0.1, device="cuda"): + """Generate well-separated clusters for testing.""" + centers = torch.randn(B, K, D, device=device) * 10 # spread apart + points = [] + for k in range(K): + cluster = centers[:, k:k+1, :] + torch.randn(B, N_per_cluster, D, device=device) * spread + points.append(cluster) + x = torch.cat(points, dim=1) # (B, K*N_per_cluster, D) + # Shuffle within each batch + for b in range(B): + perm = torch.randperm(x.shape[1], device=device) + x[b] = x[b, perm] + return x, centers + + +def test_output_shape(): + """Centroids have correct shape (B, K, D).""" + for B, N, D, K in [(1, 100, 8, 5), (4, 200, 16, 10), (32, 64, 1, 16)]: + x = torch.randn(B, N, D, device="cuda") + centroids = scalable_kmeans_pp(x, K) + assert centroids.shape == (B, K, D), f"Expected {(B, K, D)}, got {centroids.shape}" + print(" PASS: output shapes correct") + + +def test_sequential_centroids_are_data_points(): + """Sequential kmeans++ must return actual data points as centroids.""" + B, N, D, K = 2, 200, 8, 20 + x = torch.randn(B, N, D, device="cuda") + centroids = standard_kmeans_pp(x, K) # (B, K, D) + + for b in range(B): + for k in range(K): + c = centroids[b, k] + diffs = (x[b] - c.unsqueeze(0)).abs().sum(dim=-1) + min_diff = diffs.min().item() + assert min_diff < 1e-5, f"Centroid [{b},{k}] not found in data (min_diff={min_diff})" + print(" PASS: sequential centroids are actual data points") + + +def test_scalable_centroids_reasonable(): + """Scalable kmeans++ reduces candidates via weighted sequential kmeans++, + so centroids are actual candidate points. Check they're within data range.""" + B, N, D, K = 2, 200, 8, 20 + x = torch.randn(B, N, D, device="cuda") + centroids = scalable_kmeans_pp(x, K) # (B, K, D) + + assert centroids.shape == (B, K, D) + # Each centroid dimension should be within [min, max] of data (with margin) + for b in range(B): + x_min = x[b].min(dim=0).values - 0.5 + x_max = x[b].max(dim=0).values + 0.5 + assert (centroids[b] >= x_min).all(), "Centroid below data range" + assert (centroids[b] <= x_max).all(), "Centroid above data range" + print(" PASS: scalable centroids within data range") + + +def test_k1_edge_case(): + """K=1 should return a single centroid from the data.""" + B, N, D = 3, 50, 4 + x = torch.randn(B, N, D, device="cuda") + centroids = scalable_kmeans_pp(x, 1) + assert centroids.shape == (B, 1, D) + # Check it's a point from x + for b in range(B): + diffs = (x[b] - centroids[b, 0].unsqueeze(0)).abs().sum(dim=-1) + assert diffs.min().item() < 1e-5 + print(" PASS: K=1 edge case") + + +def test_quality_vs_random(): + """On well-separated clusters, kmeans++ should give better init than random.""" + torch.manual_seed(42) + B, N_per, K, D = 1, 100, 8, 4 + x, true_centers = _make_clustered_data(B, N_per, K, D, spread=0.1) + N = x.shape[1] + x_sq = (x ** 2).sum(dim=-1) + + n_trials = 20 + kpp_costs, rand_costs = [], [] + for _ in range(n_trials): + # kmeans++ init + kpp_centroids = scalable_kmeans_pp(x, K) + kpp_ids = euclid_assign_torch_native_chunked(x, kpp_centroids, x_sq) + kpp_assigned = kpp_centroids.gather(1, kpp_ids.unsqueeze(-1).expand(-1, -1, D)) + kpp_cost = ((x - kpp_assigned) ** 2).sum().item() + kpp_costs.append(kpp_cost) + + # Random init + rand_idx = torch.randint(0, N, (B, K), device=x.device) + rand_centroids = torch.gather(x, 1, rand_idx.unsqueeze(-1).expand(-1, -1, D)) + rand_ids = euclid_assign_torch_native_chunked(x, rand_centroids, x_sq) + rand_assigned = rand_centroids.gather(1, rand_ids.unsqueeze(-1).expand(-1, -1, D)) + rand_cost = ((x - rand_assigned) ** 2).sum().item() + rand_costs.append(rand_cost) + + avg_kpp = sum(kpp_costs) / n_trials + avg_rand = sum(rand_costs) / n_trials + print(f" kmeans++ avg init cost: {avg_kpp:.2f}, random avg init cost: {avg_rand:.2f}") + assert avg_kpp <= avg_rand * 1.05, ( + f"kmeans++ should be <= random on structured data, got {avg_kpp:.2f} vs {avg_rand:.2f}" + ) + print(" PASS: kmeans++ gives better/equal init quality than random") + + +def test_weighted(): + """Weighted kmeans++ should respect weights — high-weight region gets more centroids.""" + torch.manual_seed(123) + B, N, D, K = 1, 1000, 2, 10 + + # Two clusters: cluster A (first 100 pts) has 100x weight + x = torch.cat([ + torch.randn(B, 100, D, device="cuda") + 5, # cluster A at +5 + torch.randn(B, 900, D, device="cuda") - 5, # cluster B at -5 + ], dim=1) + weights = torch.ones(B, N, device="cuda") + weights[:, :100] = 100.0 # cluster A is 100x more important + + n_trials = 30 + centroids_near_A = 0 + total_centroids = 0 + for _ in range(n_trials): + centroids = scalable_kmeans_pp(x, K, weights=weights) + # Count how many centroids are near cluster A (x > 0) + near_A = (centroids[:, :, 0] > 0).sum().item() + centroids_near_A += near_A + total_centroids += K + + frac_A = centroids_near_A / total_centroids + print(f" Fraction of centroids near high-weight cluster: {frac_A:.2f}") + # With 100x weight on cluster A (10% of points), expect majority of centroids near A + assert frac_A > 0.3, f"Expected >30% centroids near high-weight cluster, got {frac_A:.2f}" + print(" PASS: weighted kmeans++ respects weights") + + +def test_batched_independence(): + """Each batch element should get independent centroids.""" + B, N, D, K = 4, 200, 8, 10 + x = torch.randn(B, N, D, device="cuda") + centroids = scalable_kmeans_pp(x, K) + + # Different batch elements should (almost certainly) have different centroids + all_same = True + for b in range(1, B): + if not torch.allclose(centroids[0], centroids[b], atol=1e-3): + all_same = False + break + assert not all_same, "All batch elements got identical centroids — not independent" + print(" PASS: batch elements have independent centroids") + + +def test_scalable_vs_sequential_quality(): + """Scalable and sequential should give comparable quality.""" + torch.manual_seed(7) + B, N_per, K, D = 1, 200, 16, 4 + x, _ = _make_clustered_data(B, N_per, K, D, spread=0.3) + x_sq = (x ** 2).sum(dim=-1) + N = x.shape[1] + + n_trials = 30 + scalable_costs, sequential_costs = [], [] + for _ in range(n_trials): + sc = scalable_kmeans_pp(x, K) + sc_ids = euclid_assign_torch_native_chunked(x, sc, x_sq) + sc_assigned = sc.gather(1, sc_ids.unsqueeze(-1).expand(-1, -1, D)) + scalable_costs.append(((x - sc_assigned) ** 2).sum().item()) + + sq = standard_kmeans_pp(x, K) + sq_ids = euclid_assign_torch_native_chunked(x, sq, x_sq) + sq_assigned = sq.gather(1, sq_ids.unsqueeze(-1).expand(-1, -1, D)) + sequential_costs.append(((x - sq_assigned) ** 2).sum().item()) + + avg_sc = sum(scalable_costs) / n_trials + avg_sq = sum(sequential_costs) / n_trials + print(f" Scalable avg cost: {avg_sc:.2f}, Sequential avg cost: {avg_sq:.2f}") + # Scalable should be within 2x of sequential (both are good) + assert avg_sc < avg_sq * 2.0, ( + f"Scalable quality unexpectedly poor: {avg_sc:.2f} vs sequential {avg_sq:.2f}" + ) + print(" PASS: scalable and sequential give comparable quality") + + +def test_speed_scalable_vs_sequential(): + """Compare scalable vs sequential speed. + + At small D (e.g. 16), kernel launch overhead dominates compute, so the + scalable approach may not be faster despite fewer data passes. The main + advantage of scalable is better init quality from oversampling + Lloyd's + refinement, not necessarily raw speed. + """ + import time + B, N, D, K = 1, 500000, 16, 1024 + x = torch.randn(B, N, D, device="cuda") + + # Warmup + scalable_kmeans_pp(x, K) + standard_kmeans_pp(x, K) + torch.cuda.synchronize() + + reps = 3 + # Time scalable + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(reps): + scalable_kmeans_pp(x, K) + torch.cuda.synchronize() + scalable_ms = (time.perf_counter() - t0) / reps * 1000 + + # Time sequential + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(reps): + standard_kmeans_pp(x, K) + torch.cuda.synchronize() + sequential_ms = (time.perf_counter() - t0) / reps * 1000 + + speedup = sequential_ms / scalable_ms + print(f" N={N}, K={K}: Scalable={scalable_ms:.1f}ms, Sequential={sequential_ms:.1f}ms, Speedup={speedup:.1f}×") + # No hard assertion — at small D, kernel launch overhead dominates. + # Scalable wins on quality (tested above), speed depends on N/D/K regime. + print(f" INFO: speed comparison (informational, no assertion)") + + +if __name__ == "__main__": + print("Testing kmeans++ initialization...\n") + + test_output_shape() + test_sequential_centroids_are_data_points() + test_scalable_centroids_reasonable() + test_k1_edge_case() + test_quality_vs_random() + test_weighted() + test_batched_independence() + test_scalable_vs_sequential_quality() + test_speed_scalable_vs_sequential() + + print("\nAll tests passed!") diff --git a/tests/test_weighted_centroid.py b/tests/test_weighted_centroid.py new file mode 100644 index 0000000..fd309b3 --- /dev/null +++ b/tests/test_weighted_centroid.py @@ -0,0 +1,96 @@ +"""Correctness tests for the weighted centroid-update kernel. + +Covers the cases the kernel must not regress on: + - power-of-two and NON-power-of-two D (the kernel streams D in BLOCK_D tiles + with masking, so D=96/192/80 must work, not just D=64/128/256), + - large D + fp32, + - all-ones weights reduce to the unweighted kernel, + - empty clusters fall back to old_centroids. +""" + +import pytest +import torch + +from flash_kmeans.centroid_update_triton import ( + triton_centroid_update_sorted_euclid, + triton_centroid_update_sorted_euclid_weighted, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" +) + + +def _torch_weighted_reference(x, cluster_ids, old_centroids, weights): + """Per-cluster weighted mean in fp32; empty clusters keep old_centroids.""" + B, N, D = x.shape + K = old_centroids.shape[1] + xf, wf = x.float(), weights.float() + out = torch.empty((B, K, D), device=x.device, dtype=torch.float32) + for b in range(B): + for k in range(K): + mask = cluster_ids[b] == k + wsum = wf[b][mask].sum() + if wsum <= 0: + out[b, k] = old_centroids[b, k].float() + else: + out[b, k] = (xf[b][mask] * wf[b][mask, None]).sum(0) / wsum + return out.to(x.dtype) + + +# (B, N, K, D, dtype) — includes non-power-of-two D and large-D/fp32. +CASES = [ + (4, 8000, 32, 64, torch.float16), + (4, 8000, 32, 128, torch.float16), + (4, 8000, 32, 96, torch.float16), # non-pow2 + (4, 8000, 32, 192, torch.float16), # non-pow2 + (2, 8000, 32, 80, torch.float16), # non-pow2 head dim + (2, 8000, 32, 256, torch.float16), + (2, 8000, 16, 1024, torch.float32), # large D + fp32 +] + + +@pytest.mark.parametrize("B,N,K,D,dtype", CASES) +def test_weighted_matches_torch_reference(B, N, K, D, dtype): + torch.manual_seed(0) + dev = "cuda" + x = torch.randn(B, N, D, device=dev, dtype=dtype) + cluster_ids = torch.randint(0, K, (B, N), device=dev) + old_c = torch.randn(B, K, D, device=dev, dtype=dtype) + weights = torch.rand(B, N, device=dev) + 0.1 + + got = triton_centroid_update_sorted_euclid_weighted(x, cluster_ids, old_c, weights) + ref = _torch_weighted_reference(x, cluster_ids, old_c, weights) + + tol = 5e-2 if dtype == torch.float16 else 1e-4 + assert (got.float() - ref.float()).abs().max().item() < tol + + +@pytest.mark.parametrize("B,N,K,D,dtype", CASES) +def test_ones_weights_match_unweighted(B, N, K, D, dtype): + torch.manual_seed(0) + dev = "cuda" + x = torch.randn(B, N, D, device=dev, dtype=dtype) + cluster_ids = torch.randint(0, K, (B, N), device=dev) + old_c = torch.randn(B, K, D, device=dev, dtype=dtype) + ones = torch.ones(B, N, device=dev) + + got_w = triton_centroid_update_sorted_euclid_weighted(x, cluster_ids, old_c, ones) + got_u = triton_centroid_update_sorted_euclid(x, cluster_ids, old_c) + + tol = 5e-2 if dtype == torch.float16 else 1e-4 + assert (got_w.float() - got_u.float()).abs().max().item() < tol + + +def test_empty_cluster_falls_back_to_old_centroids(): + torch.manual_seed(0) + dev = "cuda" + B, N, K, D = 2, 8000, 32, 96 # non-pow2 D on purpose + x = torch.randn(B, N, D, device=dev, dtype=torch.float16) + cluster_ids = torch.randint(0, K, (B, N), device=dev) + cluster_ids[cluster_ids == 0] = 1 # cluster 0 empty in every batch + old_c = torch.randn(B, K, D, device=dev, dtype=torch.float16) + weights = torch.rand(B, N, device=dev) + 0.1 + + got = triton_centroid_update_sorted_euclid_weighted(x, cluster_ids, old_c, weights) + torch.testing.assert_close(got[:, 0].float(), old_c[:, 0].float(), atol=1e-3, rtol=0) From 6e53b12c798d0acb17cd13dce2ac59214e1b8735 Mon Sep 17 00:00:00 2001 From: Zeyu Yang Date: Sat, 25 Jul 2026 16:45:05 -0700 Subject: [PATCH 3/3] Fix scalable-kmeans++ quality blow-ups: use greedy weighted reduction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit scalable_kmeans_pp (K-Means||) reduced its oversampled candidate set to K centers with the plain non-greedy _kmeanspp_sequential, while the baseline it is compared against (standard_kmeans_pp) is greedy (sklearn-style local trials). The weaker reduction occasionally picked two candidates from one cluster and left another uncovered, so on ~25% of random instances scalable init cost exceeded 2x standard, with rare 9-10x blow-ups — independent of dataset size (not a small-sample artifact). Using the greedy weighted reduction (standard_kmeans_pp with candidate weights, matching cuml's robust reduction) removes the tail entirely: across 20 seeds x {N=3200, 32000}, max ratio drops from ~10x to 1.2x, 0/20 exceed 2x, and scalable is now on par with / slightly better than standard on average. Fixes the deterministic failure of tests/test_kmeanspp.py::test_scalable_vs_sequential_quality. --- flash_kmeans/torch_fallback.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/flash_kmeans/torch_fallback.py b/flash_kmeans/torch_fallback.py index b83e5f3..1b5cf97 100644 --- a/flash_kmeans/torch_fallback.py +++ b/flash_kmeans/torch_fallback.py @@ -277,9 +277,13 @@ def scalable_kmeans_pp(x, n_clusters, x_sq=None, weights=None, ) # --- Step 5: reduce C candidates to K centroids --- - # Weighted sequential kmeans++ on candidates (matches cuml's reduction) + # Greedy weighted kmeans++ on candidates (matches cuml's reduction). The + # greedy local-trials variant is used (not plain sequential) so the final + # reduction is as robust as standard_kmeans_pp; a plain reduction + # occasionally picks two candidates from one cluster and leaves another + # uncovered, producing rare but severe quality blow-ups. cand_sq = (all_candidates.float() ** 2).sum(-1) - return _kmeanspp_sequential(all_candidates, n_clusters, cand_sq, weights=cand_weights) + return standard_kmeans_pp(all_candidates, n_clusters, cand_sq, weights=cand_weights) def euclid_assign_torch_native_chunked(x, centroids, x_sq, chunk_size_N=32768, chunk_size_K=1024):