Enable 2SM schedule for MXFP4_16 grouped GEMM (#501) - #501
Closed
ghjeong12 wants to merge 1 commit into
Closed
Conversation
|
@ghjeong12 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D101508914. |
Summary: Fix the 2SM NvF4 kernel schedule crash for MXFP4_16 (E8M0 scales with SFVecSize=16) by using non-void EpilogueElementC for 2SM tile configs. Root cause: the 2SM epilogue (PtrArrayTmaWarpSpecialized2Sm) requires non-void C type for correct stride layout configuration. With void C, the epilogue misconfigures memory access patterns, causing illegal memory access for M >= 192 per expert. NVFP4 worked because it always uses non-void C (for alpha_ptr_array global_scale support). Fix: set EpilogueElementC = ElementC (bfloat16) whenever 2SM is selected, not just for NVFP4. This adds minimal overhead (C tensor pointer is set but not used for the actual computation when global_scale is absent). Also updates MXFP4_16 tile configs to use native 2SM tiles instead of forcing 1SM fallback, enabling full 2SM performance. Reviewed By: jwfromm Differential Revision: D101508914
meta-codesync
Bot
changed the title
ghjeong12
force-pushed
the
export-D101508914
branch
from
August 25, 2026 21:36
1b6f0f5 to
d9cf4ff
Compare
|
This pull request has been merged in cd5d44e. |
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:
Fix the 2SM NvF4 kernel schedule crash for MXFP4_16 (E8M0 scales with
SFVecSize=16) by using non-void EpilogueElementC for 2SM tile configs.
Root cause: the 2SM epilogue (PtrArrayTmaWarpSpecialized2Sm) requires
non-void C type for correct stride layout configuration. With void C,
the epilogue misconfigures memory access patterns, causing illegal
memory access for M >= 192 per expert. NVFP4 worked because it always
uses non-void C (for alpha_ptr_array global_scale support).
Fix: set EpilogueElementC = ElementC (bfloat16) whenever 2SM is
selected, not just for NVFP4. This adds minimal overhead (C tensor
pointer is set but not used for the actual computation when
global_scale is absent).
Also updates MXFP4_16 tile configs to use native 2SM tiles instead
of forcing 1SM fallback, enabling full 2SM performance.
Reviewed By: jwfromm
Differential Revision: D101508914