Export Cake KMeans plan-runtime suite (124-record geomean 2.18x, 1.0x floor passes with 0 below)

[yyihuang]

Summary

This PR re-exports Cake KMeans assignment through the new plan-based launch runtime with per-signature CUDA graphs and is the direct successor of PR #15. Same reusable two-stage public API:

runtime = flashlib_cake_kmeans.init()
assignments = runtime.compute(x, centroids)

The committed artifact was generated and validated from exact merged Cake main 42070e96d0734cb580854baef60f17625ba33bb5, which includes the plan-based export runtime arc (per-signature resolved LaunchPlan + plan cache, prepacked argument buffers, TMA variant-bank pointer rebinding), the D112 a262 dispatcher winner, and the KMeans plan-runtime graph wiring (per-signature CUDA-graph capture of each record's route chain; a hot call is pointer-patch + one cuGraphLaunch). Exact provenance is embedded in manifest.json.

What changed vs PR #15

PR #15 froze the interpreted host stack into the shipped hot path and failed its 1.0x floor closed (98/124 records lost to the exact-07cf baseline hot). This export resolves every contract signature into a static launch plan at plan-construction time and captures its route chain into a CUDA graph. Publication-workflow bugs fixed along the way in Cake main (each found by this run's fail-closed gates): the correctness overlay's missing measurement_order_seed (MR 305), a stale frozen registry-identity mirror in the benchmark asset (MR 306), and an unbalanced lifecycle cold ordering (MR 307, now rank-alternated exactly 62/62 over the contract portfolio).

Correctness and scope

• candidate/baseline correctness and ordered route parity: 124/124 contract records (4 validation shards)
• hardware: 4x NVIDIA GB200 (sm_100a)
• official baseline: raw-input adapter around the exact frozen Triton implementation from 07cf2a2, same session, CUPTI timing everywhere
• frozen CUDA specializations: 358 kernels, no Cake/Loom runtime dependency

Hot steady-state synchronized E2E vs PR #15

Official metric: exact-07cf raw adapter / runtime.compute, synchronized host end-to-end (raw BF16 inputs, default output allocation, route-required internal fused norms, launch work, GPU execution and gaps, terminal synchronization).

Scope	PR #15 (interpreted runtime)	This PR (plan runtime + graphs)
All 124 records: min	0.6334x	1.0938x
All 124 records: geomean	0.8683x	2.1825x
All 124 records: median	0.8215x	2.1942x
All 124 records: p90	1.0871x	2.8423x
All 124 records: max	2.0197x	3.8773x
All 124 records: <1.0x	98/124	0/124
Publication floor (1.0x, all 124)	failed closed	passed, 0 below

Every one of the 124 contract records now beats the exact-07cf baseline hot; PR #15 lost on 98 of them.

Modeled amortized synchronized E2E

Public calls N	after-init geomean	including-init geomean
1	1.231x	0.945x
10	1.259x	0.959x
100	1.509x	1.071x
1000	1.996x	1.468x

Cold init-once median 242.8 ms (NVRTC) with first compute() median 3.13 ms; the including-init table crosses 1.0x by N=100 — unlike the KNN suites, KMeans already amortizes its init within realistic call counts. Prebuilt cubins (plan-runtime design steps 4–5) remain the queued follow-up for the N=1 cold case.

Validation contract

• All timing CUPTI, hard-checked per lane; L2 flushed; each completed record synchronized and evicted before the next (one miss + fresh-pointer hit per record; no resident multi-shape cache assumed); lifecycle cold ordering deterministically balanced 62/62 across the contract.
• Machine-readable evidence: VALIDATION.json, per-record BENCHMARK_RESULTS.json, human summary VALIDATION.md.

🤖 Generated with Claude Code

---

Mirrors yyihuang/flashlib#19 onto the FlashML-org upstream.