workaround for pytorch autotuning issue#1626
Open
dxqb wants to merge 1 commit into
Open
Conversation
…ibility torch.compile absorbs a traced @triton.autotune kernel and freezes the config benchmarked for the first shape, ignoring the autotune key (pytorch/pytorch#164124). Registering mm_8bit as an opaque custom op keeps it a single node in the compiled graph and runs its body eagerly, so Triton's autotuner selects per key. Autotune winners are cached to Triton's disk cache, and a line is printed whenever the autotuner actually benchmarks a kernel. Drafted by Claude Co-Authored-By: Claude Opus 4.8 (1M context) <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
pytorch/pytorch#164124 causes torch.compile to choose a first-fit of the triton kernel instead of a best-fit for each shape as intended
this PR defines the triton kernel a custom op, which makes it opaque to torch.compile and let's triton autotune the kernel.
first step goes from 40s to 50s in one test cold because autotuning now actually happens, but on second run the first step is still about 20s because triton autotunes are cached.
Test plan
pre-commit run --all-filespassesAI assistance