Repository navigation
[executorch][cuda] Compile an autotuned quantized GEMM's configs in parallel before benchmarking them - #23736
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23736
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 747d9c9 with merge base 20fa24e ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This was referenced Oct 13, 2026
Merged
This PR needs a
|
digantdesai
requested changes
Oct 13, 2026
digantdesai
left a comment
digantdesai
left a comment
Contributor
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
shoumikhin
approved these changes
Oct 16, 2026
meta-codesync
Bot
merged commit Oct 18, 2026
740ef01
into
gh/Gasoonjia/52/base
290 of 291 checks passed
meta-codesync Bot
pushed a commit
that referenced
this pull request
Oct 18, 2026
…pool for every split candidate Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #23787 * #23740 * #23739 * #23738 * __->__ #23737 * #23736 Builds on D124355785 (configs compiled in parallel before Triton benchmarks them). Two remaining costs of the compile-time quantized GEMM autotune: 1. **Benchmark window.** The family kernels' `triton.autotune` used Triton's default `do_bench(warmup=25, rep=100)`: about 125 ms per config whatever the kernel's length. These kernels run 10-100 us, so every config ran thousands of times; a bucket-64 shape with 288 (config, split) pairs spent ~36 s in it. They now use `quantized_gemm_utils.autotune_do_bench` (5 ms warmup, 20 ms timing), which still runs each config a few hundred times. Only the INT4/5/6/8 family kernels use it; SDPA, MoE and the rest keep Triton's default. 2. **One compile pool per shape.** `@autotune_launch_param` tries each split-K candidate in turn, and D124355785 opened one compile pool per candidate, one after another (e.g. 72, 72, 72, 48 and 24 configs at bucket 64). `@autotune_launch_param` now takes an optional `candidate_prepass`: a scope around one untimed call of every legal candidate before any is timed. The quantized GEMM launches pass `quantized_gemm_family._compile_every_candidate_at_once`, in which `launch_split_k_gemm` only queues its kernel's untuned configs on one shared pool and skips the launch. Leaving the scope waits for all compiles; timing then runs exactly as before. Each key compiled ahead is consumed once by its timed launch, which then does not open a pool of its own. Neither changes the search: the same configs and split candidates are compiled and timed, and the picks are the fastest by the same rules. **Does the shorter window still pick the best config?** For the models' linear shapes (INT4: 8448x6656, 6656x4096, 6656x19968, 256x6656, 1024x6656; INT6: 4096x5376, 5376x21504, 6656x19968, 1024x6656; INT5 and INT8: two each) at buckets 1/4/16/64 (INT5, INT8: 1/4) and every legal split, every pruned config was timed with the current window and with the new one; the union of each window's five fastest was then re-timed with a long window (100 ms warmup, 300 ms timing) as the reference. - 232 (shape, bucket, split) cases: the new window's pick is within 1% of the reference best in 199 cases (current window: 202), within 2% in 214 (213); mean regret 1.0054 (1.0050). - What ships is one split per shape: over the 44 (shape, bucket) pairs, with the split chosen as the fastest, the new window is more than 2% off the best in 2 pairs and the current window in 3. Worst: new 8.8% on INT4 256 x 6656 at M = 1 (an 8 us kernel, 0.7 us), current 6.0% on INT4 1024 x 6656 at M = 16. Both windows miss by a few percent on the shortest kernels, at the same rate. **Compile time.** One INT4 W4A16 bucket-64 shape (5376 x 4096) under `autotune_launch_params()` + `cuda_graph_timing()`, A100, empty Triton cache: 334.6 s before D124355785, 77.8 s with it, 58.0 s with this diff (26.7 s with a warm cache, was 44.8 s); split pick 4 in all, the same 292 kernels compiled, per-split times within 1%. Full Muse Glimmer 30B batching export (`export_solo_batching.py`: forward_{1,2,4,8,16,32,64} + forward_others + samplers; A100; empty Triton and Inductor caches), per method in minutes from the AOTI wrapper times: | method | before D124355785 | D124355785 | this diff | |---|---|---|---| | load + trace + forward_1 | 46.5 | 51.3 | 40.8 | | forward_2 | 21.4 | 15.3 | 14.5 | | forward_4 | 23.4 | 15.5 | 14.4 | | forward_8 | 22.3 | 12.2 | 10.3 | | forward_16 | 39.3 | 15.5 | 12.9 | | forward_32 | 52.8 | 18.8 | 14.7 | | forward_64 | 73.3 | 24.2 | 18.7 | | forward_others | 10.5 | 10.3 | 11.1 | | total | 4h55m | 2h50m | 2h25m | The export generates the same tokens as before (greedy; 512 and 4 x 2k prompts) at the same speed as with D124355785 (single_512 63.0 / 62.7 tok/s vs 62.5 / 62.9; 4 x 2k 111.3 vs 111.0). What is left per autotuned shape is mostly compiling: during forward_64 only ~5 of the 32 compile threads are busy on average, and 87% of their time is in Triton's MLIR/LLVM passes (make_llir, make_ttgir, make_ptx), which appear to serialize on the GIL. Differential Revision: D124496888 Pull Request resolved: #23737
meta-codesync Bot
pushed a commit
that referenced
this pull request
Oct 18, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #23787 * #23740 * #23739 * __->__ #23738 * #23737 * #23736 A multi-method CUDA export compiles every method against one `CudaWeightCollector`, so a storage the methods share (all of the model's weights, for a batching export) is stored under one `cuda_fqn_weight:<device>:<FQN>` key. But every method's `materialize` still copied each storage from the GPU to the host in 8 MB chunks, SHA-256'd it on one core, and wrote it to a temp file; only then did the `NamedDataStore` drop the duplicate. For an 18 GB model that is ~90 s per method (measured on 2 GiB: copy 1.9 s, SHA-256 5.9 s, write 1.3 s), all of it thrown away after the first method. The collector now remembers a fingerprint of each CUDA storage it has written, by storage key. A later method's storage under the same key is fingerprinted on the GPU first; if it matches, its write is skipped and the method's manifest points at the stored entry. Otherwise it is written and added exactly as before, so a changed storage still fails with the store's "Duplicate key ... with different data". The fingerprint is the storage's size, two independent random linear hashes of its 8-byte words (int64, chunk order kept significant) and its trailing bytes: ~0.8 s for 18 GB on an A100, and a single flipped bit changes it. Library-local constants (`_tensor_constant*`, keyed per shared library), off-graph KV entries and CPU storages keep the existing path. Full Muse Glimmer 30B batching export (`export_solo_batching.py`, 10 methods; A100; empty Triton and Inductor caches), stacked on D124496888: | | before | this diff | |---|---|---| | total | 135.1 min | 125.7 min | | main thread in `cuda_weight_collector` | ~21.4 min | ~7.7 min | | per method after the first (AOTI), min | 13.2 / 13.3 / 9.7 / 12.3 / 14.5 / 19.0 / 10.6 / 2.6 | 13.6 / 12.1 / 8.3 / 11.3 / 13.2 / 17.0 / 8.9 / 1.0 | (forward_2 .. forward_64, forward_others, sample. forward_1 writes everything once, as before.) Most of what is left in later methods is the collector's `trim_host_memory` (`gc.collect` + `malloc_trim`). The weight file is byte-identical with and without this diff (`aoti_cuda_blob.ptd`, 17,856,903,936 bytes, same SHA-256), and so is every method's weight manifest; only the methods' compiled libraries differ, as between any two cold exports. The export generates the same tokens (greedy, a 512-token prompt and 4 x 2k prompts) at the same speed (62.5 vs 62.7 tok/s; 110.0 vs 111.3). Differential Revision: D124713681 Pull Request resolved: #23738
meta-codesync Bot
pushed a commit
that referenced
this pull request
Oct 18, 2026
… Triton kernel call Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #23787 * #23740 * __->__ #23739 * #23738 * #23737 * #23736 Tracing a call of a user-defined Triton kernel runs `torch._higher_order_ops.triton_kernel_wrap.identify_accessed_tensors`, which generates the kernel's TTIR and walks it to find the tensors the kernel reads and writes. PyTorch runs it in functionalization, in Inductor's mutation propagation and in Inductor's lowering, and caches nothing, so a model whose every layer launches the same kernels repeats it for every call site of every exported method. In the Muse Glimmer batching export it was the largest single item on the main thread: ~21 minutes. The analysis depends only on the kernel and on what `generate_ttir` makes of the arguments: a fake or Inductor tensor becomes a fresh tensor of its dtype, a symbolic value becomes 2, and any other value is used as is (a real tensor also by its pointer's 16-byte alignment). `backends/cuda/triton/kernels/triton_mutation_cache.py` keys calls on exactly that. The first call with a given key goes through PyTorch; repeats return a copy of its result. Calls it cannot key exactly (TMA descriptors, unhashable values) always go through PyTorch. It is installed when `executorch.backends.cuda.triton.kernels` is imported, by replacing the function on `triton_kernel_wrap`, which both of PyTorch's callers look it up on. Full Muse Glimmer 30B batching export (`export_solo_batching.py`, 10 methods; A100; empty Triton and Inductor caches), stacked on D124736179 (min): | | before | this diff | |---|---|---| | main thread in the access analysis | ~21.1 | ~8.7 | | lowering start -> forward_1 compiled | 23.3 | 19.1 | | AOTI forward_2 / 4 / 8 / 16 / 32 / 64 | 13.3 / 12.0 / 8.5 / 9.8 / 11.6 / 15.5 | 9.8 / 8.3 / 7.2 / 9.5 / 10.8 / 13.8 | | total | 123.3 | 108.6 | (The rest of the host was busier during this run: ~15 other cores vs ~11.) What is left of the analysis is mostly its first call per key, since each method has its own shapes. The weight file is byte-identical, and the export generates the same tokens (greedy; 512 and 4 x 2k prompts) at the same speed (63.0 vs 63.2 tok/s; 110.6 vs 110.4). Differential Revision: D124772365 Pull Request resolved: #23739
meta-codesync Bot
pushed a commit
that referenced
this pull request
Oct 18, 2026
…config, timed in CUDA graphs Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #23787 * __->__ #23740 * #23739 * #23738 * #23737 * #23736 A quantized GEMM's Triton autotuner picks a config for each shape while `@autotune_launch_param` times the split candidates. But `launch_split_k_gemm` then traced the autotuner itself, and for a user-defined `Autotuner` Inductor takes every one of its configs (`kernel.configs`, without our pruning) into `user_autotune`: AOTI compiled and benchmarked all of them again for the same shapes and picked its own tile config. In the Muse Glimmer batching export Inductor was handed 2,032 GEMM configs over 56 kernel variants. **One traced config.** `launch_split_k_gemm` now traces the kernel restricted to the config its autotuner already picked for the call's key: a single-config `triton.autotune(configs=[pick], key=[])(kernel.fn)`, one object per (kernel, config), so every call site of a shape traces the same kernel and PyTorch's per-kernel caches keep working. Inductor then compiles that one config and benchmarks nothing. Called eagerly it launches the same config the autotuner would. The full autotuner is still traced when the key is untuned, when it holds a symbolic size (a dynamic M the kernel tunes on, as forward_others' W4A16 tile kernels do), or when the pick has a `pre_hook`. **CUDA-graph timing for our autotune.** With that, the configs AOTI compiles in are our autotuner's picks rather than Inductor's. Inductor's timing is already CUDA-graph based during an export (`cuda_graph_timing()`), but `triton.autotune` times with `triton.testing.do_bench`, which brackets each call with CUDA events launched from a CPU that the export saturates; the readings of these short kernels become bimodal (see `autotune/cuda_graph_timing.py`). `autotune_do_bench` now times each config with `time_with_cuda_graph_us` (calls captured in one graph, each after an L2 flush), the timer the split candidates already use, on any capturable stream (the launch-parameter timing runs on one), and falls back to Triton's timing on the legacy default stream or if the capture fails. `Int4AutotunePickTest` (compile with every core busy, re-time every candidate afterwards with our spinners paused) now records our autotuner's decisions and checks the ones compiled in: with event timing they were up to 1.73x off the best (mean 1.06x, 4 of 34 decisions > 1.10x); with CUDA-graph timing, over three runs, the 8 compiled-in decisions are at most 1.023x off (mean 1.001-1.005x) and all 34 at most 1.026x. Full Muse Glimmer 30B batching export (`export_solo_batching.py`, 10 methods; A100; empty Triton and Inductor caches), measured on the stack with D124736179 (worker-process compiles) underneath, with the rest of the host equally busy (~15 cores): | | before | one traced config | + CUDA-graph timing | |---|---|---|---| | GEMM configs handed to Inductor | 2,032 (56 variants) | 56 | 56 | | Inductor benchmarking configs (min) | 3.2 | 1.0 | 1.0 | | Inductor waiting on its compile workers (min) | 16.0 | 8.8 | 8.8 | | our autotune benchmarking (min) | | 5.7 | 8.1 | | total (min) | 108.6 | 93.7 | 98.3 (81.5 without torch.export, vs 80.8) | Weights are byte-identical and the export generates the same tokens (greedy; 512 and 4 x 2k prompts) at the same speed (63.1 / 63.3 and 111.4 / 111.3 tok/s). Mixed prefill + decode step times (batch_bench defaults, T = 1 ... 4096, two interleaved rounds, CUDA graphs, no captures in any timed window): one traced config vs before, geomean time ratio 0.9997, every T within 1% except T = 9 and 16, 2.5% faster; CUDA-graph timing vs event timing, 1.0023, every T within 1.5%. Differential Revision: D124939685 Pull Request resolved: #23740
meta-codesync Bot
pushed a commit
that referenced
this pull request
Oct 18, 2026
…processes Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * __->__ #23787 * #23740 * #23739 * #23738 * #23737 * #23736 D124355785 compiles an autotuned quantized GEMM's configs on a thread pool before Triton benchmarks them. In one process, though, Triton's MLIR and LLVM passes (`make_ttgir`, `make_llir`, `make_ptx`) hold the GIL: during a bucket-64 shape only ~5 of the 32 compile threads were busy on average, and the export's main thread spent ~21 minutes waiting on that pool. Each async compile Triton submits to the pool now first runs the same `triton.compile(source, target, options, env_vars)` in a worker process, writing the kernel into the shared Triton cache; Triton's own compile call then runs as before in the parent and loads it from the cache (same source, options and cache key, so the same kernel). The workers are Inductor's `SubprocPool`: a fresh Python process that forks them, so they neither re-run the caller's main module (as `spawn` would) nor inherit its CUDA state (as `fork` would). They are started on first use and kept for later shapes, set `CUDA_VISIBLE_DEVICES=""` (compiling needs no GPU) and use the parent's Triton cache directory. If the workers cannot import the kernel (e.g. a binary without it) or die, compiles go back to the parent process for good, with a warning; a config that fails to compile in a worker is compiled in the parent too, where the benchmark skips it as before. The search, the timing and the picks are unchanged. One INT4 W4A16 bucket-64 shape (5376 x 4096), empty Triton cache, A100: 58.0 s -> 43.7 s, including the workers' start; 289 of the autotuned kernels the parent loads are cache hits from the workers. Full Muse Glimmer 30B batching export (`export_solo_batching.py`, 10 methods; A100; empty Triton and Inductor caches), stacked on D124713681 (min): | | before | this diff | |---|---|---| | main thread waiting on the compile pool | 20.9 | 12.2 | | AOTI forward_16 / forward_32 / forward_64 | 11.3 / 13.2 / 17.0 | 9.8 / 11.6 / 15.5 | | total | 125.7 | 123.3 | The pool wait drops by 8.7 min, but the end-to-end gain was 2.4 min in this run: other phases (Inductor's own compile wait, lowering, the first method) took a few minutes longer, while the rest of the host was ~3 cores busier than during the run before. The weight file is byte-identical, and the export generates the same tokens (greedy; 512 and 4 x 2k prompts) at the same speed (63.2 vs 62.5 tok/s; 110.4 vs 110.0). The quantized GEMM test suites, which compile many configs cold, also run faster: INT5/INT6/INT8 (+ CUDA graph timing) 30:05 -> 6:47, INT4/family/dispatch/launch params 26:40 -> 18:26 in OSS. Differential Revision: D124736179 Pull Request resolved: #23787
Gasoonjia
added a commit
that referenced
this pull request
Oct 18, 2026
…arallel before benchmarking them (#23948) 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](https://our.internmc.facebook.com/intern/diff/D124355785/) @diff-train-skip-merge Co-authored-by: gasoonjia <gasoonjia@icloud.com>
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
3 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.
Stack from ghstack (oldest at bottom):
At AOTI compile time every QuantizedGemmFamily kernel autotunes its configs per shape (and per split-K candidate, via
@autotune_launch_param).triton.autotunecompiles 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_gemmnow 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 ownAsyncCompileMode+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.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.ignore_errors=True).triton.runtime._async_compilekeeps 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 (
export_solo_batching.py: 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: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_benchitself (a fixed ~25 ms warmup + 100 ms per config), which this diff does not change.Differential Revision: D124355785