[Tutorial] Remove incorrect caching from softmax tutorial #5162
Conversation
This file contains 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
The fused softmax implementation in the tutorial precompiles the kernel to query the register usage of the kernel, based on the parameters used to specialize the kernel. On top of this, it implements a simple caching system for this step based on just the block size. As noted in triton-lang#4739, this caching is incorrect, because it's also not keyed on the `num_stages` constexpr argument or the shapes of the tensors. Since triton already has its own JIT compilation cache, and this caching bit is not really relevant to the tutorial, just remove it to get rid of the footgun.
peterbell10
reviewed
Nov 15, 2024
| occupancy = min(occupancy, SIZE_SMEM // size_smem) | ||
| num_programs = NUM_SM * occupancy | ||
| kernels[BLOCK_SIZE] = (kernel, num_programs) | ||
| kernel = softmax_kernel.warmup(y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE, |
There was a problem hiding this comment.
I think there is still a bug here because warmup uses MockTensor:
triton/python/triton/runtime/jit.py
Line 729 in 1cf06c5
And MockTensor doesn't respect the real pointer's alignment:
triton/python/triton/runtime/jit.py
Lines 879 to 881 in 1cf06c5
This seems questionable though. I'm not sure why warmup couldn't just operate on the real tensors since it doesn't actually run any device code.
There was a problem hiding this comment.
Nice catch! Let me try digging into this
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.
The fused softmax implementation in the tutorial precompiles the kernel to query the register usage of the kernel, based on the parameters used to specialize the kernel. On top of this, it implements a simple caching system for this step based on just the block size.
As noted in #4739, this caching is incorrect, because it's also not keyed on the
num_stagesconstexpr argument or the shapes of the tensors. Since triton already has its own JIT compilation cache, and this caching bit is not really relevant to the tutorial, just remove it to get rid of the footgun.