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

[pggPL]

Description

Adds transformer_engine.pytorch.models, a namespace for model-specific layers built from TE modules, with a DeepSeek-V3 transformer layer analogous to TransformerLayer.

• MultiLatentAttention: low-rank q/kv latents with RMSNorm fused into the up-projections, decoupled RoPE/NoPE heads, DotProductAttention with asymmetric head dims (cuDNN fused attention), fused Triton MLA RoPE kernels, optional YaRN scaling, TP.
• DeepSeekV3MoE: fused sigmoid router with aux-loss-free bias and grouped top-k, experts as te.ops.Sequential(GroupedLinear, ScaledSwiGLU, GroupedLinear) (auto-fuses into the CuTe grouped MLP where available), local routing or expert parallelism over NCCL EP, optional shared expert.
• DeepSeekV3Layer: pre-RMSNorm + MLA, then dense LayerNormMLP or DeepSeekV3MoE, 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

• transformer_engine/pytorch/models/deepseek_v3/: MultiLatentAttention, DeepSeekV3MoE, DeepSeekV3Layer, MLA RoPE kernels; exported via transformer_engine.pytorch.models
• Docs: Model-specific layers section in docs/api/pytorch.rst
• Tests: tests/pytorch/test_models.py (RoPE, YaRN, MoE vs dense reference), test_sanity.py::test_sanity_deepseek_v3_layer (all recipes), tests/pytorch/distributed/test_models.py (EP vs all-experts-local, L1)
• examples/pytorch/deepseek_v3/: runnable EP example with a plain-PyTorch MoE baseline and a README with measurements
• Follow-up: optional seq-wise aux loss

Performance

DeepSeekV3Layer with DeepSeek-V3 dims, 8 experts per rank, top-k 8, 4096 tokens per rank, fwd+bwd, bf16 unless noted. GB300, 4 GPUs per node, one NVLink domain. naive = torch all_to_all_single + Python loop of per-expert MLPs, naive_grouped = same all_to_all with a TE grouped GEMM, te = DeepSeekV3MoE. Attention and norms identical in all three.

ms / iteration (median of 3 runs)	4 GPUs	8 GPUs
naive	27.0	27.4
naive_grouped	16.5	16.9
te	12.4	14.0
te, MXFP8	10.3	12.1
te, MXFP8 fused grouped MLP	9.5	10.5

Where the time goes in the optimized variant (te, MXFP8 fused, 8 GPUs), per GPU and iteration, from nsys stats --report cuda_gpu_kern_sum (kernel time 10.8 ms, iteration 10.5 ms without the profiler; no memsets or copies left):

Group	ms	Kernels
NCCL EP all-to-all	2.3	nccl_ep_jit_ht_dispatch_kernel (1.05), nccl_ep_jit_ht_combine_kernel (1.29), each twice per iteration (fwd + bwd)
NCCL EP local permute	0.6	local_permute_dup/reduce: staging buffer to expert-major layout, zero-filled padding
rank skew wait	0.5	ncclDevKernel_AllGather_RING_LL (routing-map all-gather in prepare, ~0.06 ms of transfer); the first collective of the layer absorbs load imbalance between ranks
fused grouped MLP (cuDNN, MXFP8)	2.4	fc1+SwiGLU fwd (0.53), fc2 fwd (0.82), dGLU bwd (0.31), wgrad (0.78)
MXFP8 quantization	1.1	group_quantize_mxfp8 on the recv buffer (0.4), quantize_mxfp8_kernel_cast_only for dense GEMM inputs (0.7)
dense MXFP8 GEMMs (MLA projections, shared expert)	1.4	nvjet_sm103_qqtst_*
attention (cuDNN SDPA)	0.9	flash fprop (0.18) + bprop (0.51) + dq / dO helpers
RMSNorm, RoPE, adds	0.8	rmsnorm_fwd/bwd (0.39), rotary_*_kv (0.18), residual adds (0.2)

For comparison the naive variant spends 3.4 ms in all_to_all, 6.5 ms in 8 separate GEMM pairs, and ~12 ms in indexing, zero-fills, copies and adds; naive_grouped removes the loop (about 10.5 ms) and NCCL EP the remaining all_to_all plus sort / gather / scatter overhead (2.9 ms on 8 GPUs). Full breakdowns in examples/pytorch/deepseek_v3/README.md.

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

[greptile-apps]

Confidence Score: 4/5

[Medium risk] Adds DeepSeek-V3 transformer layer models and MoE components.

The PR does not appear safe to merge until EP model construction works before bootstrap; the outstanding benchmark-provenance concern is non-blocking.

Findings

1.  Model construction requires bootstrap ▶
2.  Benchmark provenance removed ▶

Summary

The PR adds DeepSeek-V3 PyTorch model layers, MLA RoPE support, expert-parallel MoE composition, documentation, examples, and tests. Changes since the previous review also add EP reference tests, NVFP4 amax kernels, and optional FlashAttention import handling.

_{Reviews (8) · Last reviewed commit: "Compose DeepSeek expert parallel path wi..."}

[pggPL]

Hi @vthumbe1503 @phu0ngng this is PR with DeepSeek layer which uses fused MLP and EP, can you have a look?

[cyanguwa]

I wonder if this file should sit together with our other RoPE implementations, i.e. in transformer_engine/pytorch/attention? I feel this models/ directory should contain only the model-level/higher-level implementations.

[vthumbe1503]

I think it might make sense to use the te.Sequential to define the entire MOE block after this PR gets merged.
#3503
Would be more cleaner from API usage standpoint

[greptile-apps]

Benchmark provenance removed

The benchmark results no longer identify when they were measured or which commit was used. That makes the reported timings harder to reproduce or compare with later changes. Please keep the measurement date and commit alongside the results.

Suggested change

	Every number is the median of three independent runs (spread within 0.3 ms).
	Every number is the median of three independent runs (spread within 0.3 ms).
	Measured on September 30, 2026, at commit
	[`939c9db3`](https://github.com/NVIDIA/TransformerEngine/commit/939c9db36afcdb2617392bab57b4249c7cd4bcb0),
	before the subsequent merge of `main`.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Model construction requires bootstrap

If a caller builds a DeepSeekV3MoE with ep_group before calling ep_bootstrap(), this new check raises during construction. Previously, bootstrap was required when creating the EP buffer for forward, so callers could build the model first and bootstrap before execution. That workflow can no longer instantiate the layer.