Skip to content

fix(cutile-rs): keep every launch-site specialization - #1521

Open
rdeaton wants to merge 1 commit into
NVIDIA:mainfrom
rdeaton:fix/launch-site-specializations
Open

rdeaton wants to merge 1 commit into
NVIDIA:mainfrom
rdeaton:fix/launch-site-specializations

Conversation

@rdeaton

@rdeaton rdeaton commented Oct 15, 2026

Copy link
Copy Markdown

Summary

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. This PR adds a larger

In the host_overhead bench on my machine(Ryzen 9 9900X3D, RTX 5080) 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).

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

  • The site now keeps one entry per probe it resolved. A lookup compares against the most recently stored entry first to keep the common path fast, then looks the probe up by its FxHash.
  • I've new-typed the CacheEpoch to prevent inadvertent misuse
  • LaunchSite, SiteResolution and the new CacheEpoch are now #[doc(hidden)] launcher support for generated code
  • 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. There's probably some follow-up work to clean up how tests interact with a shared cache like this, but I left it out of scope here.

Testing

  • fmt and clippy;
  • run_cpu_tests.sh;
  • [x]the full run_gpu_tests.sh (750/750 on each branch);
  • [x]Ran the host_overhead benchmarks, verified that changes to the existing launch are within noise on my hardware (Ryzen 9 9900X3D, RTX 5080).

Checklist

  • All commits signed off (git commit -s)
  • SPDX headers on new source files

Questions about the review? Ping us in #contributors on Discord.

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>
@rdeaton
rdeaton requested a review from elibol as a code owner October 15, 2026 19:29
@copy-pr-bot

copy-pr-bot Bot commented Oct 15, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant