Snehalv dsv4 conv utility#4078
Draft
snehalv2002 wants to merge 22 commits into
Draft
Conversation
…near) Implement architectural core primitives required for DeepSeek-V4 integration into MaxText: - DeepSeekV4RMSNorm & DeepSeekV4UnweightedRMSNorm: RMS normalization layers utilizing float32 variance pooling. Includes unweighted scale-free variants that avoid allocating or synchronizing trainable weight parameters. - DeepSeekGroupedLinear: Block-diagonal grouped linear projection layer supporting parallel group projection via einsum broadcasting ([B, S, hc_mult, D] -> [B, S, D]). - DeepSeekV4RotaryEmbedding: Interleaved partial rotary positional embedding pairing consecutive even/odd channels. - Unit test suite (deepseek_v4_vs_reference_test.py) validating numerical parity against PyTorch reference implementations at atol=1e-5, rtol=1e-5.
…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.
…ghtningIndexer) Implement compressed attention mechanisms and indexer modules for DeepSeek-V4 integration into MaxText: - CSACompressor & HCACompressor: Long-range attention compressors supporting causal block bias and YaRN frequency scaling decoupling. - LightningIndexer: Memory-efficient indexer module implementing sentinel masking and dynamic RoPE scaling. - Configuration: Register attention compression hyperparameters (compress_ratios, index_head_dim, sliding_window) in types.py and base.yml. - Parity verification: Extended unit test suite (deepseek_v4_vs_reference_test.py) validating attention compression parity against PyTorch reference implementations at atol=1e-5, rtol=1e-5.
…ear projection layers from mainline
…variable tags to satisfy XLA compile rules
…ry expansion parameters for Trillium parity
…rporate DeepSeek-V4 top-k hash routers
…ked reference numerical parity
… across automated runner suites
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
2 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
Standalone and conversion util support for DeepSeek V4 Scanned.
FIXES: b/509930555
Tests
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.