Optimize Blackwell k-means assignment reduction

[henrylhtsang]

Summary

This PR optimizes the score reduction in the Blackwell CuTeDSL k-means assignment kernel. The tcgen05 MMA and existing load/MMA/consumer pipeline remain unchanged; the work focuses on loading score tiles from TMEM, applying centroid norms, and selecting the winning centroid.

The new paths preserve the original numerical behavior and lowest-index tie breaking. All 33 existing B200 k-means tests pass, every sampled benchmark output matched the reference, and adversarial tests covering cross-tile ties, all-equal centroids, NaNs, infinities, and signed zero produced zero mismatches.

Key ideas and contributions

Kernel / area	Idea	Description
BlackwellFlashKmeansAssignWSILP2	Two-chain argmin	Splits the serial centroid scan into independent even/odd reduction chains to expose instruction-level parallelism.
BlackwellFlashKmeansAssignWSILP2	Deterministic winner merge	Merges both chains by score and centroid index, preserving lowest-index tie breaking.
BlackwellFlashKmeansAssignWSNormBroadcast	Register-staged norms	Each warp cooperatively loads its centroid-norm tile directly into registers instead of shared memory.
BlackwellFlashKmeansAssignWSNormBroadcast	Warp-shuffle broadcast	Uses shuffle_sync to distribute each registered norm across the warp.
BlackwellFlashKmeansAssignWSNormBroadcast	Barrier elimination	Removes the shared norm buffer and two norm-staging barriers per centroid tile.
WS accumulator pipeline	Safe asynchronous TMEM loads	Fences each TMEM-to-register load before releasing its accumulator stage, following FlashAttention’s SM100 pattern.

Performance

These are checkout-to-checkout production measurements: main at 81e5dea versus this PR at e4ce568a. Both revisions ran the same external benchmark with the same 13 affected shapes, three rotating input sets, 10 warmup iterations, and the median of 30 timed iterations. No benchmark-only kernel selector was added to the repository.

N	D	K	Main path	PR path	Improvement
65,536	64	512	WS	ILP2	0.1%
65,536	64	4,096	WS	ILP2	2.7%
65,536	64	16,384	WS	ILP2	5.3%
65,536	128	512	WS	Norm broadcast	0.2%
65,536	128	4,096	WS	Norm broadcast	5.9%
262,144	128	4,096	WS	Norm broadcast	9.2%
262,144	128	16,384	WS	Norm broadcast	11.8%
1,048,576	128	4,096	X-resident	Norm broadcast	13.7%
65,536	256	512	WS	ILP2	2.9%
65,536	256	4,096	WS	ILP2	11.5%
262,144	256	4,096	WS	ILP2	14.0%
262,144	256	16,384	WS	ILP2	9.2%
1,048,576	256	4,096	WS	ILP2	16.2%

All 13 affected shapes improved, with a 7.77% geometric-mean speedup. Every row matched the numerical reference.

Development environment

• Modal NVIDIA B200
• PyTorch 2.13.0+cu130
• NVIDIA CUTLASS / CuTe DSL 4.5.3
• Independent synchronization and lifetime review with Codex GPT-5.6 Sol

[andy-yang-1]

@henrylhtsang Thanks for your contribution! This pr looks good to me. Can you add the comparison of nightly branch? I want to see which one is faster.

[henrylhtsang]

@andy-yang-1 sg here it is (below generated by AI):

flashlib PR #22: nightly comparison on B200

Date: 2026-08-29
Hardware: Modal NVIDIA B200 (sm_100)
PR: #22
Requested comparison: #22 (comment)

Result

PR #22 is faster than the current origin/nightly head on all 13 affected
shapes. Across the shape set, the PR is 28.05% faster by geometric mean
(geometric-mean latency 0.2641 ms versus 0.3381 ms). Per-shape gains range
from 10.14% to 64.90%. Every sampled benchmark output matched the mathematical
reference.

Revisions and environment

• Nightly: 033e0f3c43a59c18df69e55cda3a42e3e69640b3
  (origin/nightly as fetched on 2026-08-29)
• PR head: e4ce568a536e4563b15beab299879bfa4d5d4fd3
• GPU: NVIDIA B200, compute capability 10.0
• PyTorch: 2.13.0+cu130
• CUDA runtime: 13.0
• NVIDIA CUTLASS / CuTe DSL: 4.5.3

Both benchmark revisions were clean detached worktrees. Each ran on a Modal
B200 with the same cached package image and the same external benchmark script.

Benchmark methodology

• Same 13 affected production shapes used in the PR description.
• Three independently seeded BF16 input/centroid sets: 17, 29, and 43.
• Inputs rotate between kernel calls to avoid repeatedly timing only the prior
  invocation's cache state.
• 10 warmup calls and 30 timed calls per shape.
• CUDA events with synchronization; reported value is the median latency.
• Correctness check: first 2,048 assignments from every seeded input set were
  compared with the FP32 score reference.
• Improvement is (nightly latency / PR latency - 1) * 100; positive values
  favor the PR.

Benchmark results

N	D	K	Nightly path	PR path	Nightly (ms)	PR (ms)	PR faster
65,536	64	512	WS	ILP2	0.0727	0.0441	64.90%
65,536	64	4,096	WS	ILP2	0.1332	0.1022	30.37%
65,536	64	16,384	WS	ILP2	0.3462	0.2963	16.84%
65,536	128	512	WS	Norm broadcast	0.0722	0.0453	59.43%
65,536	128	4,096	WS	Norm broadcast	0.1399	0.1059	32.05%
262,144	128	4,096	WS	Norm broadcast	0.3444	0.2844	21.11%
262,144	128	16,384	WS	Norm broadcast	1.1653	1.0047	15.98%
1,048,576	128	4,096	X-resident	Norm broadcast	1.1978	1.0148	18.04%
65,536	256	512	WS	ILP2	0.0871	0.0587	48.35%
65,536	256	4,096	WS	ILP2	0.2111	0.1661	27.11%
262,144	256	4,096	WS	ILP2	0.5640	0.4631	21.78%
262,144	256	16,384	WS	ILP2	2.0881	1.8615	12.18%
1,048,576	256	4,096	WS	ILP2	2.1431	1.9458	10.14%

Geometric-mean latency is computed across the 13 rows, not by averaging the
percentage column. All rows had a minimum sampled-reference match rate of 1.0.