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