[DeepSeek-V4] Implement MoE routing primitives (HashRouter, TopKRouter, RoutedMoE)#3871
Open
parambole wants to merge 1 commit into
Open
[DeepSeek-V4] Implement MoE routing primitives (HashRouter, TopKRouter, RoutedMoE)#3871parambole wants to merge 1 commit into
parambole wants to merge 1 commit into
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
parambole
force-pushed
the
dsv4-moe-routing-primitives
branch
from
37ee811 to
31329c5
Compare
parambole
force-pushed
the
deepseek_v4_core_primitives
branch
from
1ab79e5 to
c025463
Compare
parambole
force-pushed
the
dsv4-moe-routing-primitives
branch
from
31329c5 to
22a57ff
Compare
parambole
force-pushed
the
deepseek_v4_core_primitives
branch
from
c025463 to
68c44a6
Compare
parambole
force-pushed
the
dsv4-moe-routing-primitives
branch
from
22a57ff to
32869e5
Compare
parambole
force-pushed
the
deepseek_v4_core_primitives
branch
2 times, most recently
from
72a92a7 to
e81f52d
Compare
…outer, RoutedMoE) Implement Mixture of Experts routing gates and execution layers for DeepSeek-V4 integration into MaxText: - HashRouter: Token routing mechanism utilizing MD5 hash projections for deterministic expert assignment. - TopKRouter: Gated top-k router implementing sigmoid scaling and score normalization. - RoutedMoE & RoutedAndSharedMoE: Execution layers supporting layer_idx routing and FP32 expert summation parity. - Parity verification: Extended unit test suite (deepseek_v4_vs_reference_test.py) validating routing parity against PyTorch reference implementations at atol=1e-5, rtol=1e-5.
parambole
force-pushed
the
dsv4-moe-routing-primitives
branch
from
32869e5 to
c92f2e0
Compare
parambole
changed the title
parambole
changed the title
|
🤖 Hi @parambole, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
parambole
requested review from
A9isha,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jesselu-google,
khatwanimohit,
michelle-yooh,
richjames0,
shralex,
suexu1025 and
vipannalla
as code owners
parambole
requested review from
NicoGrande,
NuojCheng,
aireenmei,
dipannita08,
igorts-git and
jiangjy1982
as code owners
parambole
assigned
entrpn,
RissyRan,
shuningjin,
gagika,
snehalv2002,
Rohan-Bierneni and
bvandermoon
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
8 participants
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
Implement Mixture of Experts (MoE) routing gates and execution layers required for DeepSeek-V4 integration into MaxText:
HashRouter: Token routing mechanism utilizing MD5 hash projections for deterministic expert assignment without auxiliary loss.TopKRouter: Gated top-k router implementing sigmoid scaling and score normalization across selected experts.RoutedMoE&RoutedAndSharedMoE: Execution layers supportinglayer_idxrouting, gate clamping, and FP32 expert summation parity.tests/unit/deepseek_v4_vs_reference_test.py) validating MoE routing parity against PyTorch reference implementations atatol=1e-5, rtol=1e-5.Tests
Tested on CPU:
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.