[PyTorch] DeepSeekV3Layer: full MoE transformer layer (MLA + DeepSeek MoE)

[pggPL]

Description

Adds transformer_engine.pytorch.models — a namespace for model-specific layers composed from TE modules — with a DeepSeekV3 transformer layer analogous to TransformerLayer, exercising TE's MoE and MLA features end to end.

MultiLatentAttention (models/deepseek_v3/multi_latent_attention.py): low-rank q/kv latents with RMSNorm fused into the up-projections (LayerNormLinear(normalization="RMSNorm")), decoupled RoPE/NoPE head split with a single shared key rope head (fused RoPE via apply_rotary_pos_emb(fused=True)), core attention through DotProductAttention with asymmetric head dims kv_channels=(qk_nope+qk_rope, v) so the cuDNN fused attention backend is used where supported (SM100+ for 192/128 training). TP via column/row parallel projections.

DeepSeekV3MoE (models/deepseek_v3/moe.py): fused sigmoid router with aux-loss-free expert bias and node-limited grouped top-k (fused_topk_with_score_function) plus an update_expert_bias() helper; routed experts as te.ops.Sequential(GroupedLinear, ScaledSwiGLU(32), GroupedLinear) which auto-fuses into the CuTe grouped-GEMM MLP on supported hardware and runs the identical unfused path elsewhere; routing probs applied per-token inside the activation, so merge is plain accumulation. Token routing either local (moe_permute_with_probs/moe_unpermute) or expert-parallel over NCCL (ep_dispatch/ep_combine with EpBuffer). Optional shared expert (te.ops SwiGLU MLP).

DeepSeekV3Layer (models/deepseek_v3/transformer_layer.py): pre-RMSNorm + MLA, then dense LayerNormMLP (RMSNorm+swiglu, first dense layers) or DeepSeekV3MoE, with residual connections.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• New transformer_engine/pytorch/models/ namespace with models/deepseek_v3/ subpackage (MultiLatentAttention, DeepSeekV3MoE, DeepSeekV3Layer), exported via transformer_engine.pytorch.models
• Docs: Model-specific layers section in docs/api/pytorch.rst
• Tests: tests/pytorch/test_deepseek.py (MLA fwd/bwd, MoE variants incl. grouped routing + shared expert, numeric check against a dense reference, dense/MoE layer fwd/bwd) — 8 passed on RTX Ada; distributed suite skips cleanly there (SM89, 1 GPU)
• Distributed EP test: tests/pytorch/distributed/test_deepseek_ep.py + run_deepseek_ep.py (EP vs all-local numeric equivalence incl. all-reduced expert wgrads, full-layer EP smoke; verified on 4x GB300 (lyris, arm64, SM103): all ranks pass)
• TODO (follow-ups): CUDA-graph/torch.compile coverage, YaRN rope scaling, optional seq-wise aux loss

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code