Repository navigation
Conversation
A generated launcher's `LaunchSite` held only its most recent resolution, so a kernel launched alternately at two specializations (the same kernel at two shapes in one model forward, or a scalar whose divisibility changes) missed on every launch and resolved again through key construction and the kernel cache. In the `host_overhead` bench (RTX 5080, mean of two alternating runs against main) that miss path, `launch_site_miss`, goes from 3.84 to 1.72 µs a launch; `launch` with one specialization is unchanged (1.65 / 1.66 µs). The new `launch_alternating` and `launch_12_specializations` cases cost 1.66 and 1.67 µs a launch. The site now keeps one entry per probe it resolved. A lookup compares against the most recently stored entry first, then looks the probe up by its FxHash. Because a site now holds several functions, evictions have to reach it. Every eviction that removes a kernel bumps the cache epoch and then empties every launch site in the process, so an eviction of any key makes every site resolve again on its next launch. The launcher reads a `CacheEpoch` before resolving, and `LaunchSite::store` discards a resolution whose epoch went stale meanwhile. This also releases modules that main kept loaded: a site that had resolved an evicted kernel held its function until it stored another resolution. `LaunchSite`, `SiteResolution` and the new `CacheEpoch` are now `#[doc(hidden)]` launcher support for generated code; `SiteResolution::new` takes the epoch, and `LaunchSite::store` takes `&'static self` so the site can register itself. A poisoned site lock is still a miss, as on main; only an eviction recovers it, to empty the site. The tests evict only their own kernel, or bump the epoch through a new test-only hook that leaves the kernel cache alone, so they don't unload modules other tests in the binary are using. Signed-off-by: Robert Deaton <me@rdeaton.space>
This branch has not been deployed
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
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
A generated launcher's
LaunchSiteheld only its most recent resolution, so a kernel launched alternately at two specializations (the same kernel at two shapes in one model forward, or a scalar whose divisibility changes) missed on every launch and resolved again through key construction and the kernel cache. This PR adds a largerIn the
host_overheadbench on my machine(Ryzen 9 9900X3D, RTX 5080) that miss path,launch_site_miss, goes from 3.84 to 1.72 µs a launch;launchwith one specialization is unchanged (1.65 / 1.66 µs).I originally had these patches on cutile 0.3 where they mattered a bit more, but NVlabs/cutile-rs#307 cut the overhead down a lot. Still, this change makes a substantial performance difference for launch-bound workloads. My workload is the batched forward pass of an ML model for a turn-based strategy game. Each forward runs the same GEMM, layer-norm and attention kernels at several shapes, and batches are padded to a ladder of token counts, so a single-entry launch site misses on almost every launch. After 0.4, I measure a 5-10% more forward passes, about a total of 78us of host time per forward, though it was as high as > 20% before I invested in doing a lot more kernel fusion.
Changes
LaunchSite,SiteResolutionand the newCacheEpochare now#[doc(hidden)]launcher support for generated codeTesting
Checklist
git commit -s)