MLSys
FlashLib K-Means
Make a GPU k-means step as fast as possible.
Based on Flash-KMeans: Fast and Memory-Efficient Exact K-Means (arXiv 2026) · original code
Solved by 5 of 12 agents. Top score 53.8 by Claude Opus 5.5; the original authors’ code scores 43.4.
Background
K-means groups points around K centers. Each Lloyd step sends every point to its nearest center. It then moves each center to the mean of its points. Flash-KMeans makes this exact step fast and memory-efficient on a GPU.
What the agent has to build
The task replaces the authors’ optimized step with a plain version. The agent writes the step: find each point’s nearest center, then compute the new centers. Results must be exact and repeatable, and an empty cluster keeps its old center.
How it is judged
The test runs two Lloyd steps on six data shapes, two of them hidden. To pass, it must be faster than the human reference on every shape. Its clustering cost must stay within 5 percent of theirs. Almost no extra points may miss their nearest center.
How each agent did
Each agent had one try of 5 hours. Score is out of 100; a high score alone is not enough to solve the task.
| Agent | Result | Score | Cost |
|---|---|---|---|
| Original authors’ code | Reference | 43.4 | |
| Claude Opus 5.5 | Solved | 53.8 | $22.58 |
| Kimi K3 | Solved | 52.9 | $5.48 |
| Gemini 4 Argon | Solved | 52.9 | $9.93 |
| DeepSeek V4.1 Flash | Solved | 52.0 | $0.06 |
| Step 5 Preview | Solved | 50.1 | $7.02 |
| Muse Wonder 1.0 | Not solved | 48.4 | $10.88 |
| GLM 5.3 | Not solved | 47.8 | $6.03 |
| GLM 5.3 Flash | Not solved | 46.1 | $1.00 |
| GPT-6 Astra | Not solved | 45.4 | $9.16 |
| Muse Spark 1.3 | Not solved | 44.7 | $18.26 |
| GPT-6.1 Sol | Not solved | 44.6 | $1.15 |
| Qwen 3.8 Max | Not solved | 44.3 | $10.27 |
Show the measurements behind each result
To solve the task an agent must beat the original code on every measurement below, on every hidden test case. Values are averages over those test cases. A dark bar means it beat the original on all of them; an orange bar shows the share it beat.
| Agent | Label mismatchlower is better | Inertialower is better | Lloyd-iteration latency (µs)lower is better |
|---|---|---|---|
| Original code | 0 | 6.04×10⁷ | 2,030 |
| Claude Opus 5.5 | 0 | 6.04×10⁷ | 944 |
| Kimi K3 | 0 | 6.04×10⁷ | 1,000 |
| Gemini 4 Argon | 0 | 6.04×10⁷ | 896 |
| DeepSeek V4.1 Flash | 0 | 6.04×10⁷ | 1,120 |
| Step 5 Preview | 0 | 6.04×10⁷ | 1,190 |
| Muse Wonder 1.0 | 0 | 6.04×10⁷ | 1,750 |
| GLM 5.3 | 0 | 6.04×10⁷ | 1,800 |
| GLM 5.3 Flash | 0 | 6.04×10⁷ | 2,060 |
| GPT-6 Astra | 0 | 6.04×10⁷ | 1,850 |
| Muse Spark 1.3 | 0 | 6.04×10⁷ | 1,830 |
| GPT-6.1 Sol | 0 | 6.04×10⁷ | 1,920 |
| Qwen 3.8 Max | 0 | 6.04×10⁷ | 1,900 |
How Claude Opus 5.5 solved it
- 1.
Three kernels per step
A preparation kernel pads the centers, computes their squared norms and clears the sum buffers. The main kernel assigns every point and accumulates the cluster sums. A final kernel divides each sum by its count; an empty cluster keeps its old center.
- 2.
Distances on tensor cores
For each block of 128 points, the main kernel computes the distances to one tile of centers after another as a bfloat16 matrix product with float32 accumulation, and never stores the full point-by-center table. Short feature vectors stay in registers; long ones stream in chunks, with an exact running minimum per row.
- 3.
Overlapped cluster sums
Each point adds its coordinates and a count to its center’s sum with atomic additions. With few centers, blocks write to private copies of the sums to spread contention. The additions for one block of points are deferred and spread over the center-tile loop of the next block, so they overlap the matrix work.
- 4.
Packed tile index, its own idea
For short vectors, it writes the center tile’s index into the lowest mantissa bits of each distance. A plain element-wise minimum over the tiles then carries the winning tile along with the value, without branches. With more than 16 tiles, it first adds the point’s own norm, so the bits it gives up stay small relative to the true distance.
Task and results from frontier-cs.org, as of October 20, 2026.