Repository navigation
[executorch][cuda] Compile an autotuned quantized GEMM's configs in parallel before benchmarking them - #23948
Merged
Merged
Conversation
…arallel before benchmarking them Pull Request resolved: #23736 At AOTI compile time every QuantizedGemmFamily kernel autotunes its configs per shape (and per split-K candidate, via `@autotune_launch_param`). `triton.autotune` compiles each config on its first benchmark call, one after another, so on a cold Triton cache the compile time dominates: for one INT4 W4A16 bucket-64 shape (5376 x 4096) it is 335 s, of which ~290 s is compiling 292 kernels serially. `launch_split_k_gemm` now compiles every pruned config of the kernel it is about to launch on a thread pool (up to 32 threads) before Triton benchmarks them, using Triton's own `AsyncCompileMode` + `Autotuner.warmup`. Triton's benchmark then times kernels that are already compiled. Nothing about the search changes: the same configs are pruned, compiled and timed, and the pick is the same. - Only for plain CUDA tensors and a tuning key the autotuner has not seen yet (it mirrors `Autotuner.run`'s key). Tensor subclasses (fake and functional tensors while tracing or under torch.compile), single-config kernels and an already-active async-compile scope are left alone. - A config that fails to compile fails again in the benchmark, which skips it as before (`ignore_errors=True`). - A Triton without `triton.runtime._async_compile` keeps the serial compile. Measured (A100, empty Triton cache, INT4 5376 x 4096 at bucket 64 under `autotune_launch_params()` + `cuda_graph_timing()`, the compile-time setup): 334.6 s -> 77.8 s; the same 292 kernels compiled; split-K pick 4 in both, per-candidate times within 1%. Full Muse Glimmer 30B batching export (forward_{1,2,4,8,16,32,64} + forward_others + samplers; A100; empty Triton and Inductor caches), per method from the AOTI wrapper times, the batching stack's INT4/INT6 large-M buckets on top: | method | before (serial compile) | this diff | |---|---|---| | load + trace + forward_1 | 46.5 min | 51.3 min | | forward_2 | 21.4 | 15.3 | | forward_4 | 23.4 | 15.5 | | forward_8 | 22.3 | 12.2 | | forward_16 | 39.3 | 15.5 | | forward_32 | 52.8 | 18.8 | | forward_64 | 73.3 | 24.2 | | forward_others | 10.5 | 10.3 | | total | 4h55m | 2h50m | The export generates the same tokens as before (greedy, 512 / 4 x 2k prompts) at equal or better speed. What remains per autotuned shape is Triton's `do_bench` itself (a fixed ~25 ms warmup + 100 ms per config), which this diff does not change. ghstack-source-id: 451021037 @exported-using-ghexport Differential Revision: [D124355785](https://our.internmc.facebook.com/intern/diff/D124355785/)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23948
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Gasoonjia
added this pull request to stack #23954
October 18, 2026 07:47
This PR needs a
|
Gasoonjia
approved these changes
Oct 18, 2026
This branch was successfully 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
2 participants
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.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #23736 by @Gasoonjia
^ Please use this as the source of truth for the PR details, comments, and reviews
ghstack PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/52/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/52/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/52/orig
Differential Revision: D124355785
@diff-train-skip-merge