[BugFix][MoE] Guard fused mapping edge tiles#20
Open
JayceSu98 wants to merge 1 commit into
Open
Conversation
Full H100/CuTeDSL coverage exercises MoE routing shapes where the generated work loops are rounded up to alignment or warp-size boundaries. The logical tensors remain sized by the real hidden dimension, scaling-factor dimension, and token-topk count. The expand-to-fused kernel previously iterated over aligned hidden and scaling-factor extents and wrote every lane into the destination tensors. For tail dimensions this can address columns outside the real tensor shape, especially with TMA-aligned column-major scaling-factor layouts. The fused mapping kernel also used a Select expression whose inactive arm could still expose an out-of-range topk_idx load to code generation, and its fallback allocation estimate rounded the per-expert alignment upper bound down to an alignment multiple. Guard aligned expand stores with the real hidden/scaling-factor extents, initialize the shared cumsum buffer for all threads, avoid the speculative out-of-range topk load, and reserve a rounded-up per-expert alignment upper bound before trimming with the synchronized expert counts. Co-authored-by: dingsg <shengge.ding@enflame-tech.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.
Context
Full H100/CuTeDSL coverage exercises MoE routing shapes where the generated loops are rounded up to alignment or warp-size boundaries while the logical tensors keep the real token, hidden, and scaling-factor extents.
That exposed three TileKernels-side assumptions:
expand_to_fusedwrote all lanes inhidden_aligned/hidden_sf_aligned, even when the destination tensor only hashidden/hidden_sfcolumns.get_fused_mappingused aSelectexpression for the aligned tail. The inactive branch still exposed atopk_idx_1d[i]expression to lowering wheni >= numel.num_expanded_tokens == 0used an aligned expression that could round the per-expert alignment upper bound down.Fix
expand_to_fusedwrites with the real hidden and scaling-factor extents.T.cumsum.Selectload with an explicit guarded load into a localexpert_idxinitialized to-1.Why this belongs in TileKernels
The kernels intentionally schedule aligned work, but the memory operations still need to respect the logical tensor shape. These guards and allocation sizing changes make that relationship explicit and are independent of a particular CuTeDSL version.
Validation
Validated as part of the H100 PCIe CuTeDSL full correctness and benchmark run after the matching TileLang CuTeDSL fixes were applied.
Local sanity checks for this branch:
git diff --check origin/main..2dfff4eruff checkon the changed TileKernels Python filesJayceSu98 jayce.su@enflame-tech.com authored and validated this patch.
Co-authored-by: dingsg shengge.ding@enflame-tech.com