Optimize softmax #69
Closed
Optimize softmax #69
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
1. ensure that decorator cascading is working as expected, i.e. inner decorator can use arguments provided by outer decorator 2. ensure that grid function can use all the arguments provided by decorators(Autotuner & Heuristics) 3. simply LibEntry, extract captured constant arguments from CompiledKernel, instead of traversing layers of decorator.
…e ONE_TILE_PER_CTA static condition(to decide whether to load only one tile per cta.
…e ONE_TILE_PER_CTA static condition(to decide whether to load only one tile per cta.
|
need to rebase after merging pr68 |
StrongSpoon
reviewed
Jun 18, 2024
| @@ -535,8 +535,8 @@ def _torch_rms_norm(x, residual, weight, eps): | |||
|
|
|||
| @pytest.mark.parametrize("shape", REDUCTION_SHAPES) | |||
| @pytest.mark.parametrize("dtype", FLOAT_DTYPES) | |||
| def test_accuracy_softmax(shape, dtype): | |||
| dim = 1 | |||
| @pytest.mark.parametrize("dim", [0, 1]) | |||
There was a problem hiding this comment.
"dim" passed to the function here is an integer.
Do you mean we should add test for non-inner reduction in general?
I didn't resolve the conflicts properly. So I opened a new pull request #76 . |
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.
a. inner: for reduction the last dim(and the input is preprocessed to be contiguous)
a. inner: for reduce along other dimensions(and the input is preprocessed to be contiguous)
ONE_TILE_PER_CTAstatic conditiona. when
ONE_TILE_PER_CTAis True, load only one tile per cta without looping over reduction dimb. when
ONE_TILE_PER_CTAis False, use online softmax normalizer algorithm to save one swipe over the input.We can leave other optimizations(optimize accordign to input layout or two-pass-reduction) for future PRs