Fix incorrect initialization of GenericFusedScaleMaskSoftmax #1586
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
GenericFusedScaleMaskSoftmax incorrectly sets self.scaled_mask_softmax_fusion to the generic_scaled_mask_softmax function. Its forward function calls self.fused_softmax_function however which means generic_scaled_mask_softmax is never called and standard FusedScaleMaskSoftmax will be used instead. Also fixed wrong import in GenericScaledMaskedSoftmax.
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.
Summary
Fixes an error in the initialization of GenericFusedScaleMaskSoftmax which leads to
generic_scaled_masked_softmaxnever actually being called.Fixes import error in
GenericScaledMaskedSoftmax.backwardwhich had also gone unnoticed.Also adds a test for GenericFusedScaleMaskSoftmax similar to the one for FusedScaleMaskSoftmax to prevent future errors.
Details
Currently when GenericFusedScaleMaskSoftmax is initialized we set
self.scaled_masked_softmax_fusionto to thegeneric_scaled_masked_softmaxfunction. I believe this is done in error as the function called by the fused forward isself.fused_softmax_funcand notself.scaled_masked_softmax_fusion(which is supposed to be the flag that indicates whether to use the fused version).In practice this means that when GenericFusedScaleMaskSoftmax is initialized it actually does not use
generic_scaled_masked_softmaxat all andis_kernel_availablewill return True even when initialized withscaled_masked_softmax_fusion = False.