Export Cake KNN search plan-runtime suite (full-198 geomean 3.11x, 0 rows below 1.0x)#5
Open
yyihuang wants to merge 1 commit into
Open
Export Cake KNN search plan-runtime suite (full-198 geomean 3.11x, 0 rows below 1.0x)#5yyihuang wants to merge 1 commit into
yyihuang wants to merge 1 commit into
Conversation
Plan-based launch runtime (per-signature LaunchPlan cache, TMA variant-bank hot path) exported and validated on 4x NVIDIA GB200: 198/198 correctness, 198/198 ordered route parity, hot public synchronized-E2E speedup vs live flashlib.flash_knn min 1.5043 / geomean 3.1134 across the full 198-shape denominator, 0 rows below 1.0x, declared 11-shape publication floor min 2.0693 / geomean 2.9272 at the 1.0x floor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR re-exports Cake KNN search through the new plan-based launch runtime and is the direct successor of PR #16. Same reusable two-stage public API:
The artifact was generated and validated from exact merged Cake main
0b4375909f78358fd1eb3d1d3a804e31353a18b9, which includes the full plan-based export runtime arc (per-signature resolvedLaunchPlan+ plan cache, prepacked argument buffers, TMA variant-bank pointer rebinding — Cake MRs 293/299/300/301). Exact provenance is embedded inmanifest.json.What changed vs PR #16
PR #16 froze the evolve-time interpreted host stack (guard cascade, per-argument ctypes marshalling, per-call TMA descriptor encode) into the shipped hot path; the kernels won but ~0.1–0.2 ms/call of Python host work ate the wins. This export resolves every one of the 198 contract signatures into a static launch plan at export time: a hot call is signature lookup → pointer patch → submit. The generic dispatcher remains only as the slow-path plan constructor, and route parity against
expected_routes.jsonstays a hard gate.Correctness and inventory
sm_100a), 4 validation shardsflashlib.flash_knn, same session, CUPTI timing everywhere (no wall-clock or event timing)Hot steady-state synchronized E2E vs PR #16
The official metric is
flashlib.flash_knn / runtime.computesynchronized host end-to-end latency (raw inputs, default output allocation, slot lookup, launch work, GPU execution and gaps, terminal synchronization).<1.0xEvery one of the 198 diagnostic shapes now beats FlashLib hot; PR #16 lost on 53 of them. The declared 11-shape publication floor (1.0x) passes fail-closed with
below_floor_count: 0.Modeled amortized synchronized E2E
after-init(N) = (F + (N-1) x H) / N,including-init(N) = (I + F + (N-1) x H) / N, evaluated per shape from measured initI, first-callF, and hot-medianH.Cold init-once synchronized E2E median 423.8 ms and first
compute()median 52.3 ms remain NVRTC-dominated; prebuilt cubins (plan-runtime design steps 4–5) are the queued follow-up, so the including-init table still loses at small N by design.Validation contract
timing_backend: "cupti", hard-checked per lane); L2 flushed; each completed shape synchronized and evicted before the next (one miss + fresh-pointer hit per shape; no pre-populated resident cache assumed).VALIDATION.json, per-shapeBENCHMARK_RESULTS.json, human summaryVALIDATION.md.🤖 Generated with Claude Code
Mirrors yyihuang/flashlib#17 onto the FlashML-org upstream.