Add bf16/fp32 token-per-expert to the MoE aux loss kernel #2162
Conversation
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
…fusion Signed-off-by: tongliu <[email protected]>
Autumn1998
force-pushed
the
tongliu/router_fusion
branch
from
5708f20 to
de887b4
Compare
for more information, see https://pre-commit.ci
|
/te-ci |
yaox12
changed the title
yaox12
approved these changes
Sep 9, 2025
vthumbe1503
added a commit
to vthumbe1503/TransformerEngine
that referenced
this pull request
Sep 19, 2025
Signed-off-by: vthumbe1503 <[email protected]> Add bf16/fp32 token-per-expert to the MoE aux loss kernel (NVIDIA#2162) * add bf16/fp32 token-per-expert on the moe-loss-computation on router fusion Signed-off-by: tongliu <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: tongliu <[email protected]> Co-authored-by: tongliu <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> [JAX] Scale swizzling via JAX transpose op (NVIDIA#2163) * add swizzle in jax Signed-off-by: Phuong Nguyen <[email protected]> * added outer_impl Signed-off-by: Phuong Nguyen <[email protected]> * clean up FFI Signed-off-by: Phuong Nguyen <[email protected]> --------- Signed-off-by: Phuong Nguyen <[email protected]> Extract cpp distributed tests into a separate project (NVIDIA#2165) * Extract cpp distributed tests into a separate project Signed-off-by: Vladimir Cherepanov <[email protected]> * Remove obsolete exclusion Signed-off-by: Vladimir Cherepanov <[email protected]> * Run L1_cpp_distributed tests if at least 4 GPUs Signed-off-by: Vladimir Cherepanov <[email protected]> --------- Signed-off-by: Vladimir Cherepanov <[email protected]> Adds context parallelism utilities: moving cp shards to diff ranks and pad sequence to divisibility factory (NVIDIA#2129) * test - adds unit test for cp utilities and the utilites Signed-off-by: Jonathan Mitchell <[email protected]> * assert line change Signed-off-by: Jonathan Mitchell <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jonathan Mitchell <[email protected]> Co-authored-by: Jonathan Mitchell <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sudhakar Singh <[email protected]>
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.
Description
Add the new datatype(1. bf16, 2. fp32) support on the token-per-expert on the moe-aux-loss-computation kernel.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: