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

docs/api/pytorch.rst

@@ -6,6 +6,11 @@
 PyTorch
 =======
 
+.. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None)
+
+Standard layers
+---------------
+
 .. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, bias=True, **kwargs)
   :members: forward, set_tensor_parallel_group
 
@@ -34,20 +39,34 @@ PyTorch
 .. autoapiclass:: transformer_engine.pytorch.TransformerLayer(hidden_size, ffn_hidden_size, num_attention_heads, **kwargs)
   :members: forward, set_context_parallel_group, set_tensor_parallel_group
 
+Model-specific layers
+---------------------
+
+DeepSeek-V3
+^^^^^^^^^^^
+
+.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(hidden_size, num_attention_heads, **kwargs)
+  :members: forward
+
+.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(hidden_size, moe_ffn_hidden_size, num_experts, **kwargs)
+  :members: forward, update_expert_bias
+
+.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(hidden_size, num_attention_heads, **kwargs)
+  :members: forward
+
+Other
+-----
+
 .. autoapiclass:: transformer_engine.pytorch.dot_product_attention.inference.InferenceParams(max_batch_size, max_sequence_length)
   :members: reset, allocate_memory, pre_step, get_seqlens_pre_step, convert_paged_to_nonpaged, step
 
 .. autoapiclass:: transformer_engine.pytorch.CudaRNGStatesTracker()
   :members: reset, get_states, set_states, add, fork
 
-
-.. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None)
-
 .. autoapifunction:: transformer_engine.pytorch.quantized_model_init
 
 .. autoapifunction:: transformer_engine.pytorch.checkpoint
 
-
 .. autoapifunction:: transformer_engine.pytorch.make_graphed_callables
 
 .. autoapifunction:: transformer_engine.pytorch.get_cpu_offload_context

qa/L0_pytorch_unittest/test.sh

@@ -58,6 +58,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_overrid
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cross_entropy.xml $TE_PATH/tests/pytorch/test_cross_entropy.py || test_fail "test_cross_entropy.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_models.xml $TE_PATH/tests/pytorch/test_models.py || test_fail "test_models.py"
 NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py"

qa/L1_pytorch_distributed_unittest/test.sh

@@ -56,6 +56,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_ep.xml $TE_PATH/tests/pytorch/distributed/test_ep.py || test_fail "test_ep.py"
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_models.xml $TE_PATH/tests/pytorch/distributed/test_models.py || test_fail "distributed/test_models.py"
 
 
 # debug tests

tests/pytorch/attention/test_linear_mxfp8_attention.py

@@ -36,7 +36,11 @@
 _current_file = pathlib.Path(__file__).resolve()
 sys.path = [str(_current_file.parent.parent)] + sys.path
 from utils import ModelConfig, compare_and_assert, get_available_attention_backends
-from mla_rope_utils import apply_mla_rope, build_rope_tables
+from transformer_engine.pytorch.models.deepseek_v3.mla_rope import (
+    apply_mla_rope_kv,
+    apply_mla_rope_q,
+    build_rope_tables,
+)
 
 
 try:
@@ -183,14 +187,21 @@ def _run_projections(
     return q_flat, kv_flat, q, kv, k_pos_emb
 
 
+def _apply_rope(q, kv, k_pos_emb, rope_tables):
+    cos, sin = rope_tables
+    q = apply_mla_rope_q(q, cos, sin, HEAD_DIM_NOPE, HEAD_DIM_ROPE)
+    k, v = apply_mla_rope_kv(kv, k_pos_emb, cos, sin, HEAD_DIM_NOPE, HEAD_DIM_ROPE, HEAD_DIM_V)
+    return q, k, v
+
+
 def _run_forward_bf16(
     modules: tuple,
     x: torch.Tensor,
     rope_tables: tuple[torch.Tensor, torch.Tensor],
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     q_proj, kv_proj, dpa, out_linear = modules
     _, _, q, kv, k_pos_emb = _run_projections(q_proj, kv_proj, x)
-    q, k, v = apply_mla_rope(q, kv, k_pos_emb, cos_table=rope_tables[0], sin_table=rope_tables[1])
+    q, k, v = _apply_rope(q, kv, k_pos_emb, rope_tables)
     attn_out = dpa(q, k, v, qkv_format="sbhd")
     return q, k, v, out_linear(attn_out.view(x.shape[0], x.shape[1], HIDDEN_SIZE))
 
@@ -212,13 +223,7 @@ def _run_forward_mxfp8(
             x,
             is_first_microbatch,
         )
-        q, k, v = apply_mla_rope(
-            q,
-            kv,
-            k_pos_emb,
-            cos_table=rope_tables[0],
-            sin_table=rope_tables[1],
-        )
+        q, k, v = _apply_rope(q, kv, k_pos_emb, rope_tables)
         attn_out = dpa(q, k, v, qkv_format="sbhd")
         out = out_linear(
             attn_out.view(x.shape[0], x.shape[1], HIDDEN_SIZE),
@@ -292,7 +297,7 @@ def test_accuracy(self, batch_size: int, seq_len: int) -> None:
         _set_seed()
         baseline_modules, mxfp8_modules = _build_modules()
         x = torch.randn(seq_len, batch_size, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
-        rope_tables = build_rope_tables(seq_len, device=x.device)
+        rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device)
 
         q_bf16, k_bf16, v_bf16, out_bf16 = _run_forward_bf16(baseline_modules, x, rope_tables)
         q_mxfp8, k_mxfp8, v_mxfp8, out_mxfp8 = _run_forward_mxfp8(
@@ -378,7 +383,7 @@ def test_backward(self, batch_size: int, seq_len: int) -> None:
             device="cuda",
             requires_grad=True,
         )
-        rope_tables = build_rope_tables(seq_len, device=x.device)
+        rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device)
 
         *_, out_mxfp8 = _run_forward_mxfp8(mxfp8_modules, x, fp8_recipe, rope_tables)
         out_mxfp8.sum().backward()
@@ -412,7 +417,7 @@ def test_performance(self, batch_size: int, seq_len: int) -> None:
             device="cuda",
             requires_grad=True,
         )
-        rope_tables = build_rope_tables(seq_len, device=x.device)
+        rope_tables = build_rope_tables(seq_len, HEAD_DIM_ROPE, device=x.device)
 
         mxfp8_fprop_ms, mxfp8_bprop_ms = _benchmark_training_step(
             _run_forward_mxfp8, mxfp8_modules, x, fp8_recipe, rope_tables

tests/pytorch/distributed/run_models.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""Multi-process tests for model-specific layers (te.models), launched via torchrun."""
+
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+from transformer_engine.pytorch.ep import ep_bootstrap, ep_finalize, release_symm_mem_pool
+from transformer_engine.pytorch.models import DeepSeekV3Layer
+
+HIDDEN = 256
+MOE_FFN = 128
+SHARED_FFN = 128
+NUM_LOCAL_EXPERTS = 2
+TOP_K = 2
+TOKENS_PER_RANK = 64
+HEADS = 4
+DTYPE = torch.bfloat16
+
+MLA_KWARGS = dict(
+    q_lora_rank=96,
+    kv_lora_rank=64,
+    qk_nope_head_dim=64,
+    qk_rope_head_dim=32,
+    v_head_dim=64,
+)
+
+
+def _device_sm() -> int:
+    major, minor = torch.cuda.get_device_capability()
+    return major * 10 + minor
+
+
+def _recv_capacity(ep_size: int) -> int:
+    cap = ep_size * TOKENS_PER_RANK * TOP_K + NUM_LOCAL_EXPERTS * 128
+    return -(-cap // 128) * 128
+
+
+def _broadcast_params(module: torch.nn.Module) -> None:
+    for t in list(module.parameters()) + list(module.buffers()):
+        dist.broadcast(t.detach(), src=0)
+
+
+def _make_layer(ep_group, ep_size: int, num_experts: int) -> DeepSeekV3Layer:
+    ep = ep_group is not None
+    return DeepSeekV3Layer(
+        HIDDEN,
+        HEADS,
+        num_experts=num_experts,
+        moe_ffn_hidden_size=MOE_FFN,
+        topk=TOP_K,
+        shared_expert_ffn_hidden_size=SHARED_FFN,
+        params_dtype=DTYPE,
+        ep_group=ep_group,
+        ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None,
+        **MLA_KWARGS,
+    )
+
+
+def _copy_weights(ep_layer: DeepSeekV3Layer, ref: DeepSeekV3Layer, rank: int) -> None:
+    ref_params = dict(ref.named_parameters())
+    ref_bufs = dict(ref.named_buffers())
+    with torch.no_grad():
+        for name, p in ep_layer.named_parameters():
+            if not name.startswith("mlp.experts."):
+                p.copy_(ref_params[name])
+        for name, b in ep_layer.named_buffers():
+            if name in ref_bufs and b.shape == ref_bufs[name].shape:
+                b.copy_(ref_bufs[name])
+        ep_fc1, _, ep_fc2 = ep_layer.mlp.experts
+        ref_fc1, _, ref_fc2 = ref.mlp.experts
+        for local_e in range(NUM_LOCAL_EXPERTS):
+            global_e = rank * NUM_LOCAL_EXPERTS + local_e
+            getattr(ep_fc1, f"weight{local_e}").copy_(getattr(ref_fc1, f"weight{global_e}"))
+            getattr(ep_fc2, f"weight{local_e}").copy_(getattr(ref_fc2, f"weight{global_e}"))
+
+
+def test_layer_ep_matches_local(rank: int, ep_size: int, ep_group) -> None:
+    """Full DeepSeekV3Layer with EP must match the all-experts-local layer numerically."""
+    num_experts = NUM_LOCAL_EXPERTS * ep_size
+    torch.manual_seed(0)
+    ref = _make_layer(None, ep_size, num_experts)
+    _broadcast_params(ref)
+    ep_layer = _make_layer(ep_group, ep_size, num_experts)
+    _copy_weights(ep_layer, ref, rank)
+
+    torch.manual_seed(1234 + rank)
+    x = torch.randn(TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda")
+    x_ep = x.clone().requires_grad_(True)
+    x_ref = x.clone().requires_grad_(True)
+
+    out_ep = ep_layer(x_ep)
+    out_ref = ref(x_ref)
+    assert out_ep.shape == x.shape
+    torch.testing.assert_close(out_ep, out_ref, rtol=0.05, atol=0.05)
+
+    grad_out = torch.randn_like(out_ep)
+    out_ep.backward(grad_out)
+    out_ref.backward(grad_out)
+    torch.testing.assert_close(x_ep.grad, x_ref.grad, rtol=0.05, atol=0.05)
+
+    ref_params = dict(ref.named_parameters())
+    for name, p in ep_layer.named_parameters():
+        if name.startswith("mlp.experts.") or p.grad is None:
+            continue
+        torch.testing.assert_close(p.grad, ref_params[name].grad, rtol=0.1, atol=0.1, msg=name)
+
+    # A local expert's wgrad on its owner rank equals the sum of the
+    # reference wgrads over all ranks. all_reduce is collective, so every
+    # rank must reduce every expert's grad (in the same order).
+    ep_fc1, _, ep_fc2 = ep_layer.mlp.experts
+    ref_fc1, _, ref_fc2 = ref.mlp.experts
+    for ep_fc, ref_fc in ((ep_fc1, ref_fc1), (ep_fc2, ref_fc2)):
+        ref_grads = [getattr(ref_fc, f"weight{e}").grad.float().clone() for e in range(num_experts)]
+        for g in ref_grads:
+            dist.all_reduce(g)
+        for local_e in range(NUM_LOCAL_EXPERTS):
+            global_e = rank * NUM_LOCAL_EXPERTS + local_e
+            ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float()
+            torch.testing.assert_close(ep_grad, ref_grads[global_e], rtol=0.1, atol=0.1)
+
+    counts = ep_layer.mlp._last_tokens_per_expert.clone()
+    dist.all_reduce(counts)
+    assert counts.sum().item() == ep_size * TOKENS_PER_RANK * TOP_K
+
+    ep_layer.mlp.update_expert_bias()
+    assert torch.isfinite(ep_layer.mlp.expert_bias).all()
+
+
+def main() -> int:
+    dist.init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    from torch.distributed import _symmetric_memory as _symm_mem
+
+    _symm_mem.set_backend("NCCL")
+
+    rank = dist.get_rank()
+    ep_size = dist.get_world_size()
+    if _device_sm() < 90:
+        if rank == 0:
+            print(f"NCCL EP requires SM>=90 (got SM{_device_sm()}); skipping.")
+        dist.destroy_process_group()
+        return 0
+
+    ep_group = dist.new_group(ranks=list(range(ep_size)), backend="nccl")
+    ep_bootstrap(
+        ep_group,
+        num_experts=NUM_LOCAL_EXPERTS * ep_size,
+        max_tokens_per_rank=TOKENS_PER_RANK,
+        hidden_dim=HIDDEN,
+        num_topk=TOP_K,
+        recv_capacity_per_rank=_recv_capacity(ep_size),
+    )
+    test_layer_ep_matches_local(rank, ep_size, ep_group)
+    print(f"[rank {rank}] PASSED")
+
+    dist.barrier()
+    ep_finalize()
+    release_symm_mem_pool()
+    dist.destroy_process_group()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/pytorch/distributed/test_models.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import os
+import subprocess
+from pathlib import Path
+
+import pytest
+import torch
+
+TEST_ROOT = Path(__file__).parent.resolve()
+NUM_PROCS = min(8, torch.cuda.device_count())
+LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"]
+
+
+def _has_nvlink() -> bool:
+    # NCCL EP falls back to the network transport and deadlocks on PCIe-only nodes.
+    out = subprocess.run(
+        ["nvidia-smi", "nvlink", "--status"], capture_output=True, text=True, check=False
+    ).stdout
+    return "GB/s" in out
+
+
+@pytest.mark.skipif(NUM_PROCS < 2, reason="EP requires >= 2 GPUs")
+@pytest.mark.skipif(not _has_nvlink(), reason="NCCL EP requires NVLink")
+def test_deepseek_layer_ep():
+    result = subprocess.run(
+        LAUNCH_CMD + [str(TEST_ROOT / "run_models.py")], env=os.environ, check=False, timeout=300
+    )
+    assert result.returncode == 0