Flatten assign/update launch grids to remove the 65535 batch limit

[zeyuyang8]

Summary

Splits the flattened-launch-grid fix out of #21 into its own PR, since it touches the shared assign/update paths that all current users hit (as suggested by @andy-yang-1).

The non-split assign kernels (euclid/cosine) and the shared _centroid_update_chunk_kernel launched a 2D 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 Triton Error [CUDA]: invalid argument.

This flattens them to a 1D grid (B * n_tiles,) and decodes 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 identical to the old 2D grid (CUDA enumerates grid.x fastest), so per-program work and atomic contention are unchanged.

The split-D assign kernels decode a 2D grid and keep their 2D launch — this change only touches the non-split kernels.

Benchmarks (H100 80GB, fp16)

benchmarks/grid_fix/ is included so this is reproducible.

No throughput regression on the non-split paths — AFTER/BEFORE latency ratio (main vs this branch), < 1.0 = faster:

shape (B,N,K,D)	euclid_assign	cosine_assign	euclid_update	cosine_update
(1, 1M, 256, 128)	0.99x	0.99x	1.00x	1.00x
(8, 131K, 256, 128)	1.00x	0.97x	1.02x	1.02x
(64, 16K, 256, 128)	1.01x	0.99x	1.01x	1.01x
(1, 4M, 1024, 64)	1.02x	1.01x	1.00x	1.00x
(256, 8K, 128, 128)	0.99x	0.98x	1.01x	1.01x
(4, 262K, 512, 256)	1.00x	1.01x	1.01x	1.01x

Worst case 1.02x, within run-to-run noise.

Capability (bench_large_b.py, B=70000 > 65535):

	euclid_assign	cosine_assign	euclid_update
main	❌ invalid argument	❌ invalid argument	❌ invalid argument
this branch	✅	✅	✅

How to reproduce

git checkout main
python benchmarks/grid_fix/bench_grid.py --out /tmp/grid_main.json
git checkout fix/flattened-launch-grid
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
python benchmarks/grid_fix/bench_large_b.py   # run on both branches

Numbers above are on an H100; the same scripts will produce H200 numbers if you'd like to confirm on that arch.

[zeyuyang8]

Split out from #21 per review. Grid-only change; no throughput regression (worst 1.02x on H100) and fixes a real CUDA: invalid argument launch failure at B > 65535. Feature work (weighted k-means, k-means++) is stacked on this in #21.

[andy-yang-1]

@zeyuyang8 Thanks for your contribution! Can you also fix the split-D path? I can merge it once split-D is fixed

[zeyuyang8]

Thanks @andy-yang-1! Done — the split-D assign kernels (euclid + cosine) now decode the same flattened 1D grid (B * n_tiles,), so they launch for B > 65535 too. Pushed in 6906cee.

Verified on H100:

• Capability: B=70000 split-D (D=768) euclid + cosine now launch (both failed with CUDA: invalid argument before); bench_large_b.py covers this case.
• Correctness: batched split-D output is bit-identical to a per-batch (B=1) loop, so the (b, tile) decode is exact; euclid matches the torch reference at 99.7–100% across fp16/bf16/fp32 for D ∈ {513, 768, 1024}.

The block→(b, tile) linearization matches the old 2D grid (CUDA enumerates grid.x fastest), so per-program work and atomic contention are unchanged.

[andy-yang-1]

@zeyuyang8 Great work!