From 45ed767a7af61a64dc41e29be1b6c0cc14c21d9c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 12:37:21 +0200 Subject: [PATCH 01/43] [PyTorch] Add DeepSeekV3Layer skeleton (MLA + MoE) Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/deepseek/__init__.py | 11 ++++++++ transformer_engine/pytorch/deepseek/moe.py | 27 +++++++++++++++++++ .../deepseek/multi_latent_attention.py | 24 +++++++++++++++++ .../pytorch/deepseek/transformer_layer.py | 25 +++++++++++++++++ 4 files changed, 87 insertions(+) create mode 100644 transformer_engine/pytorch/deepseek/__init__.py create mode 100644 transformer_engine/pytorch/deepseek/moe.py create mode 100644 transformer_engine/pytorch/deepseek/multi_latent_attention.py create mode 100644 transformer_engine/pytorch/deepseek/transformer_layer.py diff --git a/transformer_engine/pytorch/deepseek/__init__.py b/transformer_engine/pytorch/deepseek/__init__.py new file mode 100644 index 0000000000..5dafdf1fff --- /dev/null +++ b/transformer_engine/pytorch/deepseek/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 transformer layer built from Transformer Engine MoE building blocks.""" + +from transformer_engine.pytorch.deepseek.multi_latent_attention import MultiLatentAttention +from transformer_engine.pytorch.deepseek.moe import DeepSeekV3MoE +from transformer_engine.pytorch.deepseek.transformer_layer import DeepSeekV3Layer + +__all__ = ["DeepSeekV3Layer", "DeepSeekV3MoE", "MultiLatentAttention"] diff --git a/transformer_engine/pytorch/deepseek/moe.py b/transformer_engine/pytorch/deepseek/moe.py new file mode 100644 index 0000000000..f4f787743b --- /dev/null +++ b/transformer_engine/pytorch/deepseek/moe.py @@ -0,0 +1,27 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 MoE block: sigmoid router with aux-loss-free bias, shared + +routed experts.""" + +import torch + +__all__ = ["DeepSeekV3MoE"] + + +class DeepSeekV3MoE(torch.nn.Module): + """ + DeepSeekV3-style Mixture of Experts block composed from TE MoE + primitives: ``fused_topk_with_score_function`` (sigmoid score function, + expert bias, grouped top-k), ``moe_permute_with_probs``/``moe_unpermute``, + :class:`GroupedLinear` routed experts, a shared expert + (:class:`LayerNormMLP`), ``Fp8Padding``/``Fp8Unpadding`` and optional + expert parallelism via ``ep_dispatch``/``ep_combine``. + + .. warning:: Work in progress, not functional yet. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + raise NotImplementedError("DeepSeekV3MoE is under development") diff --git a/transformer_engine/pytorch/deepseek/multi_latent_attention.py b/transformer_engine/pytorch/deepseek/multi_latent_attention.py new file mode 100644 index 0000000000..6c2bb7420b --- /dev/null +++ b/transformer_engine/pytorch/deepseek/multi_latent_attention.py @@ -0,0 +1,24 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Multi-Latent Attention (MLA) block as used in DeepSeekV3.""" + +import torch + +__all__ = ["MultiLatentAttention"] + + +class MultiLatentAttention(torch.nn.Module): + """ + Multi-Latent Attention with low-rank Q/KV down-projections and a + decoupled RoPE/NoPE head split, composed from :class:`Linear`, + :class:`LayerNormLinear` and :class:`DotProductAttention` + (``kv_channels=(head_dim_qk, head_dim_v)``). + + .. warning:: Work in progress, not functional yet. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + raise NotImplementedError("MultiLatentAttention is under development") diff --git a/transformer_engine/pytorch/deepseek/transformer_layer.py b/transformer_engine/pytorch/deepseek/transformer_layer.py new file mode 100644 index 0000000000..2a28a6ceb3 --- /dev/null +++ b/transformer_engine/pytorch/deepseek/transformer_layer.py @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DeepSeekV3 transformer layer.""" + +import torch + +__all__ = ["DeepSeekV3Layer"] + + +class DeepSeekV3Layer(torch.nn.Module): + """ + A full DeepSeekV3 transformer layer, analogous to + :class:`TransformerLayer`: :class:`MultiLatentAttention` followed by + either a dense :class:`LayerNormMLP` (first layers) or + :class:`DeepSeekV3MoE`, with the same residual and fused + bias-dropout-add plumbing as :class:`TransformerLayer`. + + .. warning:: Work in progress, not functional yet. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + raise NotImplementedError("DeepSeekV3Layer is under development") From c306c6f840bfe187cbdc3ecad39ad5aa517b669e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 13:56:21 +0200 Subject: [PATCH 02/43] Move DeepSeekV3 skeleton to models/deepseek_v3 subpackage Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/models/__init__.py | 13 +++++++++++++ .../{deepseek => models/deepseek_v3}/__init__.py | 8 +++++--- .../pytorch/{deepseek => models/deepseek_v3}/moe.py | 0 .../deepseek_v3}/multi_latent_attention.py | 0 .../deepseek_v3}/transformer_layer.py | 0 5 files changed, 18 insertions(+), 3 deletions(-) create mode 100644 transformer_engine/pytorch/models/__init__.py rename transformer_engine/pytorch/{deepseek => models/deepseek_v3}/__init__.py (50%) rename transformer_engine/pytorch/{deepseek => models/deepseek_v3}/moe.py (100%) rename transformer_engine/pytorch/{deepseek => models/deepseek_v3}/multi_latent_attention.py (100%) rename transformer_engine/pytorch/{deepseek => models/deepseek_v3}/transformer_layer.py (100%) diff --git a/transformer_engine/pytorch/models/__init__.py b/transformer_engine/pytorch/models/__init__.py new file mode 100644 index 0000000000..bee5474c81 --- /dev/null +++ b/transformer_engine/pytorch/models/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Model-specific transformer layers composed from Transformer Engine modules.""" + +from transformer_engine.pytorch.models.deepseek_v3 import ( + DeepSeekV3Layer, + DeepSeekV3MoE, + MultiLatentAttention, +) + +__all__ = ["DeepSeekV3Layer", "DeepSeekV3MoE", "MultiLatentAttention"] diff --git a/transformer_engine/pytorch/deepseek/__init__.py b/transformer_engine/pytorch/models/deepseek_v3/__init__.py similarity index 50% rename from transformer_engine/pytorch/deepseek/__init__.py rename to transformer_engine/pytorch/models/deepseek_v3/__init__.py index 5dafdf1fff..a7cbb50ae2 100644 --- a/transformer_engine/pytorch/deepseek/__init__.py +++ b/transformer_engine/pytorch/models/deepseek_v3/__init__.py @@ -4,8 +4,10 @@ """DeepSeekV3 transformer layer built from Transformer Engine MoE building blocks.""" -from transformer_engine.pytorch.deepseek.multi_latent_attention import MultiLatentAttention -from transformer_engine.pytorch.deepseek.moe import DeepSeekV3MoE -from transformer_engine.pytorch.deepseek.transformer_layer import DeepSeekV3Layer +from transformer_engine.pytorch.models.deepseek_v3.multi_latent_attention import ( + MultiLatentAttention, +) +from transformer_engine.pytorch.models.deepseek_v3.moe import DeepSeekV3MoE +from transformer_engine.pytorch.models.deepseek_v3.transformer_layer import DeepSeekV3Layer __all__ = ["DeepSeekV3Layer", "DeepSeekV3MoE", "MultiLatentAttention"] diff --git a/transformer_engine/pytorch/deepseek/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py similarity index 100% rename from transformer_engine/pytorch/deepseek/moe.py rename to transformer_engine/pytorch/models/deepseek_v3/moe.py diff --git a/transformer_engine/pytorch/deepseek/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py similarity index 100% rename from transformer_engine/pytorch/deepseek/multi_latent_attention.py rename to transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py diff --git a/transformer_engine/pytorch/deepseek/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py similarity index 100% rename from transformer_engine/pytorch/deepseek/transformer_layer.py rename to transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py From f73d04edaf0e16740f428b537c78d70d0c9c1ec1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 14:00:07 +0200 Subject: [PATCH 03/43] Add DeepSeekV3 layer entries to PyTorch API docs Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 5fac0a89a6..bd3099b590 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -59,6 +59,15 @@ PyTorch .. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor +Model-specific layers +--------------------- + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(**kwargs) + +.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(**kwargs) + +.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(**kwargs) + Data types ---------- From 09f28a9a3903b85ea28acff8ef63149738f38ea8 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 14:13:44 +0200 Subject: [PATCH 04/43] [PyTorch] Implement DeepSeekV3Layer: MLA + MoE from TE building blocks MultiLatentAttention: low-rank q/kv latents (RMSNorm fused into LayerNormLinear up-projections), decoupled RoPE/NoPE head split with a shared key rope head, DotProductAttention with kv_channels=(qk, v) for the cuDNN fused backend. DeepSeekV3MoE: fused sigmoid router with aux-loss-free expert bias and grouped top-k, routed experts as te.ops GroupedLinear+ScaledSwiGLU+ GroupedLinear (CuTe fused grouped MLP on supported HW), probs applied per-token in the activation, local permute/unpermute or NCCL expert parallelism via ep_dispatch/ep_combine, optional shared expert. DeepSeekV3Layer: pre-RMSNorm + MLA and dense LayerNormMLP (RMSNorm, swiglu) or MoE with residual connections. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_deepseek.py | 124 +++++++++ transformer_engine/pytorch/__init__.py | 1 + .../pytorch/models/deepseek_v3/moe.py | 245 +++++++++++++++++- .../deepseek_v3/multi_latent_attention.py | 193 +++++++++++++- .../models/deepseek_v3/transformer_layer.py | 163 +++++++++++- 5 files changed, 702 insertions(+), 24 deletions(-) create mode 100644 tests/pytorch/test_deepseek.py diff --git a/tests/pytorch/test_deepseek.py b/tests/pytorch/test_deepseek.py new file mode 100644 index 0000000000..7778d0448c --- /dev/null +++ b/tests/pytorch/test_deepseek.py @@ -0,0 +1,124 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pytest +import torch + +from transformer_engine.pytorch.utils import deinterleave_glu_tensor +from transformer_engine.pytorch.models import ( + DeepSeekV3Layer, + DeepSeekV3MoE, + MultiLatentAttention, +) + +SEQ_LEN = 128 +BATCH = 2 +HIDDEN = 256 +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 _input(requires_grad=True): + torch.manual_seed(1234) + return torch.randn( + SEQ_LEN, BATCH, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=requires_grad + ) + + +def test_mla_forward_backward(): + torch.manual_seed(0) + mla = MultiLatentAttention(HIDDEN, HEADS, params_dtype=DTYPE, **MLA_KWARGS) + x = _input() + out = mla(x) + assert out.shape == x.shape + out.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + +@pytest.mark.parametrize("shared", [False, True], ids=["no_shared", "shared"]) +@pytest.mark.parametrize("grouped", [False, True], ids=["ungrouped", "grouped"]) +def test_moe_forward_backward(shared, grouped): + torch.manual_seed(0) + moe = DeepSeekV3MoE( + HIDDEN, + moe_ffn_hidden_size=128, + num_experts=8, + topk=2, + num_groups=4 if grouped else None, + group_topk=2 if grouped else None, + shared_expert_ffn_hidden_size=128 if shared else None, + params_dtype=DTYPE, + ) + x = _input() + out = moe(x) + assert out.shape == x.shape + out.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + counts = moe._last_tokens_per_expert + assert counts.sum().item() == SEQ_LEN * BATCH * 2 + bias_before = moe.expert_bias.clone() + moe.update_expert_bias() + assert not torch.equal(bias_before, moe.expert_bias) + + +def test_moe_matches_dense_reference(): + """topk == num_experts with uniform probs must reduce to a sum of expert MLPs.""" + torch.manual_seed(0) + num_experts = 4 + moe = DeepSeekV3MoE( + HIDDEN, + moe_ffn_hidden_size=128, + num_experts=num_experts, + topk=num_experts, + routed_scaling_factor=1.0, + params_dtype=DTYPE, + ) + x = _input(requires_grad=False) + out = moe(x) + + tokens = x.reshape(-1, HIDDEN) + probs, _ = moe._route(moe.gate(tokens).float()) + fc1, _, fc2 = moe.experts + ref = torch.zeros_like(tokens) + for e in range(num_experts): + w1 = deinterleave_glu_tensor(getattr(fc1, f"weight{e}"), 32) + w2 = getattr(fc2, f"weight{e}") + gate_part, lin_part = (tokens @ w1.t()).chunk(2, dim=-1) + act = torch.nn.functional.silu(gate_part.float()) * lin_part.float() + ref += (act.to(DTYPE) * probs[:, e : e + 1].to(DTYPE)) @ w2.t() + torch.testing.assert_close(out.reshape(-1, HIDDEN), ref, rtol=0.05, atol=0.05) + + +@pytest.mark.parametrize("num_experts", [None, 8], ids=["dense", "moe"]) +def test_layer_forward_backward(num_experts): + torch.manual_seed(0) + layer = ( + DeepSeekV3Layer( + HIDDEN, + HEADS, + ffn_hidden_size=512, + num_experts=num_experts, + moe_ffn_hidden_size=128 if num_experts else None, + topk=2 if num_experts else None, + shared_expert_ffn_hidden_size=128 if num_experts else None, + params_dtype=DTYPE, + **MLA_KWARGS, + ) + if num_experts + else DeepSeekV3Layer(HIDDEN, HEADS, ffn_hidden_size=512, params_dtype=DTYPE, **MLA_KWARGS) + ) + x = _input() + out = layer(x) + assert out.shape == x.shape + out.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 2b1803bfb2..fae4d973e5 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -34,6 +34,7 @@ from transformer_engine.pytorch.attention import InferenceParams from transformer_engine.pytorch.attention import RotaryPositionEmbedding from transformer_engine.pytorch.transformer import TransformerLayer +from transformer_engine.pytorch import models from transformer_engine.pytorch.permutation import ( moe_permute, moe_permute_with_probs, diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index f4f787743b..5a1c8d650c 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -5,23 +5,248 @@ """DeepSeekV3 MoE block: sigmoid router with aux-loss-free bias, shared + routed experts.""" +from typing import Optional, Union + import torch +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.pytorch.router import fused_topk_with_score_function +from transformer_engine.pytorch.permutation import moe_permute_with_probs, moe_unpermute + __all__ = ["DeepSeekV3MoE"] +def _make_expert_mlp(num_experts, hidden_size, ffn_hidden_size, dtype, device): + # GroupedLinear + ScaledSwiGLU + GroupedLinear fuses into a single CuTe + # grouped MLP on supported hardware; elsewhere it runs as three ops with + # the same API and checkpoint layout. + return te_ops.Sequential( + te_ops.GroupedLinear( + num_experts, hidden_size, 2 * ffn_hidden_size, bias=False, dtype=dtype, device=device + ), + te_ops.ScaledSwiGLU(glu_interleave_size=32), + te_ops.GroupedLinear( + num_experts, ffn_hidden_size, hidden_size, bias=False, dtype=dtype, device=device + ), + ) + + class DeepSeekV3MoE(torch.nn.Module): """ - DeepSeekV3-style Mixture of Experts block composed from TE MoE - primitives: ``fused_topk_with_score_function`` (sigmoid score function, - expert bias, grouped top-k), ``moe_permute_with_probs``/``moe_unpermute``, - :class:`GroupedLinear` routed experts, a shared expert - (:class:`LayerNormMLP`), ``Fp8Padding``/``Fp8Unpadding`` and optional - expert parallelism via ``ep_dispatch``/``ep_combine``. - - .. warning:: Work in progress, not functional yet. + DeepSeekV3-style Mixture of Experts block. + + Routing uses the fused sigmoid router with aux-loss-free expert bias and + node-limited (grouped) top-k (``fused_topk_with_score_function``). Routed + experts run as a grouped SwiGLU MLP built from ``te.ops`` (fusable into a + single CuTe grouped-GEMM kernel); routing probabilities are applied + per-token inside the expert MLP, so unpermute/combine is a plain + accumulation. Token routing is either local + (``moe_permute_with_probs``/``moe_unpermute``) or, when ``ep_group`` is + given, expert-parallel over NCCL (``ep_dispatch``/``ep_combine``). + + When expert parallelism is used, ``transformer_engine.pytorch.ep.ep_bootstrap`` + must be called once per process before the first forward, and inputs must + be bfloat16. + + Parameters + ---------- + hidden_size : int + size of each input sample. + moe_ffn_hidden_size : int + ffn size of each routed expert. + num_experts : int + total number of routed experts. + topk : int, default = 8 + number of experts per token. + num_groups : int, optional + number of expert groups for node-limited routing. + group_topk : int, optional + number of groups each token is limited to. + routed_scaling_factor : float, default = 2.5 + scaling applied to the routing probabilities. + shared_expert_ffn_hidden_size : int, optional + ffn size of the shared expert; ``None`` + disables the shared expert. + expert_bias_update_rate : float, default = 1e-3 + step size of the aux-loss-free bias update + (see :meth:`update_expert_bias`). + params_dtype : torch.dtype, optional + dtype of module parameters. + ep_group : ProcessGroup, optional + expert-parallel process group; enables the NCCL EP path. + ep_max_tokens_per_rank : int, optional + max local tokens per forward (required with EP). + ep_recv_capacity_per_rank : int, optional + receive-buffer capacity; defaults to + ``ep_size * ep_max_tokens_per_rank * topk``. + ep_alignment : int, default = 128 + per-expert row alignment of the EP receive buffer. """ - def __init__(self, *args, **kwargs): + def __init__( + self, + hidden_size: int, + moe_ffn_hidden_size: int, + num_experts: int, + topk: int = 8, + num_groups: Optional[int] = None, + group_topk: Optional[int] = None, + routed_scaling_factor: float = 2.5, + shared_expert_ffn_hidden_size: Optional[int] = None, + expert_bias_update_rate: float = 1e-3, + params_dtype: Optional[torch.dtype] = None, + device: Union[torch.device, str] = "cuda", + ep_group: Optional[torch.distributed.ProcessGroup] = None, + ep_max_tokens_per_rank: Optional[int] = None, + ep_recv_capacity_per_rank: Optional[int] = None, + ep_alignment: int = 128, + ) -> None: super().__init__() - raise NotImplementedError("DeepSeekV3MoE is under development") + + dtype = params_dtype if params_dtype is not None else torch.get_default_dtype() + self.hidden_size = hidden_size + self.num_experts = num_experts + self.topk = topk + self.num_groups = num_groups + self.group_topk = group_topk + self.routed_scaling_factor = routed_scaling_factor + self.expert_bias_update_rate = expert_bias_update_rate + + self.gate = torch.nn.Linear( + hidden_size, num_experts, bias=False, dtype=dtype, device=device + ) + self.register_buffer( + "expert_bias", torch.zeros(num_experts, dtype=torch.float32, device=device) + ) + self._last_tokens_per_expert: Optional[torch.Tensor] = None + + self.ep_group = ep_group + self.ep_size = 1 if ep_group is None else torch.distributed.get_world_size(ep_group) + assert num_experts % self.ep_size == 0 + num_local_experts = num_experts // self.ep_size + + self.experts = _make_expert_mlp( + num_local_experts, hidden_size, moe_ffn_hidden_size, dtype, device + ) + + self.shared_expert = None + if shared_expert_ffn_hidden_size is not None: + self.shared_expert = te_ops.Sequential( + te_ops.Linear( + hidden_size, + 2 * shared_expert_ffn_hidden_size, + bias=False, + dtype=dtype, + device=device, + ), + te_ops.SwiGLU(), + te_ops.Linear( + shared_expert_ffn_hidden_size, + hidden_size, + bias=False, + dtype=dtype, + device=device, + ), + ) + + self.ep_buffer = None + if ep_group is not None: + from transformer_engine.pytorch.ep import EpBuffer + + assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." + if ep_recv_capacity_per_rank is None: + ep_recv_capacity_per_rank = self.ep_size * ep_max_tokens_per_rank * topk + self.ep_buffer = EpBuffer( + top_k=topk, + max_tokens_per_rank=ep_max_tokens_per_rank, + hidden_dim=hidden_size, + num_local_experts=num_local_experts, + recv_capacity_per_rank=ep_recv_capacity_per_rank, + alignment=ep_alignment, + device=device, + ) + + def _route(self, logits: torch.Tensor, topk_indices: Optional[torch.Tensor] = None): + return fused_topk_with_score_function( + logits=logits, + topk=self.topk, + use_pre_softmax=False, + num_groups=self.num_groups, + group_topk=self.group_topk, + scaling_factor=self.routed_scaling_factor, + score_function="sigmoid", + expert_bias=self.expert_bias, + topk_indices=topk_indices, + ) + + def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: + probs, routing_map = self._route(self.gate(tokens).float()) + tokens_per_expert = routing_map.sum(dim=0) + self._last_tokens_per_expert = tokens_per_expert.detach() + + num_out = tokens.shape[0] * self.topk + permuted, permuted_probs, row_id_map = moe_permute_with_probs( + tokens, probs, routing_map, num_out_tokens=num_out + ) + + # The fused grouped MLP requires the total row count to be a multiple + # of 128; rows beyond sum(tokens_per_expert) fall outside every group. + pad = (-num_out) % 128 + if pad: + permuted = torch.nn.functional.pad(permuted, (0, 0, 0, pad)) + permuted_probs = torch.nn.functional.pad(permuted_probs, (0, pad)) + + out = self.experts( + permuted, tokens_per_expert, permuted_probs.to(tokens.dtype), tokens_per_expert + ) + return moe_unpermute(out[:num_out], row_id_map, restore_shape=tokens.shape) + + def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: + from transformer_engine.pytorch.ep import ep_dispatch, ep_combine + + assert tokens.dtype == torch.bfloat16, "The EP path requires bfloat16 inputs." + topk_idx = torch.empty( + (tokens.shape[0], self.topk), dtype=torch.int64, device=tokens.device + ) + probs, topk_idx = self._route(self.gate(tokens).float(), topk_indices=topk_idx) + self._last_tokens_per_expert = torch.bincount( + topk_idx.flatten(), minlength=self.num_experts + ) + topk_weights = probs.gather(1, topk_idx).float() + + recv_tokens, recv_weights, tokens_per_expert = ep_dispatch( + self.ep_buffer, tokens, topk_idx, topk_weights + ) + expert_out = self.experts( + recv_tokens, tokens_per_expert, recv_weights.to(tokens.dtype), tokens_per_expert + ) + return ep_combine(self.ep_buffer, expert_out, num_local_tokens=tokens.shape[0]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[..., hidden_size]``. + """ + tokens = hidden_states.reshape(-1, self.hidden_size) + if self.ep_group is not None: + out = self._forward_ep(tokens) + else: + out = self._forward_local(tokens) + if self.shared_expert is not None: + out = out + self.shared_expert(tokens) + return out.view_as(hidden_states) + + @torch.no_grad() + def update_expert_bias(self) -> None: + """Aux-loss-free bias update from the last forward's routing counts. + + With data/expert parallelism, all-reduce ``_last_tokens_per_expert`` + across ranks before calling (or call on identically-routed ranks). + """ + counts = self._last_tokens_per_expert + if counts is None: + return + err = counts.float().mean() - counts.float() + self.expert_bias += self.expert_bias_update_rate * torch.sign(err) diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index 6c2bb7420b..a36075f2c5 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -4,21 +4,200 @@ """Multi-Latent Attention (MLA) block as used in DeepSeekV3.""" +from typing import Optional, Union + import torch +from transformer_engine.pytorch.module import Linear, LayerNormLinear +from transformer_engine.pytorch.attention import DotProductAttention, RotaryPositionEmbedding +from transformer_engine.pytorch.attention.rope import apply_rotary_pos_emb + __all__ = ["MultiLatentAttention"] class MultiLatentAttention(torch.nn.Module): """ - Multi-Latent Attention with low-rank Q/KV down-projections and a - decoupled RoPE/NoPE head split, composed from :class:`Linear`, - :class:`LayerNormLinear` and :class:`DotProductAttention` - (``kv_channels=(head_dim_qk, head_dim_v)``). + Multi-Latent Attention as used in DeepSeekV3. - .. warning:: Work in progress, not functional yet. + Queries and key-values are projected through low-rank latents + (``q_lora_rank``, ``kv_lora_rank``); RMSNorm on each latent is fused into + the up-projection (:class:`LayerNormLinear` with RMSNorm). Each query/key + head is split into a ``qk_nope_head_dim`` part and a ``qk_rope_head_dim`` + part; RoPE is applied only to the rope part, and the key rope part comes + from a single shared head broadcast to all heads. Attention runs through + :class:`DotProductAttention` with asymmetric head dims + ``kv_channels=(qk_nope_head_dim + qk_rope_head_dim, v_head_dim)``, which + supports the cuDNN fused attention backend. + + Parameters + ---------- + hidden_size : int + size of each input sample. + num_attention_heads : int + number of attention heads. + q_lora_rank : int, default = 1536 + rank of the query latent. + kv_lora_rank : int, default = 512 + rank of the key-value latent. + qk_nope_head_dim : int, default = 128 + per-head dim of the non-rotary query/key part. + qk_rope_head_dim : int, default = 64 + per-head dim of the rotary query/key part. + v_head_dim : int, default = 128 + per-head dim of the values. + attention_dropout : float, default = 0.0 + dropout probability on attention scores. + attn_mask_type : str, default = "causal" + attention mask type passed to :class:`DotProductAttention`. + rotary_base : float, default = 10000.0 + RoPE base. + softmax_scale : float, optional + softmax scale; defaults to ``1/sqrt(qk head dim)`` inside + :class:`DotProductAttention`. + qkv_format : str, default = "sbhd" + layout of the input/output tensors. + params_dtype : torch.dtype, optional + dtype of module parameters. + tp_group : ProcessGroup, optional + tensor-parallel process group for the up/output projections. + tp_size : int, default = 1 + tensor-parallel world size. """ - def __init__(self, *args, **kwargs): + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + q_lora_rank: int = 1536, + kv_lora_rank: int = 512, + qk_nope_head_dim: int = 128, + qk_rope_head_dim: int = 64, + v_head_dim: int = 128, + attention_dropout: float = 0.0, + attn_mask_type: str = "causal", + rotary_base: float = 10000.0, + softmax_scale: Optional[float] = None, + qkv_format: str = "sbhd", + params_dtype: Optional[torch.dtype] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + tp_size: int = 1, + device: Union[torch.device, str] = "cuda", + ) -> None: super().__init__() - raise NotImplementedError("MultiLatentAttention is under development") + + assert qkv_format in ("sbhd", "bshd"), "MultiLatentAttention supports sbhd/bshd formats." + assert num_attention_heads % tp_size == 0 + + self.qkv_format = qkv_format + self.num_attention_heads = num_attention_heads + self.num_attention_heads_per_partition = num_attention_heads // tp_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.kv_lora_rank = kv_lora_rank + + common = {"bias": False, "params_dtype": params_dtype, "device": device} + tp = {"tp_group": tp_group, "tp_size": tp_size} + + self.q_down_proj = Linear(hidden_size, q_lora_rank, **common) + self.q_up_proj = LayerNormLinear( + q_lora_rank, + num_attention_heads * self.qk_head_dim, + normalization="RMSNorm", + parallel_mode="column" if tp_size > 1 else None, + **tp, + **common, + ) + self.kv_down_proj = Linear(hidden_size, kv_lora_rank + qk_rope_head_dim, **common) + self.kv_up_proj = LayerNormLinear( + kv_lora_rank, + num_attention_heads * (qk_nope_head_dim + v_head_dim), + normalization="RMSNorm", + parallel_mode="column" if tp_size > 1 else None, + **tp, + **common, + ) + self.out_proj = Linear( + num_attention_heads * v_head_dim, + hidden_size, + parallel_mode="row" if tp_size > 1 else None, + **tp, + **common, + ) + + self.rope = RotaryPositionEmbedding(qk_rope_head_dim, rotary_base=rotary_base) + self._rope_freqs: Optional[torch.Tensor] = None + + self.core_attention = DotProductAttention( + num_attention_heads, + kv_channels=(self.qk_head_dim, v_head_dim), + attention_dropout=attention_dropout, + qkv_format=qkv_format, + attn_mask_type=attn_mask_type, + softmax_scale=softmax_scale, + tp_group=tp_group, + tp_size=tp_size, + ) + + def _rope_freqs_for(self, seq_len: int, device: torch.device) -> torch.Tensor: + if self._rope_freqs is None or self._rope_freqs.shape[0] < seq_len: + self._rope_freqs = self.rope(seq_len).to(device) + return self._rope_freqs[:seq_len] + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + attn_mask_type: Optional[str] = None, + checkpoint_core_attention: bool = False, + ) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[sq, b, h]`` (sbhd) or ``[b, sq, h]`` (bshd). + attention_mask : torch.Tensor, optional + boolean mask passed to :class:`DotProductAttention`. + attn_mask_type : str, optional + override of the constructor's mask type. + checkpoint_core_attention : bool, default = False + checkpoint the core attention computation. + """ + seq_dim = 0 if self.qkv_format == "sbhd" else 1 + seq_len = hidden_states.shape[seq_dim] + heads = self.num_attention_heads_per_partition + + q = self.q_up_proj(self.q_down_proj(hidden_states)) + q = q.view(*q.shape[:-1], heads, self.qk_head_dim) + + kv_down = self.kv_down_proj(hidden_states) + kv_latent, k_pos = torch.split(kv_down, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + kv = self.kv_up_proj(kv_latent) + kv = kv.view(*kv.shape[:-1], heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + freqs = self._rope_freqs_for(seq_len, hidden_states.device) + q_rope = apply_rotary_pos_emb( + q[..., self.qk_nope_head_dim :].contiguous(), + freqs, + tensor_format=self.qkv_format, + fused=True, + ) + k_rope = apply_rotary_pos_emb( + k_pos.unsqueeze(-2), freqs, tensor_format=self.qkv_format, fused=True + ) + + q = torch.cat([q[..., : self.qk_nope_head_dim], q_rope], dim=-1) + k = torch.cat([k_nope, k_rope.expand(*k_nope.shape[:-1], -1)], dim=-1) + + context = self.core_attention( + q, + k, + v.contiguous(), + attention_mask=attention_mask, + qkv_format=self.qkv_format, + attn_mask_type=attn_mask_type, + checkpoint_core_attention=checkpoint_core_attention, + ) + return self.out_proj(context) diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index 2a28a6ceb3..af1eeb1a95 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -4,22 +4,171 @@ """DeepSeekV3 transformer layer.""" +from typing import Optional, Union + import torch +from transformer_engine.pytorch.module import LayerNormMLP, RMSNorm +from transformer_engine.pytorch.models.deepseek_v3.multi_latent_attention import ( + MultiLatentAttention, +) +from transformer_engine.pytorch.models.deepseek_v3.moe import DeepSeekV3MoE + __all__ = ["DeepSeekV3Layer"] class DeepSeekV3Layer(torch.nn.Module): """ A full DeepSeekV3 transformer layer, analogous to - :class:`TransformerLayer`: :class:`MultiLatentAttention` followed by - either a dense :class:`LayerNormMLP` (first layers) or - :class:`DeepSeekV3MoE`, with the same residual and fused - bias-dropout-add plumbing as :class:`TransformerLayer`. + :class:`TransformerLayer`: pre-RMSNorm + :class:`MultiLatentAttention`, + then either a dense SwiGLU MLP (:class:`LayerNormMLP` with RMSNorm, used + for the first dense layers of DeepSeekV3) or :class:`DeepSeekV3MoE`, each + with a residual connection. - .. warning:: Work in progress, not functional yet. + Parameters + ---------- + hidden_size : int + size of each input sample. + num_attention_heads : int + number of attention heads. + ffn_hidden_size : int + ffn size of the dense MLP (used when ``num_experts`` is + ``None``). + num_experts : int, optional + number of routed experts; ``None`` makes this a dense layer. + moe_ffn_hidden_size : int, optional + ffn size of each routed expert (required with MoE). + hidden_dropout : float, default = 0.0 + dropout probability on the residual branches. + kwargs common to the submodules (``q_lora_rank``, ``kv_lora_rank``, + ``qk_nope_head_dim``, ``qk_rope_head_dim``, ``v_head_dim``, + ``attention_dropout``, ``attn_mask_type``, ``qkv_format``, ``topk``, + ``num_groups``, ``group_topk``, ``routed_scaling_factor``, + ``shared_expert_ffn_hidden_size``, EP options, ...) are forwarded to + :class:`MultiLatentAttention` and :class:`DeepSeekV3MoE`. """ - def __init__(self, *args, **kwargs): + _MLA_KWARGS = frozenset( + { + "q_lora_rank", + "kv_lora_rank", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "attention_dropout", + "attn_mask_type", + "rotary_base", + "softmax_scale", + "qkv_format", + "tp_group", + "tp_size", + } + ) + _MOE_KWARGS = frozenset( + { + "topk", + "num_groups", + "group_topk", + "routed_scaling_factor", + "shared_expert_ffn_hidden_size", + "expert_bias_update_rate", + "ep_group", + "ep_max_tokens_per_rank", + "ep_recv_capacity_per_rank", + "ep_alignment", + } + ) + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + ffn_hidden_size: Optional[int] = None, + num_experts: Optional[int] = None, + moe_ffn_hidden_size: Optional[int] = None, + hidden_dropout: float = 0.0, + layernorm_epsilon: float = 1e-5, + params_dtype: Optional[torch.dtype] = None, + device: Union[torch.device, str] = "cuda", + **kwargs, + ) -> None: super().__init__() - raise NotImplementedError("DeepSeekV3Layer is under development") + + unknown = set(kwargs) - self._MLA_KWARGS - self._MOE_KWARGS + if unknown: + raise TypeError(f"Unexpected keyword arguments: {sorted(unknown)}") + mla_kwargs = {k: v for k, v in kwargs.items() if k in self._MLA_KWARGS} + moe_kwargs = {k: v for k, v in kwargs.items() if k in self._MOE_KWARGS} + + self.hidden_dropout = hidden_dropout + + self.input_layernorm = RMSNorm( + hidden_size, eps=layernorm_epsilon, device=device, dtype=params_dtype + ) + self.self_attention = MultiLatentAttention( + hidden_size, + num_attention_heads, + params_dtype=params_dtype, + device=device, + **mla_kwargs, + ) + + if num_experts is None: + assert ffn_hidden_size is not None, "Dense layers require ffn_hidden_size." + self.pre_mlp_layernorm = None + self.mlp = LayerNormMLP( + hidden_size, + ffn_hidden_size, + eps=layernorm_epsilon, + normalization="RMSNorm", + activation="swiglu", + bias=False, + params_dtype=params_dtype, + device=device, + ) + else: + assert moe_ffn_hidden_size is not None, "MoE layers require moe_ffn_hidden_size." + self.pre_mlp_layernorm = RMSNorm( + hidden_size, eps=layernorm_epsilon, device=device, dtype=params_dtype + ) + self.mlp = DeepSeekV3MoE( + hidden_size, + moe_ffn_hidden_size, + num_experts, + params_dtype=params_dtype, + device=device, + **moe_kwargs, + ) + + def _residual_add(self, out: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + out = torch.nn.functional.dropout(out, p=self.hidden_dropout, training=self.training) + return residual + out + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + checkpoint_core_attention: bool = False, + ) -> torch.Tensor: + """ + Parameters + ---------- + hidden_states : torch.Tensor + input of shape ``[sq, b, h]`` (sbhd) or ``[b, sq, h]`` (bshd). + attention_mask : torch.Tensor, optional + boolean attention mask. + checkpoint_core_attention : bool, default = False + checkpoint the core attention computation. + """ + attention_out = self.self_attention( + self.input_layernorm(hidden_states), + attention_mask=attention_mask, + checkpoint_core_attention=checkpoint_core_attention, + ) + hidden_states = self._residual_add(attention_out, hidden_states) + + if self.pre_mlp_layernorm is not None: + mlp_out = self.mlp(self.pre_mlp_layernorm(hidden_states)) + else: + mlp_out = self.mlp(hidden_states) + return self._residual_add(mlp_out, hidden_states) From e23100b73cef8e7f7c955af8f626f83177aef06e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 14:28:46 +0200 Subject: [PATCH 05/43] Add distributed EP test for DeepSeekV3 MoE/layer run_deepseek_ep.py checks the EP path against the all-experts-local path numerically (forward, input/gate grads, all-reduced expert wgrads) and smoke-tests the full layer with EP. Also size the default EP recv capacity for per-expert alignment padding and the fused grouped MLP's row-count requirement. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_deepseek_ep.py | 185 ++++++++++++++++++ .../distributed/run_test_deepseek_ep.sh | 52 +++++ tests/pytorch/distributed/test_deepseek_ep.py | 26 +++ .../pytorch/models/deepseek_v3/moe.py | 6 +- 4 files changed, 268 insertions(+), 1 deletion(-) create mode 100644 tests/pytorch/distributed/run_deepseek_ep.py create mode 100644 tests/pytorch/distributed/run_test_deepseek_ep.sh create mode 100644 tests/pytorch/distributed/test_deepseek_ep.py diff --git a/tests/pytorch/distributed/run_deepseek_ep.py b/tests/pytorch/distributed/run_deepseek_ep.py new file mode 100644 index 0000000000..0edae05961 --- /dev/null +++ b/tests/pytorch/distributed/run_deepseek_ep.py @@ -0,0 +1,185 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Multi-process DeepSeekV3 MoE/layer EP tests, launched via torchrun.""" + +import os +import sys +import unittest + +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, DeepSeekV3MoE + +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) + + +class TestDeepSeekEP(unittest.TestCase): + @classmethod + def setUpClass(cls): + if _device_sm() < 90: + raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{_device_sm()})") + cls.rank = dist.get_rank() + cls.ep_size = dist.get_world_size() + cls.num_experts = NUM_LOCAL_EXPERTS * cls.ep_size + world_pg = dist.distributed_c10d._get_default_group() + cls.ep_group = dist.new_group(ranks=list(range(world_pg.size())), backend="nccl") + ep_bootstrap( + cls.ep_group, + num_experts=cls.num_experts, + max_tokens_per_rank=TOKENS_PER_RANK, + hidden_dim=HIDDEN, + num_topk=TOP_K, + recv_capacity_per_rank=_recv_capacity(cls.ep_size), + ) + + def _make_moe(self, ep: bool, shared: bool = True) -> DeepSeekV3MoE: + return DeepSeekV3MoE( + HIDDEN, + moe_ffn_hidden_size=MOE_FFN, + num_experts=self.num_experts, + topk=TOP_K, + shared_expert_ffn_hidden_size=SHARED_FFN if shared else None, + params_dtype=DTYPE, + ep_group=self.ep_group if ep else None, + ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, + ep_recv_capacity_per_rank=_recv_capacity(self.ep_size) if ep else None, + ) + + def _copy_local_expert_weights(self, ep_moe: DeepSeekV3MoE, ref: DeepSeekV3MoE) -> None: + with torch.no_grad(): + ep_moe.gate.weight.copy_(ref.gate.weight) + if ref.shared_expert is not None: + for dst, src in zip( + ep_moe.shared_expert.parameters(), ref.shared_expert.parameters() + ): + dst.copy_(src) + ep_fc1, _, ep_fc2 = ep_moe.experts + ref_fc1, _, ref_fc2 = ref.experts + for local_e in range(NUM_LOCAL_EXPERTS): + global_e = self.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_moe_ep_matches_local(self): + """EP MoE must match the single-GPU (all-experts-local) path numerically.""" + torch.manual_seed(0) + ref = self._make_moe(ep=False) + _broadcast_params(ref) + ep_moe = self._make_moe(ep=True) + self._copy_local_expert_weights(ep_moe, ref) + + torch.manual_seed(1234 + self.rank) + x = torch.randn(TOKENS_PER_RANK, HIDDEN, dtype=DTYPE, device="cuda") + x_ep = x.clone().requires_grad_(True) + x_ref = x.clone().requires_grad_(True) + + out_ep = ep_moe(x_ep) + out_ref = ref(x_ref) + 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) + torch.testing.assert_close( + ep_moe.gate.weight.grad, ref.gate.weight.grad, rtol=0.1, atol=0.1 + ) + + # A local expert's wgrad on its owner rank equals the sum of the + # reference wgrads over all ranks. + ep_fc1, _, ep_fc2 = ep_moe.experts + ref_fc1, _, ref_fc2 = ref.experts + for ep_fc, ref_fc in ((ep_fc1, ref_fc1), (ep_fc2, ref_fc2)): + for local_e in range(NUM_LOCAL_EXPERTS): + global_e = self.rank * NUM_LOCAL_EXPERTS + local_e + ref_grad = getattr(ref_fc, f"weight{global_e}").grad.float() + dist.all_reduce(ref_grad) + ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float() + torch.testing.assert_close(ep_grad, ref_grad, rtol=0.1, atol=0.1) + + counts = ep_moe._last_tokens_per_expert.clone() + dist.all_reduce(counts) + self.assertEqual(counts.sum().item(), self.ep_size * TOKENS_PER_RANK * TOP_K) + + def test_layer_ep_forward_backward(self): + """Full DeepSeekV3Layer smoke test with an EP MoE block.""" + torch.manual_seed(10 + self.rank) + layer = DeepSeekV3Layer( + HIDDEN, + HEADS, + num_experts=self.num_experts, + moe_ffn_hidden_size=MOE_FFN, + topk=TOP_K, + shared_expert_ffn_hidden_size=SHARED_FFN, + params_dtype=DTYPE, + ep_group=self.ep_group, + ep_max_tokens_per_rank=TOKENS_PER_RANK, + ep_recv_capacity_per_rank=_recv_capacity(self.ep_size), + **MLA_KWARGS, + ) + x = torch.randn( + TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=True + ) + out = layer(x) + self.assertEqual(out.shape, x.shape) + out.sum().backward() + self.assertIsNotNone(x.grad) + self.assertTrue(torch.isfinite(x.grad).all()) + + layer.mlp.update_expert_bias() + self.assertTrue(torch.isfinite(layer.mlp.expert_bias).all()) + + +def _init_distributed(): + dist.init_process_group(backend="nccl") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + try: + from torch.distributed import _symmetric_memory as _symm_mem + + _symm_mem.set_backend("NCCL") + except (ImportError, RuntimeError): + pass + + +if __name__ == "__main__": + _init_distributed() + suite = unittest.TestLoader().loadTestsFromTestCase(TestDeepSeekEP) + result = unittest.TextTestRunner(stream=sys.stdout, verbosity=2).run(suite) + dist.barrier() + ep_finalize() + release_symm_mem_pool() + dist.destroy_process_group() + sys.exit(0 if result.wasSuccessful() else 1) diff --git a/tests/pytorch/distributed/run_test_deepseek_ep.sh b/tests/pytorch/distributed/run_test_deepseek_ep.sh new file mode 100644 index 0000000000..8c0bbbc5b9 --- /dev/null +++ b/tests/pytorch/distributed/run_test_deepseek_ep.sh @@ -0,0 +1,52 @@ +#!/bin/bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Launcher for tests/pytorch/distributed/run_deepseek_ep.py. Auto-detects GPU count. + +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +if [ "${DETECTED_GPUS}" -lt 2 ]; then + echo "DeepSeek EP test requires >= 2 GPUs (found ${DETECTED_GPUS}); SKIPPING." + exit 0 +fi + +# NCCL EP requires active NVLink P2P among ranks on the node. +if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then + echo "No NVLink between GPUs (PCIe-only fabric); NCCL EP is unsupported here. SKIPPING." + exit 0 +fi + +NUM_RANKS="${NVTE_TEST_EP_NUM_RANKS:-${DETECTED_GPUS}}" +if [ "${NUM_RANKS}" -gt 8 ]; then NUM_RANKS=8; fi + +TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" + +: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} +export NCCL_EP_JIT_CACHE_DIR +mkdir -p "$NCCL_EP_JIT_CACHE_DIR" + +SCRIPT="${SCRIPT_DIR}/run_deepseek_ep.py" +LOG="stdout_deepseek_ep.txt" + +echo "=== Running ${SCRIPT} on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" +setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ + torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ + "${SCRIPT}" 2>&1 | tee "${LOG}" +RC=${PIPESTATUS[0]} +pkill -9 -f "tests/pytorch/distributed/run_deepseek_ep.py" 2>/dev/null || true + +RET=0 +if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; RET=1; fi +if grep -qE "(^|]:)FAILED|(^|]:)Traceback" "${LOG}"; then RET=1; fi +if ! grep -qE "Ran [0-9]+ test|^OK$" "${LOG}"; then + echo "ERROR: no test summary — likely hang or early crash" + RET=1 +fi +if [ -z "${KEEP_EP_LOGS:-}" ]; then rm -f "${LOG}"; fi + +exit $RET diff --git a/tests/pytorch/distributed/test_deepseek_ep.py b/tests/pytorch/distributed/test_deepseek_ep.py new file mode 100644 index 0000000000..4a4d9a8dea --- /dev/null +++ b/tests/pytorch/distributed/test_deepseek_ep.py @@ -0,0 +1,26 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Pytest driver — spawns run_deepseek_ep.py under torchrun and asserts it passed.""" + +import os +import subprocess +from pathlib import Path + +import pytest +import torch + +TEST_ROOT = Path(__file__).parent.resolve() +LAUNCHER = TEST_ROOT / "run_test_deepseek_ep.sh" + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="DeepSeek EP requires >= 2 GPUs") +def test_multi_process_deepseek_ep(): + timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) + proc = subprocess.run( + ["bash", str(LAUNCHER)], + env={**os.environ, "KEEP_EP_LOGS": "1", "TEST_TIMEOUT_S": str(timeout_s)}, + timeout=timeout_s + 30, + check=False, + ) + assert proc.returncode == 0, f"DeepSeek EP test suite failed (rc={proc.returncode})" diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 5a1c8d650c..f413221bd1 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -155,7 +155,11 @@ def __init__( assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." if ep_recv_capacity_per_rank is None: - ep_recv_capacity_per_rank = self.ep_size * ep_max_tokens_per_rank * topk + # Worst case plus per-expert alignment padding, rounded up to + # the multiple of 128 required by the fused grouped MLP. + cap = self.ep_size * ep_max_tokens_per_rank * topk + cap += num_local_experts * max(ep_alignment, 1) + ep_recv_capacity_per_rank = -(-cap // 128) * 128 self.ep_buffer = EpBuffer( top_k=topk, max_tokens_per_rank=ep_max_tokens_per_rank, From 4c6e1e8aff62def1cd0bd72ce9bfb960a4aab035 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 15:49:08 +0200 Subject: [PATCH 06/43] Fix EP wgrad test collective + zero EP recv/grad buffers The per-expert wgrad check called all_reduce on different tensors per rank (rank-local experts), corrupting the reference grads; reduce every expert's grad on every rank instead. Also pass zero-filled recv/grad buffers to ep_dispatch/ep_combine so alignment-padding rows inside the grouped-GEMM m_splits can never poison expert wgrads. Verified on lyris (4x GB300, arm64): run_test_deepseek_ep.sh passes on all ranks (EP forward/dgrad/gate-grad/expert-wgrad match the all-local reference; full-layer EP smoke passes). Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_deepseek_ep.py | 12 +++++++---- .../pytorch/models/deepseek_v3/moe.py | 20 +++++++++++++++++-- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/distributed/run_deepseek_ep.py b/tests/pytorch/distributed/run_deepseek_ep.py index 0edae05961..bf756b69ad 100644 --- a/tests/pytorch/distributed/run_deepseek_ep.py +++ b/tests/pytorch/distributed/run_deepseek_ep.py @@ -119,16 +119,20 @@ def test_moe_ep_matches_local(self): ) # A local expert's wgrad on its owner rank equals the sum of the - # reference wgrads over all ranks. + # 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_moe.experts ref_fc1, _, ref_fc2 = ref.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(self.num_experts) + ] + for g in ref_grads: + dist.all_reduce(g) for local_e in range(NUM_LOCAL_EXPERTS): global_e = self.rank * NUM_LOCAL_EXPERTS + local_e - ref_grad = getattr(ref_fc, f"weight{global_e}").grad.float() - dist.all_reduce(ref_grad) ep_grad = getattr(ep_fc, f"weight{local_e}").grad.float() - torch.testing.assert_close(ep_grad, ref_grad, rtol=0.1, atol=0.1) + torch.testing.assert_close(ep_grad, ref_grads[global_e], rtol=0.1, atol=0.1) counts = ep_moe._last_tokens_per_expert.clone() dist.all_reduce(counts) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index f413221bd1..3c182b4405 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -218,13 +218,29 @@ def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: ) topk_weights = probs.gather(1, topk_idx).float() + # Zero-filled recv/grad buffers: per-expert alignment padding lands + # inside the grouped-GEMM m_splits, so uninitialized rows would poison + # the expert wgrads. + cap = self.ep_buffer.recv_capacity_per_rank recv_tokens, recv_weights, tokens_per_expert = ep_dispatch( - self.ep_buffer, tokens, topk_idx, topk_weights + self.ep_buffer, + tokens, + topk_idx, + topk_weights, + recv_tokens=torch.zeros( + (cap, self.hidden_size), dtype=tokens.dtype, device=tokens.device + ), + recv_topk_weights=torch.zeros((cap,), dtype=torch.float32, device=tokens.device), ) expert_out = self.experts( recv_tokens, tokens_per_expert, recv_weights.to(tokens.dtype), tokens_per_expert ) - return ep_combine(self.ep_buffer, expert_out, num_local_tokens=tokens.shape[0]) + return ep_combine( + self.ep_buffer, + expert_out, + num_local_tokens=tokens.shape[0], + grad_out=torch.zeros_like(expert_out), + ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """ From aa17c37fb9a0e8cd74c3b5d67a5d34365f4133d7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 18 Aug 2026 17:23:10 +0200 Subject: [PATCH 07/43] Use fused MLA RoPE kernels in MultiLatentAttention Move the Triton MLA RoPE kernels (Megatron-LM fused_mla_yarn_rope_apply port) from tests/pytorch/attention/ mla_rope_utils.py into models/deepseek_v3/mla_rope.py and use them in MultiLatentAttention: the q kernel rotates the rope slice in place and the kv kernel assembles key/value in a single pass, removing the torch.cat/expand/contiguous copies (~10% of layer GPU time). PyTorch fallback (same convention) covers missing Triton and bshd. Fix a latent bug from the test util: the q backward kernel assumed a contiguous incoming gradient, but cuDNN attention backward can hand over a strided one (allocator-state dependent IMA). The old test file stays as a compat shim. Add a Triton-vs-PyTorch parity test. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/mla_rope_utils.py | 652 +----------------- tests/pytorch/test_deepseek.py | 60 ++ .../pytorch/models/deepseek_v3/mla_rope.py | 495 +++++++++++++ .../deepseek_v3/multi_latent_attention.py | 55 +- 4 files changed, 601 insertions(+), 661 deletions(-) create mode 100644 transformer_engine/pytorch/models/deepseek_v3/mla_rope.py diff --git a/tests/pytorch/attention/mla_rope_utils.py b/tests/pytorch/attention/mla_rope_utils.py index 90eebfc66a..d022757886 100644 --- a/tests/pytorch/attention/mla_rope_utils.py +++ b/tests/pytorch/attention/mla_rope_utils.py @@ -2,26 +2,17 @@ # # See LICENSE for license information. -"""MLA RoPE for DSv3 671B - Triton forward and backward kernels. - -Source: Megatron-LM megatron/core/fusions/fused_mla_yarn_rope_apply.py -Falls back to pure PyTorch when Triton is unavailable. - -Note: DSv3 uses YaRN-scaled RoPE for long-context extrapolation. This test -intentionally uses plain RoPE (base=10000) because it only validates MXFP8 -attention path wiring, tensor shapes, forward/backward flow, and relative BF16 -vs MXFP8 behavior. Both reference and MXFP8 paths use the same RoPE tables. -""" +"""Compat shim: the MLA RoPE kernels moved to +``transformer_engine.pytorch.models.deepseek_v3.mla_rope``.""" import torch -try: - import triton - import triton.language as tl - - HAVE_TRITON = True -except ImportError: - HAVE_TRITON = False +from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( # noqa: F401 + HAVE_TRITON, + apply_mla_rope_kv, + apply_mla_rope_q, + build_rope_tables, +) HEAD_DIM_ROPE = 64 HEAD_DIM_NOPE = 128 @@ -29,576 +20,6 @@ ROTARY_BASE = 10000 -def build_rope_tables( - seq_len: int, - emb_dim: int = HEAD_DIM_ROPE, - base: int = ROTARY_BASE, - device: torch.device = None, -) -> tuple[torch.Tensor, torch.Tensor]: - inv_freq = 1.0 / ( - base ** (torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim) - ) - t = torch.arange(seq_len, device=device, dtype=torch.float32) - freqs = torch.outer(t, inv_freq) - freqs = torch.cat([freqs, freqs], dim=-1) - return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous() - - -if HAVE_TRITON: - - # Not used for non-packed batches; kept for THD compatibility. - @triton.jit - def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): - token_idx = -1 - this_seq_len = 0 - seq_idx = 0 - last_cum_seqlen = tl.load(cu_seqlens) // cp_size - while seq_idx < seq_num: - cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size - if token_idx == -1 and cur_cum_seqlen > pid_m: - token_idx = pid_m - last_cum_seqlen - this_seq_len = cur_cum_seqlen - last_cum_seqlen - last_cum_seqlen = cur_cum_seqlen - seq_idx += 1 - if cp_size > 1: - if token_idx < this_seq_len // 2: - token_idx = token_idx + cp_rank * this_seq_len // 2 - else: - token_idx = (token_idx - this_seq_len // 2) + ( - 2 * cp_size - cp_rank - 1 - ) * this_seq_len // 2 - return token_idx - - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "head_num"], - restore_value=["Q"], - ) - @triton.jit - def rotary_fwd_q_kernel( - Q, - COS, - SIN, - qk_head_dim, - emb_dim: tl.constexpr, - head_num: tl.constexpr, - batch_size, - seq_num, - cu_seqlens_q, - stride_x_seq, - stride_x_nheads, - cp_rank, - cp_size, - BLOCK_H: tl.constexpr, - ): - pid_m = tl.program_id(axis=0) - pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: - token_idx = pid_m // batch_size - else: - token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - Q = Q + pid_m * stride_x_seq - x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim - mask = head_offsets[:, None] < head_num - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 - x_1 = tl.load(Q + x_1_off, mask=mask) - x_2 = tl.load(Q + x_2_off, mask=mask) - x_left = x_1 * cos_left - x_2 * sin_left - x_right = x_2 * cos_right + x_1 * sin_right - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - tl.store(Q + x_left_off, x_left, mask=mask) - tl.store(Q + x_right_off, x_right, mask=mask) - - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "head_num"], - restore_value=["DO"], - ) - @triton.jit - def rotary_bwd_q_kernel( - DO, - COS, - SIN, - qk_head_dim, - emb_dim: tl.constexpr, - head_num: tl.constexpr, - batch_size, - seq_num, - cu_seqlens_q, - stride_x_seq, - stride_x_nheads, - cp_rank, - cp_size, - BLOCK_H: tl.constexpr, - ): - pid_m = tl.program_id(axis=0) - pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: - token_idx = pid_m // batch_size - else: - token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - DO = DO + pid_m * stride_x_seq - x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim - mask = head_offsets[:, None] < head_num - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(DO + x_left_off, mask=mask) - x_right = tl.load(DO + x_right_off, mask=mask) - x_1 = x_left * cos_left + x_right * sin_right - x_2 = -x_left * sin_left + x_right * cos_right - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 - tl.store(DO + x_1_off, x_1, mask=mask) - tl.store(DO + x_2_off, x_2, mask=mask) - - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "k_dim", "v_dim", "head_num"], - ) - @triton.jit - def rotary_fwd_kv_kernel( - KV, - K_POS_EMB, - O_KEY, - O_VALUE, - COS, - SIN, - emb_dim: tl.constexpr, - k_dim: tl.constexpr, - v_dim: tl.constexpr, - head_num: tl.constexpr, - batch_size, - seq_num, - cu_seqlens_kv, - stride_kv_seq, - stride_kv_nheads, - stride_emb_seq, - stride_k_seq, - stride_k_nheads, - stride_v_seq, - stride_v_nheads, - cp_rank, - cp_size, - BLOCK_H: tl.constexpr, - ): - pid_m = tl.program_id(axis=0) - pid_head = tl.program_id(axis=1) - if cu_seqlens_kv is None: - token_idx = pid_m // batch_size - else: - token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - KV_ptr = KV + pid_m * stride_kv_seq - kv_off = head_offsets[:, None] * stride_kv_nheads - mask = head_offsets[:, None] < head_num - k_in_off = kv_off + tl.arange(0, k_dim)[None, :] - v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] - k = tl.load(KV_ptr + k_in_off, mask=mask) - v = tl.load(KV_ptr + v_in_off, mask=mask) - K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads - V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads - k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] - v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] - tl.store(K_ptr + k_out_off, k, mask=mask) - tl.store(V_ptr + v_out_off, v, mask=mask) - EMB = K_POS_EMB + pid_m * stride_emb_seq - x_1 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2) - x_2 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2 + 1) - x_left = x_1 * cos_left - x_2 * sin_left - x_right = x_2 * cos_right + x_1 * sin_right - x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_left_off = ( - tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads - + k_dim - + tl.arange(0, emb_dim // 2)[None, :] - ) - x_right_off = x_left_off + emb_dim // 2 - tl.store(K_ptr + x_left_off, x_left, mask=mask) - tl.store(K_ptr + x_right_off, x_right, mask=mask) - - @triton.autotune( - configs=[ - triton.Config({"BLOCK_H": 1}), - triton.Config({"BLOCK_H": 2}), - triton.Config({"BLOCK_H": 4}), - triton.Config({"BLOCK_H": 8}), - triton.Config({"BLOCK_H": 16}), - triton.Config({"BLOCK_H": 32}), - triton.Config({"BLOCK_H": 64}), - triton.Config({"BLOCK_H": 128}), - ], - key=["emb_dim", "k_dim", "v_dim", "head_num"], - ) - @triton.jit - def rotary_bwd_kv_kernel( - dK, - dV, - dKV, - dEMB, - COS, - SIN, - emb_dim: tl.constexpr, - k_dim: tl.constexpr, - v_dim: tl.constexpr, - head_num: tl.constexpr, - batch_size, - seq_num, - cu_seqlens_kv, - stride_dk_seq, - stride_dk_nheads, - stride_dv_seq, - stride_dv_nheads, - stride_dkv_seq, - stride_dkv_nheads, - stride_demb_seq, - cp_rank, - cp_size, - BLOCK_H: tl.constexpr, - ): - pid_m = tl.program_id(axis=0) - pid_head = tl.program_id(axis=1) - if cu_seqlens_kv is None: - token_idx = pid_m // batch_size - else: - token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) - head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) - dKV_ptr = dKV + pid_m * stride_dkv_seq - dkv_off = head_offsets[:, None] * stride_dkv_nheads - mask = head_offsets[:, None] < head_num - dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] - dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] - dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads - dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads - dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :] - dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :] - dk = tl.load(dK_ptr + dk_in_off, mask=mask) - dv = tl.load(dV_ptr + dv_in_off, mask=mask) - tl.store(dKV_ptr + dk_out_off, dk, mask=mask) - tl.store(dKV_ptr + dv_out_off, dv, mask=mask) - if pid_head == 0: - x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) - x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) - for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): - head_offsets_i = i * BLOCK_H + tl.arange(0, BLOCK_H) - dK_ptr_i = dK + pid_m * stride_dk_seq - x_off = head_offsets_i[:, None] * stride_dk_nheads + k_dim - mask_i = head_offsets_i[:, None] < head_num - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left_accum += tl.load(dK_ptr_i + x_left_off, mask=mask_i) - x_right_accum += tl.load(dK_ptr_i + x_right_off, mask=mask_i) - x_left_accum = tl.sum(x_left_accum, axis=0) - x_right_accum = tl.sum(x_right_accum, axis=0) - x_left_accum = x_left_accum.to(dEMB.dtype.element_ty) - x_right_accum = x_right_accum.to(dEMB.dtype.element_ty) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load( - COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) - ) - sin_right = tl.load( - SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) - ) - x_1 = x_left_accum * cos_left + x_right_accum * sin_right - x_2 = -x_left_accum * sin_left + x_right_accum * cos_right - dEMB_ptr = dEMB + pid_m * stride_demb_seq - tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2, x_1) - tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) - - def _flattened_token_stride(tensor: torch.Tensor) -> int: - if tensor.dim() == 4: - return tensor.stride(1) - return tensor.stride(0) - - class _MLARoPEQTriton(torch.autograd.Function): - @staticmethod - def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): - s, b, nheads, _ = q.shape - total = s * b - - grid_q = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_q_kernel[grid_q]( - q, - cos, - sin, - head_dim_nope, - head_dim_rope, - nheads, - b, - None, - None, - _flattened_token_stride(q), - q.stride(2), - 0, - 1, - ) - - ctx.save_for_backward(cos, sin) - ctx.head_dim_nope = head_dim_nope - ctx.head_dim_rope = head_dim_rope - ctx.nheads = nheads - ctx.s = s - ctx.b = b - return q - - @staticmethod - def backward(ctx, dq): - cos, sin = ctx.saved_tensors - s, b, nheads = ctx.s, ctx.b, ctx.nheads - total = s * b - - grid_q = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_q_kernel[grid_q]( - dq, - cos, - sin, - ctx.head_dim_nope, - ctx.head_dim_rope, - nheads, - b, - None, - None, - _flattened_token_stride(dq), - dq.stride(2), - 0, - 1, - ) - return dq, None, None, None, None - - class _MLARoPEKVTriton(torch.autograd.Function): - @staticmethod - def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim_v): - s, b, nheads, _ = kv.shape - total = s * b - - o_key = kv.new_empty(s, b, nheads, head_dim_nope + head_dim_rope) - o_value = kv.new_empty(s, b, nheads, head_dim_v) - grid_kv = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_kv_kernel[grid_kv]( - kv, - k_pos_emb, - o_key, - o_value, - cos, - sin, - head_dim_rope, - head_dim_nope, - head_dim_v, - nheads, - b, - None, - None, - _flattened_token_stride(kv), - kv.stride(2), - _flattened_token_stride(k_pos_emb), - _flattened_token_stride(o_key), - o_key.stride(2), - _flattened_token_stride(o_value), - o_value.stride(2), - 0, - 1, - ) - - ctx.save_for_backward(cos, sin) - ctx.head_dim_nope = head_dim_nope - ctx.head_dim_rope = head_dim_rope - ctx.head_dim_v = head_dim_v - ctx.nheads = nheads - ctx.s = s - ctx.b = b - return o_key, o_value - - @staticmethod - def backward(ctx, dk_out, dv_out): - cos, sin = ctx.saved_tensors - s, b, nheads = ctx.s, ctx.b, ctx.nheads - ndp, ndr, ndv = ctx.head_dim_nope, ctx.head_dim_rope, ctx.head_dim_v - total = s * b - - d_kv = dk_out.new_empty(s, b, nheads, ndp + ndv) - d_emb = dk_out.new_empty(s, b, 1, ndr) - grid_kv = lambda META: (total, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_kv_kernel[grid_kv]( - dk_out, - dv_out, - d_kv, - d_emb, - cos, - sin, - ndr, - ndp, - ndv, - nheads, - b, - None, - None, - _flattened_token_stride(dk_out), - dk_out.stride(2), - _flattened_token_stride(dv_out), - dv_out.stride(2), - _flattened_token_stride(d_kv), - d_kv.stride(2), - _flattened_token_stride(d_emb), - 0, - 1, - ) - return d_kv, d_emb, None, None, None, None, None - - -def _apply_mla_rope_q_with_tables( - q: torch.Tensor, - cos_table: torch.Tensor, - sin_table: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, -) -> torch.Tensor: - if HAVE_TRITON: - return _MLARoPEQTriton.apply( - q, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - ) - return _apply_pytorch_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope) - - -def _apply_mla_rope_kv_with_tables( - kv: torch.Tensor, - k_pos_emb: torch.Tensor, - cos_table: torch.Tensor, - sin_table: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, -) -> tuple[torch.Tensor, torch.Tensor]: - if HAVE_TRITON: - return _MLARoPEKVTriton.apply( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - return _apply_pytorch_kv( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - - -def apply_mla_rope_q( - q: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -) -> torch.Tensor: - if cos_table is None or sin_table is None: - s = q.shape[0] - cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=q.device, - ) - return _apply_mla_rope_q_with_tables( - q, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - ) - - -def apply_mla_rope_kv( - kv: torch.Tensor, - k_pos_emb: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - if cos_table is None or sin_table is None: - s = kv.shape[0] - cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=kv.device, - ) - return _apply_mla_rope_kv_with_tables( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, - ) - - def apply_mla_rope( q: torch.Tensor, kv: torch.Tensor, @@ -609,60 +30,13 @@ def apply_mla_rope( base: int = ROTARY_BASE, cos_table: torch.Tensor | None = None, sin_table: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +): if cos_table is None or sin_table is None: - s = q.shape[0] cos_table, sin_table = build_rope_tables( - s, - emb_dim=head_dim_rope, - base=base, - device=q.device, + q.shape[0], head_dim_rope, base=base, device=q.device ) - q = _apply_mla_rope_q_with_tables(q, cos_table, sin_table, head_dim_nope, head_dim_rope) - k, v = _apply_mla_rope_kv_with_tables( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, + q = apply_mla_rope_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope) + k, v = apply_mla_rope_kv( + kv, k_pos_emb, cos_table, sin_table, head_dim_nope, head_dim_rope, head_dim_v ) return q, k, v - - -def _rotate_interleaved_to_neox( - x: torch.Tensor, cos_table: torch.Tensor, sin_table: torch.Tensor -) -> torch.Tensor: - cos_ = cos_table[:, None, None, :].to(x.dtype) - sin_ = sin_table[:, None, None, :].to(x.dtype) - half_dim = x.shape[-1] // 2 - x_1 = x[..., 0::2] - x_2 = x[..., 1::2] - x_left = x_1 * cos_[..., :half_dim] - x_2 * sin_[..., :half_dim] - x_right = x_2 * cos_[..., half_dim:] + x_1 * sin_[..., half_dim:] - return torch.cat((x_left, x_right), dim=-1) - - -def _apply_pytorch_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope): - q_nope = q[..., :head_dim_nope] - q_rope = q[..., head_dim_nope : head_dim_nope + head_dim_rope] - q_rope = _rotate_interleaved_to_neox(q_rope, cos_table, sin_table) - return torch.cat((q_nope, q_rope), dim=-1) - - -def _apply_pytorch_kv( - kv, - k_pos_emb, - cos_table, - sin_table, - head_dim_nope, - head_dim_rope, - head_dim_v, -): - k_nope = kv[..., :head_dim_nope] - v = kv[..., head_dim_nope : head_dim_nope + head_dim_v] - k_rope = _rotate_interleaved_to_neox(k_pos_emb, cos_table, sin_table).expand( - -1, -1, kv.shape[2], -1 - ) - return torch.cat((k_nope, k_rope), dim=-1), v diff --git a/tests/pytorch/test_deepseek.py b/tests/pytorch/test_deepseek.py index 7778d0448c..4c4aea0a92 100644 --- a/tests/pytorch/test_deepseek.py +++ b/tests/pytorch/test_deepseek.py @@ -34,6 +34,66 @@ def _input(requires_grad=True): ) +def test_mla_rope_triton_matches_pytorch(): + from transformer_engine.pytorch.models.deepseek_v3 import mla_rope + + if not mla_rope.HAVE_TRITON: + pytest.skip("Triton unavailable") + s, b, h = 64, 2, 4 + nope, rope, vdim = 64, 32, 64 + cos, sin = mla_rope.build_rope_tables(s, rope, device="cuda") + + torch.manual_seed(0) + q_leaf = torch.randn(s, b, h, nope + rope, device="cuda", requires_grad=True) + kv_leaf = torch.randn(s, b, h, nope + vdim, device="cuda", requires_grad=True) + pos_leaf = torch.randn(s, b, 1, rope, device="cuda", requires_grad=True) + grad_q = torch.randn(s, b, h, nope + rope, device="cuda") + grad_k = torch.randn(s, b, h, nope + rope, device="cuda") + grad_v = torch.randn(s, b, h, vdim, device="cuda") + + def run(fmt): + # non-leaf copies: the Triton q kernel rotates in place + q, kv, pos = q_leaf * 1.0, kv_leaf * 1.0, pos_leaf * 1.0 + q_out = mla_rope.apply_mla_rope_q(q, cos, sin, nope, rope, fmt) + k_out, v_out = mla_rope.apply_mla_rope_kv(kv, pos, cos, sin, nope, rope, vdim, fmt) + # fresh grad clones: the Triton q backward modifies its input grad in place + torch.autograd.backward( + [q_out, k_out, v_out], [grad_q.clone(), grad_k.clone(), grad_v.clone()] + ) + grads = (q_leaf.grad.clone(), kv_leaf.grad.clone(), pos_leaf.grad.clone()) + q_leaf.grad = kv_leaf.grad = pos_leaf.grad = None + return (q_out.clone(), k_out, v_out), grads + + (q_t, k_t, v_t), grads_t = run("sbhd") + + seq_dim = 0 + q_ref = torch.cat( + ( + (q_leaf * 1.0)[..., :nope], + mla_rope._rotate_interleaved_to_neox((q_leaf * 1.0)[..., nope:], cos, sin, seq_dim), + ), + dim=-1, + ) + k_ref = torch.cat( + ( + (kv_leaf * 1.0)[..., :nope], + mla_rope._rotate_interleaved_to_neox(pos_leaf * 1.0, cos, sin, seq_dim).expand( + s, b, h, rope + ), + ), + dim=-1, + ) + v_ref = (kv_leaf * 1.0)[..., nope:] + torch.autograd.backward([q_ref, k_ref, v_ref], [grad_q.clone(), grad_k.clone(), grad_v.clone()]) + + torch.testing.assert_close(q_t, q_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(k_t, k_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(v_t, v_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[0], q_leaf.grad, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[1], kv_leaf.grad, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(grads_t[2], pos_leaf.grad, rtol=1e-5, atol=1e-5) + + def test_mla_forward_backward(): torch.manual_seed(0) mla = MultiLatentAttention(HIDDEN, HEADS, params_dtype=DTYPE, **MLA_KWARGS) diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py new file mode 100644 index 0000000000..350bedb69b --- /dev/null +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -0,0 +1,495 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused MLA RoPE kernels (DeepSeekV3-style decoupled RoPE/NoPE). + +Triton forward/backward kernels adapted from Megatron-LM +``megatron/core/fusions/fused_mla_yarn_rope_apply.py``. The query kernel +rotates the trailing ``head_dim_rope`` slice in place (no concat); the KV +kernel builds the final key (nope | broadcast-rotated shared rope head) and +value tensors in a single pass. Falls back to pure PyTorch when Triton is +unavailable or for the ``bshd`` layout (the Triton path is ``sbhd``-only). + +Rotation convention: the rope slice is read interleaved (as stored in +HF/Megatron DeepSeekV3 checkpoints) and written in NeoX half-split layout, +matching the Megatron fused kernel semantics. +""" + +from typing import Optional, Tuple + +import torch + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: + HAVE_TRITON = False + +__all__ = ["build_rope_tables", "apply_mla_rope_q", "apply_mla_rope_kv"] + + +def build_rope_tables( + seq_len: int, + emb_dim: int, + base: float = 10000.0, + device: Optional[torch.device] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """cos/sin tables of shape ``[seq_len, emb_dim]`` (fp32, NeoX duplicated halves).""" + inv_freq = 1.0 / ( + base ** (torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim) + ) + t = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + freqs = torch.cat([freqs, freqs], dim=-1) + return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous() + + +if HAVE_TRITON: + + # Not used for non-packed batches; kept for THD compatibility. + @triton.jit + def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): + token_idx = -1 + this_seq_len = 0 + seq_idx = 0 + last_cum_seqlen = tl.load(cu_seqlens) // cp_size + while seq_idx < seq_num: + cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size + if token_idx == -1 and cur_cum_seqlen > pid_m: + token_idx = pid_m - last_cum_seqlen + this_seq_len = cur_cum_seqlen - last_cum_seqlen + last_cum_seqlen = cur_cum_seqlen + seq_idx += 1 + if cp_size > 1: + if token_idx < this_seq_len // 2: + token_idx = token_idx + cp_rank * this_seq_len // 2 + else: + token_idx = (token_idx - this_seq_len // 2) + ( + 2 * cp_size - cp_rank - 1 + ) * this_seq_len // 2 + return token_idx + + _AUTOTUNE_CONFIGS = [triton.Config({"BLOCK_H": h}) for h in (1, 2, 4, 8, 16, 32, 64, 128)] + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "head_num"], restore_value=["Q"]) + @triton.jit + def rotary_fwd_q_kernel( + Q, + COS, + SIN, + qk_head_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + stride_x_seq, + stride_x_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + Q = Q + pid_m * stride_x_seq + x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim + mask = head_offsets[:, None] < head_num + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_1 = tl.load(Q + x_1_off, mask=mask) + x_2 = tl.load(Q + x_2_off, mask=mask) + x_left = x_1 * cos_left - x_2 * sin_left + x_right = x_2 * cos_right + x_1 * sin_right + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + tl.store(Q + x_left_off, x_left, mask=mask) + tl.store(Q + x_right_off, x_right, mask=mask) + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "head_num"], restore_value=["DO"]) + @triton.jit + def rotary_bwd_q_kernel( + DO, + COS, + SIN, + qk_head_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + stride_x_seq, + stride_x_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + DO = DO + pid_m * stride_x_seq + x_off = head_offsets[:, None] * stride_x_nheads + qk_head_dim + mask = head_offsets[:, None] < head_num + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + x_1 = x_left * cos_left + x_right * sin_right + x_2 = -x_left * sin_left + x_right * cos_right + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + tl.store(DO + x_1_off, x_1, mask=mask) + tl.store(DO + x_2_off, x_2, mask=mask) + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "k_dim", "v_dim", "head_num"]) + @triton.jit + def rotary_fwd_kv_kernel( + KV, + K_POS_EMB, + O_KEY, + O_VALUE, + COS, + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_kv, + stride_kv_seq, + stride_kv_nheads, + stride_emb_seq, + stride_k_seq, + stride_k_nheads, + stride_v_seq, + stride_v_nheads, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_kv is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + KV_ptr = KV + pid_m * stride_kv_seq + kv_off = head_offsets[:, None] * stride_kv_nheads + mask = head_offsets[:, None] < head_num + k_in_off = kv_off + tl.arange(0, k_dim)[None, :] + v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] + k = tl.load(KV_ptr + k_in_off, mask=mask) + v = tl.load(KV_ptr + v_in_off, mask=mask) + K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads + V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads + k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] + v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] + tl.store(K_ptr + k_out_off, k, mask=mask) + tl.store(V_ptr + v_out_off, v, mask=mask) + EMB = K_POS_EMB + pid_m * stride_emb_seq + x_1 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2) + x_2 = tl.load(EMB + tl.arange(0, emb_dim // 2) * 2 + 1) + x_left = x_1 * cos_left - x_2 * sin_left + x_right = x_2 * cos_right + x_1 * sin_right + x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_left_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 + tl.store(K_ptr + x_left_off, x_left, mask=mask) + tl.store(K_ptr + x_right_off, x_right, mask=mask) + + @triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["emb_dim", "k_dim", "v_dim", "head_num"]) + @triton.jit + def rotary_bwd_kv_kernel( + dK, + dV, + dKV, + dEMB, + COS, + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_kv, + stride_dk_seq, + stride_dk_nheads, + stride_dv_seq, + stride_dv_nheads, + stride_dkv_seq, + stride_dkv_nheads, + stride_demb_seq, + cp_rank, + cp_size, + BLOCK_H: tl.constexpr, + ): + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + if cu_seqlens_kv is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + head_offsets = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + dKV_ptr = dKV + pid_m * stride_dkv_seq + dkv_off = head_offsets[:, None] * stride_dkv_nheads + mask = head_offsets[:, None] < head_num + dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] + dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] + dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads + dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads + dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :] + dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :] + dk = tl.load(dK_ptr + dk_in_off, mask=mask) + dv = tl.load(dV_ptr + dv_in_off, mask=mask) + tl.store(dKV_ptr + dk_out_off, dk, mask=mask) + tl.store(dKV_ptr + dv_out_off, dv, mask=mask) + if pid_head == 0: + x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): + head_offsets_i = i * BLOCK_H + tl.arange(0, BLOCK_H) + dK_ptr_i = dK + pid_m * stride_dk_seq + x_off = head_offsets_i[:, None] * stride_dk_nheads + k_dim + mask_i = head_offsets_i[:, None] < head_num + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left_accum += tl.load(dK_ptr_i + x_left_off, mask=mask_i) + x_right_accum += tl.load(dK_ptr_i + x_right_off, mask=mask_i) + x_left_accum = tl.sum(x_left_accum, axis=0) + x_right_accum = tl.sum(x_right_accum, axis=0) + x_left_accum = x_left_accum.to(dEMB.dtype.element_ty) + x_right_accum = x_right_accum.to(dEMB.dtype.element_ty) + cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + x_1 = x_left_accum * cos_left + x_right_accum * sin_right + x_2 = -x_left_accum * sin_left + x_right_accum * cos_right + dEMB_ptr = dEMB + pid_m * stride_demb_seq + tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2, x_1) + tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) + + def _token_stride(tensor: torch.Tensor) -> int: + return tensor.stride(1) if tensor.dim() == 4 else tensor.stride(0) + + class _MLARoPEQTriton(torch.autograd.Function): + """In-place RoPE on the trailing rope slice of q [s, b, h, nope+rope].""" + + @staticmethod + def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): + if not q.is_contiguous(): + q = q.contiguous() + s, b, nheads, _ = q.shape + grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_fwd_q_kernel[grid]( + q, + cos, + sin, + head_dim_nope, + head_dim_rope, + nheads, + b, + None, + None, + _token_stride(q), + q.stride(2), + 0, + 1, + ) + ctx.save_for_backward(cos, sin) + ctx.dims = (s, b, nheads, head_dim_nope, head_dim_rope) + return q + + @staticmethod + def backward(ctx, dq): + cos, sin = ctx.saved_tensors + # attention backward may hand over a strided grad; the kernel + # assumes a contiguous [s, b, h, d] layout + dq = dq.contiguous() + s, b, nheads, head_dim_nope, head_dim_rope = ctx.dims + grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_bwd_q_kernel[grid]( + dq, + cos, + sin, + head_dim_nope, + head_dim_rope, + nheads, + b, + None, + None, + _token_stride(dq), + dq.stride(2), + 0, + 1, + ) + return dq, None, None, None, None + + class _MLARoPEKVTriton(torch.autograd.Function): + """kv [s, b, h, nope+v] + shared rope head [s, b, 1, rope] -> (k, v).""" + + @staticmethod + def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim_v): + if not kv.is_contiguous(): + kv = kv.contiguous() + s, b, nheads, _ = kv.shape + o_key = kv.new_empty(s, b, nheads, head_dim_nope + head_dim_rope) + o_value = kv.new_empty(s, b, nheads, head_dim_v) + grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_fwd_kv_kernel[grid]( + kv, + k_pos_emb, + o_key, + o_value, + cos, + sin, + head_dim_rope, + head_dim_nope, + head_dim_v, + nheads, + b, + None, + None, + _token_stride(kv), + kv.stride(2), + _token_stride(k_pos_emb), + _token_stride(o_key), + o_key.stride(2), + _token_stride(o_value), + o_value.stride(2), + 0, + 1, + ) + ctx.save_for_backward(cos, sin) + ctx.dims = (s, b, nheads, head_dim_nope, head_dim_rope, head_dim_v) + return o_key, o_value + + @staticmethod + def backward(ctx, dk_out, dv_out): + cos, sin = ctx.saved_tensors + s, b, nheads, ndp, ndr, ndv = ctx.dims + dk_out = dk_out.contiguous() + dv_out = dv_out.contiguous() + d_kv = dk_out.new_empty(s, b, nheads, ndp + ndv) + d_emb = dk_out.new_empty(s, b, 1, ndr) + grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + rotary_bwd_kv_kernel[grid]( + dk_out, + dv_out, + d_kv, + d_emb, + cos, + sin, + ndr, + ndp, + ndv, + nheads, + b, + None, + None, + _token_stride(dk_out), + dk_out.stride(2), + _token_stride(dv_out), + dv_out.stride(2), + _token_stride(d_kv), + d_kv.stride(2), + _token_stride(d_emb), + 0, + 1, + ) + return d_kv, d_emb, None, None, None, None, None + + +def _rotate_interleaved_to_neox(x, cos_table, sin_table, seq_dim): + shape = [1, 1, 1, cos_table.shape[-1]] + shape[seq_dim] = cos_table.shape[0] + cos_ = cos_table.view(shape).to(x.dtype) + sin_ = sin_table.view(shape).to(x.dtype) + half = x.shape[-1] // 2 + x_1 = x[..., 0::2] + x_2 = x[..., 1::2] + x_left = x_1 * cos_[..., :half] - x_2 * sin_[..., :half] + x_right = x_2 * cos_[..., half:] + x_1 * sin_[..., half:] + return torch.cat((x_left, x_right), dim=-1) + + +def apply_mla_rope_q( + q: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + head_dim_nope: int, + head_dim_rope: int, + tensor_format: str = "sbhd", +) -> torch.Tensor: + """RoPE on the trailing ``head_dim_rope`` slice of q; in place on the Triton path.""" + if HAVE_TRITON and tensor_format == "sbhd": + return _MLARoPEQTriton.apply(q, cos_table, sin_table, head_dim_nope, head_dim_rope) + seq_dim = 0 if tensor_format == "sbhd" else 1 + q_rope = _rotate_interleaved_to_neox(q[..., head_dim_nope:], cos_table, sin_table, seq_dim) + return torch.cat((q[..., :head_dim_nope], q_rope), dim=-1) + + +def apply_mla_rope_kv( + kv: torch.Tensor, + k_pos_emb: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + head_dim_nope: int, + head_dim_rope: int, + head_dim_v: int, + tensor_format: str = "sbhd", +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build (k, v) from kv ``[.., h, nope+v]`` and the shared rope head ``[.., 1, rope]``.""" + if HAVE_TRITON and tensor_format == "sbhd": + return _MLARoPEKVTriton.apply( + kv, k_pos_emb, cos_table, sin_table, head_dim_nope, head_dim_rope, head_dim_v + ) + seq_dim = 0 if tensor_format == "sbhd" else 1 + k_nope = kv[..., :head_dim_nope] + v = kv[..., head_dim_nope : head_dim_nope + head_dim_v] + k_rope = _rotate_interleaved_to_neox(k_pos_emb, cos_table, sin_table, seq_dim) + k_rope = k_rope.expand(*k_nope.shape[:-1], -1) + return torch.cat((k_nope, k_rope), dim=-1), v.contiguous() diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index a36075f2c5..e5840a812f 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -9,8 +9,12 @@ import torch from transformer_engine.pytorch.module import Linear, LayerNormLinear -from transformer_engine.pytorch.attention import DotProductAttention, RotaryPositionEmbedding -from transformer_engine.pytorch.attention.rope import apply_rotary_pos_emb +from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( + apply_mla_rope_kv, + apply_mla_rope_q, + build_rope_tables, +) __all__ = ["MultiLatentAttention"] @@ -29,6 +33,10 @@ class MultiLatentAttention(torch.nn.Module): ``kv_channels=(qk_nope_head_dim + qk_rope_head_dim, v_head_dim)``, which supports the cuDNN fused attention backend. + RoPE uses the fused MLA kernels from :mod:`.mla_rope` (in-place on the + query rope slice, single-pass key/value assembly); the rope slice follows + the HF/Megatron DeepSeekV3 convention (interleaved weights, NeoX output). + Parameters ---------- hidden_size : int @@ -126,8 +134,8 @@ def __init__( **common, ) - self.rope = RotaryPositionEmbedding(qk_rope_head_dim, rotary_base=rotary_base) - self._rope_freqs: Optional[torch.Tensor] = None + self.rotary_base = rotary_base + self._rope_tables: Optional[tuple] = None self.core_attention = DotProductAttention( num_attention_heads, @@ -140,10 +148,13 @@ def __init__( tp_size=tp_size, ) - def _rope_freqs_for(self, seq_len: int, device: torch.device) -> torch.Tensor: - if self._rope_freqs is None or self._rope_freqs.shape[0] < seq_len: - self._rope_freqs = self.rope(seq_len).to(device) - return self._rope_freqs[:seq_len] + def _rope_tables_for(self, seq_len: int, device: torch.device): + if self._rope_tables is None or self._rope_tables[0].shape[0] < seq_len: + self._rope_tables = build_rope_tables( + seq_len, self.qk_rope_head_dim, base=self.rotary_base, device=device + ) + cos, sin = self._rope_tables + return cos[:seq_len], sin[:seq_len] def forward( self, @@ -175,26 +186,26 @@ def forward( kv_latent, k_pos = torch.split(kv_down, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) kv = self.kv_up_proj(kv_latent) kv = kv.view(*kv.shape[:-1], heads, self.qk_nope_head_dim + self.v_head_dim) - k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - - freqs = self._rope_freqs_for(seq_len, hidden_states.device) - q_rope = apply_rotary_pos_emb( - q[..., self.qk_nope_head_dim :].contiguous(), - freqs, - tensor_format=self.qkv_format, - fused=True, + + cos, sin = self._rope_tables_for(seq_len, hidden_states.device) + q = apply_mla_rope_q( + q, cos, sin, self.qk_nope_head_dim, self.qk_rope_head_dim, self.qkv_format ) - k_rope = apply_rotary_pos_emb( - k_pos.unsqueeze(-2), freqs, tensor_format=self.qkv_format, fused=True + k, v = apply_mla_rope_kv( + kv, + k_pos.unsqueeze(-2), + cos, + sin, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.qkv_format, ) - q = torch.cat([q[..., : self.qk_nope_head_dim], q_rope], dim=-1) - k = torch.cat([k_nope, k_rope.expand(*k_nope.shape[:-1], -1)], dim=-1) - context = self.core_attention( q, k, - v.contiguous(), + v, attention_mask=attention_mask, qkv_format=self.qkv_format, attn_mask_type=attn_mask_type, From c713af7bf06fca793e48c86be6fd203b4836bf14 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 11:42:40 +0200 Subject: [PATCH 08/43] Add HF transformers numeric reference test for DeepSeekV3Layer Maps HF DeepseekV3DecoderLayer weights into DeepSeekV3Layer (GLU interleave for routed experts, fused latent norms) and checks forward and input grads match within bf16 tolerance. Expose layernorm_epsilon on MultiLatentAttention (HF latent RMSNorms use 1e-6). Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_deepseek_hf.py | 144 ++++++++++++++++++ .../deepseek_v3/multi_latent_attention.py | 5 + 2 files changed, 149 insertions(+) create mode 100644 tests/pytorch/test_deepseek_hf.py diff --git a/tests/pytorch/test_deepseek_hf.py b/tests/pytorch/test_deepseek_hf.py new file mode 100644 index 0000000000..08f1c06ccd --- /dev/null +++ b/tests/pytorch/test_deepseek_hf.py @@ -0,0 +1,144 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Numeric comparison of DeepSeekV3Layer against the HF transformers reference.""" + +import pytest +import torch + +transformers = pytest.importorskip("transformers") +from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config +from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3DecoderLayer, + DeepseekV3RotaryEmbedding, +) + +from transformer_engine.pytorch.models import DeepSeekV3Layer +from transformer_engine.pytorch.utils import interleave_glu_tensor + +SEQ, BATCH = 64, 2 +HIDDEN, HEADS = 256, 4 +Q_LORA, KV_LORA = 96, 64 +NOPE, ROPE, VDIM = 64, 32, 64 +NUM_EXPERTS, TOPK, N_GROUP, TOPK_GROUP = 16, 4, 4, 2 +MOE_FFN, N_SHARED = 128, 1 +DTYPE = torch.bfloat16 + + +def _hf_config(): + return DeepseekV3Config( + hidden_size=HIDDEN, + intermediate_size=4 * HIDDEN, + moe_intermediate_size=MOE_FFN, + num_hidden_layers=1, + num_attention_heads=HEADS, + num_key_value_heads=HEADS, + n_shared_experts=N_SHARED, + n_routed_experts=NUM_EXPERTS, + routed_scaling_factor=2.5, + kv_lora_rank=KV_LORA, + q_lora_rank=Q_LORA, + qk_rope_head_dim=ROPE, + v_head_dim=VDIM, + qk_nope_head_dim=NOPE, + n_group=N_GROUP, + topk_group=TOPK_GROUP, + num_experts_per_tok=TOPK, + first_k_dense_replace=0, + norm_topk_prob=True, + rms_norm_eps=1e-5, + attention_bias=False, + attention_dropout=0.0, + rope_interleave=True, + _attn_implementation="eager", + ) + + +def _init_hf_layer(config): + torch.manual_seed(0) + layer = DeepseekV3DecoderLayer(config, layer_idx=0).to(device="cuda", dtype=DTYPE) + with torch.no_grad(): + for name, p in layer.named_parameters(): + if "layernorm" in name or "norm" in name: + p.copy_(1.0 + 0.1 * torch.randn_like(p)) + else: + p.normal_(0.0, 0.02) + bias = layer.mlp.gate.e_score_correction_bias + bias.copy_(0.1 * torch.randn_like(bias)) + return layer + + +def _build_te_layer(hf): + te_layer = DeepSeekV3Layer( + HIDDEN, + HEADS, + num_experts=NUM_EXPERTS, + moe_ffn_hidden_size=MOE_FFN, + topk=TOPK, + num_groups=N_GROUP, + group_topk=TOPK_GROUP, + routed_scaling_factor=2.5, + shared_expert_ffn_hidden_size=MOE_FFN * N_SHARED, + q_lora_rank=Q_LORA, + kv_lora_rank=KV_LORA, + qk_nope_head_dim=NOPE, + qk_rope_head_dim=ROPE, + v_head_dim=VDIM, + params_dtype=DTYPE, + ) + attn, mla = hf.self_attn, te_layer.self_attention + with torch.no_grad(): + te_layer.input_layernorm.weight.copy_(hf.input_layernorm.weight) + te_layer.pre_mlp_layernorm.weight.copy_(hf.post_attention_layernorm.weight) + + mla.q_down_proj.weight.copy_(attn.q_a_proj.weight) + mla.q_up_proj.layer_norm_weight.copy_(attn.q_a_layernorm.weight) + mla.q_up_proj.weight.copy_(attn.q_b_proj.weight) + mla.kv_down_proj.weight.copy_(attn.kv_a_proj_with_mqa.weight) + mla.kv_up_proj.layer_norm_weight.copy_(attn.kv_a_layernorm.weight) + mla.kv_up_proj.weight.copy_(attn.kv_b_proj.weight) + mla.out_proj.weight.copy_(attn.o_proj.weight) + + moe = te_layer.mlp + moe.gate.weight.copy_(hf.mlp.gate.weight) + moe.expert_bias.copy_(hf.mlp.gate.e_score_correction_bias) + fc1, _, fc2 = moe.experts + for e in range(NUM_EXPERTS): + getattr(fc1, f"weight{e}").copy_( + interleave_glu_tensor(hf.mlp.experts.gate_up_proj[e], 32) + ) + getattr(fc2, f"weight{e}").copy_(hf.mlp.experts.down_proj[e]) + shared = hf.mlp.shared_experts + moe.shared_expert[0].weight.copy_( + torch.cat([shared.gate_proj.weight, shared.up_proj.weight], dim=0) + ) + moe.shared_expert[2].weight.copy_(shared.down_proj.weight) + return te_layer + + +def test_layer_matches_hf(): + config = _hf_config() + hf = _init_hf_layer(config) + te_layer = _build_te_layer(hf) + + torch.manual_seed(1) + x = torch.randn(BATCH, SEQ, HIDDEN, dtype=DTYPE, device="cuda") + x_hf = x.clone().requires_grad_(True) + x_te = x.transpose(0, 1).contiguous().requires_grad_(True) # sbhd + + rotary = DeepseekV3RotaryEmbedding(config).to("cuda") + position_ids = torch.arange(SEQ, device="cuda").unsqueeze(0).expand(BATCH, -1) + cos, sin = rotary(x_hf, position_ids) + causal = torch.full((SEQ, SEQ), float("-inf"), device="cuda", dtype=DTYPE).triu(1) + causal = causal[None, None].expand(BATCH, 1, SEQ, SEQ) + + out_hf = hf(x_hf, attention_mask=causal, position_embeddings=(cos, sin)) + out_te = te_layer(x_te) + + torch.testing.assert_close(out_te.transpose(0, 1), out_hf, rtol=5e-2, atol=5e-2) + + grad = torch.randn_like(out_hf) + out_hf.backward(grad) + out_te.backward(grad.transpose(0, 1).contiguous()) + torch.testing.assert_close(x_te.grad.transpose(0, 1), x_hf.grad, rtol=5e-2, atol=5e-2) diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index e5840a812f..0ddf2f75b2 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -57,6 +57,8 @@ class MultiLatentAttention(torch.nn.Module): dropout probability on attention scores. attn_mask_type : str, default = "causal" attention mask type passed to :class:`DotProductAttention`. + layernorm_epsilon : float, default = 1e-6 + epsilon of the latent RMSNorms (matches DeepSeekV3). rotary_base : float, default = 10000.0 RoPE base. softmax_scale : float, optional @@ -83,6 +85,7 @@ def __init__( v_head_dim: int = 128, attention_dropout: float = 0.0, attn_mask_type: str = "causal", + layernorm_epsilon: float = 1e-6, rotary_base: float = 10000.0, softmax_scale: Optional[float] = None, qkv_format: str = "sbhd", @@ -113,6 +116,7 @@ def __init__( q_lora_rank, num_attention_heads * self.qk_head_dim, normalization="RMSNorm", + eps=layernorm_epsilon, parallel_mode="column" if tp_size > 1 else None, **tp, **common, @@ -122,6 +126,7 @@ def __init__( kv_lora_rank, num_attention_heads * (qk_nope_head_dim + v_head_dim), normalization="RMSNorm", + eps=layernorm_epsilon, parallel_mode="column" if tp_size > 1 else None, **tp, **common, From 88028a94f061e0ce075cfd788cf01e580d571ad4 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 12:08:12 +0200 Subject: [PATCH 09/43] Docstring cleanups for lint and docs build Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../pytorch/models/deepseek_v3/mla_rope.py | 28 ++++++++++++++++--- .../models/deepseek_v3/transformer_layer.py | 13 +++++---- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py index 350bedb69b..f8c01f75d6 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -92,6 +92,7 @@ def rotary_fwd_q_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """In-place RoPE fwd on the trailing rope slice of q.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_q is None: @@ -139,6 +140,7 @@ def rotary_bwd_q_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """In-place RoPE bwd on the trailing rope slice of dq.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_q is None: @@ -195,6 +197,7 @@ def rotary_fwd_kv_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """Fwd: build (key, value) from kv and the shared rotated rope head.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_kv is None: @@ -262,6 +265,7 @@ def rotary_bwd_kv_kernel( cp_size, BLOCK_H: tl.constexpr, ): + """Bwd: scatter (dk, dv) into dkv and reduce rope-slice grads into demb.""" pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) if cu_seqlens_kv is None: @@ -320,10 +324,14 @@ class _MLARoPEQTriton(torch.autograd.Function): @staticmethod def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): + """Rotate the rope slice of q in place.""" if not q.is_contiguous(): q = q.contiguous() s, b, nheads, _ = q.shape - grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + rotary_fwd_q_kernel[grid]( q, cos, @@ -345,12 +353,16 @@ def forward(ctx, q, cos, sin, head_dim_nope, head_dim_rope): @staticmethod def backward(ctx, dq): + """Counter-rotate the rope slice of dq (in place on the copy).""" cos, sin = ctx.saved_tensors # attention backward may hand over a strided grad; the kernel # assumes a contiguous [s, b, h, d] layout dq = dq.contiguous() s, b, nheads, head_dim_nope, head_dim_rope = ctx.dims - grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + rotary_bwd_q_kernel[grid]( dq, cos, @@ -373,12 +385,16 @@ class _MLARoPEKVTriton(torch.autograd.Function): @staticmethod def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim_v): + """Build (k, v) from kv and the shared rope head.""" if not kv.is_contiguous(): kv = kv.contiguous() s, b, nheads, _ = kv.shape o_key = kv.new_empty(s, b, nheads, head_dim_nope + head_dim_rope) o_value = kv.new_empty(s, b, nheads, head_dim_v) - grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + rotary_fwd_kv_kernel[grid]( kv, k_pos_emb, @@ -409,13 +425,17 @@ def forward(ctx, kv, k_pos_emb, cos, sin, head_dim_nope, head_dim_rope, head_dim @staticmethod def backward(ctx, dk_out, dv_out): + """Gradients for (kv, k_pos_emb) from (dk, dv).""" cos, sin = ctx.saved_tensors s, b, nheads, ndp, ndr, ndv = ctx.dims dk_out = dk_out.contiguous() dv_out = dv_out.contiguous() d_kv = dk_out.new_empty(s, b, nheads, ndp + ndv) d_emb = dk_out.new_empty(s, b, 1, ndr) - grid = lambda META: (s * b, triton.cdiv(nheads, META["BLOCK_H"])) + + def grid(meta): + return (s * b, triton.cdiv(nheads, meta["BLOCK_H"])) + rotary_bwd_kv_kernel[grid]( dk_out, dv_out, diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index af1eeb1a95..f41ec0061c 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -40,12 +40,13 @@ class DeepSeekV3Layer(torch.nn.Module): ffn size of each routed expert (required with MoE). hidden_dropout : float, default = 0.0 dropout probability on the residual branches. - kwargs common to the submodules (``q_lora_rank``, ``kv_lora_rank``, - ``qk_nope_head_dim``, ``qk_rope_head_dim``, ``v_head_dim``, - ``attention_dropout``, ``attn_mask_type``, ``qkv_format``, ``topk``, - ``num_groups``, ``group_topk``, ``routed_scaling_factor``, - ``shared_expert_ffn_hidden_size``, EP options, ...) are forwarded to - :class:`MultiLatentAttention` and :class:`DeepSeekV3MoE`. + **kwargs + kwargs common to the submodules (``q_lora_rank``, ``kv_lora_rank``, + ``qk_nope_head_dim``, ``qk_rope_head_dim``, ``v_head_dim``, + ``attention_dropout``, ``attn_mask_type``, ``qkv_format``, ``topk``, + ``num_groups``, ``group_topk``, ``routed_scaling_factor``, + ``shared_expert_ffn_hidden_size``, EP options, ...), forwarded to + :class:`MultiLatentAttention` and :class:`DeepSeekV3MoE`. """ _MLA_KWARGS = frozenset( From 5f68c9bbdaa984abb35ba071f44c0271a8fce986 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 14:09:24 +0200 Subject: [PATCH 10/43] Move model-specific layers to a dedicated docs page docs/api/pytorch_models.rst: usage (local and EP), fused-path notes, HF checkpoint weight mapping, and the class API; linked from the PyTorch API page via a toctree entry. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 8 ++- docs/api/pytorch_models.rst | 113 ++++++++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+), 3 deletions(-) create mode 100644 docs/api/pytorch_models.rst diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index bd3099b590..4fa279cfbc 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -62,11 +62,13 @@ PyTorch Model-specific layers --------------------- -.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(**kwargs) +Full transformer layers for specific model families live in +``transformer_engine.pytorch.models``: -.. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3MoE(**kwargs) +.. toctree:: + :maxdepth: 1 -.. autoapiclass:: transformer_engine.pytorch.models.MultiLatentAttention(**kwargs) + pytorch_models Data types ---------- diff --git a/docs/api/pytorch_models.rst b/docs/api/pytorch_models.rst new file mode 100644 index 0000000000..0534acb282 --- /dev/null +++ b/docs/api/pytorch_models.rst @@ -0,0 +1,113 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Model-specific layers (te.models) +================================= + +The ``transformer_engine.pytorch.models`` namespace holds full transformer +layers for specific model families, composed from Transformer Engine modules +and fused kernels. Each family lives in its own subpackage. + +DeepSeek-V3 +----------- + +A DeepSeek-V3 transformer layer analogous to +:class:`transformer_engine.pytorch.TransformerLayer`: Multi-Latent Attention +(MLA) with low-rank q/kv latents and decoupled RoPE/NoPE heads, plus a +DeepSeek-style Mixture of Experts block (fused sigmoid router with +aux-loss-free expert bias and node-limited grouped top-k, grouped-GEMM SwiGLU +experts, optional shared expert). The same architecture is used by other +model families (e.g. GLM-5, Kimi K2), which can reuse these modules. + +Basic usage (single GPU, all experts local): + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + layer = te.models.DeepSeekV3Layer( + hidden_size=7168, + num_attention_heads=128, + num_experts=64, + moe_ffn_hidden_size=2048, + topk=8, + shared_expert_ffn_hidden_size=2048, + params_dtype=torch.bfloat16, + ) + x = torch.randn(seq_len, batch, 7168, dtype=torch.bfloat16, device="cuda") + y = layer(x) # sbhd layout + +Expert parallelism routes tokens between GPUs with the NCCL EP backend +(``transformer_engine.pytorch.ep``). Call ``ep_bootstrap`` once per process +before the first forward; EP requires bfloat16 inputs and NCCL >= 2.30.4: + +.. code-block:: python + + from transformer_engine.pytorch.ep import ep_bootstrap + + ep_bootstrap(ep_group, num_experts=64, max_tokens_per_rank=tokens, + hidden_dim=7168, num_topk=8, recv_capacity_per_rank=capacity) + layer = te.models.DeepSeekV3Layer( + ..., + ep_group=ep_group, + ep_max_tokens_per_rank=tokens, + ) + +On SM100-class GPUs the routed experts fuse into a single CuTe grouped-GEMM +MLP when running under ``te.autocast`` with an MXFP8/NVFP4 recipe and +``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``; elsewhere the same modules run unfused +with an identical checkpoint layout. + +Loading HuggingFace checkpoints +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The layer follows the HF/Megatron DeepSeek-V3 conventions (interleaved rope +weights, sigmoid router bias used for selection only). Weights map from +``transformers`` ``DeepseekV3DecoderLayer`` as follows (latent RMSNorms are +fused into the up-projections): + +.. list-table:: + :header-rows: 1 + + * - Transformer Engine + - HuggingFace + * - ``input_layernorm.weight`` + - ``input_layernorm.weight`` + * - ``pre_mlp_layernorm.weight`` + - ``post_attention_layernorm.weight`` + * - ``self_attention.q_down_proj.weight`` + - ``self_attn.q_a_proj.weight`` + * - ``self_attention.q_up_proj.{layer_norm_weight, weight}`` + - ``self_attn.{q_a_layernorm, q_b_proj}.weight`` + * - ``self_attention.kv_down_proj.weight`` + - ``self_attn.kv_a_proj_with_mqa.weight`` + * - ``self_attention.kv_up_proj.{layer_norm_weight, weight}`` + - ``self_attn.{kv_a_layernorm, kv_b_proj}.weight`` + * - ``self_attention.out_proj.weight`` + - ``self_attn.o_proj.weight`` + * - ``mlp.gate.weight`` / ``mlp.expert_bias`` + - ``mlp.gate.weight`` / ``mlp.gate.e_score_correction_bias`` + * - ``mlp.experts[0].weight{i}`` + - ``interleave_glu_tensor(cat([gate_proj, up_proj]), 32)`` of expert *i* + * - ``mlp.experts[2].weight{i}`` + - ``mlp.experts.down_proj[i]`` + * - ``mlp.shared_expert[0].weight`` / ``[2].weight`` + - ``cat([gate_proj, up_proj])`` / ``down_proj`` of ``shared_experts`` + +See ``tests/pytorch/test_deepseek_hf.py`` for a complete, numerically +verified mapping. + +API +^^^ + +.. 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 From 9db495a6577397f898f8d7a812fff557e1102e6e Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 14:53:21 +0200 Subject: [PATCH 11/43] Drop HF-transformers comparison test from the repo Keep the verified weight-mapping table in the docs; the comparison itself stays as an out-of-tree script. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- docs/api/pytorch_models.rst | 4 +- tests/pytorch/test_deepseek_hf.py | 144 ------------------------------ 2 files changed, 2 insertions(+), 146 deletions(-) delete mode 100644 tests/pytorch/test_deepseek_hf.py diff --git a/docs/api/pytorch_models.rst b/docs/api/pytorch_models.rst index 0534acb282..665d547cc8 100644 --- a/docs/api/pytorch_models.rst +++ b/docs/api/pytorch_models.rst @@ -97,8 +97,8 @@ fused into the up-projections): * - ``mlp.shared_expert[0].weight`` / ``[2].weight`` - ``cat([gate_proj, up_proj])`` / ``down_proj`` of ``shared_experts`` -See ``tests/pytorch/test_deepseek_hf.py`` for a complete, numerically -verified mapping. +The routed-expert fc1 layout can be produced with +:func:`transformer_engine.pytorch.interleave_glu_tensor`. API ^^^ diff --git a/tests/pytorch/test_deepseek_hf.py b/tests/pytorch/test_deepseek_hf.py deleted file mode 100644 index 08f1c06ccd..0000000000 --- a/tests/pytorch/test_deepseek_hf.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Numeric comparison of DeepSeekV3Layer against the HF transformers reference.""" - -import pytest -import torch - -transformers = pytest.importorskip("transformers") -from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config -from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( - DeepseekV3DecoderLayer, - DeepseekV3RotaryEmbedding, -) - -from transformer_engine.pytorch.models import DeepSeekV3Layer -from transformer_engine.pytorch.utils import interleave_glu_tensor - -SEQ, BATCH = 64, 2 -HIDDEN, HEADS = 256, 4 -Q_LORA, KV_LORA = 96, 64 -NOPE, ROPE, VDIM = 64, 32, 64 -NUM_EXPERTS, TOPK, N_GROUP, TOPK_GROUP = 16, 4, 4, 2 -MOE_FFN, N_SHARED = 128, 1 -DTYPE = torch.bfloat16 - - -def _hf_config(): - return DeepseekV3Config( - hidden_size=HIDDEN, - intermediate_size=4 * HIDDEN, - moe_intermediate_size=MOE_FFN, - num_hidden_layers=1, - num_attention_heads=HEADS, - num_key_value_heads=HEADS, - n_shared_experts=N_SHARED, - n_routed_experts=NUM_EXPERTS, - routed_scaling_factor=2.5, - kv_lora_rank=KV_LORA, - q_lora_rank=Q_LORA, - qk_rope_head_dim=ROPE, - v_head_dim=VDIM, - qk_nope_head_dim=NOPE, - n_group=N_GROUP, - topk_group=TOPK_GROUP, - num_experts_per_tok=TOPK, - first_k_dense_replace=0, - norm_topk_prob=True, - rms_norm_eps=1e-5, - attention_bias=False, - attention_dropout=0.0, - rope_interleave=True, - _attn_implementation="eager", - ) - - -def _init_hf_layer(config): - torch.manual_seed(0) - layer = DeepseekV3DecoderLayer(config, layer_idx=0).to(device="cuda", dtype=DTYPE) - with torch.no_grad(): - for name, p in layer.named_parameters(): - if "layernorm" in name or "norm" in name: - p.copy_(1.0 + 0.1 * torch.randn_like(p)) - else: - p.normal_(0.0, 0.02) - bias = layer.mlp.gate.e_score_correction_bias - bias.copy_(0.1 * torch.randn_like(bias)) - return layer - - -def _build_te_layer(hf): - te_layer = DeepSeekV3Layer( - HIDDEN, - HEADS, - num_experts=NUM_EXPERTS, - moe_ffn_hidden_size=MOE_FFN, - topk=TOPK, - num_groups=N_GROUP, - group_topk=TOPK_GROUP, - routed_scaling_factor=2.5, - shared_expert_ffn_hidden_size=MOE_FFN * N_SHARED, - q_lora_rank=Q_LORA, - kv_lora_rank=KV_LORA, - qk_nope_head_dim=NOPE, - qk_rope_head_dim=ROPE, - v_head_dim=VDIM, - params_dtype=DTYPE, - ) - attn, mla = hf.self_attn, te_layer.self_attention - with torch.no_grad(): - te_layer.input_layernorm.weight.copy_(hf.input_layernorm.weight) - te_layer.pre_mlp_layernorm.weight.copy_(hf.post_attention_layernorm.weight) - - mla.q_down_proj.weight.copy_(attn.q_a_proj.weight) - mla.q_up_proj.layer_norm_weight.copy_(attn.q_a_layernorm.weight) - mla.q_up_proj.weight.copy_(attn.q_b_proj.weight) - mla.kv_down_proj.weight.copy_(attn.kv_a_proj_with_mqa.weight) - mla.kv_up_proj.layer_norm_weight.copy_(attn.kv_a_layernorm.weight) - mla.kv_up_proj.weight.copy_(attn.kv_b_proj.weight) - mla.out_proj.weight.copy_(attn.o_proj.weight) - - moe = te_layer.mlp - moe.gate.weight.copy_(hf.mlp.gate.weight) - moe.expert_bias.copy_(hf.mlp.gate.e_score_correction_bias) - fc1, _, fc2 = moe.experts - for e in range(NUM_EXPERTS): - getattr(fc1, f"weight{e}").copy_( - interleave_glu_tensor(hf.mlp.experts.gate_up_proj[e], 32) - ) - getattr(fc2, f"weight{e}").copy_(hf.mlp.experts.down_proj[e]) - shared = hf.mlp.shared_experts - moe.shared_expert[0].weight.copy_( - torch.cat([shared.gate_proj.weight, shared.up_proj.weight], dim=0) - ) - moe.shared_expert[2].weight.copy_(shared.down_proj.weight) - return te_layer - - -def test_layer_matches_hf(): - config = _hf_config() - hf = _init_hf_layer(config) - te_layer = _build_te_layer(hf) - - torch.manual_seed(1) - x = torch.randn(BATCH, SEQ, HIDDEN, dtype=DTYPE, device="cuda") - x_hf = x.clone().requires_grad_(True) - x_te = x.transpose(0, 1).contiguous().requires_grad_(True) # sbhd - - rotary = DeepseekV3RotaryEmbedding(config).to("cuda") - position_ids = torch.arange(SEQ, device="cuda").unsqueeze(0).expand(BATCH, -1) - cos, sin = rotary(x_hf, position_ids) - causal = torch.full((SEQ, SEQ), float("-inf"), device="cuda", dtype=DTYPE).triu(1) - causal = causal[None, None].expand(BATCH, 1, SEQ, SEQ) - - out_hf = hf(x_hf, attention_mask=causal, position_embeddings=(cos, sin)) - out_te = te_layer(x_te) - - torch.testing.assert_close(out_te.transpose(0, 1), out_hf, rtol=5e-2, atol=5e-2) - - grad = torch.randn_like(out_hf) - out_hf.backward(grad) - out_te.backward(grad.transpose(0, 1).contiguous()) - torch.testing.assert_close(x_te.grad.transpose(0, 1), x_hf.grad, rtol=5e-2, atol=5e-2) From 4883b1728218c11cba4c14d77f46f98a66d57f5a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Fri, 21 Aug 2026 14:54:35 +0200 Subject: [PATCH 12/43] Docs: reduce models page to a plain API listing Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 7 +-- docs/api/pytorch_models.rst | 98 +------------------------------------ 2 files changed, 4 insertions(+), 101 deletions(-) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 4fa279cfbc..1c39469f2c 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -59,11 +59,8 @@ PyTorch .. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor -Model-specific layers ---------------------- - -Full transformer layers for specific model families live in -``transformer_engine.pytorch.models``: +Models +------ .. toctree:: :maxdepth: 1 diff --git a/docs/api/pytorch_models.rst b/docs/api/pytorch_models.rst index 665d547cc8..2cde879ffb 100644 --- a/docs/api/pytorch_models.rst +++ b/docs/api/pytorch_models.rst @@ -3,106 +3,12 @@ See LICENSE for license information. -Model-specific layers (te.models) -================================= - -The ``transformer_engine.pytorch.models`` namespace holds full transformer -layers for specific model families, composed from Transformer Engine modules -and fused kernels. Each family lives in its own subpackage. +Models +====== DeepSeek-V3 ----------- -A DeepSeek-V3 transformer layer analogous to -:class:`transformer_engine.pytorch.TransformerLayer`: Multi-Latent Attention -(MLA) with low-rank q/kv latents and decoupled RoPE/NoPE heads, plus a -DeepSeek-style Mixture of Experts block (fused sigmoid router with -aux-loss-free expert bias and node-limited grouped top-k, grouped-GEMM SwiGLU -experts, optional shared expert). The same architecture is used by other -model families (e.g. GLM-5, Kimi K2), which can reuse these modules. - -Basic usage (single GPU, all experts local): - -.. code-block:: python - - import torch - import transformer_engine.pytorch as te - - layer = te.models.DeepSeekV3Layer( - hidden_size=7168, - num_attention_heads=128, - num_experts=64, - moe_ffn_hidden_size=2048, - topk=8, - shared_expert_ffn_hidden_size=2048, - params_dtype=torch.bfloat16, - ) - x = torch.randn(seq_len, batch, 7168, dtype=torch.bfloat16, device="cuda") - y = layer(x) # sbhd layout - -Expert parallelism routes tokens between GPUs with the NCCL EP backend -(``transformer_engine.pytorch.ep``). Call ``ep_bootstrap`` once per process -before the first forward; EP requires bfloat16 inputs and NCCL >= 2.30.4: - -.. code-block:: python - - from transformer_engine.pytorch.ep import ep_bootstrap - - ep_bootstrap(ep_group, num_experts=64, max_tokens_per_rank=tokens, - hidden_dim=7168, num_topk=8, recv_capacity_per_rank=capacity) - layer = te.models.DeepSeekV3Layer( - ..., - ep_group=ep_group, - ep_max_tokens_per_rank=tokens, - ) - -On SM100-class GPUs the routed experts fuse into a single CuTe grouped-GEMM -MLP when running under ``te.autocast`` with an MXFP8/NVFP4 recipe and -``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1``; elsewhere the same modules run unfused -with an identical checkpoint layout. - -Loading HuggingFace checkpoints -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -The layer follows the HF/Megatron DeepSeek-V3 conventions (interleaved rope -weights, sigmoid router bias used for selection only). Weights map from -``transformers`` ``DeepseekV3DecoderLayer`` as follows (latent RMSNorms are -fused into the up-projections): - -.. list-table:: - :header-rows: 1 - - * - Transformer Engine - - HuggingFace - * - ``input_layernorm.weight`` - - ``input_layernorm.weight`` - * - ``pre_mlp_layernorm.weight`` - - ``post_attention_layernorm.weight`` - * - ``self_attention.q_down_proj.weight`` - - ``self_attn.q_a_proj.weight`` - * - ``self_attention.q_up_proj.{layer_norm_weight, weight}`` - - ``self_attn.{q_a_layernorm, q_b_proj}.weight`` - * - ``self_attention.kv_down_proj.weight`` - - ``self_attn.kv_a_proj_with_mqa.weight`` - * - ``self_attention.kv_up_proj.{layer_norm_weight, weight}`` - - ``self_attn.{kv_a_layernorm, kv_b_proj}.weight`` - * - ``self_attention.out_proj.weight`` - - ``self_attn.o_proj.weight`` - * - ``mlp.gate.weight`` / ``mlp.expert_bias`` - - ``mlp.gate.weight`` / ``mlp.gate.e_score_correction_bias`` - * - ``mlp.experts[0].weight{i}`` - - ``interleave_glu_tensor(cat([gate_proj, up_proj]), 32)`` of expert *i* - * - ``mlp.experts[2].weight{i}`` - - ``mlp.experts.down_proj[i]`` - * - ``mlp.shared_expert[0].weight`` / ``[2].weight`` - - ``cat([gate_proj, up_proj])`` / ``down_proj`` of ``shared_experts`` - -The routed-expert fc1 layout can be produced with -:func:`transformer_engine.pytorch.interleave_glu_tensor`. - -API -^^^ - .. autoapiclass:: transformer_engine.pytorch.models.DeepSeekV3Layer(hidden_size, num_attention_heads, **kwargs) :members: forward From 2f52af137544575d6aedac704ae6cf49989a72bc Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 11:24:45 +0200 Subject: [PATCH 13/43] Rename distributed DeepSeek EP tests to generic test_models Signed-off-by: Pawel Gadzinski --- .../distributed/{run_deepseek_ep.py => run_models.py} | 2 +- .../{run_test_deepseek_ep.sh => run_test_models.sh} | 8 ++++---- .../distributed/{test_deepseek_ep.py => test_models.py} | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) rename tests/pytorch/distributed/{run_deepseek_ep.py => run_models.py} (98%) rename tests/pytorch/distributed/{run_test_deepseek_ep.sh => run_test_models.sh} (86%) rename tests/pytorch/distributed/{test_deepseek_ep.py => test_models.py} (84%) diff --git a/tests/pytorch/distributed/run_deepseek_ep.py b/tests/pytorch/distributed/run_models.py similarity index 98% rename from tests/pytorch/distributed/run_deepseek_ep.py rename to tests/pytorch/distributed/run_models.py index bf756b69ad..ff5e233503 100644 --- a/tests/pytorch/distributed/run_deepseek_ep.py +++ b/tests/pytorch/distributed/run_models.py @@ -1,7 +1,7 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Multi-process DeepSeekV3 MoE/layer EP tests, launched via torchrun.""" +"""Multi-process tests for model-specific layers (te.models), launched via torchrun.""" import os import sys diff --git a/tests/pytorch/distributed/run_test_deepseek_ep.sh b/tests/pytorch/distributed/run_test_models.sh similarity index 86% rename from tests/pytorch/distributed/run_test_deepseek_ep.sh rename to tests/pytorch/distributed/run_test_models.sh index 8c0bbbc5b9..e3411ee7d9 100644 --- a/tests/pytorch/distributed/run_test_deepseek_ep.sh +++ b/tests/pytorch/distributed/run_test_models.sh @@ -3,7 +3,7 @@ # # See LICENSE for license information. # -# Launcher for tests/pytorch/distributed/run_deepseek_ep.py. Auto-detects GPU count. +# Launcher for tests/pytorch/distributed/run_models.py (model-specific layers). Auto-detects GPU count. set -uo pipefail @@ -30,15 +30,15 @@ TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" export NCCL_EP_JIT_CACHE_DIR mkdir -p "$NCCL_EP_JIT_CACHE_DIR" -SCRIPT="${SCRIPT_DIR}/run_deepseek_ep.py" -LOG="stdout_deepseek_ep.txt" +SCRIPT="${SCRIPT_DIR}/run_models.py" +LOG="stdout_models.txt" echo "=== Running ${SCRIPT} on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ "${SCRIPT}" 2>&1 | tee "${LOG}" RC=${PIPESTATUS[0]} -pkill -9 -f "tests/pytorch/distributed/run_deepseek_ep.py" 2>/dev/null || true +pkill -9 -f "tests/pytorch/distributed/run_models.py" 2>/dev/null || true RET=0 if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; RET=1; fi diff --git a/tests/pytorch/distributed/test_deepseek_ep.py b/tests/pytorch/distributed/test_models.py similarity index 84% rename from tests/pytorch/distributed/test_deepseek_ep.py rename to tests/pytorch/distributed/test_models.py index 4a4d9a8dea..83213bd6a8 100644 --- a/tests/pytorch/distributed/test_deepseek_ep.py +++ b/tests/pytorch/distributed/test_models.py @@ -1,7 +1,7 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Pytest driver — spawns run_deepseek_ep.py under torchrun and asserts it passed.""" +"""Pytest driver — spawns run_models.py (model-specific layers, multi-GPU) under torchrun.""" import os import subprocess @@ -11,7 +11,7 @@ import torch TEST_ROOT = Path(__file__).parent.resolve() -LAUNCHER = TEST_ROOT / "run_test_deepseek_ep.sh" +LAUNCHER = TEST_ROOT / "run_test_models.sh" @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="DeepSeek EP requires >= 2 GPUs") From e841f966ea345ac11c5c612d3f311dfc0b6bf7f7 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 11:26:27 +0200 Subject: [PATCH 14/43] Rename test_deepseek.py to test_models.py and add models tests to QA scripts Signed-off-by: Pawel Gadzinski --- qa/L0_pytorch_unittest/test.sh | 1 + qa/L1_pytorch_distributed_unittest/test.sh | 1 + tests/pytorch/{test_deepseek.py => test_models.py} | 0 3 files changed, 2 insertions(+) rename tests/pytorch/{test_deepseek.py => test_models.py} (100%) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 14a5f4fe3d..ae4e183cda 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/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_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_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" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index ec19492ee7..6773055b19 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -54,6 +54,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 diff --git a/tests/pytorch/test_deepseek.py b/tests/pytorch/test_models.py similarity index 100% rename from tests/pytorch/test_deepseek.py rename to tests/pytorch/test_models.py From 935e475c1e4c267f92974111715036b96c949443 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:08:14 +0200 Subject: [PATCH 15/43] Add YaRN RoPE scaling to DeepSeek V3 MLA Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_models.py | 47 +++++++++++++ .../pytorch/models/deepseek_v3/mla_rope.py | 68 +++++++++++++++++-- .../deepseek_v3/multi_latent_attention.py | 43 +++++++++++- .../models/deepseek_v3/transformer_layer.py | 6 ++ 4 files changed, 155 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_models.py b/tests/pytorch/test_models.py index 4c4aea0a92..678c104bb8 100644 --- a/tests/pytorch/test_models.py +++ b/tests/pytorch/test_models.py @@ -2,6 +2,8 @@ # # See LICENSE for license information. +import math + import pytest import torch @@ -104,6 +106,51 @@ def test_mla_forward_backward(): assert x.grad is not None and torch.isfinite(x.grad).all() +def test_rope_tables_yarn(): + from transformer_engine.pytorch.models.deepseek_v3 import mla_rope + + s, rope = 8192, 64 + cos, sin = mla_rope.build_rope_tables(s, rope, device="cuda") + cos_none, sin_none = mla_rope.build_rope_tables(s, rope, device="cuda", scaling_factor=None) + assert torch.equal(cos, cos_none) and torch.equal(sin, sin_none) + + yarn = dict(scaling_factor=40.0, original_max_position_embeddings=4096) + cos_y, sin_y = mla_rope.build_rope_tables(s, rope, device="cuda", **yarn) + factor = mla_rope.yarn_concentration_factor(40.0, 1.0, 0.0) + assert factor == pytest.approx(0.1 * math.log(40.0) + 1.0) + # amplitude scaled by the concentration factor + torch.testing.assert_close(cos_y**2 + sin_y**2, torch.full_like(cos_y, factor**2)) + # high-frequency dims untouched, low-frequency dims interpolated by 1/scaling_factor + torch.testing.assert_close(cos_y[:, 0] / factor, cos[:, 0]) + angle_y = torch.atan2(sin_y[:, rope // 2 - 1], cos_y[:, rope // 2 - 1]) + angle = torch.atan2(sin[:, rope // 2 - 1], cos[:, rope // 2 - 1]) + torch.testing.assert_close(angle_y[:64], angle[:64] / 40.0, atol=1e-4, rtol=0) + + +@pytest.mark.parametrize("mscale_all_dim", [0.0, 1.0]) +def test_mla_yarn_forward_backward(mscale_all_dim): + torch.manual_seed(0) + mla = MultiLatentAttention( + HIDDEN, + HEADS, + params_dtype=DTYPE, + rope_scaling_factor=40.0, + original_max_position_embeddings=64, + mscale_all_dim=mscale_all_dim, + **MLA_KWARGS, + ) + m = 0.1 * mscale_all_dim * math.log(40.0) + 1.0 + qk_head_dim = MLA_KWARGS["qk_nope_head_dim"] + MLA_KWARGS["qk_rope_head_dim"] + assert mla.core_attention.unfused_attention.softmax_scale == pytest.approx( + m * m / math.sqrt(qk_head_dim) + ) + x = _input() + out = mla(x) + assert out.shape == x.shape + out.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + @pytest.mark.parametrize("shared", [False, True], ids=["no_shared", "shared"]) @pytest.mark.parametrize("grouped", [False, True], ids=["ungrouped", "grouped"]) def test_moe_forward_backward(shared, grouped): diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py index f8c01f75d6..341850fddb 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -16,6 +16,7 @@ matching the Megatron fused kernel semantics. """ +import math from typing import Optional, Tuple import torch @@ -28,7 +29,44 @@ except ImportError: HAVE_TRITON = False -__all__ = ["build_rope_tables", "apply_mla_rope_q", "apply_mla_rope_kv"] +__all__ = [ + "build_rope_tables", + "apply_mla_rope_q", + "apply_mla_rope_kv", + "yarn_mscale", + "yarn_concentration_factor", +] + + +def _yarn_correction_dim(num_rotations, dim, base, max_pos): + return (dim * math.log(max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) + + +def _yarn_correction_range(beta_fast, beta_slow, dim, base, max_pos, round_to_int=True): + low = _yarn_correction_dim(beta_fast, dim, base, max_pos) + high = _yarn_correction_dim(beta_slow, dim, base, max_pos) + if round_to_int: + low, high = math.floor(low), math.ceil(high) + return max(low, 0), min(high, dim - 1) + + +def _yarn_linear_ramp(low, high, dim, device): + if low == high: + high += 0.001 + ramp = (torch.arange(dim, dtype=torch.float32, device=device) - low) / (high - low) + return torch.clamp(ramp, 0, 1) + + +def yarn_mscale(scale: float, mscale: float = 1.0) -> float: + """YaRN attention temperature factor ``0.1 * mscale * ln(scale) + 1`` (1 for scale <= 1).""" + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def yarn_concentration_factor(scaling_factor: float, mscale: float, mscale_all_dim: float) -> float: + """Factor multiplied into cos/sin tables (as in Megatron-Core).""" + return yarn_mscale(scaling_factor, mscale) / yarn_mscale(scaling_factor, mscale_all_dim) def build_rope_tables( @@ -36,15 +74,33 @@ def build_rope_tables( emb_dim: int, base: float = 10000.0, device: Optional[torch.device] = None, + scaling_factor: Optional[float] = None, + original_max_position_embeddings: int = 4096, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, ) -> Tuple[torch.Tensor, torch.Tensor]: - """cos/sin tables of shape ``[seq_len, emb_dim]`` (fp32, NeoX duplicated halves).""" - inv_freq = 1.0 / ( - base ** (torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim) - ) + """cos/sin tables of shape ``[seq_len, emb_dim]`` (fp32, NeoX duplicated halves). + + With ``scaling_factor`` set, frequencies follow YaRN (NTK-by-parts ramp between + ``beta_fast``/``beta_slow`` rotations over ``original_max_position_embeddings``) and the + tables are scaled by the YaRN concentration factor. + """ + exponent = torch.arange(0, emb_dim, 2, dtype=torch.float32, device=device) / emb_dim + inv_freq = 1.0 / (base**exponent) + factor = 1.0 + if scaling_factor is not None: + low, high = _yarn_correction_range( + beta_fast, beta_slow, emb_dim, base, original_max_position_embeddings + ) + extra_mask = 1.0 - _yarn_linear_ramp(low, high, emb_dim // 2, device) + inv_freq = (inv_freq / scaling_factor) * (1 - extra_mask) + inv_freq * extra_mask + factor = yarn_concentration_factor(scaling_factor, mscale, mscale_all_dim) t = torch.arange(seq_len, device=device, dtype=torch.float32) freqs = torch.outer(t, inv_freq) freqs = torch.cat([freqs, freqs], dim=-1) - return torch.cos(freqs).contiguous(), torch.sin(freqs).contiguous() + return (torch.cos(freqs) * factor).contiguous(), (torch.sin(freqs) * factor).contiguous() if HAVE_TRITON: diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index 0ddf2f75b2..ca9c5ed3a7 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -4,6 +4,7 @@ """Multi-Latent Attention (MLA) block as used in DeepSeekV3.""" +import math from typing import Optional, Union import torch @@ -14,6 +15,7 @@ apply_mla_rope_kv, apply_mla_rope_q, build_rope_tables, + yarn_mscale, ) __all__ = ["MultiLatentAttention"] @@ -61,9 +63,22 @@ class MultiLatentAttention(torch.nn.Module): epsilon of the latent RMSNorms (matches DeepSeekV3). rotary_base : float, default = 10000.0 RoPE base. + rope_scaling_factor : float, optional + YaRN context-extension factor; ``None`` disables YaRN. + original_max_position_embeddings : int, default = 4096 + pre-extension context length (YaRN). + beta_fast : float, default = 32.0 + YaRN high-frequency rotation bound. + beta_slow : float, default = 1.0 + YaRN low-frequency rotation bound. + mscale : float, default = 1.0 + YaRN mscale of the rope part. + mscale_all_dim : float, default = 0.0 + YaRN mscale of all dims; sets the default softmax scale to + ``m**2 / sqrt(qk head dim)`` with ``m = 0.1 * mscale_all_dim * ln(factor) + 1``. softmax_scale : float, optional - softmax scale; defaults to ``1/sqrt(qk head dim)`` inside - :class:`DotProductAttention`. + softmax scale; defaults to ``1/sqrt(qk head dim)`` (times the YaRN + ``m**2`` when YaRN is enabled). qkv_format : str, default = "sbhd" layout of the input/output tensors. params_dtype : torch.dtype, optional @@ -87,6 +102,12 @@ def __init__( attn_mask_type: str = "causal", layernorm_epsilon: float = 1e-6, rotary_base: float = 10000.0, + rope_scaling_factor: Optional[float] = None, + original_max_position_embeddings: int = 4096, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, softmax_scale: Optional[float] = None, qkv_format: str = "sbhd", params_dtype: Optional[torch.dtype] = None, @@ -140,8 +161,20 @@ def __init__( ) self.rotary_base = rotary_base + self._yarn_kwargs = dict( + scaling_factor=rope_scaling_factor, + original_max_position_embeddings=original_max_position_embeddings, + beta_fast=beta_fast, + beta_slow=beta_slow, + mscale=mscale, + mscale_all_dim=mscale_all_dim, + ) self._rope_tables: Optional[tuple] = None + if softmax_scale is None and rope_scaling_factor is not None: + m = yarn_mscale(rope_scaling_factor, mscale_all_dim) + softmax_scale = m * m / math.sqrt(self.qk_head_dim) + self.core_attention = DotProductAttention( num_attention_heads, kv_channels=(self.qk_head_dim, v_head_dim), @@ -156,7 +189,11 @@ def __init__( def _rope_tables_for(self, seq_len: int, device: torch.device): if self._rope_tables is None or self._rope_tables[0].shape[0] < seq_len: self._rope_tables = build_rope_tables( - seq_len, self.qk_rope_head_dim, base=self.rotary_base, device=device + seq_len, + self.qk_rope_head_dim, + base=self.rotary_base, + device=device, + **self._yarn_kwargs, ) cos, sin = self._rope_tables return cos[:seq_len], sin[:seq_len] diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index f41ec0061c..aa2fab232d 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -59,6 +59,12 @@ class DeepSeekV3Layer(torch.nn.Module): "attention_dropout", "attn_mask_type", "rotary_base", + "rope_scaling_factor", + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", "softmax_scale", "qkv_format", "tp_group", From 7eaefd97fc0790f0a1805b0d3f098fe911c45dc5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:10:54 +0200 Subject: [PATCH 16/43] Drop tests/pytorch/attention/mla_rope_utils.py shim; use models.deepseek_v3.mla_rope directly Signed-off-by: Pawel Gadzinski --- tests/pytorch/attention/mla_rope_utils.py | 42 ------------------- .../attention/test_linear_mxfp8_attention.py | 29 +++++++------ 2 files changed, 17 insertions(+), 54 deletions(-) delete mode 100644 tests/pytorch/attention/mla_rope_utils.py diff --git a/tests/pytorch/attention/mla_rope_utils.py b/tests/pytorch/attention/mla_rope_utils.py deleted file mode 100644 index d022757886..0000000000 --- a/tests/pytorch/attention/mla_rope_utils.py +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Compat shim: the MLA RoPE kernels moved to -``transformer_engine.pytorch.models.deepseek_v3.mla_rope``.""" - -import torch - -from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( # noqa: F401 - HAVE_TRITON, - apply_mla_rope_kv, - apply_mla_rope_q, - build_rope_tables, -) - -HEAD_DIM_ROPE = 64 -HEAD_DIM_NOPE = 128 -HEAD_DIM_V = 128 -ROTARY_BASE = 10000 - - -def apply_mla_rope( - q: torch.Tensor, - kv: torch.Tensor, - k_pos_emb: torch.Tensor, - head_dim_nope: int = HEAD_DIM_NOPE, - head_dim_rope: int = HEAD_DIM_ROPE, - head_dim_v: int = HEAD_DIM_V, - base: int = ROTARY_BASE, - cos_table: torch.Tensor | None = None, - sin_table: torch.Tensor | None = None, -): - if cos_table is None or sin_table is None: - cos_table, sin_table = build_rope_tables( - q.shape[0], head_dim_rope, base=base, device=q.device - ) - q = apply_mla_rope_q(q, cos_table, sin_table, head_dim_nope, head_dim_rope) - k, v = apply_mla_rope_kv( - kv, k_pos_emb, cos_table, sin_table, head_dim_nope, head_dim_rope, head_dim_v - ) - return q, k, v diff --git a/tests/pytorch/attention/test_linear_mxfp8_attention.py b/tests/pytorch/attention/test_linear_mxfp8_attention.py index f1bba7bc9a..95770a6bb8 100644 --- a/tests/pytorch/attention/test_linear_mxfp8_attention.py +++ b/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,6 +187,13 @@ 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, @@ -190,7 +201,7 @@ def _run_forward_bf16( ) -> 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 From cef2e39607767a03954824de5f257ee02b7bcb0b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:26:37 +0200 Subject: [PATCH 17/43] Distributed models test: single full DeepSeekV3Layer EP-vs-local numerical comparison Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 91 ++++++++++--------------- 1 file changed, 37 insertions(+), 54 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index ff5e233503..8bfdd7bfb3 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -11,7 +11,7 @@ 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, DeepSeekV3MoE +from transformer_engine.pytorch.models import DeepSeekV3Layer HIDDEN = 256 MOE_FFN = 128 @@ -65,64 +65,72 @@ def setUpClass(cls): recv_capacity_per_rank=_recv_capacity(cls.ep_size), ) - def _make_moe(self, ep: bool, shared: bool = True) -> DeepSeekV3MoE: - return DeepSeekV3MoE( + def _make_layer(self, ep: bool) -> DeepSeekV3Layer: + return DeepSeekV3Layer( HIDDEN, - moe_ffn_hidden_size=MOE_FFN, + HEADS, num_experts=self.num_experts, + moe_ffn_hidden_size=MOE_FFN, topk=TOP_K, - shared_expert_ffn_hidden_size=SHARED_FFN if shared else None, + shared_expert_ffn_hidden_size=SHARED_FFN, params_dtype=DTYPE, ep_group=self.ep_group if ep else None, ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, ep_recv_capacity_per_rank=_recv_capacity(self.ep_size) if ep else None, + **MLA_KWARGS, ) - def _copy_local_expert_weights(self, ep_moe: DeepSeekV3MoE, ref: DeepSeekV3MoE) -> None: + def _copy_weights(self, ep_layer: DeepSeekV3Layer, ref: DeepSeekV3Layer) -> None: + ref_params = dict(ref.named_parameters()) + ref_bufs = dict(ref.named_buffers()) with torch.no_grad(): - ep_moe.gate.weight.copy_(ref.gate.weight) - if ref.shared_expert is not None: - for dst, src in zip( - ep_moe.shared_expert.parameters(), ref.shared_expert.parameters() - ): - dst.copy_(src) - ep_fc1, _, ep_fc2 = ep_moe.experts - ref_fc1, _, ref_fc2 = ref.experts + 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 = self.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_moe_ep_matches_local(self): - """EP MoE must match the single-GPU (all-experts-local) path numerically.""" + def test_layer_ep_matches_local(self): + """Full DeepSeekV3Layer with EP must match the all-experts-local layer numerically.""" torch.manual_seed(0) - ref = self._make_moe(ep=False) + ref = self._make_layer(ep=False) _broadcast_params(ref) - ep_moe = self._make_moe(ep=True) - self._copy_local_expert_weights(ep_moe, ref) + ep_layer = self._make_layer(ep=True) + self._copy_weights(ep_layer, ref) torch.manual_seed(1234 + self.rank) - x = torch.randn(TOKENS_PER_RANK, HIDDEN, dtype=DTYPE, device="cuda") + 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_moe(x_ep) + out_ep = ep_layer(x_ep) out_ref = ref(x_ref) + self.assertEqual(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) - torch.testing.assert_close( - ep_moe.gate.weight.grad, ref.gate.weight.grad, rtol=0.1, atol=0.1 - ) + + 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_moe.experts - ref_fc1, _, ref_fc2 = ref.experts + 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(self.num_experts) @@ -134,37 +142,12 @@ def test_moe_ep_matches_local(self): 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_moe._last_tokens_per_expert.clone() + counts = ep_layer.mlp._last_tokens_per_expert.clone() dist.all_reduce(counts) self.assertEqual(counts.sum().item(), self.ep_size * TOKENS_PER_RANK * TOP_K) - def test_layer_ep_forward_backward(self): - """Full DeepSeekV3Layer smoke test with an EP MoE block.""" - torch.manual_seed(10 + self.rank) - layer = DeepSeekV3Layer( - HIDDEN, - HEADS, - num_experts=self.num_experts, - moe_ffn_hidden_size=MOE_FFN, - topk=TOP_K, - shared_expert_ffn_hidden_size=SHARED_FFN, - params_dtype=DTYPE, - ep_group=self.ep_group, - ep_max_tokens_per_rank=TOKENS_PER_RANK, - ep_recv_capacity_per_rank=_recv_capacity(self.ep_size), - **MLA_KWARGS, - ) - x = torch.randn( - TOKENS_PER_RANK // 2, 2, HIDDEN, dtype=DTYPE, device="cuda", requires_grad=True - ) - out = layer(x) - self.assertEqual(out.shape, x.shape) - out.sum().backward() - self.assertIsNotNone(x.grad) - self.assertTrue(torch.isfinite(x.grad).all()) - - layer.mlp.update_expert_bias() - self.assertTrue(torch.isfinite(layer.mlp.expert_bias).all()) + ep_layer.mlp.update_expert_bias() + self.assertTrue(torch.isfinite(ep_layer.mlp.expert_bias).all()) def _init_distributed(): From 5453fea6f0f760d6e3c05e4cbe653a412889ae96 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:32:05 +0200 Subject: [PATCH 18/43] run_models.py: plain main() instead of unittest, simplify launcher Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 224 ++++++++++--------- tests/pytorch/distributed/run_test_models.sh | 19 +- tests/pytorch/distributed/test_models.py | 2 +- 3 files changed, 118 insertions(+), 127 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index 8bfdd7bfb3..ed41f020f9 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -5,7 +5,6 @@ import os import sys -import unittest import torch import torch.distributed as dist @@ -46,111 +45,94 @@ def _broadcast_params(module: torch.nn.Module) -> None: dist.broadcast(t.detach(), src=0) -class TestDeepSeekEP(unittest.TestCase): - @classmethod - def setUpClass(cls): - if _device_sm() < 90: - raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{_device_sm()})") - cls.rank = dist.get_rank() - cls.ep_size = dist.get_world_size() - cls.num_experts = NUM_LOCAL_EXPERTS * cls.ep_size - world_pg = dist.distributed_c10d._get_default_group() - cls.ep_group = dist.new_group(ranks=list(range(world_pg.size())), backend="nccl") - ep_bootstrap( - cls.ep_group, - num_experts=cls.num_experts, - max_tokens_per_rank=TOKENS_PER_RANK, - hidden_dim=HIDDEN, - num_topk=TOP_K, - recv_capacity_per_rank=_recv_capacity(cls.ep_size), - ) - - def _make_layer(self, ep: bool) -> DeepSeekV3Layer: - return DeepSeekV3Layer( - HIDDEN, - HEADS, - num_experts=self.num_experts, - moe_ffn_hidden_size=MOE_FFN, - topk=TOP_K, - shared_expert_ffn_hidden_size=SHARED_FFN, - params_dtype=DTYPE, - ep_group=self.ep_group if ep else None, - ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, - ep_recv_capacity_per_rank=_recv_capacity(self.ep_size) if ep else None, - **MLA_KWARGS, - ) - - def _copy_weights(self, ep_layer: DeepSeekV3Layer, ref: DeepSeekV3Layer) -> 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 = self.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(self): - """Full DeepSeekV3Layer with EP must match the all-experts-local layer numerically.""" - torch.manual_seed(0) - ref = self._make_layer(ep=False) - _broadcast_params(ref) - ep_layer = self._make_layer(ep=True) - self._copy_weights(ep_layer, ref) - - torch.manual_seed(1234 + self.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) - self.assertEqual(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()) +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, + ep_recv_capacity_per_rank=_recv_capacity(ep_size) 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 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). + 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 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(self.num_experts) - ] - for g in ref_grads: - dist.all_reduce(g) - for local_e in range(NUM_LOCAL_EXPERTS): - global_e = self.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) - self.assertEqual(counts.sum().item(), self.ep_size * TOKENS_PER_RANK * TOP_K) - - ep_layer.mlp.update_expert_bias() - self.assertTrue(torch.isfinite(ep_layer.mlp.expert_bias).all()) - - -def _init_distributed(): + 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"])) try: @@ -160,13 +142,33 @@ def _init_distributed(): except (ImportError, RuntimeError): pass + 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), + ) + try: + test_layer_ep_matches_local(rank, ep_size, ep_group) + print(f"[rank {rank}] PASSED") + finally: + dist.barrier() + ep_finalize() + release_symm_mem_pool() + dist.destroy_process_group() + return 0 + if __name__ == "__main__": - _init_distributed() - suite = unittest.TestLoader().loadTestsFromTestCase(TestDeepSeekEP) - result = unittest.TextTestRunner(stream=sys.stdout, verbosity=2).run(suite) - dist.barrier() - ep_finalize() - release_symm_mem_pool() - dist.destroy_process_group() - sys.exit(0 if result.wasSuccessful() else 1) + sys.exit(main()) diff --git a/tests/pytorch/distributed/run_test_models.sh b/tests/pytorch/distributed/run_test_models.sh index e3411ee7d9..b0230b10ac 100644 --- a/tests/pytorch/distributed/run_test_models.sh +++ b/tests/pytorch/distributed/run_test_models.sh @@ -31,22 +31,11 @@ export NCCL_EP_JIT_CACHE_DIR mkdir -p "$NCCL_EP_JIT_CACHE_DIR" SCRIPT="${SCRIPT_DIR}/run_models.py" -LOG="stdout_models.txt" echo "=== Running ${SCRIPT} on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ - torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ - "${SCRIPT}" 2>&1 | tee "${LOG}" -RC=${PIPESTATUS[0]} + torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" "${SCRIPT}" +RC=$? pkill -9 -f "tests/pytorch/distributed/run_models.py" 2>/dev/null || true - -RET=0 -if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; RET=1; fi -if grep -qE "(^|]:)FAILED|(^|]:)Traceback" "${LOG}"; then RET=1; fi -if ! grep -qE "Ran [0-9]+ test|^OK$" "${LOG}"; then - echo "ERROR: no test summary — likely hang or early crash" - RET=1 -fi -if [ -z "${KEEP_EP_LOGS:-}" ]; then rm -f "${LOG}"; fi - -exit $RET +if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; fi +exit $RC diff --git a/tests/pytorch/distributed/test_models.py b/tests/pytorch/distributed/test_models.py index 83213bd6a8..d7aa3a478e 100644 --- a/tests/pytorch/distributed/test_models.py +++ b/tests/pytorch/distributed/test_models.py @@ -19,7 +19,7 @@ def test_multi_process_deepseek_ep(): timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) proc = subprocess.run( ["bash", str(LAUNCHER)], - env={**os.environ, "KEEP_EP_LOGS": "1", "TEST_TIMEOUT_S": str(timeout_s)}, + env={**os.environ, "TEST_TIMEOUT_S": str(timeout_s)}, timeout=timeout_s + 30, check=False, ) From f24835b5ab7072665fc45e8c2261db2b6fe3469d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:37:56 +0200 Subject: [PATCH 19/43] run_models.py: fail hard instead of swallowing symm-mem/cleanup errors Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index ed41f020f9..a77b18e205 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -135,12 +135,9 @@ def test_layer_ep_matches_local(rank: int, ep_size: int, ep_group) -> None: def main() -> int: dist.init_process_group(backend="nccl") torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) - try: - from torch.distributed import _symmetric_memory as _symm_mem + from torch.distributed import _symmetric_memory as _symm_mem - _symm_mem.set_backend("NCCL") - except (ImportError, RuntimeError): - pass + _symm_mem.set_backend("NCCL") rank = dist.get_rank() ep_size = dist.get_world_size() @@ -159,14 +156,13 @@ def main() -> int: num_topk=TOP_K, recv_capacity_per_rank=_recv_capacity(ep_size), ) - try: - test_layer_ep_matches_local(rank, ep_size, ep_group) - print(f"[rank {rank}] PASSED") - finally: - dist.barrier() - ep_finalize() - release_symm_mem_pool() - dist.destroy_process_group() + 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 From d66fc1c872d50d31bc82cf9f2a7d3e27fe41ef6d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:39:30 +0200 Subject: [PATCH 20/43] DeepSeekV3MoE docstring: ep_bootstrap must precede construction Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/models/deepseek_v3/moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 3c182b4405..caa2359e5a 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -45,7 +45,8 @@ class DeepSeekV3MoE(torch.nn.Module): given, expert-parallel over NCCL (``ep_dispatch``/``ep_combine``). When expert parallelism is used, ``transformer_engine.pytorch.ep.ep_bootstrap`` - must be called once per process before the first forward, and inputs must + must be called once per process before constructing the module (it allocates + the ``EpBuffer`` in ``__init__``), and inputs must be bfloat16. Parameters From 80041fadbef3c9d39a673f59a0c972bab672b413 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 12:43:28 +0200 Subject: [PATCH 21/43] Distributed models test: launch torchrun directly from pytest, drop shell launcher Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_test_models.sh | 41 -------------------- tests/pytorch/distributed/test_models.py | 27 +++++++------ 2 files changed, 16 insertions(+), 52 deletions(-) delete mode 100644 tests/pytorch/distributed/run_test_models.sh diff --git a/tests/pytorch/distributed/run_test_models.sh b/tests/pytorch/distributed/run_test_models.sh deleted file mode 100644 index b0230b10ac..0000000000 --- a/tests/pytorch/distributed/run_test_models.sh +++ /dev/null @@ -1,41 +0,0 @@ -#!/bin/bash -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. -# -# Launcher for tests/pytorch/distributed/run_models.py (model-specific layers). Auto-detects GPU count. - -set -uo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) -if [ "${DETECTED_GPUS}" -lt 2 ]; then - echo "DeepSeek EP test requires >= 2 GPUs (found ${DETECTED_GPUS}); SKIPPING." - exit 0 -fi - -# NCCL EP requires active NVLink P2P among ranks on the node. -if ! nvidia-smi nvlink --status 2>/dev/null | grep -qE 'Link [0-9]+:.*GB/s'; then - echo "No NVLink between GPUs (PCIe-only fabric); NCCL EP is unsupported here. SKIPPING." - exit 0 -fi - -NUM_RANKS="${NVTE_TEST_EP_NUM_RANKS:-${DETECTED_GPUS}}" -if [ "${NUM_RANKS}" -gt 8 ]; then NUM_RANKS=8; fi - -TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" - -: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} -export NCCL_EP_JIT_CACHE_DIR -mkdir -p "$NCCL_EP_JIT_CACHE_DIR" - -SCRIPT="${SCRIPT_DIR}/run_models.py" - -echo "=== Running ${SCRIPT} on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" -setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ - torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" "${SCRIPT}" -RC=$? -pkill -9 -f "tests/pytorch/distributed/run_models.py" 2>/dev/null || true -if [ "${RC}" -ne 0 ]; then echo "torchrun exited with ${RC}"; fi -exit $RC diff --git a/tests/pytorch/distributed/test_models.py b/tests/pytorch/distributed/test_models.py index d7aa3a478e..1b96eae2aa 100644 --- a/tests/pytorch/distributed/test_models.py +++ b/tests/pytorch/distributed/test_models.py @@ -1,7 +1,6 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Pytest driver — spawns run_models.py (model-specific layers, multi-GPU) under torchrun.""" import os import subprocess @@ -11,16 +10,22 @@ import torch TEST_ROOT = Path(__file__).parent.resolve() -LAUNCHER = TEST_ROOT / "run_test_models.sh" +NUM_PROCS = min(8, torch.cuda.device_count()) +LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="DeepSeek EP requires >= 2 GPUs") -def test_multi_process_deepseek_ep(): - timeout_s = int(os.environ.get("NVTE_TEST_EP_TIMEOUT_S", "180")) - proc = subprocess.run( - ["bash", str(LAUNCHER)], - env={**os.environ, "TEST_TIMEOUT_S": str(timeout_s)}, - timeout=timeout_s + 30, - check=False, +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 proc.returncode == 0, f"DeepSeek EP test suite failed (rc={proc.returncode})" + assert result.returncode == 0 From 2475a9113d988f784ec3b649799431224a899223 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:04:23 +0200 Subject: [PATCH 22/43] Add DeepSeekV3Layer to test_sanity; pad per-expert rows for quantized grouped GEMM in local MoE path Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_sanity.py | 36 +++++++++++++++++++ .../pytorch/models/deepseek_v3/moe.py | 34 ++++++++++++++---- 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index c9b620fa1e..835813d5c6 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -35,6 +35,7 @@ is_bf16_available, ) from transformer_engine.common import recipe +from transformer_engine.pytorch.models import DeepSeekV3Layer from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.pytorch.tensor.utils import replace_raw_data from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported @@ -736,6 +737,41 @@ def test_sanity_layernorm_mlp( _test_sanity_common(block, dtype, config, fp8_recipe, skip_wgrad, skip_dgrad, microbatching) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) +@pytest.mark.parametrize("model", ["small"]) +@pytest.mark.parametrize("skip_wgrad", all_boolean) +@pytest.mark.parametrize("moe", all_boolean) +def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, model, skip_wgrad, moe): + config = model_configs[model] + + if fp8_recipe is not None: + if not is_fp8_supported(config): + pytest.skip("Model config does not support FP8") + if fp8_recipe.nvfp4() and dtype == torch.float16: + pytest.skip("FP16 output for NVFP4 not supported") + + mlp_kwargs = ( + dict(num_experts=4, topk=2, moe_ffn_hidden_size=32, shared_expert_ffn_hidden_size=32) + if moe + else dict(ffn_hidden_size=4 * config.hidden_size) + ) + block = DeepSeekV3Layer( + config.hidden_size, + config.num_heads, + q_lora_rank=16, + kv_lora_rank=16, + qk_nope_head_dim=16, + qk_rope_head_dim=16, + v_head_dim=16, + params_dtype=dtype, + device="cuda", + **mlp_kwargs, + ) + + _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad) + + @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index caa2359e5a..60080d8816 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -11,7 +11,15 @@ import transformer_engine.pytorch.ops as te_ops from transformer_engine.pytorch.router import fused_topk_with_score_function -from transformer_engine.pytorch.permutation import moe_permute_with_probs, moe_unpermute +from transformer_engine.pytorch.permutation import ( + moe_permute_and_pad_with_probs, + moe_permute_with_probs, + moe_unpermute, +) +from transformer_engine.pytorch.quantization import ( + FP8GlobalStateManager, + get_align_size_for_quantization, +) __all__ = ["DeepSeekV3MoE"] @@ -189,14 +197,24 @@ def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: tokens_per_expert = routing_map.sum(dim=0) self._last_tokens_per_expert = tokens_per_expert.detach() - num_out = tokens.shape[0] * self.topk - permuted, permuted_probs, row_id_map = moe_permute_with_probs( - tokens, probs, routing_map, num_out_tokens=num_out - ) + # Quantized grouped GEMMs need every expert's row count aligned. + align = 1 + if FP8GlobalStateManager.is_fp8_enabled(): + align = get_align_size_for_quantization(FP8GlobalStateManager.get_fp8_recipe()) + if align > 1: + permuted, permuted_probs, row_id_map, pad_offsets, tokens_per_expert = ( + moe_permute_and_pad_with_probs(tokens, probs, routing_map, tokens_per_expert, align) + ) + else: + permuted, permuted_probs, row_id_map = moe_permute_with_probs( + tokens, probs, routing_map, num_out_tokens=tokens.shape[0] * self.topk + ) + pad_offsets = None # The fused grouped MLP requires the total row count to be a multiple # of 128; rows beyond sum(tokens_per_expert) fall outside every group. - pad = (-num_out) % 128 + num_rows = permuted.shape[0] + pad = (-num_rows) % 128 if pad: permuted = torch.nn.functional.pad(permuted, (0, 0, 0, pad)) permuted_probs = torch.nn.functional.pad(permuted_probs, (0, pad)) @@ -204,7 +222,9 @@ def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: out = self.experts( permuted, tokens_per_expert, permuted_probs.to(tokens.dtype), tokens_per_expert ) - return moe_unpermute(out[:num_out], row_id_map, restore_shape=tokens.shape) + return moe_unpermute( + out[:num_rows], row_id_map, restore_shape=tokens.shape, pad_offsets=pad_offsets + ) def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: from transformer_engine.pytorch.ep import ep_dispatch, ep_combine From babc5e7102b5251f978a83a312a31ee2c0a9a060 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:21:08 +0200 Subject: [PATCH 23/43] Tests: drop fwd/bwd smoke tests covered by sanity, trim sanity combos; expose MLA softmax_scale; clean docstrings Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_models.py | 101 ++++-------------- tests/pytorch/test_sanity.py | 10 +- .../pytorch/models/deepseek_v3/mla_rope.py | 20 ++-- .../deepseek_v3/multi_latent_attention.py | 3 +- 4 files changed, 37 insertions(+), 97 deletions(-) diff --git a/tests/pytorch/test_models.py b/tests/pytorch/test_models.py index 678c104bb8..cd1903e79f 100644 --- a/tests/pytorch/test_models.py +++ b/tests/pytorch/test_models.py @@ -8,11 +8,7 @@ import torch from transformer_engine.pytorch.utils import deinterleave_glu_tensor -from transformer_engine.pytorch.models import ( - DeepSeekV3Layer, - DeepSeekV3MoE, - MultiLatentAttention, -) +from transformer_engine.pytorch.models import DeepSeekV3MoE, MultiLatentAttention SEQ_LEN = 128 BATCH = 2 @@ -96,16 +92,6 @@ def run(fmt): torch.testing.assert_close(grads_t[2], pos_leaf.grad, rtol=1e-5, atol=1e-5) -def test_mla_forward_backward(): - torch.manual_seed(0) - mla = MultiLatentAttention(HIDDEN, HEADS, params_dtype=DTYPE, **MLA_KWARGS) - x = _input() - out = mla(x) - assert out.shape == x.shape - out.sum().backward() - assert x.grad is not None and torch.isfinite(x.grad).all() - - def test_rope_tables_yarn(): from transformer_engine.pytorch.models.deepseek_v3 import mla_rope @@ -128,8 +114,7 @@ def test_rope_tables_yarn(): @pytest.mark.parametrize("mscale_all_dim", [0.0, 1.0]) -def test_mla_yarn_forward_backward(mscale_all_dim): - torch.manual_seed(0) +def test_mla_yarn_softmax_scale(mscale_all_dim): mla = MultiLatentAttention( HIDDEN, HEADS, @@ -141,27 +126,23 @@ def test_mla_yarn_forward_backward(mscale_all_dim): ) m = 0.1 * mscale_all_dim * math.log(40.0) + 1.0 qk_head_dim = MLA_KWARGS["qk_nope_head_dim"] + MLA_KWARGS["qk_rope_head_dim"] - assert mla.core_attention.unfused_attention.softmax_scale == pytest.approx( - m * m / math.sqrt(qk_head_dim) - ) - x = _input() - out = mla(x) - assert out.shape == x.shape - out.sum().backward() - assert x.grad is not None and torch.isfinite(x.grad).all() + assert mla.softmax_scale == pytest.approx(m * m / math.sqrt(qk_head_dim)) @pytest.mark.parametrize("shared", [False, True], ids=["no_shared", "shared"]) @pytest.mark.parametrize("grouped", [False, True], ids=["ungrouped", "grouped"]) -def test_moe_forward_backward(shared, grouped): +@pytest.mark.parametrize("topk", [2, 4]) +def test_moe_matches_dense_reference(shared, grouped, topk): + """Routed output must equal the prob-weighted sum of the selected expert MLPs.""" torch.manual_seed(0) + num_experts = 4 moe = DeepSeekV3MoE( HIDDEN, moe_ffn_hidden_size=128, - num_experts=8, - topk=2, - num_groups=4 if grouped else None, - group_topk=2 if grouped else None, + num_experts=num_experts, + topk=topk, + num_groups=2 if grouped else None, + group_topk=topk // 2 if grouped else None, shared_expert_ffn_hidden_size=128 if shared else None, params_dtype=DTYPE, ) @@ -169,32 +150,13 @@ def test_moe_forward_backward(shared, grouped): out = moe(x) assert out.shape == x.shape out.sum().backward() - assert x.grad is not None and torch.isfinite(x.grad).all() - - counts = moe._last_tokens_per_expert - assert counts.sum().item() == SEQ_LEN * BATCH * 2 - bias_before = moe.expert_bias.clone() - moe.update_expert_bias() - assert not torch.equal(bias_before, moe.expert_bias) - - -def test_moe_matches_dense_reference(): - """topk == num_experts with uniform probs must reduce to a sum of expert MLPs.""" - torch.manual_seed(0) - num_experts = 4 - moe = DeepSeekV3MoE( - HIDDEN, - moe_ffn_hidden_size=128, - num_experts=num_experts, - topk=num_experts, - routed_scaling_factor=1.0, - params_dtype=DTYPE, - ) - x = _input(requires_grad=False) - out = moe(x) + assert torch.isfinite(x.grad).all() - tokens = x.reshape(-1, HIDDEN) + tokens = x.detach().reshape(-1, HIDDEN) probs, _ = moe._route(moe.gate(tokens).float()) + assert (probs > 0).sum(dim=1).eq(topk).all() + assert moe._last_tokens_per_expert.sum().item() == tokens.shape[0] * topk + fc1, _, fc2 = moe.experts ref = torch.zeros_like(tokens) for e in range(num_experts): @@ -203,29 +165,12 @@ def test_moe_matches_dense_reference(): gate_part, lin_part = (tokens @ w1.t()).chunk(2, dim=-1) act = torch.nn.functional.silu(gate_part.float()) * lin_part.float() ref += (act.to(DTYPE) * probs[:, e : e + 1].to(DTYPE)) @ w2.t() + if shared: + ref += moe.shared_expert(tokens) torch.testing.assert_close(out.reshape(-1, HIDDEN), ref, rtol=0.05, atol=0.05) - -@pytest.mark.parametrize("num_experts", [None, 8], ids=["dense", "moe"]) -def test_layer_forward_backward(num_experts): - torch.manual_seed(0) - layer = ( - DeepSeekV3Layer( - HIDDEN, - HEADS, - ffn_hidden_size=512, - num_experts=num_experts, - moe_ffn_hidden_size=128 if num_experts else None, - topk=2 if num_experts else None, - shared_expert_ffn_hidden_size=128 if num_experts else None, - params_dtype=DTYPE, - **MLA_KWARGS, - ) - if num_experts - else DeepSeekV3Layer(HIDDEN, HEADS, ffn_hidden_size=512, params_dtype=DTYPE, **MLA_KWARGS) - ) - x = _input() - out = layer(x) - assert out.shape == x.shape - out.sum().backward() - assert x.grad is not None and torch.isfinite(x.grad).all() + bias_before = moe.expert_bias.clone() + moe.update_expert_bias() + assert torch.isfinite(moe.expert_bias).all() + if topk < num_experts: + assert not torch.equal(bias_before, moe.expert_bias) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 835813d5c6..60d17e39af 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -739,11 +739,9 @@ def test_sanity_layernorm_mlp( @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) -@pytest.mark.parametrize("model", ["small"]) -@pytest.mark.parametrize("skip_wgrad", all_boolean) -@pytest.mark.parametrize("moe", all_boolean) -def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, model, skip_wgrad, moe): - config = model_configs[model] +@pytest.mark.parametrize("moe", all_boolean, ids=["dense", "moe"]) +def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, moe): + config = model_configs["small"] if fp8_recipe is not None: if not is_fp8_supported(config): @@ -769,7 +767,7 @@ def test_sanity_deepseek_v3_layer(dtype, fp8_recipe, model, skip_wgrad, moe): **mlp_kwargs, ) - _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad) + _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad=False) @pytest.mark.parametrize("dtype", param_types) diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py index 341850fddb..0aea284afd 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -4,17 +4,13 @@ """Fused MLA RoPE kernels (DeepSeekV3-style decoupled RoPE/NoPE). -Triton forward/backward kernels adapted from Megatron-LM -``megatron/core/fusions/fused_mla_yarn_rope_apply.py``. The query kernel -rotates the trailing ``head_dim_rope`` slice in place (no concat); the KV -kernel builds the final key (nope | broadcast-rotated shared rope head) and -value tensors in a single pass. Falls back to pure PyTorch when Triton is -unavailable or for the ``bshd`` layout (the Triton path is ``sbhd``-only). - -Rotation convention: the rope slice is read interleaved (as stored in -HF/Megatron DeepSeekV3 checkpoints) and written in NeoX half-split layout, -matching the Megatron fused kernel semantics. -""" +The query kernel rotates the trailing ``head_dim_rope`` slice in place; the KV +kernel builds the key (nope | broadcast-rotated shared rope head) and value +tensors in a single pass. Falls back to pure PyTorch when Triton is unavailable +or for the ``bshd`` layout. + +The rope slice is read interleaved (checkpoint layout) and written in NeoX +half-split layout.""" import math from typing import Optional, Tuple @@ -65,7 +61,7 @@ def yarn_mscale(scale: float, mscale: float = 1.0) -> float: def yarn_concentration_factor(scaling_factor: float, mscale: float, mscale_all_dim: float) -> float: - """Factor multiplied into cos/sin tables (as in Megatron-Core).""" + """Factor multiplied into cos/sin tables.""" return yarn_mscale(scaling_factor, mscale) / yarn_mscale(scaling_factor, mscale_all_dim) diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index ca9c5ed3a7..4b4b1c9926 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -37,7 +37,7 @@ class MultiLatentAttention(torch.nn.Module): RoPE uses the fused MLA kernels from :mod:`.mla_rope` (in-place on the query rope slice, single-pass key/value assembly); the rope slice follows - the HF/Megatron DeepSeekV3 convention (interleaved weights, NeoX output). + the DeepSeekV3 checkpoint convention (interleaved weights, NeoX output). Parameters ---------- @@ -174,6 +174,7 @@ def __init__( if softmax_scale is None and rope_scaling_factor is not None: m = yarn_mscale(rope_scaling_factor, mscale_all_dim) softmax_scale = m * m / math.sqrt(self.qk_head_dim) + self.softmax_scale = softmax_scale self.core_attention = DotProductAttention( num_attention_heads, From 33549066ccd65e623e103975bd44664aae7c3a02 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:55:17 +0200 Subject: [PATCH 24/43] Rewrite DeepSeekV3MoE class docstring Signed-off-by: Pawel Gadzinski --- .../pytorch/models/deepseek_v3/moe.py | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 60080d8816..0b26933c4c 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -41,21 +41,23 @@ def _make_expert_mlp(num_experts, hidden_size, ffn_hidden_size, dtype, device): class DeepSeekV3MoE(torch.nn.Module): """ - DeepSeekV3-style Mixture of Experts block. + DeepSeekV3 Mixture-of-Experts block. - Routing uses the fused sigmoid router with aux-loss-free expert bias and - node-limited (grouped) top-k (``fused_topk_with_score_function``). Routed - experts run as a grouped SwiGLU MLP built from ``te.ops`` (fusable into a - single CuTe grouped-GEMM kernel); routing probabilities are applied - per-token inside the expert MLP, so unpermute/combine is a plain - accumulation. Token routing is either local - (``moe_permute_with_probs``/``moe_unpermute``) or, when ``ep_group`` is - given, expert-parallel over NCCL (``ep_dispatch``/``ep_combine``). + Each token is scored by a sigmoid router with a non-trainable expert bias + updated by ``update_expert_bias()`` (aux-loss-free load balancing) and, + optionally, group-limited routing: experts are split into ``num_groups`` + groups, the top ``group_topk`` groups are selected by their summed scores, + and the final ``topk`` experts are chosen only from those groups. Selected + tokens run through the routed experts, a SwiGLU MLP shared across experts + as a grouped GEMM, with the routing probability applied inside the MLP. An + optional shared expert (dense SwiGLU MLP) is added to every token. On + hardware that supports it the expert MLP runs as a single fused + grouped-GEMM kernel. - When expert parallelism is used, ``transformer_engine.pytorch.ep.ep_bootstrap`` - must be called once per process before constructing the module (it allocates - the ``EpBuffer`` in ``__init__``), and inputs must - be bfloat16. + Without ``ep_group`` all experts live on the local device. With + ``ep_group`` the experts are split across the group and tokens are + exchanged over NCCL; this requires ``ep_bootstrap`` to be called once per + process before constructing the module, and bfloat16 inputs. Parameters ---------- From 397733e75f8acb6f72bc6509b1baa5b9341aed9b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 13:59:02 +0200 Subject: [PATCH 25/43] DeepSeekV3MoE: drop ep_recv_capacity_per_rank and ep_alignment parameters Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 1 - .../pytorch/models/deepseek_v3/moe.py | 25 ++++++++----------- .../models/deepseek_v3/transformer_layer.py | 2 -- 3 files changed, 10 insertions(+), 18 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index a77b18e205..9561d20117 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -57,7 +57,6 @@ def _make_layer(ep_group, ep_size: int, num_experts: int) -> DeepSeekV3Layer: params_dtype=DTYPE, ep_group=ep_group, ep_max_tokens_per_rank=TOKENS_PER_RANK if ep else None, - ep_recv_capacity_per_rank=_recv_capacity(ep_size) if ep else None, **MLA_KWARGS, ) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 0b26933c4c..b35cedab0f 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -24,6 +24,9 @@ __all__ = ["DeepSeekV3MoE"] +_EP_ALIGNMENT = 128 + + def _make_expert_mlp(num_experts, hidden_size, ffn_hidden_size, dtype, device): # GroupedLinear + ScaledSwiGLU + GroupedLinear fuses into a single CuTe # grouped MLP on supported hardware; elsewhere it runs as three ops with @@ -87,11 +90,6 @@ class DeepSeekV3MoE(torch.nn.Module): expert-parallel process group; enables the NCCL EP path. ep_max_tokens_per_rank : int, optional max local tokens per forward (required with EP). - ep_recv_capacity_per_rank : int, optional - receive-buffer capacity; defaults to - ``ep_size * ep_max_tokens_per_rank * topk``. - ep_alignment : int, default = 128 - per-expert row alignment of the EP receive buffer. """ def __init__( @@ -109,8 +107,6 @@ def __init__( device: Union[torch.device, str] = "cuda", ep_group: Optional[torch.distributed.ProcessGroup] = None, ep_max_tokens_per_rank: Optional[int] = None, - ep_recv_capacity_per_rank: Optional[int] = None, - ep_alignment: int = 128, ) -> None: super().__init__() @@ -165,19 +161,18 @@ def __init__( from transformer_engine.pytorch.ep import EpBuffer assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." - if ep_recv_capacity_per_rank is None: - # Worst case plus per-expert alignment padding, rounded up to - # the multiple of 128 required by the fused grouped MLP. - cap = self.ep_size * ep_max_tokens_per_rank * topk - cap += num_local_experts * max(ep_alignment, 1) - ep_recv_capacity_per_rank = -(-cap // 128) * 128 + # Worst case plus per-expert alignment padding, rounded up to + # the multiple of 128 required by the fused grouped MLP. + cap = self.ep_size * ep_max_tokens_per_rank * topk + cap += num_local_experts * _EP_ALIGNMENT + cap = -(-cap // _EP_ALIGNMENT) * _EP_ALIGNMENT self.ep_buffer = EpBuffer( top_k=topk, max_tokens_per_rank=ep_max_tokens_per_rank, hidden_dim=hidden_size, num_local_experts=num_local_experts, - recv_capacity_per_rank=ep_recv_capacity_per_rank, - alignment=ep_alignment, + recv_capacity_per_rank=cap, + alignment=_EP_ALIGNMENT, device=device, ) diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index aa2fab232d..ab11fe6394 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -81,8 +81,6 @@ class DeepSeekV3Layer(torch.nn.Module): "expert_bias_update_rate", "ep_group", "ep_max_tokens_per_rank", - "ep_recv_capacity_per_rank", - "ep_alignment", } ) From cc10354bf59d6342e411aadf38c74a8452a68ff1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 14:01:30 +0200 Subject: [PATCH 26/43] DeepSeekV3MoE: build shared expert with the same SwiGLU MLP helper as routed experts Signed-off-by: Pawel Gadzinski --- .../pytorch/models/deepseek_v3/moe.py | 46 ++++++++----------- 1 file changed, 18 insertions(+), 28 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index b35cedab0f..e78e94fc2a 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -27,18 +27,22 @@ _EP_ALIGNMENT = 128 -def _make_expert_mlp(num_experts, hidden_size, ffn_hidden_size, dtype, device): - # GroupedLinear + ScaledSwiGLU + GroupedLinear fuses into a single CuTe - # grouped MLP on supported hardware; elsewhere it runs as three ops with - # the same API and checkpoint layout. +def _make_swiglu_mlp(hidden_size, ffn_hidden_size, dtype, device, num_experts=None): + """Dense SwiGLU MLP, or a grouped one (probs applied inside the activation) per expert. + + The grouped variant fuses into a single CuTe grouped MLP on supported hardware. + """ + common = dict(bias=False, dtype=dtype, device=device) + if num_experts is None: + return te_ops.Sequential( + te_ops.Linear(hidden_size, 2 * ffn_hidden_size, **common), + te_ops.SwiGLU(), + te_ops.Linear(ffn_hidden_size, hidden_size, **common), + ) return te_ops.Sequential( - te_ops.GroupedLinear( - num_experts, hidden_size, 2 * ffn_hidden_size, bias=False, dtype=dtype, device=device - ), + te_ops.GroupedLinear(num_experts, hidden_size, 2 * ffn_hidden_size, **common), te_ops.ScaledSwiGLU(glu_interleave_size=32), - te_ops.GroupedLinear( - num_experts, ffn_hidden_size, hidden_size, bias=False, dtype=dtype, device=device - ), + te_ops.GroupedLinear(num_experts, ffn_hidden_size, hidden_size, **common), ) @@ -132,28 +136,14 @@ def __init__( assert num_experts % self.ep_size == 0 num_local_experts = num_experts // self.ep_size - self.experts = _make_expert_mlp( - num_local_experts, hidden_size, moe_ffn_hidden_size, dtype, device + self.experts = _make_swiglu_mlp( + hidden_size, moe_ffn_hidden_size, dtype, device, num_experts=num_local_experts ) self.shared_expert = None if shared_expert_ffn_hidden_size is not None: - self.shared_expert = te_ops.Sequential( - te_ops.Linear( - hidden_size, - 2 * shared_expert_ffn_hidden_size, - bias=False, - dtype=dtype, - device=device, - ), - te_ops.SwiGLU(), - te_ops.Linear( - shared_expert_ffn_hidden_size, - hidden_size, - bias=False, - dtype=dtype, - device=device, - ), + self.shared_expert = _make_swiglu_mlp( + hidden_size, shared_expert_ffn_hidden_size, dtype, device ) self.ep_buffer = None From 2c1cd4c5fb8df0bc33e4efd59582a2f08245a7af Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 14:15:35 +0200 Subject: [PATCH 27/43] DeepSeekV3MoE EP path: count tokens per expert with scatter_add instead of syncing bincount Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/models/deepseek_v3/moe.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index e78e94fc2a..2bb47d49ce 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -221,10 +221,11 @@ def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: (tokens.shape[0], self.topk), dtype=torch.int64, device=tokens.device ) probs, topk_idx = self._route(self.gate(tokens).float(), topk_indices=topk_idx) - self._last_tokens_per_expert = torch.bincount( - topk_idx.flatten(), minlength=self.num_experts - ) - topk_weights = probs.gather(1, topk_idx).float() + flat_idx = topk_idx.flatten() + self._last_tokens_per_expert = torch.zeros( + self.num_experts, dtype=torch.long, device=tokens.device + ).scatter_add_(0, flat_idx, torch.ones_like(flat_idx)) + topk_weights = probs.gather(1, topk_idx) # Zero-filled recv/grad buffers: per-expert alignment padding lands # inside the grouped-GEMM m_splits, so uninitialized rows would poison From 6bae1ba180081748fbc54bcd1d69387519135dfa Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 15:18:05 +0200 Subject: [PATCH 28/43] Docs: list model-specific layers inline on the PyTorch API page; group standard layers, autocast and other utilities Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 35 +++++++++++++++++++++++------------ docs/api/pytorch_models.rst | 19 ------------------- 2 files changed, 23 insertions(+), 31 deletions(-) delete mode 100644 docs/api/pytorch_models.rst diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 8b2b742372..497f414aee 100644 --- a/docs/api/pytorch.rst +++ b/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 @@ -62,14 +81,6 @@ PyTorch .. autoapifunction:: transformer_engine.pytorch.deinterleave_glu_tensor -Models ------- - -.. toctree:: - :maxdepth: 1 - - pytorch_models - Data types ---------- diff --git a/docs/api/pytorch_models.rst b/docs/api/pytorch_models.rst deleted file mode 100644 index 2cde879ffb..0000000000 --- a/docs/api/pytorch_models.rst +++ /dev/null @@ -1,19 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -Models -====== - -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 From 610a1e236b10c0f275dc18f6e81d19b96be1958d Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 3 Sep 2026 15:56:42 +0200 Subject: [PATCH 29/43] Lint: use dict literals in models.deepseek_v3 Signed-off-by: Pawel Gadzinski --- .../pytorch/models/deepseek_v3/moe.py | 2 +- .../models/deepseek_v3/multi_latent_attention.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 2bb47d49ce..42f5048e6e 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -32,7 +32,7 @@ def _make_swiglu_mlp(hidden_size, ffn_hidden_size, dtype, device, num_experts=No The grouped variant fuses into a single CuTe grouped MLP on supported hardware. """ - common = dict(bias=False, dtype=dtype, device=device) + common = {"bias": False, "dtype": dtype, "device": device} if num_experts is None: return te_ops.Sequential( te_ops.Linear(hidden_size, 2 * ffn_hidden_size, **common), diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index 4b4b1c9926..56b4d3d0d1 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -161,14 +161,14 @@ def __init__( ) self.rotary_base = rotary_base - self._yarn_kwargs = dict( - scaling_factor=rope_scaling_factor, - original_max_position_embeddings=original_max_position_embeddings, - beta_fast=beta_fast, - beta_slow=beta_slow, - mscale=mscale, - mscale_all_dim=mscale_all_dim, - ) + self._yarn_kwargs = { + "scaling_factor": rope_scaling_factor, + "original_max_position_embeddings": original_max_position_embeddings, + "beta_fast": beta_fast, + "beta_slow": beta_slow, + "mscale": mscale, + "mscale_all_dim": mscale_all_dim, + } self._rope_tables: Optional[tuple] = None if softmax_scale is None and rope_scaling_factor is not None: From 93a40dfd47dca7e76b07e800ed6b39a3502d9c45 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 15:53:32 +0200 Subject: [PATCH 30/43] [PyTorch] DeepSeekV3MoE EP: drop full-buffer zeroing of recv/grad buffers NCCL EP zero-fills the alignment padding between experts itself, so the recv and grad buffers are allocated uninitialized instead of zeroing ~3.8 GB per buffer every step. Only a small margin past the received tokens is zeroed (device-side offsets, no host sync) because the fused grouped MLP reads a tile past the last expert. Expert zones are aligned to 256 rows: the fused MXFP8 grouped MLP gives wrong, nondeterministic weight gradients for splits that are not a multiple of 256. DeepSeekV3MoE.ep_recv_capacity() gives the recv capacity for ep_bootstrap; the distributed test uses it. On 2x4 GB300 (4096 tokens/rank, hidden 7168, 64 experts, top-k 8): bf16 11.7 -> 10.3 ms/iter, MXFP8 fused 8.8 -> 7.3 ms/iter. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 5 +-- .../pytorch/models/deepseek_v3/moe.py | 45 ++++++++++++------- 2 files changed, 32 insertions(+), 18 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index 9561d20117..f1f7d5e58f 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -10,7 +10,7 @@ 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 +from transformer_engine.pytorch.models import DeepSeekV3Layer, DeepSeekV3MoE HIDDEN = 256 MOE_FFN = 128 @@ -36,8 +36,7 @@ def _device_sm() -> int: def _recv_capacity(ep_size: int) -> int: - cap = ep_size * TOKENS_PER_RANK * TOP_K + NUM_LOCAL_EXPERTS * 128 - return -(-cap // 128) * 128 + return DeepSeekV3MoE.ep_recv_capacity(ep_size, TOKENS_PER_RANK, TOP_K, NUM_LOCAL_EXPERTS) def _broadcast_params(module: torch.nn.Module) -> None: diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 42f5048e6e..aa1ff9408d 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -24,7 +24,9 @@ __all__ = ["DeepSeekV3MoE"] -_EP_ALIGNMENT = 128 +_EP_ALIGNMENT = 256 +_FUSED_MLP_ROWS = 256 +_FUSED_MLP_MARGIN = 1024 def _make_swiglu_mlp(hidden_size, ffn_hidden_size, dtype, device, num_experts=None): @@ -151,11 +153,9 @@ def __init__( from transformer_engine.pytorch.ep import EpBuffer assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." - # Worst case plus per-expert alignment padding, rounded up to - # the multiple of 128 required by the fused grouped MLP. - cap = self.ep_size * ep_max_tokens_per_rank * topk - cap += num_local_experts * _EP_ALIGNMENT - cap = -(-cap // _EP_ALIGNMENT) * _EP_ALIGNMENT + cap = self.ep_recv_capacity( + self.ep_size, ep_max_tokens_per_rank, topk, num_local_experts + ) self.ep_buffer = EpBuffer( top_k=topk, max_tokens_per_rank=ep_max_tokens_per_rank, @@ -166,6 +166,16 @@ def __init__( device=device, ) + @staticmethod + def ep_recv_capacity( + ep_size: int, max_tokens_per_rank: int, topk: int, num_local_experts: int + ) -> int: + """Recv rows per rank for ``ep_bootstrap``: worst-case routing plus per-expert + alignment padding and the fused grouped MLP margin, rounded to its row multiple.""" + cap = ep_size * max_tokens_per_rank * topk + cap += num_local_experts * _EP_ALIGNMENT + _FUSED_MLP_MARGIN + return -(-cap // _FUSED_MLP_ROWS) * _FUSED_MLP_ROWS + def _route(self, logits: torch.Tensor, topk_indices: Optional[torch.Tensor] = None): return fused_topk_with_score_function( logits=logits, @@ -227,28 +237,33 @@ def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: ).scatter_add_(0, flat_idx, torch.ones_like(flat_idx)) topk_weights = probs.gather(1, topk_idx) - # Zero-filled recv/grad buffers: per-expert alignment padding lands - # inside the grouped-GEMM m_splits, so uninitialized rows would poison - # the expert wgrads. + # NCCL EP zero-fills the alignment padding between experts itself, so + # the recv/grad buffers can stay uninitialized. The fused grouped MLP + # reads up to a tile past the last expert, so zero a margin there + # (offsets stay on device: no host sync). cap = self.ep_buffer.recv_capacity_per_rank recv_tokens, recv_weights, tokens_per_expert = ep_dispatch( self.ep_buffer, tokens, topk_idx, topk_weights, - recv_tokens=torch.zeros( + recv_tokens=torch.empty( (cap, self.hidden_size), dtype=tokens.dtype, device=tokens.device ), - recv_topk_weights=torch.zeros((cap,), dtype=torch.float32, device=tokens.device), + recv_topk_weights=torch.empty((cap,), dtype=torch.float32, device=tokens.device), ) + grad_out = torch.empty((cap, self.hidden_size), dtype=tokens.dtype, device=tokens.device) + with torch.no_grad(): + margin = ( + torch.arange(_FUSED_MLP_MARGIN, device=tokens.device) + tokens_per_expert.sum() + ).clamp_(max=cap - 1) + for buf in (recv_tokens, recv_weights, grad_out): + buf.detach().index_fill_(0, margin, 0) expert_out = self.experts( recv_tokens, tokens_per_expert, recv_weights.to(tokens.dtype), tokens_per_expert ) return ep_combine( - self.ep_buffer, - expert_out, - num_local_tokens=tokens.shape[0], - grad_out=torch.zeros_like(expert_out), + self.ep_buffer, expert_out, num_local_tokens=tokens.shape[0], grad_out=grad_out ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: From 5e36ef1a1f7c742148b3f7a00dca3f29cfa69b67 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 15:53:32 +0200 Subject: [PATCH 31/43] [PyTorch] Add DeepSeekV3Layer expert-parallel example torchrun script running one DeepSeekV3Layer with the routed experts sharded over all ranks via NCCL EP, a launcher for local GPUs with an nsys option, and a README with the configuration, measured timings on GB300 and a kernel-level breakdown. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 125 ++++++++++++++ .../deepseek_v3/deepseek_v3_layer_ep.py | 153 ++++++++++++++++++ .../deepseek_v3/run_deepseek_v3_layer_ep.sh | 35 ++++ 3 files changed, 313 insertions(+) create mode 100644 examples/pytorch/deepseek_v3/README.md create mode 100644 examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py create mode 100755 examples/pytorch/deepseek_v3/run_deepseek_v3_layer_ep.sh diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md new file mode 100644 index 0000000000..5e4b28ed70 --- /dev/null +++ b/examples/pytorch/deepseek_v3/README.md @@ -0,0 +1,125 @@ +# DeepSeekV3Layer with expert parallelism + +`deepseek_v3_layer_ep.py` runs one `transformer_engine.pytorch.models.DeepSeekV3Layer` +(RMSNorm, multi-latent attention, DeepSeek MoE with a shared expert) with the routed experts +sharded over all ranks. Tokens are exchanged with NCCL EP (`transformer_engine.pytorch.ep`). +Each iteration is a forward and a backward pass on random data; the script reports the time +per iteration and tokens per second and checks that the output is finite. + +## Requirements + +- SM90 or newer GPUs connected with NVLink (NCCL EP falls back to the network transport and + deadlocks on PCIe-only nodes). +- NCCL >= 2.30.4, PyTorch >= 2.11 (symmetric memory), Transformer Engine built with the + `3rdparty/nccl-extensions` submodule. +- `NVTE_CUTEDSL_FUSED_GROUPED_MLP=1` additionally needs SM100-class GPUs and the CuTe DSL + (`nvidia-cutlass-dsl`) for the fused MXFP8 grouped MLP. + +## Running + +Single node, all local GPUs: + +```bash +bash run_deepseek_v3_layer_ep.sh # small dims, bf16 +bash run_deepseek_v3_layer_ep.sh --dsv3 # DeepSeek-V3 layer dims, bf16 +bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 # MXFP8 experts (unfused grouped GEMM) +NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 +``` + +Multi-node: launch `torchrun` yourself, EP spans every rank: + +```bash +torchrun --nnodes=2 --nproc-per-node=4 --rdzv-backend=c10d --rdzv-endpoint=:29500 \ + deepseek_v3_layer_ep.py --dsv3 --recipe mxfp8 +``` + +Every rank owns `--num-local-experts` experts (default 8), so the expert count is +`8 * world_size`. Other knobs: `--tokens-per-rank`, `--topk`, `--hidden`, `--num-heads`, +`--moe-ffn`, the MLA dims (`--q-lora-rank`, `--kv-lora-rank`, `--qk-nope-head-dim`, +`--qk-rope-head-dim`, `--v-head-dim`), `--warmup`, `--iters`. + +## Profiling with nsys + +The timed iterations run inside a `torch.cuda.profiler.start()` / `stop()` window, so +`-c cudaProfilerApi` records only them, one NVTX range per iteration: + +```bash +NSYS=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 +# -> results/deepseek_v3_layer_ep_.nsys-rep +nsys stats --report cuda_gpu_kern_sum results/deepseek_v3_layer_ep_.nsys-rep +``` + +The launcher wraps `torchrun`, so all local ranks land in one report. For multi-node runs put +the same `nsys profile ... -o _%q{SLURM_NODEID}` in front of `torchrun` on each node. + +## Configuration used below + +| | value | +|---|---| +| hardware | GB300 (SM103), 4 GPUs per node, both nodes in one NVLink domain (MNNVL) | +| software | CUDA 13.3, NCCL 2.30.7, PyTorch 2.13 nightly, cuDNN 9.24 | +| `--dsv3` dims | hidden 7168, 128 heads, MLA q_lora 1536 / kv_lora 512 / nope 128 / rope 64 / v 128, expert ffn 2048, shared expert ffn 2048 | +| MoE | 8 local experts per rank (32 on 4 GPUs, 64 on 8), top-k 8, 4096 tokens per rank | +| precision | bf16 params and activations; `--recipe mxfp8` = MXFP8 block scaling for the expert GEMMs and dense projections | +| timing | fwd + bwd, 5 warmup, 10 timed iterations, no CUDA graphs | + +## Results + +Time per iteration is forward + backward of one layer on 4096 tokens per rank; throughput +counts tokens over all ranks. `fused` means `NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`. + +1 node, 4 GPUs, default dims (hidden 2048, 16 heads, expert ffn 1024, 32 experts): + +| precision | ms / iter | Mtok / s | +|---|---|---| +| bf16 | 6.17 | 2.65 | +| mxfp8 fused | 8.4 | 1.95 | + +1 node, 4 GPUs, `--dsv3` (32 experts): + +| precision | ms / iter | Mtok / s | +|---|---|---| +| bf16 | 13.09 | 1.25 | +| mxfp8 | 11.02 | 1.49 | +| mxfp8 fused | 10.19 | 1.61 | + +2 nodes, 8 GPUs, `--dsv3` (64 experts): + +| precision | ms / iter | Mtok / s | +|---|---|---| +| bf16 | 15.06 | 2.18 | +| mxfp8 | 12.88 | 2.54 | +| mxfp8 fused | 11.32 | 2.89 | + +At the small default dims MXFP8 is slower than bf16: the fused path launches many small +quantization kernels and the layer becomes CPU-launch-bound. Running under `nsys` adds +about 1.5 ms per iteration to these numbers. + +## Where the time goes (8 GPUs, `--dsv3`, mxfp8 fused) + +Per GPU and iteration, from `nsys stats --report cuda_gpu_kern_sum` on one node +(kernel time 11.7 ms of an 11.3 ms iteration: the GPU is busy back to back): + +| group | ms | kernels | +|---|---|---| +| NCCL EP all-to-all | 2.3 | `nccl_ep_jit_ht_dispatch_kernel` (0.99), `nccl_ep_jit_ht_combine_kernel` (1.31), each twice per iteration (fwd + bwd) | +| NCCL EP local permute + routing all-gather | 1.3 | `local_permute_dup/reduce` (0.62), `ncclDevKernel_AllGather_RING_LL` (0.65, includes waiting for slower ranks) | +| fused grouped MLP (cuDNN, MXFP8) | 2.7 | fc1+SwiGLU fwd (0.58), fc2 fwd (0.9), dGLU bwd (0.35), wgrad (0.85) | +| 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.3 | `nvjet_sm103_qqtst_*` | +| attention (cuDNN SDPA) | 0.7 | flash fprop (0.19) + bprop (0.51) | +| elementwise | 1.1 | residual/shared-expert adds (0.73), RMSNorm fwd+bwd (0.39) | + +Both nodes sit in one NVLink domain, so dispatch and combine move roughly 470 MB per GPU per +call over NVLink at close to link bandwidth; on an InfiniBand-connected pair of nodes the +all-to-all share would be much larger. + +## Notes on the EP path + +- `ep_bootstrap` must be given the same recv capacity the layer uses: + `DeepSeekV3MoE.ep_recv_capacity(ep_size, tokens_per_rank, topk, num_local_experts)`. +- Per-expert zones in the recv buffer are aligned to 256 rows; the fused grouped MLP requires + that alignment and reads a little past the last expert, which is why the layer keeps a + zeroed margin after the received tokens. +- The recv and grad buffers are allocated uninitialized: NCCL EP zero-fills the alignment + padding between experts itself. diff --git a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py new file mode 100644 index 0000000000..69086a4548 --- /dev/null +++ b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py @@ -0,0 +1,153 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""DeepSeekV3Layer with expert parallelism over all ranks: forward + backward timing. + +One process per GPU, launched via run_deepseek_v3_layer_ep.sh (torchrun). Every rank +holds ``--num-local-experts`` routed experts; tokens are exchanged with NCCL EP. +Timed iterations run inside a ``torch.cuda.profiler`` window, so +``nsys profile -c cudaProfilerApi --capture-range-end=stop torchrun ...`` records +only them. +""" + +import argparse +import os +import sys +import time +from contextlib import nullcontext + +import torch +import torch.distributed as dist + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe as te_recipe +from transformer_engine.pytorch.ep import ep_bootstrap, ep_finalize, release_symm_mem_pool +from transformer_engine.pytorch.models import DeepSeekV3Layer, DeepSeekV3MoE + + +def _parse_args(): + p = argparse.ArgumentParser(description="DeepSeekV3Layer EP example (fwd + bwd)") + p.add_argument("--tokens-per-rank", type=int, default=4096) + p.add_argument("--hidden", type=int, default=2048) + p.add_argument("--num-heads", type=int, default=16) + p.add_argument("--moe-ffn", type=int, default=1024) + p.add_argument("--num-local-experts", type=int, default=8) + p.add_argument("--topk", type=int, default=8) + p.add_argument("--q-lora-rank", type=int, default=512) + p.add_argument("--kv-lora-rank", type=int, default=256) + p.add_argument("--qk-nope-head-dim", type=int, default=64) + p.add_argument("--qk-rope-head-dim", type=int, default=32) + p.add_argument("--v-head-dim", type=int, default=64) + p.add_argument( + "--dsv3", + action="store_true", + help=( + "DeepSeek-V3 layer dims (hidden 7168, 128 heads, MLA 1536/512/128/64/128, expert ffn" + " 2048)." + ), + ) + p.add_argument("--recipe", choices=["none", "mxfp8"], default="none") + p.add_argument("--warmup", type=int, default=5) + p.add_argument("--iters", type=int, default=10) + args = p.parse_args() + if args.dsv3: + args.hidden, args.num_heads, args.moe_ffn = 7168, 128, 2048 + args.q_lora_rank, args.kv_lora_rank = 1536, 512 + args.qk_nope_head_dim, args.qk_rope_head_dim, args.v_head_dim = 128, 64, 128 + return args + + +def _autocast(name): + if name == "none": + return nullcontext() + return te.autocast(enabled=True, recipe=te_recipe.MXFP8BlockScaling()) + + +def main(): + """Build the layer, run warmup + timed fwd/bwd iterations, print throughput on rank 0.""" + args = _parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank)) + rank, world_size = dist.get_rank(), dist.get_world_size() + + major, minor = torch.cuda.get_device_capability() + if major * 10 + minor < 90: + if rank == 0: + print(f"SKIPPED: NCCL EP requires SM>=90 (got SM{major}{minor})") + dist.destroy_process_group() + return 0 + + ep_group = dist.new_group(ranks=list(range(world_size)), backend="nccl") + num_experts = args.num_local_experts * world_size + ep_bootstrap( + ep_group, + num_experts=num_experts, + max_tokens_per_rank=args.tokens_per_rank, + hidden_dim=args.hidden, + num_topk=args.topk, + recv_capacity_per_rank=DeepSeekV3MoE.ep_recv_capacity( + world_size, args.tokens_per_rank, args.topk, args.num_local_experts + ), + ) + + torch.manual_seed(0) + layer = DeepSeekV3Layer( + args.hidden, + args.num_heads, + num_experts=num_experts, + moe_ffn_hidden_size=args.moe_ffn, + shared_expert_ffn_hidden_size=args.moe_ffn, + topk=args.topk, + params_dtype=torch.bfloat16, + ep_group=ep_group, + ep_max_tokens_per_rank=args.tokens_per_rank, + q_lora_rank=args.q_lora_rank, + kv_lora_rank=args.kv_lora_rank, + qk_nope_head_dim=args.qk_nope_head_dim, + qk_rope_head_dim=args.qk_rope_head_dim, + v_head_dim=args.v_head_dim, + ) + seq = args.tokens_per_rank // 4 + x = torch.randn(seq, 4, args.hidden, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + def step(): + with _autocast(args.recipe): + out = layer(x) + out.backward(torch.ones_like(out)) + x.grad = None + return out + + for _ in range(args.warmup): + out = step() + finite = bool(torch.isfinite(out).all()) + torch.cuda.synchronize() + dist.barrier() + + torch.cuda.profiler.start() + start = time.perf_counter() + for i in range(args.iters): + with torch.cuda.nvtx.range(f"iter{i}"): + step() + torch.cuda.synchronize() + ms = (time.perf_counter() - start) / args.iters * 1e3 + torch.cuda.profiler.stop() + dist.barrier() + + if rank == 0: + tok_s = args.tokens_per_rank * world_size / (ms / 1e3) + print( + f"DeepSeekV3Layer EP: ranks={world_size} experts={num_experts} topk={args.topk} " + f"tokens/rank={args.tokens_per_rank} hidden={args.hidden} recipe={args.recipe} " + f"fused_mlp={os.environ.get('NVTE_CUTEDSL_FUSED_GROUPED_MLP', '0')} " + f"fwd+bwd {ms:.3f} ms/iter ({tok_s / 1e6:.2f} Mtok/s) finite={finite}", + flush=True, + ) + ep_finalize() + release_symm_mem_pool() + dist.destroy_process_group() + return 0 if finite else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/pytorch/deepseek_v3/run_deepseek_v3_layer_ep.sh b/examples/pytorch/deepseek_v3/run_deepseek_v3_layer_ep.sh new file mode 100755 index 0000000000..9341b7f14a --- /dev/null +++ b/examples/pytorch/deepseek_v3/run_deepseek_v3_layer_ep.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Launcher for deepseek_v3_layer_ep.py on all local GPUs. Extra args go to the script: +# bash run_deepseek_v3_layer_ep.sh # bf16, small dims +# bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 # DeepSeek-V3 dims, MXFP8 experts +# NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 +# NSYS=1 bash run_deepseek_v3_layer_ep.sh --dsv3 # nsys report in results/ +# Multi-node: run torchrun yourself with --nnodes/--rdzv-endpoint; EP spans all ranks. + +set -uo pipefail + +DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}" +if [ "${NUM_GPUS}" -lt 2 ]; then + echo "EP requires >= 2 GPUs (found ${NUM_GPUS}); SKIPPING." + exit 0 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} +export NCCL_EP_JIT_CACHE_DIR +mkdir -p "$NCCL_EP_JIT_CACHE_DIR" + +PREFIX=() +if [ "${NSYS:-0}" = "1" ]; then + mkdir -p "${SCRIPT_DIR}/results" + PREFIX=(nsys profile -t cuda,nvtx,nccl -c cudaProfilerApi --capture-range-end=stop + --cuda-graph-trace=node -o "${SCRIPT_DIR}/results/deepseek_v3_layer_ep_%h") +fi + +"${PREFIX[@]}" torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_GPUS}" \ + "${SCRIPT_DIR}/deepseek_v3_layer_ep.py" "$@" From 361287989266ea0ed5ba4c09f45f3891aa6f66ee Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 18:41:58 +0200 Subject: [PATCH 32/43] [PyTorch] DeepSeekV3Layer example: naive PyTorch MoE and dense baselines --impl naive is an expert-parallel MoE written with plain PyTorch (all_to_all_single, a loop of dense per-expert MLPs); --impl dense swaps the MoE for the dense SwiGLU MLP of DeepSeek-V3's first layers. --dsv3 selects the DeepSeek-V3 layer dims. README gets the comparison and the kernel breakdown of the naive variant. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 48 ++++++ .../deepseek_v3/deepseek_v3_layer_ep.py | 140 +++++++++++++++--- 2 files changed, 167 insertions(+), 21 deletions(-) diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md index 5e4b28ed70..dc4b10c30b 100644 --- a/examples/pytorch/deepseek_v3/README.md +++ b/examples/pytorch/deepseek_v3/README.md @@ -24,8 +24,21 @@ bash run_deepseek_v3_layer_ep.sh # small dims, bf16 bash run_deepseek_v3_layer_ep.sh --dsv3 # DeepSeek-V3 layer dims, bf16 bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 # MXFP8 experts (unfused grouped GEMM) NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 +bash run_deepseek_v3_layer_ep.sh --dsv3 --impl naive # plain PyTorch MoE baseline +bash run_deepseek_v3_layer_ep.sh --dsv3 --impl dense # dense SwiGLU MLP instead of MoE ``` +`--impl` selects the MLP block inside the same layer (attention and norms are identical): + +- `te` (default): `DeepSeekV3MoE`, NCCL EP dispatch/combine, experts as one grouped GEMM. +- `naive`: MoE written with plain PyTorch, no TE MoE code: sigmoid top-k router with expert + bias, `all_to_all_single` dispatch and combine (two host syncs per layer for the split sizes), + a Python loop over the local experts with dense `F.linear` SwiGLU MLPs, `index_copy` / + `index_add` to place results, and a shared expert. This is what an EP MoE looks like before + any fused kernels. +- `dense`: no MoE, the dense SwiGLU MLP used in DeepSeek-V3's first three layers (`--dense-ffn`, + default 18432). Gives the cost of a non-MoE layer of the same model for reference. + Multi-node: launch `torchrun` yourself, EP spans every rank: ```bash @@ -95,6 +108,41 @@ At the small default dims MXFP8 is slower than bf16: the fused path launches man quantization kernels and the layer becomes CPU-launch-bound. Running under `nsys` adds about 1.5 ms per iteration to these numbers. +## TE MoE vs. plain PyTorch MoE vs. dense layer + +Same layer, same dims (`--dsv3`), bf16 unless noted: + +| | 4 GPUs (32 experts) | 8 GPUs (64 experts) | +|---|---|---| +| `--impl te`, bf16 | 13.09 ms | 15.06 ms | +| `--impl te`, mxfp8 fused | 10.19 ms | 11.32 ms | +| `--impl naive`, bf16 | 26.83 ms | 27.34 ms | +| `--impl dense`, bf16 (ffn 18432) | 9.99 ms | 10.04 ms | +| `--impl dense`, mxfp8 | 7.49 ms | 7.58 ms | + +At the small default dims the gap is similar: `naive` 10.08 ms vs `te` 6.17 ms on 4 GPUs. + +Per GPU and iteration, the naive MoE spends (8 GPUs, kernel time 25.9 ms of a 28.5 ms +iteration): + +| group | ms | what | +|---|---|---| +| `ncclDevKernel_SendRecv` | 5.2 | 7 all_to_all launches per iteration (tokens fwd/bwd, results fwd/bwd, counts, indices, probs) | +| expert and dense GEMMs (`nvjet_*`) | 6.4 | 8 separate GEMM pairs per rank instead of one grouped GEMM, plus the MLA projections | +| elementwise adds | 4.1 | `index_add` and its backward, residuals | +| `FillFunctor` (zeros) | 2.5 | `zeros_like` for the per-expert output buffer and `index_add` targets | +| device-to-device copies | 2.3 | `index_copy` and gathers materialising per-expert slices | +| indexing kernels | 3.0 | `x[tok]`, `nonzero` masks, `index_copy`, `indexing_backward` | +| attention, norms | 1.1 | same as in the TE variant | + +The TE variant replaces all of the communication and indexing rows with two NCCL EP kernels +per direction (dispatch, combine) writing directly into the expert-major layout, and the +per-expert GEMMs with one grouped GEMM, which is where the roughly 2x comes from. The dense +layer is faster than either MoE variant here because with 8 experts per rank and top-k 8 the +MoE moves 8 activations per token across ranks while the dense MLP does the equivalent FLOPs +locally; the MoE wins only once the expert count grows past what a dense layer of equal +per-token FLOPs can hold. + ## Where the time goes (8 GPUs, `--dsv3`, mxfp8 fused) Per GPU and iteration, from `nsys stats --report cuda_gpu_kern_sum` on one node diff --git a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py index 69086a4548..d42797dfe0 100644 --- a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py +++ b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py @@ -4,7 +4,10 @@ """DeepSeekV3Layer with expert parallelism over all ranks: forward + backward timing. One process per GPU, launched via run_deepseek_v3_layer_ep.sh (torchrun). Every rank -holds ``--num-local-experts`` routed experts; tokens are exchanged with NCCL EP. +holds ``--num-local-experts`` routed experts. ``--impl te`` exchanges tokens with NCCL EP +and runs the experts as one grouped GEMM (DeepSeekV3MoE); ``--impl naive`` is a plain +PyTorch MoE (all_to_all_single + a Python loop over experts) dropped into the same layer; +``--impl dense`` replaces the MoE with the dense SwiGLU MLP of DeepSeek-V3's first layers. Timed iterations run inside a ``torch.cuda.profiler`` window, so ``nsys profile -c cudaProfilerApi --capture-range-end=stop torchrun ...`` records only them. @@ -18,6 +21,8 @@ import torch import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed.nn.functional import all_to_all_single import transformer_engine.pytorch as te from transformer_engine.common import recipe as te_recipe @@ -46,6 +51,8 @@ def _parse_args(): " 2048)." ), ) + p.add_argument("--impl", choices=["te", "naive", "dense"], default="te") + p.add_argument("--dense-ffn", type=int, default=18432, help="ffn size for --impl dense") p.add_argument("--recipe", choices=["none", "mxfp8"], default="none") p.add_argument("--warmup", type=int, default=5) p.add_argument("--iters", type=int, default=10) @@ -63,6 +70,77 @@ def _autocast(name): return te.autocast(enabled=True, recipe=te_recipe.MXFP8BlockScaling()) +class NaiveMoE(torch.nn.Module): + """DeepSeek-style MoE without TE: sigmoid top-k router with expert bias, torch all_to_all + dispatch/combine, a Python loop of dense SwiGLU experts, and a shared expert.""" + + def __init__(self, hidden, ffn, num_experts, topk, ep_group, shared_ffn, dtype): + super().__init__() + self.hidden, self.topk, self.group = hidden, topk, ep_group + self.ws, self.rank = dist.get_world_size(ep_group), dist.get_rank(ep_group) + self.num_experts, self.local = num_experts, num_experts // self.ws + self.gate = torch.nn.Linear(hidden, num_experts, bias=False, dtype=dtype, device="cuda") + self.register_buffer("expert_bias", torch.zeros(num_experts, device="cuda")) + std = hidden**-0.5 + self.w1 = torch.nn.Parameter( + torch.randn(self.local, 2 * ffn, hidden, dtype=dtype, device="cuda") * std + ) + self.w2 = torch.nn.Parameter( + torch.randn(self.local, hidden, ffn, dtype=dtype, device="cuda") * ffn**-0.5 + ) + self.shared_w1 = torch.nn.Linear( + hidden, 2 * shared_ffn, bias=False, dtype=dtype, device="cuda" + ) + self.shared_w2 = torch.nn.Linear(shared_ffn, hidden, bias=False, dtype=dtype, device="cuda") + + @staticmethod + def _swiglu(h): + a, g = h.chunk(2, dim=-1) + return F.silu(a) * g + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, self.hidden) + scores = torch.sigmoid(self.gate(x).float()) + _, idx = torch.topk(scores + self.expert_bias, self.topk, dim=-1) + probs = scores.gather(1, idx) + probs = probs / probs.sum(-1, keepdim=True) * 2.5 + # Dispatch: sort (token, expert) pairs by destination rank, exchange counts, all_to_all. + flat_e, flat_p = idx.reshape(-1), probs.reshape(-1) + tok = torch.arange(x.shape[0], device=x.device).repeat_interleave(self.topk) + order = torch.argsort(flat_e // self.local, stable=True) + flat_e, flat_p, tok = flat_e[order], flat_p[order], tok[order] + send = torch.bincount(flat_e // self.local, minlength=self.ws) + recv = torch.empty_like(send) + dist.all_to_all_single(recv, send, group=self.group) + send, recv = send.tolist(), recv.tolist() + n_recv = sum(recv) + x_recv = all_to_all_single( + torch.empty(n_recv, self.hidden, dtype=x.dtype, device=x.device), + x[tok], + recv, + send, + group=self.group, + ) + e_recv = torch.empty(n_recv, dtype=flat_e.dtype, device=x.device) + p_recv = torch.empty(n_recv, dtype=flat_p.dtype, device=x.device) + dist.all_to_all_single(e_recv, flat_e.contiguous(), recv, send, group=self.group) + dist.all_to_all_single(p_recv, flat_p.contiguous(), recv, send, group=self.group) + # Experts: one dense SwiGLU MLP per local expert. + local_e = e_recv - self.rank * self.local + y_recv = torch.zeros_like(x_recv) + for e in range(self.local): + sel = (local_e == e).nonzero().squeeze(1) + if sel.numel() == 0: + continue + h = self._swiglu(F.linear(x_recv[sel], self.w1[e])) * p_recv[sel, None].to(x.dtype) + y_recv = y_recv.index_copy(0, sel, F.linear(h, self.w2[e])) + # Combine: reverse all_to_all, sum the top-k contributions per token. + y = all_to_all_single(torch.empty_like(x[tok]), y_recv, send, recv, group=self.group) + out = torch.zeros_like(x).index_add(0, tok, y) + out = out + self.shared_w2(self._swiglu(self.shared_w1(x))) + return out.view_as(hidden_states) + + def main(): """Build the layer, run warmup + timed fwd/bwd iterations, print throughput on rank 0.""" args = _parse_args() @@ -79,35 +157,55 @@ def main(): return 0 ep_group = dist.new_group(ranks=list(range(world_size)), backend="nccl") + dist.all_reduce(torch.zeros(1, device="cuda"), group=ep_group) num_experts = args.num_local_experts * world_size - ep_bootstrap( - ep_group, - num_experts=num_experts, - max_tokens_per_rank=args.tokens_per_rank, - hidden_dim=args.hidden, - num_topk=args.topk, - recv_capacity_per_rank=DeepSeekV3MoE.ep_recv_capacity( - world_size, args.tokens_per_rank, args.topk, args.num_local_experts - ), - ) + if args.impl == "te": + ep_bootstrap( + ep_group, + num_experts=num_experts, + max_tokens_per_rank=args.tokens_per_rank, + hidden_dim=args.hidden, + num_topk=args.topk, + recv_capacity_per_rank=DeepSeekV3MoE.ep_recv_capacity( + world_size, args.tokens_per_rank, args.topk, args.num_local_experts + ), + ) torch.manual_seed(0) - layer = DeepSeekV3Layer( - args.hidden, - args.num_heads, + mlp_kwargs = dict( num_experts=num_experts, moe_ffn_hidden_size=args.moe_ffn, shared_expert_ffn_hidden_size=args.moe_ffn, topk=args.topk, - params_dtype=torch.bfloat16, ep_group=ep_group, ep_max_tokens_per_rank=args.tokens_per_rank, + ) + if args.impl == "dense": + mlp_kwargs = dict(ffn_hidden_size=args.dense_ffn) + elif args.impl == "naive": + # Build the TE MoE without EP (replaced below); keeps the pre-MLP RMSNorm. + mlp_kwargs.pop("ep_group"), mlp_kwargs.pop("ep_max_tokens_per_rank") + layer = DeepSeekV3Layer( + args.hidden, + args.num_heads, + params_dtype=torch.bfloat16, + **mlp_kwargs, q_lora_rank=args.q_lora_rank, kv_lora_rank=args.kv_lora_rank, qk_nope_head_dim=args.qk_nope_head_dim, qk_rope_head_dim=args.qk_rope_head_dim, v_head_dim=args.v_head_dim, ) + if args.impl == "naive": + layer.mlp = NaiveMoE( + args.hidden, + args.moe_ffn, + num_experts, + args.topk, + ep_group, + args.moe_ffn, + torch.bfloat16, + ) seq = args.tokens_per_rank // 4 x = torch.randn(seq, 4, args.hidden, dtype=torch.bfloat16, device="cuda", requires_grad=True) @@ -137,14 +235,14 @@ def step(): if rank == 0: tok_s = args.tokens_per_rank * world_size / (ms / 1e3) print( - f"DeepSeekV3Layer EP: ranks={world_size} experts={num_experts} topk={args.topk} " - f"tokens/rank={args.tokens_per_rank} hidden={args.hidden} recipe={args.recipe} " - f"fused_mlp={os.environ.get('NVTE_CUTEDSL_FUSED_GROUPED_MLP', '0')} " - f"fwd+bwd {ms:.3f} ms/iter ({tok_s / 1e6:.2f} Mtok/s) finite={finite}", + f"DeepSeekV3Layer impl={args.impl}:" + f" ranks={world_size} experts={num_experts} topk={args.topk} tokens/rank={args.tokens_per_rank} hidden={args.hidden} recipe={args.recipe} fused_mlp={os.environ.get('NVTE_CUTEDSL_FUSED_GROUPED_MLP', '0')} fwd+bwd" + f" {ms:.3f} ms/iter ({tok_s / 1e6:.2f} Mtok/s) finite={finite}", flush=True, ) - ep_finalize() - release_symm_mem_pool() + if args.impl == "te": + ep_finalize() + release_symm_mem_pool() dist.destroy_process_group() return 0 if finite else 1 From eb36d7d28b60f986b04cbe328abe40d86d9d3dae Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 18:54:13 +0200 Subject: [PATCH 33/43] [PyTorch] DeepSeekV3Layer example: two-level plain PyTorch MoE baseline --impl naive: torch all_to_all dispatch/combine and a Python loop of dense per-expert MLPs. --impl naive_grouped: the same all_to_all path with the experts as one TE grouped GEMM stack. Drops the dense variant. README compares both against the NCCL EP + grouped GEMM layer and adds the kernel breakdown of each baseline. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 53 +++++++++------ .../deepseek_v3/deepseek_v3_layer_ep.py | 64 ++++++++++++------- 2 files changed, 73 insertions(+), 44 deletions(-) diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md index dc4b10c30b..00e0f9ceb9 100644 --- a/examples/pytorch/deepseek_v3/README.md +++ b/examples/pytorch/deepseek_v3/README.md @@ -24,11 +24,11 @@ bash run_deepseek_v3_layer_ep.sh # small dims, bf16 bash run_deepseek_v3_layer_ep.sh --dsv3 # DeepSeek-V3 layer dims, bf16 bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 # MXFP8 experts (unfused grouped GEMM) NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 -bash run_deepseek_v3_layer_ep.sh --dsv3 --impl naive # plain PyTorch MoE baseline -bash run_deepseek_v3_layer_ep.sh --dsv3 --impl dense # dense SwiGLU MLP instead of MoE +bash run_deepseek_v3_layer_ep.sh --dsv3 --impl naive # all_to_all + loop over experts +bash run_deepseek_v3_layer_ep.sh --dsv3 --impl naive_grouped # all_to_all + TE grouped GEMM ``` -`--impl` selects the MLP block inside the same layer (attention and norms are identical): +`--impl` selects the MoE block inside the same layer (attention and norms are identical): - `te` (default): `DeepSeekV3MoE`, NCCL EP dispatch/combine, experts as one grouped GEMM. - `naive`: MoE written with plain PyTorch, no TE MoE code: sigmoid top-k router with expert @@ -36,8 +36,9 @@ bash run_deepseek_v3_layer_ep.sh --dsv3 --impl dense # dense SwiGLU MLP inst a Python loop over the local experts with dense `F.linear` SwiGLU MLPs, `index_copy` / `index_add` to place results, and a shared expert. This is what an EP MoE looks like before any fused kernels. -- `dense`: no MoE, the dense SwiGLU MLP used in DeepSeek-V3's first three layers (`--dense-ffn`, - default 18432). Gives the cost of a non-MoE layer of the same model for reference. +- `naive_grouped`: the same all_to_all dispatch and combine, but the received rows are sorted by + local expert and run through one `te.ops.GroupedLinear` / `ScaledSwiGLU` / `GroupedLinear` + stack. Isolates the cost of the Python loop from the cost of the communication path. Multi-node: launch `torchrun` yourself, EP spans every rank: @@ -108,21 +109,21 @@ At the small default dims MXFP8 is slower than bf16: the fused path launches man quantization kernels and the layer becomes CPU-launch-bound. Running under `nsys` adds about 1.5 ms per iteration to these numbers. -## TE MoE vs. plain PyTorch MoE vs. dense layer +## TE MoE vs. plain PyTorch MoE Same layer, same dims (`--dsv3`), bf16 unless noted: | | 4 GPUs (32 experts) | 8 GPUs (64 experts) | |---|---|---| -| `--impl te`, bf16 | 13.09 ms | 15.06 ms | -| `--impl te`, mxfp8 fused | 10.19 ms | 11.32 ms | -| `--impl naive`, bf16 | 26.83 ms | 27.34 ms | -| `--impl dense`, bf16 (ffn 18432) | 9.99 ms | 10.04 ms | -| `--impl dense`, mxfp8 | 7.49 ms | 7.58 ms | +| `naive`: all_to_all + loop over experts | 26.83 ms | 27.34 ms | +| `naive_grouped`: all_to_all + TE grouped GEMM | 16.84 ms | 17.20 ms | +| `te`: NCCL EP + grouped GEMM | 13.09 ms | 15.06 ms | +| `te`, mxfp8 (unfused grouped GEMM) | 11.02 ms | 12.88 ms | +| `te`, mxfp8 fused | 10.19 ms | 11.32 ms | -At the small default dims the gap is similar: `naive` 10.08 ms vs `te` 6.17 ms on 4 GPUs. +At the small default dims (4 GPUs): `naive` 10.08 ms, `naive_grouped` 6.80 ms, `te` 6.17 ms. -Per GPU and iteration, the naive MoE spends (8 GPUs, kernel time 25.9 ms of a 28.5 ms +Per GPU and iteration, the `naive` MoE spends (8 GPUs, kernel time 25.9 ms of a 28.5 ms iteration): | group | ms | what | @@ -135,13 +136,25 @@ iteration): | indexing kernels | 3.0 | `x[tok]`, `nonzero` masks, `index_copy`, `indexing_backward` | | attention, norms | 1.1 | same as in the TE variant | -The TE variant replaces all of the communication and indexing rows with two NCCL EP kernels -per direction (dispatch, combine) writing directly into the expert-major layout, and the -per-expert GEMMs with one grouped GEMM, which is where the roughly 2x comes from. The dense -layer is faster than either MoE variant here because with 8 experts per rank and top-k 8 the -MoE moves 8 activations per token across ranks while the dense MLP does the equivalent FLOPs -locally; the MoE wins only once the expert count grows past what a dense layer of equal -per-token FLOPs can hold. +`naive_grouped` (8 GPUs, kernel time 17.1 ms of a 17.2 ms iteration): + +| group | ms | what | +|---|---|---| +| `ncclDevKernel_SendRecv` | 3.8 | the same 7 all_to_all launches, less time because the GPU is no longer stalled between them | +| grouped GEMMs (`nvjet_*_ptrGroup_*`) | 4.1 | fc1 / fc2 forward, dgrad, wgrad as grouped GEMMs, same as in `te` | +| sorting rows by expert and back | 2.6 | `argsort`, gathers (`x[tok]`, `x_recv[by_expert]`), `index_copy`, `indexing_backward` | +| dense GEMMs (MLA projections, shared expert) | 2.4 | same as in `te` | +| elementwise adds | 0.7 | `index_add`, residuals | +| attention, norms | 1.1 | same as in `te` | + +Reading the three rows of the table together: the Python loop over experts costs about +10 ms per iteration (`naive` -> `naive_grouped`, 8 separate GEMM pairs, `nonzero` masks, +zero-filled buffers, copies); replacing torch `all_to_all` plus the surrounding sort / gather / +scatter with NCCL EP dispatch and combine, which write straight into the expert-major layout +and zero-fill the padding, saves another 2 ms (`naive_grouped` -> `te`). Both `naive` +variants also synchronise with the host twice per layer to learn the all_to_all split sizes; +`te` does not. `--recipe mxfp8` is only supported by `te`: the naive variants would need the +per-expert row counts padded to the MXFP8 block size. ## Where the time goes (8 GPUs, `--dsv3`, mxfp8 fused) diff --git a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py index d42797dfe0..880ca5c582 100644 --- a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py +++ b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py @@ -7,7 +7,7 @@ holds ``--num-local-experts`` routed experts. ``--impl te`` exchanges tokens with NCCL EP and runs the experts as one grouped GEMM (DeepSeekV3MoE); ``--impl naive`` is a plain PyTorch MoE (all_to_all_single + a Python loop over experts) dropped into the same layer; -``--impl dense`` replaces the MoE with the dense SwiGLU MLP of DeepSeek-V3's first layers. +``--impl naive_grouped`` keeps the all_to_all but runs the experts as one TE grouped GEMM. Timed iterations run inside a ``torch.cuda.profiler`` window, so ``nsys profile -c cudaProfilerApi --capture-range-end=stop torchrun ...`` records only them. @@ -51,8 +51,7 @@ def _parse_args(): " 2048)." ), ) - p.add_argument("--impl", choices=["te", "naive", "dense"], default="te") - p.add_argument("--dense-ffn", type=int, default=18432, help="ffn size for --impl dense") + p.add_argument("--impl", choices=["te", "naive", "naive_grouped"], default="te") p.add_argument("--recipe", choices=["none", "mxfp8"], default="none") p.add_argument("--warmup", type=int, default=5) p.add_argument("--iters", type=int, default=10) @@ -71,23 +70,32 @@ def _autocast(name): class NaiveMoE(torch.nn.Module): - """DeepSeek-style MoE without TE: sigmoid top-k router with expert bias, torch all_to_all - dispatch/combine, a Python loop of dense SwiGLU experts, and a shared expert.""" + """DeepSeek-style MoE with torch all_to_all dispatch/combine: sigmoid top-k router with + expert bias, a shared expert, and experts either as a Python loop of dense SwiGLU MLPs or + (``grouped=True``) as one TE grouped GEMM stack.""" - def __init__(self, hidden, ffn, num_experts, topk, ep_group, shared_ffn, dtype): + def __init__(self, hidden, ffn, num_experts, topk, ep_group, shared_ffn, dtype, grouped=False): super().__init__() + self.grouped = grouped self.hidden, self.topk, self.group = hidden, topk, ep_group self.ws, self.rank = dist.get_world_size(ep_group), dist.get_rank(ep_group) self.num_experts, self.local = num_experts, num_experts // self.ws self.gate = torch.nn.Linear(hidden, num_experts, bias=False, dtype=dtype, device="cuda") self.register_buffer("expert_bias", torch.zeros(num_experts, device="cuda")) std = hidden**-0.5 - self.w1 = torch.nn.Parameter( - torch.randn(self.local, 2 * ffn, hidden, dtype=dtype, device="cuda") * std - ) - self.w2 = torch.nn.Parameter( - torch.randn(self.local, hidden, ffn, dtype=dtype, device="cuda") * ffn**-0.5 - ) + if grouped: + self.experts = te.ops.Sequential( + te.ops.GroupedLinear(self.local, hidden, 2 * ffn, bias=False, dtype=dtype), + te.ops.ScaledSwiGLU(glu_interleave_size=32), + te.ops.GroupedLinear(self.local, ffn, hidden, bias=False, dtype=dtype), + ) + else: + self.w1 = torch.nn.Parameter( + torch.randn(self.local, 2 * ffn, hidden, dtype=dtype, device="cuda") * std + ) + self.w2 = torch.nn.Parameter( + torch.randn(self.local, hidden, ffn, dtype=dtype, device="cuda") * ffn**-0.5 + ) self.shared_w1 = torch.nn.Linear( hidden, 2 * shared_ffn, bias=False, dtype=dtype, device="cuda" ) @@ -125,15 +133,24 @@ def forward(self, hidden_states): p_recv = torch.empty(n_recv, dtype=flat_p.dtype, device=x.device) dist.all_to_all_single(e_recv, flat_e.contiguous(), recv, send, group=self.group) dist.all_to_all_single(p_recv, flat_p.contiguous(), recv, send, group=self.group) - # Experts: one dense SwiGLU MLP per local expert. local_e = e_recv - self.rank * self.local - y_recv = torch.zeros_like(x_recv) - for e in range(self.local): - sel = (local_e == e).nonzero().squeeze(1) - if sel.numel() == 0: - continue - h = self._swiglu(F.linear(x_recv[sel], self.w1[e])) * p_recv[sel, None].to(x.dtype) - y_recv = y_recv.index_copy(0, sel, F.linear(h, self.w2[e])) + if self.grouped: + # Experts: sort received rows by local expert, one grouped GEMM stack. + by_expert = torch.argsort(local_e, stable=True) + counts = torch.bincount(local_e, minlength=self.local) + y_sorted = self.experts( + x_recv[by_expert], counts, p_recv[by_expert].to(x.dtype), counts + ) + y_recv = torch.empty_like(x_recv).index_copy(0, by_expert, y_sorted) + else: + # Experts: one dense SwiGLU MLP per local expert. + y_recv = torch.zeros_like(x_recv) + for e in range(self.local): + sel = (local_e == e).nonzero().squeeze(1) + if sel.numel() == 0: + continue + h = self._swiglu(F.linear(x_recv[sel], self.w1[e])) * p_recv[sel, None].to(x.dtype) + y_recv = y_recv.index_copy(0, sel, F.linear(h, self.w2[e])) # Combine: reverse all_to_all, sum the top-k contributions per token. y = all_to_all_single(torch.empty_like(x[tok]), y_recv, send, recv, group=self.group) out = torch.zeros_like(x).index_add(0, tok, y) @@ -180,9 +197,7 @@ def main(): ep_group=ep_group, ep_max_tokens_per_rank=args.tokens_per_rank, ) - if args.impl == "dense": - mlp_kwargs = dict(ffn_hidden_size=args.dense_ffn) - elif args.impl == "naive": + if args.impl != "te": # Build the TE MoE without EP (replaced below); keeps the pre-MLP RMSNorm. mlp_kwargs.pop("ep_group"), mlp_kwargs.pop("ep_max_tokens_per_rank") layer = DeepSeekV3Layer( @@ -196,7 +211,7 @@ def main(): qk_rope_head_dim=args.qk_rope_head_dim, v_head_dim=args.v_head_dim, ) - if args.impl == "naive": + if args.impl != "te": layer.mlp = NaiveMoE( args.hidden, args.moe_ffn, @@ -205,6 +220,7 @@ def main(): ep_group, args.moe_ffn, torch.bfloat16, + grouped=args.impl == "naive_grouped", ) seq = args.tokens_per_rank // 4 x = torch.randn(seq, 4, args.hidden, dtype=torch.bfloat16, device="cuda", requires_grad=True) From 7a1693a8a007e2c454dbbd7fdb408855ce1a85e9 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 19:19:34 +0200 Subject: [PATCH 34/43] [PyTorch] DeepSeekV3: MLA TP gradient reduction, RoPE guards, per-forward EP buffer MultiLatentAttention: all-reduce the shared rope-key gradient across TP ranks, derive tp_size from tp_group, refresh RoPE tables on device change. MLA RoPE: use the Triton kernels only for power-of-two head dims. DeepSeekV3MoE: create the EpBuffer per forward so several in-flight microbatches keep their own routing state until backward. DeepSeekV3Layer: pass layernorm_epsilon to the MoE pre-norm. Tests: MLA TP vs unsharded reference, EP with 1 and 3 in-flight microbatches, RoPE fallback dims, NaiveMoE gradient check. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 98 +++++++++++++++---- tests/pytorch/distributed/test_models.py | 11 +++ tests/pytorch/test_models.py | 68 ++++++++++++- .../pytorch/models/deepseek_v3/mla_rope.py | 18 +++- .../pytorch/models/deepseek_v3/moe.py | 32 +++--- .../deepseek_v3/multi_latent_attention.py | 28 +++++- .../models/deepseek_v3/transformer_layer.py | 1 + 7 files changed, 216 insertions(+), 40 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index f1f7d5e58f..302da451c5 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -10,7 +10,7 @@ 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, DeepSeekV3MoE +from transformer_engine.pytorch.models import DeepSeekV3Layer, DeepSeekV3MoE, MultiLatentAttention HIDDEN = 256 MOE_FFN = 128 @@ -78,7 +78,9 @@ def _copy_weights(ep_layer: DeepSeekV3Layer, ref: DeepSeekV3Layer, rank: int) -> 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: +def test_layer_ep_matches_local( + rank: int, ep_size: int, ep_group, num_microbatches: int = 1 +) -> None: """Full DeepSeekV3Layer with EP must match the all-experts-local layer numerically.""" num_experts = NUM_LOCAL_EXPERTS * ep_size torch.manual_seed(0) @@ -88,19 +90,24 @@ def test_layer_ep_matches_local(rank: int, ep_size: int, ep_group) -> None: _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) + microbatches = [] + for mb in range(num_microbatches): + seq_len = TOKENS_PER_RANK // 2 - 8 * mb + x = torch.randn(seq_len, 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) + microbatches.append((x_ep, x_ref, out_ep, out_ref)) + + # All microbatches must retain their routing until their own backward. + for x_ep, x_ref, out_ep, out_ref in microbatches: + grad_out = torch.randn_like(out_ep) + out_ep.backward(grad_out.clone()) + out_ref.backward(grad_out.clone()) + 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(): @@ -124,15 +131,71 @@ def test_layer_ep_matches_local(rank: int, ep_size: int, ep_group) -> None: counts = ep_layer.mlp._last_tokens_per_expert.clone() dist.all_reduce(counts) - assert counts.sum().item() == ep_size * TOKENS_PER_RANK * TOP_K + last_num_tokens = microbatches[-1][0].numel() // HIDDEN + assert counts.sum().item() == ep_size * last_num_tokens * TOP_K ep_layer.mlp.update_expert_bias() assert torch.isfinite(ep_layer.mlp.expert_bias).all() +def test_mla_tp_matches_local(rank: int, tp_size: int, tp_group) -> None: + """Compare sharded MLA outputs and gradients with an unsharded reference.""" + torch.backends.cuda.matmul.allow_tf32 = False + for fmt in ("sbhd", "bshd"): + for explicit_size in (False, True): + torch.manual_seed(0) + common = dict(params_dtype=torch.float32, qkv_format=fmt, **MLA_KWARGS) + ref = MultiLatentAttention(HIDDEN, 2 * tp_size, **common) + _broadcast_params(ref) + tp = MultiLatentAttention( + HIDDEN, + 2 * tp_size, + tp_group=tp_group, + **({"tp_size": tp_size} if explicit_size else {}), + **common, + ) + ref_params = dict(ref.named_parameters()) + + def shard(name, tensor): + if name in ("q_up_proj.weight", "kv_up_proj.weight"): + return tensor.chunk(tp_size, dim=0)[rank] + if name == "out_proj.weight": + return tensor.chunk(tp_size, dim=1)[rank] + return tensor + + with torch.no_grad(): + for name, param in tp.named_parameters(): + param.copy_(shard(name, ref_params[name])) + + torch.manual_seed(1234) + shape = (16, 2, HIDDEN) if fmt == "sbhd" else (2, 16, HIDDEN) + x = torch.randn(shape, device="cuda", requires_grad=True) + x_ref = x.detach().clone().requires_grad_() + out = tp(x) + out_ref = ref(x_ref) + torch.testing.assert_close(out, out_ref, rtol=1e-3, atol=1e-3) + grad = torch.randn_like(out) + out.backward(grad.clone()) + out_ref.backward(grad.clone()) + torch.testing.assert_close(x.grad, x_ref.grad, rtol=1e-3, atol=1e-3) + for name, param in tp.named_parameters(): + torch.testing.assert_close( + param.grad, + shard(name, ref_params[name].grad), + rtol=1e-3, + atol=1e-3, + msg=name, + ) + + def main() -> int: dist.init_process_group(backend="nccl") torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + if "--tp" in sys.argv: + test_mla_tp_matches_local(dist.get_rank(), dist.get_world_size(), dist.group.WORLD) + print(f"[rank {dist.get_rank()}] TP PASSED") + dist.destroy_process_group() + return 0 from torch.distributed import _symmetric_memory as _symm_mem _symm_mem.set_backend("NCCL") @@ -154,7 +217,8 @@ def main() -> int: num_topk=TOP_K, recv_capacity_per_rank=_recv_capacity(ep_size), ) - test_layer_ep_matches_local(rank, ep_size, ep_group) + for num_microbatches in (1, 3): + test_layer_ep_matches_local(rank, ep_size, ep_group, num_microbatches) print(f"[rank {rank}] PASSED") dist.barrier() diff --git a/tests/pytorch/distributed/test_models.py b/tests/pytorch/distributed/test_models.py index 1b96eae2aa..9f7ff6d6e0 100644 --- a/tests/pytorch/distributed/test_models.py +++ b/tests/pytorch/distributed/test_models.py @@ -29,3 +29,14 @@ def test_deepseek_layer_ep(): LAUNCH_CMD + [str(TEST_ROOT / "run_models.py")], env=os.environ, check=False, timeout=300 ) assert result.returncode == 0 + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="TP requires >= 2 GPUs") +def test_mla_tp(): + result = subprocess.run( + LAUNCH_CMD + [str(TEST_ROOT / "run_models.py"), "--tp"], + env=os.environ, + check=False, + timeout=300, + ) + assert result.returncode == 0 diff --git a/tests/pytorch/test_models.py b/tests/pytorch/test_models.py index cd1903e79f..41f0833b20 100644 --- a/tests/pytorch/test_models.py +++ b/tests/pytorch/test_models.py @@ -3,9 +3,12 @@ # See LICENSE for license information. import math +import runpy +from pathlib import Path import pytest import torch +import torch.distributed as dist from transformer_engine.pytorch.utils import deinterleave_glu_tensor from transformer_engine.pytorch.models import DeepSeekV3MoE, MultiLatentAttention @@ -32,13 +35,13 @@ def _input(requires_grad=True): ) -def test_mla_rope_triton_matches_pytorch(): +@pytest.mark.parametrize("nope,rope,vdim", [(64, 32, 64), (48, 32, 64), (64, 48, 64), (64, 32, 48)]) +def test_mla_rope_matches_pytorch(nope, rope, vdim): from transformer_engine.pytorch.models.deepseek_v3 import mla_rope if not mla_rope.HAVE_TRITON: pytest.skip("Triton unavailable") s, b, h = 64, 2, 4 - nope, rope, vdim = 64, 32, 64 cos, sin = mla_rope.build_rope_tables(s, rope, device="cuda") torch.manual_seed(0) @@ -174,3 +177,64 @@ def test_moe_matches_dense_reference(shared, grouped, topk): assert torch.isfinite(moe.expert_bias).all() if topk < num_experts: assert not torch.equal(bias_before, moe.expert_bias) + + +@pytest.fixture(scope="module") +def deepseek_example(): + path = ( + Path(__file__).resolve().parents[2] / "examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py" + ) + return runpy.run_path(str(path)) + + +@pytest.fixture +def single_rank_group(tmp_path): + if dist.is_initialized(): + pytest.skip("Requires an isolated process group") + dist.init_process_group("nccl", init_method=(tmp_path / "store").as_uri(), rank=0, world_size=1) + try: + yield dist.group.WORLD + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("grouped", [False, True]) +def test_naive_moe_gradients(deepseek_example, single_rank_group, grouped): + torch.manual_seed(123) + moe = deepseek_example["NaiveMoE"]( + HIDDEN, 128, 4, 2, single_rank_group, 128, torch.float32, grouped + ) + x = torch.randn(64, HIDDEN, device="cuda", requires_grad=True) + x_ref = x.detach().clone().requires_grad_() + out = moe(x) + + scores = torch.sigmoid(moe.gate(x_ref)) + idx = torch.topk(scores + moe.expert_bias, moe.topk, dim=-1).indices + selected = scores.gather(1, idx) + selected = selected / selected.sum(-1, keepdim=True) * 2.5 + probs = torch.zeros_like(scores).scatter(1, idx, selected) + ref = torch.zeros_like(x_ref) + for e in range(moe.local): + if grouped: + fc1, _, fc2 = moe.experts + w1 = deinterleave_glu_tensor(getattr(fc1, f"weight{e}"), 32) + w2 = getattr(fc2, f"weight{e}") + else: + w1, w2 = moe.w1[e], moe.w2[e] + act = moe._swiglu(torch.nn.functional.linear(x_ref, w1)) + ref = ref + torch.nn.functional.linear(act * probs[:, e : e + 1], w2) + ref = ref + moe.shared_w2(moe._swiglu(moe.shared_w1(x_ref))) + + grad = torch.randn_like(out) + actual_grads = torch.autograd.grad(out, (x, moe.gate.weight), grad) + ref_grads = torch.autograd.grad(ref, (x_ref, moe.gate.weight), grad) + torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3) + for actual, expected in zip(actual_grads, ref_grads): + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + assert actual_grads[1].abs().max() > 0 + + +@pytest.mark.parametrize("value", [1.0, float("nan"), float("inf"), None]) +def test_example_finite_check(deepseek_example, single_rank_group, value): + tensor = None if value is None else torch.tensor(value, device="cuda") + assert deepseek_example["_check_finite"]([tensor], torch.device("cuda")) == (value == 1.0) diff --git a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py index 0aea284afd..7c47a02848 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py +++ b/transformer_engine/pytorch/models/deepseek_v3/mla_rope.py @@ -537,7 +537,12 @@ def apply_mla_rope_q( tensor_format: str = "sbhd", ) -> torch.Tensor: """RoPE on the trailing ``head_dim_rope`` slice of q; in place on the Triton path.""" - if HAVE_TRITON and tensor_format == "sbhd": + if ( + HAVE_TRITON + and tensor_format == "sbhd" + and head_dim_rope >= 2 + and _is_power_of_two(head_dim_rope) + ): return _MLARoPEQTriton.apply(q, cos_table, sin_table, head_dim_nope, head_dim_rope) seq_dim = 0 if tensor_format == "sbhd" else 1 q_rope = _rotate_interleaved_to_neox(q[..., head_dim_nope:], cos_table, sin_table, seq_dim) @@ -555,7 +560,12 @@ def apply_mla_rope_kv( tensor_format: str = "sbhd", ) -> Tuple[torch.Tensor, torch.Tensor]: """Build (k, v) from kv ``[.., h, nope+v]`` and the shared rope head ``[.., 1, rope]``.""" - if HAVE_TRITON and tensor_format == "sbhd": + if ( + HAVE_TRITON + and tensor_format == "sbhd" + and head_dim_rope >= 2 + and all(_is_power_of_two(dim) for dim in (head_dim_nope, head_dim_rope, head_dim_v)) + ): return _MLARoPEKVTriton.apply( kv, k_pos_emb, cos_table, sin_table, head_dim_nope, head_dim_rope, head_dim_v ) @@ -565,3 +575,7 @@ def apply_mla_rope_kv( k_rope = _rotate_interleaved_to_neox(k_pos_emb, cos_table, sin_table, seq_dim) k_rope = k_rope.expand(*k_nope.shape[:-1], -1) return torch.cat((k_nope, k_rope), dim=-1), v.contiguous() + + +def _is_power_of_two(value: int) -> bool: + return value > 0 and value & (value - 1) == 0 diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index aa1ff9408d..22be051749 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -148,23 +148,20 @@ def __init__( hidden_size, shared_expert_ffn_hidden_size, dtype, device ) - self.ep_buffer = None + self._ep_buffer_kwargs = None if ep_group is not None: - from transformer_engine.pytorch.ep import EpBuffer - assert ep_max_tokens_per_rank is not None, "EP requires ep_max_tokens_per_rank." cap = self.ep_recv_capacity( self.ep_size, ep_max_tokens_per_rank, topk, num_local_experts ) - self.ep_buffer = EpBuffer( - top_k=topk, - max_tokens_per_rank=ep_max_tokens_per_rank, - hidden_dim=hidden_size, - num_local_experts=num_local_experts, - recv_capacity_per_rank=cap, - alignment=_EP_ALIGNMENT, - device=device, - ) + self._ep_buffer_kwargs = { + "top_k": topk, + "max_tokens_per_rank": ep_max_tokens_per_rank, + "hidden_dim": hidden_size, + "num_local_experts": num_local_experts, + "recv_capacity_per_rank": cap, + "alignment": _EP_ALIGNMENT, + } @staticmethod def ep_recv_capacity( @@ -224,9 +221,10 @@ def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: ) def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: - from transformer_engine.pytorch.ep import ep_dispatch, ep_combine + from transformer_engine.pytorch.ep import EpBuffer, ep_dispatch, ep_combine assert tokens.dtype == torch.bfloat16, "The EP path requires bfloat16 inputs." + buffer = EpBuffer(**self._ep_buffer_kwargs, device=tokens.device) topk_idx = torch.empty( (tokens.shape[0], self.topk), dtype=torch.int64, device=tokens.device ) @@ -241,9 +239,9 @@ def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: # the recv/grad buffers can stay uninitialized. The fused grouped MLP # reads up to a tile past the last expert, so zero a margin there # (offsets stay on device: no host sync). - cap = self.ep_buffer.recv_capacity_per_rank + cap = buffer.recv_capacity_per_rank recv_tokens, recv_weights, tokens_per_expert = ep_dispatch( - self.ep_buffer, + buffer, tokens, topk_idx, topk_weights, @@ -262,9 +260,7 @@ def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: expert_out = self.experts( recv_tokens, tokens_per_expert, recv_weights.to(tokens.dtype), tokens_per_expert ) - return ep_combine( - self.ep_buffer, expert_out, num_local_tokens=tokens.shape[0], grad_out=grad_out - ) + return ep_combine(buffer, expert_out, num_local_tokens=tokens.shape[0], grad_out=grad_out) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """ diff --git a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py index 56b4d3d0d1..f7f9053205 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py +++ b/transformer_engine/pytorch/models/deepseek_v3/multi_latent_attention.py @@ -11,6 +11,7 @@ from transformer_engine.pytorch.module import Linear, LayerNormLinear from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.distributed import allreduce, get_distributed_world_size from transformer_engine.pytorch.models.deepseek_v3.mla_rope import ( apply_mla_rope_kv, apply_mla_rope_q, @@ -21,6 +22,23 @@ __all__ = ["MultiLatentAttention"] +class _ReduceGrad(torch.autograd.Function): + """Replicate an input across TP ranks and sum its gradients.""" + + @staticmethod + def forward(ctx, inp, tp_group): + """Return the replicated input unchanged.""" + ctx.tp_group = tp_group + return inp + + @staticmethod + def backward(ctx, grad_output): + """Sum gradients from each rank's attention heads.""" + grad_input = grad_output.clone(memory_format=torch.contiguous_format) + grad_input, _ = allreduce(grad_input, ctx.tp_group) + return grad_input, None + + class MultiLatentAttention(torch.nn.Module): """ Multi-Latent Attention as used in DeepSeekV3. @@ -117,6 +135,8 @@ def __init__( ) -> None: super().__init__() + if tp_group is not None: + tp_size = get_distributed_world_size(tp_group) assert qkv_format in ("sbhd", "bshd"), "MultiLatentAttention supports sbhd/bshd formats." assert num_attention_heads % tp_size == 0 @@ -188,7 +208,11 @@ def __init__( ) def _rope_tables_for(self, seq_len: int, device: torch.device): - if self._rope_tables is None or self._rope_tables[0].shape[0] < seq_len: + if ( + self._rope_tables is None + or self._rope_tables[0].shape[0] < seq_len + or self._rope_tables[0].device != device + ): self._rope_tables = build_rope_tables( seq_len, self.qk_rope_head_dim, @@ -227,6 +251,8 @@ def forward( kv_down = self.kv_down_proj(hidden_states) kv_latent, k_pos = torch.split(kv_down, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + if self.kv_up_proj.tp_size > 1: + k_pos = _ReduceGrad.apply(k_pos, self.kv_up_proj.tp_group) kv = self.kv_up_proj(kv_latent) kv = kv.view(*kv.shape[:-1], heads, self.qk_nope_head_dim + self.v_head_dim) diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index ab11fe6394..741b235bd9 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -113,6 +113,7 @@ def __init__( self.self_attention = MultiLatentAttention( hidden_size, num_attention_heads, + layernorm_epsilon=layernorm_epsilon, params_dtype=params_dtype, device=device, **mla_kwargs, From 49ecaffde81b4113cdd676de064f006b5fc4cf40 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 19:19:34 +0200 Subject: [PATCH 35/43] [PyTorch] DeepSeekV3Layer example: argument validation, gradient checks, differentiable naive probs Validate CLI args, check output and all gradients (including the router) for finiteness across ranks after timing, clear parameter grads each step. The naive MoE now sends routing probabilities through the differentiable all_to_all so the router gets its gradient. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 13 +++++- .../deepseek_v3/deepseek_v3_layer_ep.py | 40 +++++++++++++++---- 2 files changed, 44 insertions(+), 9 deletions(-) diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md index 00e0f9ceb9..9d2608b92b 100644 --- a/examples/pytorch/deepseek_v3/README.md +++ b/examples/pytorch/deepseek_v3/README.md @@ -4,10 +4,14 @@ (RMSNorm, multi-latent attention, DeepSeek MoE with a shared expert) with the routed experts sharded over all ranks. Tokens are exchanged with NCCL EP (`transformer_engine.pytorch.ep`). Each iteration is a forward and a backward pass on random data; the script reports the time -per iteration and tokens per second and checks that the output is finite. +per iteration and tokens per second. After timing, it checks the final output and gradients +on every rank, including the presence of input and router gradients. ## Requirements +The following NCCL EP requirements apply to `--impl te`. The naive variants use +ordinary NCCL collectives and can also run on older GPUs without NVLink. + - SM90 or newer GPUs connected with NVLink (NCCL EP falls back to the network transport and deadlocks on PCIe-only nodes). - NCCL >= 2.30.4, PyTorch >= 2.11 (symmetric memory), Transformer Engine built with the @@ -51,6 +55,8 @@ Every rank owns `--num-local-experts` experts (default 8), so the expert count i `8 * world_size`. Other knobs: `--tokens-per-rank`, `--topk`, `--hidden`, `--num-heads`, `--moe-ffn`, the MLA dims (`--q-lora-rank`, `--kv-lora-rank`, `--qk-nope-head-dim`, `--qk-rope-head-dim`, `--v-head-dim`), `--warmup`, `--iters`. +`--warmup 0` is supported; `--iters` must be positive and `--tokens-per-rank` must be a +positive multiple of four. MXFP8 requires `--impl te`. ## Profiling with nsys @@ -111,6 +117,11 @@ about 1.5 ms per iteration to these numbers. ## TE MoE vs. plain PyTorch MoE +The naive timings and profiles below predate the routing-probability autograd fix: +they omit router backward and probability-gradient communication. They are historical +measurements and must be rerun before drawing training-speedup conclusions. The current +benchmark also clears parameter gradients before each step. + Same layer, same dims (`--dsv3`), bf16 unless noted: | | 4 GPUs (32 experts) | 8 GPUs (64 experts) | diff --git a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py index 880ca5c582..ca7144b332 100644 --- a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py +++ b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py @@ -30,7 +30,7 @@ from transformer_engine.pytorch.models import DeepSeekV3Layer, DeepSeekV3MoE -def _parse_args(): +def _parse_args(argv=None): p = argparse.ArgumentParser(description="DeepSeekV3Layer EP example (fwd + bwd)") p.add_argument("--tokens-per-rank", type=int, default=4096) p.add_argument("--hidden", type=int, default=2048) @@ -55,7 +55,15 @@ def _parse_args(): p.add_argument("--recipe", choices=["none", "mxfp8"], default="none") p.add_argument("--warmup", type=int, default=5) p.add_argument("--iters", type=int, default=10) - args = p.parse_args() + args = p.parse_args(argv) + if args.warmup < 0: + p.error("--warmup must be non-negative") + if args.iters <= 0: + p.error("--iters must be positive") + if args.tokens_per_rank <= 0 or args.tokens_per_rank % 4: + p.error("--tokens-per-rank must be a positive multiple of 4") + if args.impl != "te" and args.recipe != "none": + p.error("--recipe mxfp8 is only supported with --impl te") if args.dsv3: args.hidden, args.num_heads, args.moe_ffn = 7168, 128, 2048 args.q_lora_rank, args.kv_lora_rank = 1536, 512 @@ -69,6 +77,17 @@ def _autocast(name): return te.autocast(enabled=True, recipe=te_recipe.MXFP8BlockScaling()) +def _check_finite(tensors, device): + finite = torch.ones((), dtype=torch.int32, device=device) + for tensor in tensors: + if tensor is None: + finite.zero_() + else: + finite.mul_(torch.isfinite(tensor).all()) + dist.all_reduce(finite, op=dist.ReduceOp.MIN) + return bool(finite.item()) + + class NaiveMoE(torch.nn.Module): """DeepSeek-style MoE with torch all_to_all dispatch/combine: sigmoid top-k router with expert bias, a shared expert, and experts either as a Python loop of dense SwiGLU MLPs or @@ -132,7 +151,7 @@ def forward(self, hidden_states): e_recv = torch.empty(n_recv, dtype=flat_e.dtype, device=x.device) p_recv = torch.empty(n_recv, dtype=flat_p.dtype, device=x.device) dist.all_to_all_single(e_recv, flat_e.contiguous(), recv, send, group=self.group) - dist.all_to_all_single(p_recv, flat_p.contiguous(), recv, send, group=self.group) + p_recv = all_to_all_single(p_recv, flat_p.contiguous(), recv, send, group=self.group) local_e = e_recv - self.rank * self.local if self.grouped: # Experts: sort received rows by local expert, one grouped GEMM stack. @@ -167,7 +186,7 @@ def main(): rank, world_size = dist.get_rank(), dist.get_world_size() major, minor = torch.cuda.get_device_capability() - if major * 10 + minor < 90: + if args.impl == "te" and major * 10 + minor < 90: if rank == 0: print(f"SKIPPED: NCCL EP requires SM>=90 (got SM{major}{minor})") dist.destroy_process_group() @@ -226,15 +245,15 @@ def main(): x = torch.randn(seq, 4, args.hidden, dtype=torch.bfloat16, device="cuda", requires_grad=True) def step(): + layer.zero_grad(set_to_none=True) + x.grad = None with _autocast(args.recipe): out = layer(x) out.backward(torch.ones_like(out)) - x.grad = None return out for _ in range(args.warmup): - out = step() - finite = bool(torch.isfinite(out).all()) + step() torch.cuda.synchronize() dist.barrier() @@ -242,10 +261,15 @@ def step(): start = time.perf_counter() for i in range(args.iters): with torch.cuda.nvtx.range(f"iter{i}"): - step() + out = step() torch.cuda.synchronize() ms = (time.perf_counter() - start) / args.iters * 1e3 torch.cuda.profiler.stop() + finite = _check_finite( + [out, x.grad, layer.mlp.gate.weight.grad] + + [p.grad for p in layer.parameters() if p.grad is not None], + x.device, + ) dist.barrier() if rank == 0: From a131eb3e939a82324f45386accf4c07bfc1c2e83 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 19:24:18 +0200 Subject: [PATCH 36/43] [PyTorch] Drop example-dependent tests from test_models.py Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_models.py | 64 ------------------------------------ 1 file changed, 64 deletions(-) diff --git a/tests/pytorch/test_models.py b/tests/pytorch/test_models.py index 41f0833b20..5d3a2d5531 100644 --- a/tests/pytorch/test_models.py +++ b/tests/pytorch/test_models.py @@ -3,12 +3,9 @@ # See LICENSE for license information. import math -import runpy -from pathlib import Path import pytest import torch -import torch.distributed as dist from transformer_engine.pytorch.utils import deinterleave_glu_tensor from transformer_engine.pytorch.models import DeepSeekV3MoE, MultiLatentAttention @@ -177,64 +174,3 @@ def test_moe_matches_dense_reference(shared, grouped, topk): assert torch.isfinite(moe.expert_bias).all() if topk < num_experts: assert not torch.equal(bias_before, moe.expert_bias) - - -@pytest.fixture(scope="module") -def deepseek_example(): - path = ( - Path(__file__).resolve().parents[2] / "examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py" - ) - return runpy.run_path(str(path)) - - -@pytest.fixture -def single_rank_group(tmp_path): - if dist.is_initialized(): - pytest.skip("Requires an isolated process group") - dist.init_process_group("nccl", init_method=(tmp_path / "store").as_uri(), rank=0, world_size=1) - try: - yield dist.group.WORLD - finally: - dist.destroy_process_group() - - -@pytest.mark.parametrize("grouped", [False, True]) -def test_naive_moe_gradients(deepseek_example, single_rank_group, grouped): - torch.manual_seed(123) - moe = deepseek_example["NaiveMoE"]( - HIDDEN, 128, 4, 2, single_rank_group, 128, torch.float32, grouped - ) - x = torch.randn(64, HIDDEN, device="cuda", requires_grad=True) - x_ref = x.detach().clone().requires_grad_() - out = moe(x) - - scores = torch.sigmoid(moe.gate(x_ref)) - idx = torch.topk(scores + moe.expert_bias, moe.topk, dim=-1).indices - selected = scores.gather(1, idx) - selected = selected / selected.sum(-1, keepdim=True) * 2.5 - probs = torch.zeros_like(scores).scatter(1, idx, selected) - ref = torch.zeros_like(x_ref) - for e in range(moe.local): - if grouped: - fc1, _, fc2 = moe.experts - w1 = deinterleave_glu_tensor(getattr(fc1, f"weight{e}"), 32) - w2 = getattr(fc2, f"weight{e}") - else: - w1, w2 = moe.w1[e], moe.w2[e] - act = moe._swiglu(torch.nn.functional.linear(x_ref, w1)) - ref = ref + torch.nn.functional.linear(act * probs[:, e : e + 1], w2) - ref = ref + moe.shared_w2(moe._swiglu(moe.shared_w1(x_ref))) - - grad = torch.randn_like(out) - actual_grads = torch.autograd.grad(out, (x, moe.gate.weight), grad) - ref_grads = torch.autograd.grad(ref, (x_ref, moe.gate.weight), grad) - torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3) - for actual, expected in zip(actual_grads, ref_grads): - torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) - assert actual_grads[1].abs().max() > 0 - - -@pytest.mark.parametrize("value", [1.0, float("nan"), float("inf"), None]) -def test_example_finite_check(deepseek_example, single_rank_group, value): - tensor = None if value is None else torch.tensor(value, device="cuda") - assert deepseek_example["_check_finite"]([tensor], torch.device("cuda")) == (value == 1.0) From b5605bcb346be12118a86122cece5aab5cb06b01 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 19:37:15 +0200 Subject: [PATCH 37/43] [Docs] Organize DeepSeekV3 example README around results Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 177 ++++++++++++++----------- 1 file changed, 102 insertions(+), 75 deletions(-) diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md index 9d2608b92b..b2d1181f33 100644 --- a/examples/pytorch/deepseek_v3/README.md +++ b/examples/pytorch/deepseek_v3/README.md @@ -1,13 +1,65 @@ # DeepSeekV3Layer with expert parallelism -`deepseek_v3_layer_ep.py` runs one `transformer_engine.pytorch.models.DeepSeekV3Layer` -(RMSNorm, multi-latent attention, DeepSeek MoE with a shared expert) with the routed experts -sharded over all ranks. Tokens are exchanged with NCCL EP (`transformer_engine.pytorch.ep`). -Each iteration is a forward and a backward pass on random data; the script reports the time -per iteration and tokens per second. After timing, it checks the final output and gradients -on every rank, including the presence of input and router gradients. +A forward + backward benchmark of one `DeepSeekV3Layer`: RMSNorm, multi-latent attention, +and MoE with a shared expert. Routed experts are sharded across GPUs using NCCL EP. -## Requirements +## Results + +GB300 GPUs, 4096 tokens per rank, top-k 8, 8 local experts per GPU. +Times cover one layer's forward + backward; throughput is global, in millions of tokens/s. +All rows below use `--impl te`. See [benchmark configuration](#c-benchmark-configuration) +for the full setup. + +### DeepSeek-V3 dimensions (`--dsv3`) + +| Precision | 4 GPUs · ms/iter | 4 GPUs · Mtok/s | 8 GPUs · ms/iter | 8 GPUs · Mtok/s | +|---|---:|---:|---:|---:| +| BF16 | 13.09 | 1.25 | 15.06 | 2.18 | +| MXFP8 | 11.02 | 1.49 | 12.88 | 2.54 | +| MXFP8 fused | 10.19 | 1.61 | 11.32 | 2.89 | + +4 GPUs = 1 node / 32 experts; 8 GPUs = 2 nodes / 64 experts. +Both nodes share one NVLink domain (MNNVL). + +### Small default dimensions + +1 node, 4 GPUs, hidden 2048, 16 heads, expert FFN 1024, 32 experts. + +| Precision | ms/iter | Mtok/s | +|---|---:|---:| +| BF16 | 6.17 | 2.65 | +| MXFP8 fused | 8.4 | 1.95 | + +“Fused” enables `NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`. +At the small dimensions, MXFP8 fused is slower than BF16. + +The [historical naive comparison](#f-historical-comparison-with-plain-pytorch) is in the +appendix; those measurements predate the router-backward fix. + +## Quick start + +Requires SM90+ GPUs with NVLink and an NCCL EP-enabled TE build; +see [full requirements](#a-requirements). From this directory: + +```bash +bash run_deepseek_v3_layer_ep.sh --dsv3 +``` + +For MXFP8 with the fused grouped MLP (SM100-class GPUs): + +```bash +NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 bash run_deepseek_v3_layer_ep.sh --dsv3 --recipe mxfp8 +``` + +## Appendix + +[Requirements](#a-requirements) · [Running](#b-running) · +[Configuration](#c-benchmark-configuration) · [Profiling](#d-profiling-with-nsys) · +[TE kernels](#e-te-kernel-breakdown) · +[Historical naive comparison](#f-historical-comparison-with-plain-pytorch) · +[EP internals](#g-ep-implementation-notes) + +### A. Requirements The following NCCL EP requirements apply to `--impl te`. The naive variants use ordinary NCCL collectives and can also run on older GPUs without NVLink. @@ -19,9 +71,9 @@ ordinary NCCL collectives and can also run on older GPUs without NVLink. - `NVTE_CUTEDSL_FUSED_GROUPED_MLP=1` additionally needs SM100-class GPUs and the CuTe DSL (`nvidia-cutlass-dsl`) for the fused MXFP8 grouped MLP. -## Running +### B. Running -Single node, all local GPUs: +Run these commands from this directory. Single node, all local GPUs: ```bash bash run_deepseek_v3_layer_ep.sh # small dims, bf16 @@ -58,7 +110,21 @@ Every rank owns `--num-local-experts` experts (default 8), so the expert count i `--warmup 0` is supported; `--iters` must be positive and `--tokens-per-rank` must be a positive multiple of four. MXFP8 requires `--impl te`. -## Profiling with nsys +### C. Benchmark configuration + +| Setting | Value | +|---|---| +| hardware | GB300 (SM103), 4 GPUs per node, both nodes in one NVLink domain (MNNVL) | +| software | CUDA 13.3, NCCL 2.30.7, PyTorch 2.13.0a0+9186a08b2c (NGC 26.07 build), cuDNN 9.24, Transformer Engine from this PR | +| `--dsv3` dims | hidden 7168, 128 heads, MLA q_lora 1536 / kv_lora 512 / nope 128 / rope 64 / v 128, expert ffn 2048, shared expert ffn 2048 | +| MoE | 8 local experts per rank (32 on 4 GPUs, 64 on 8), top-k 8, 4096 tokens per rank | +| precision | bf16 params and activations; `--recipe mxfp8` = MXFP8 block scaling for the expert GEMMs and dense projections | +| timing | fwd + bwd, 5 warmup, 10 timed iterations, no CUDA graphs | + +Each iteration uses random data. After timing, the script checks the final output and +gradients on every rank, including the presence of input and router gradients. + +### D. Profiling with nsys The timed iterations run inside a `torch.cuda.profiler.start()` / `stop()` window, so `-c cudaProfilerApi` records only them, one NVTX range per iteration: @@ -72,50 +138,30 @@ nsys stats --report cuda_gpu_kern_sum results/deepseek_v3_layer_ep_.ns The launcher wraps `torchrun`, so all local ranks land in one report. For multi-node runs put the same `nsys profile ... -o _%q{SLURM_NODEID}` in front of `torchrun` on each node. -## Configuration used below - -| | value | -|---|---| -| hardware | GB300 (SM103), 4 GPUs per node, both nodes in one NVLink domain (MNNVL) | -| software | CUDA 13.3, NCCL 2.30.7, PyTorch 2.13 nightly, cuDNN 9.24 | -| `--dsv3` dims | hidden 7168, 128 heads, MLA q_lora 1536 / kv_lora 512 / nope 128 / rope 64 / v 128, expert ffn 2048, shared expert ffn 2048 | -| MoE | 8 local experts per rank (32 on 4 GPUs, 64 on 8), top-k 8, 4096 tokens per rank | -| precision | bf16 params and activations; `--recipe mxfp8` = MXFP8 block scaling for the expert GEMMs and dense projections | -| timing | fwd + bwd, 5 warmup, 10 timed iterations, no CUDA graphs | - -## Results - -Time per iteration is forward + backward of one layer on 4096 tokens per rank; throughput -counts tokens over all ranks. `fused` means `NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`. - -1 node, 4 GPUs, default dims (hidden 2048, 16 heads, expert ffn 1024, 32 experts): - -| precision | ms / iter | Mtok / s | -|---|---|---| -| bf16 | 6.17 | 2.65 | -| mxfp8 fused | 8.4 | 1.95 | - -1 node, 4 GPUs, `--dsv3` (32 experts): +At the small default dims MXFP8 is slower than bf16: the fused path launches many small +quantization kernels and the layer becomes CPU-launch-bound. Running under `nsys` adds +about 1.5 ms per iteration to these numbers. -| precision | ms / iter | Mtok / s | -|---|---|---| -| bf16 | 13.09 | 1.25 | -| mxfp8 | 11.02 | 1.49 | -| mxfp8 fused | 10.19 | 1.61 | +### E. TE kernel breakdown -2 nodes, 8 GPUs, `--dsv3` (64 experts): +8 GPUs, `--dsv3`, MXFP8 fused. Per GPU and iteration, from `nsys stats --report cuda_gpu_kern_sum` on one node +(kernel time 11.7 ms of an 11.3 ms iteration: the GPU is busy back to back): -| precision | ms / iter | Mtok / s | -|---|---|---| -| bf16 | 15.06 | 2.18 | -| mxfp8 | 12.88 | 2.54 | -| mxfp8 fused | 11.32 | 2.89 | +| Group | ms | Kernels | +|---|---:|---| +| NCCL EP all-to-all | 2.3 | `nccl_ep_jit_ht_dispatch_kernel` (0.99), `nccl_ep_jit_ht_combine_kernel` (1.31), each twice per iteration (fwd + bwd) | +| NCCL EP local permute + routing all-gather | 1.3 | `local_permute_dup/reduce` (0.62), `ncclDevKernel_AllGather_RING_LL` (0.65, includes waiting for slower ranks) | +| fused grouped MLP (cuDNN, MXFP8) | 2.7 | fc1+SwiGLU fwd (0.58), fc2 fwd (0.9), dGLU bwd (0.35), wgrad (0.85) | +| 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.3 | `nvjet_sm103_qqtst_*` | +| attention (cuDNN SDPA) | 0.7 | flash fprop (0.19) + bprop (0.51) | +| elementwise | 1.1 | residual/shared-expert adds (0.73), RMSNorm fwd+bwd (0.39) | -At the small default dims MXFP8 is slower than bf16: the fused path launches many small -quantization kernels and the layer becomes CPU-launch-bound. Running under `nsys` adds -about 1.5 ms per iteration to these numbers. +Both nodes sit in one NVLink domain, so dispatch and combine move roughly 470 MB per GPU per +call over NVLink at close to link bandwidth; on an InfiniBand-connected pair of nodes the +all-to-all share would be much larger. -## TE MoE vs. plain PyTorch MoE +### F. Historical comparison with plain PyTorch The naive timings and profiles below predate the routing-probability autograd fix: they omit router backward and probability-gradient communication. They are historical @@ -124,8 +170,8 @@ benchmark also clears parameter gradients before each step. Same layer, same dims (`--dsv3`), bf16 unless noted: -| | 4 GPUs (32 experts) | 8 GPUs (64 experts) | -|---|---|---| +| MoE implementation | 4 GPUs (32 experts) | 8 GPUs (64 experts) | +|---|---:|---:| | `naive`: all_to_all + loop over experts | 26.83 ms | 27.34 ms | | `naive_grouped`: all_to_all + TE grouped GEMM | 16.84 ms | 17.20 ms | | `te`: NCCL EP + grouped GEMM | 13.09 ms | 15.06 ms | @@ -137,8 +183,8 @@ At the small default dims (4 GPUs): `naive` 10.08 ms, `naive_grouped` 6.80 ms, ` Per GPU and iteration, the `naive` MoE spends (8 GPUs, kernel time 25.9 ms of a 28.5 ms iteration): -| group | ms | what | -|---|---|---| +| Group | ms | Details | +|---|---:|---| | `ncclDevKernel_SendRecv` | 5.2 | 7 all_to_all launches per iteration (tokens fwd/bwd, results fwd/bwd, counts, indices, probs) | | expert and dense GEMMs (`nvjet_*`) | 6.4 | 8 separate GEMM pairs per rank instead of one grouped GEMM, plus the MLA projections | | elementwise adds | 4.1 | `index_add` and its backward, residuals | @@ -149,8 +195,8 @@ iteration): `naive_grouped` (8 GPUs, kernel time 17.1 ms of a 17.2 ms iteration): -| group | ms | what | -|---|---|---| +| Group | ms | Details | +|---|---:|---| | `ncclDevKernel_SendRecv` | 3.8 | the same 7 all_to_all launches, less time because the GPU is no longer stalled between them | | grouped GEMMs (`nvjet_*_ptrGroup_*`) | 4.1 | fc1 / fc2 forward, dgrad, wgrad as grouped GEMMs, same as in `te` | | sorting rows by expert and back | 2.6 | `argsort`, gathers (`x[tok]`, `x_recv[by_expert]`), `index_copy`, `indexing_backward` | @@ -167,26 +213,7 @@ variants also synchronise with the host twice per layer to learn the all_to_all `te` does not. `--recipe mxfp8` is only supported by `te`: the naive variants would need the per-expert row counts padded to the MXFP8 block size. -## Where the time goes (8 GPUs, `--dsv3`, mxfp8 fused) - -Per GPU and iteration, from `nsys stats --report cuda_gpu_kern_sum` on one node -(kernel time 11.7 ms of an 11.3 ms iteration: the GPU is busy back to back): - -| group | ms | kernels | -|---|---|---| -| NCCL EP all-to-all | 2.3 | `nccl_ep_jit_ht_dispatch_kernel` (0.99), `nccl_ep_jit_ht_combine_kernel` (1.31), each twice per iteration (fwd + bwd) | -| NCCL EP local permute + routing all-gather | 1.3 | `local_permute_dup/reduce` (0.62), `ncclDevKernel_AllGather_RING_LL` (0.65, includes waiting for slower ranks) | -| fused grouped MLP (cuDNN, MXFP8) | 2.7 | fc1+SwiGLU fwd (0.58), fc2 fwd (0.9), dGLU bwd (0.35), wgrad (0.85) | -| 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.3 | `nvjet_sm103_qqtst_*` | -| attention (cuDNN SDPA) | 0.7 | flash fprop (0.19) + bprop (0.51) | -| elementwise | 1.1 | residual/shared-expert adds (0.73), RMSNorm fwd+bwd (0.39) | - -Both nodes sit in one NVLink domain, so dispatch and combine move roughly 470 MB per GPU per -call over NVLink at close to link bandwidth; on an InfiniBand-connected pair of nodes the -all-to-all share would be much larger. - -## Notes on the EP path +### G. EP implementation notes - `ep_bootstrap` must be given the same recv capacity the layer uses: `DeepSeekV3MoE.ep_recv_capacity(ep_size, tokens_per_rank, topk, num_local_experts)`. From e9fac4f6f92e4fc4fe62050f93b405b19df56349 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 19:39:27 +0200 Subject: [PATCH 38/43] [Docs] Keep naive MoE comparison in headline results Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 43 +++++++++++++------------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md index b2d1181f33..793cec8465 100644 --- a/examples/pytorch/deepseek_v3/README.md +++ b/examples/pytorch/deepseek_v3/README.md @@ -7,10 +7,26 @@ and MoE with a shared expert. Routed experts are sharded across GPUs using NCCL GB300 GPUs, 4096 tokens per rank, top-k 8, 8 local experts per GPU. Times cover one layer's forward + backward; throughput is global, in millions of tokens/s. -All rows below use `--impl te`. See [benchmark configuration](#c-benchmark-configuration) -for the full setup. +See [benchmark configuration](#c-benchmark-configuration) for the full setup. -### DeepSeek-V3 dimensions (`--dsv3`) +### TE vs. plain PyTorch (`--dsv3`) + +Same layer and dimensions, BF16 unless noted. + +The naive measurements predate the router-backward fix and omit probability-gradient +communication; updated timings are pending. They do not yet establish training speedups. + +| MoE implementation | 4 GPUs (32 experts) | 8 GPUs (64 experts) | +|---|---:|---:| +| `naive`: all_to_all + loop over experts | 26.83 ms | 27.34 ms | +| `naive_grouped`: all_to_all + TE grouped GEMM | 16.84 ms | 17.20 ms | +| `te`: NCCL EP + grouped GEMM | 13.09 ms | 15.06 ms | +| `te`, mxfp8 (unfused grouped GEMM) | 11.02 ms | 12.88 ms | +| `te`, mxfp8 fused | 10.19 ms | 11.32 ms | + +At the small default dims (4 GPUs): `naive` 10.08 ms, `naive_grouped` 6.80 ms, `te` 6.17 ms. + +### TE precision and throughput (`--dsv3`) | Precision | 4 GPUs · ms/iter | 4 GPUs · Mtok/s | 8 GPUs · ms/iter | 8 GPUs · Mtok/s | |---|---:|---:|---:|---:| @@ -21,7 +37,7 @@ for the full setup. 4 GPUs = 1 node / 32 experts; 8 GPUs = 2 nodes / 64 experts. Both nodes share one NVLink domain (MNNVL). -### Small default dimensions +### TE at small default dimensions 1 node, 4 GPUs, hidden 2048, 16 heads, expert FFN 1024, 32 experts. @@ -33,9 +49,6 @@ Both nodes share one NVLink domain (MNNVL). “Fused” enables `NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`. At the small dimensions, MXFP8 fused is slower than BF16. -The [historical naive comparison](#f-historical-comparison-with-plain-pytorch) is in the -appendix; those measurements predate the router-backward fix. - ## Quick start Requires SM90+ GPUs with NVLink and an NCCL EP-enabled TE build; @@ -56,7 +69,7 @@ NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 bash run_deepseek_v3_layer_ep.sh --dsv3 --recip [Requirements](#a-requirements) · [Running](#b-running) · [Configuration](#c-benchmark-configuration) · [Profiling](#d-profiling-with-nsys) · [TE kernels](#e-te-kernel-breakdown) · -[Historical naive comparison](#f-historical-comparison-with-plain-pytorch) · +[Naive kernel profiles](#f-naive-kernel-profiles) · [EP internals](#g-ep-implementation-notes) ### A. Requirements @@ -161,25 +174,13 @@ Both nodes sit in one NVLink domain, so dispatch and combine move roughly 470 MB call over NVLink at close to link bandwidth; on an InfiniBand-connected pair of nodes the all-to-all share would be much larger. -### F. Historical comparison with plain PyTorch +### F. Naive kernel profiles The naive timings and profiles below predate the routing-probability autograd fix: they omit router backward and probability-gradient communication. They are historical measurements and must be rerun before drawing training-speedup conclusions. The current benchmark also clears parameter gradients before each step. -Same layer, same dims (`--dsv3`), bf16 unless noted: - -| MoE implementation | 4 GPUs (32 experts) | 8 GPUs (64 experts) | -|---|---:|---:| -| `naive`: all_to_all + loop over experts | 26.83 ms | 27.34 ms | -| `naive_grouped`: all_to_all + TE grouped GEMM | 16.84 ms | 17.20 ms | -| `te`: NCCL EP + grouped GEMM | 13.09 ms | 15.06 ms | -| `te`, mxfp8 (unfused grouped GEMM) | 11.02 ms | 12.88 ms | -| `te`, mxfp8 fused | 10.19 ms | 11.32 ms | - -At the small default dims (4 GPUs): `naive` 10.08 ms, `naive_grouped` 6.80 ms, `te` 6.17 ms. - Per GPU and iteration, the `naive` MoE spends (8 GPUs, kernel time 25.9 ms of a 28.5 ms iteration): From 00bcbd9252cdf9dfa293eb46150583b8f9393a3a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 19:41:29 +0200 Subject: [PATCH 39/43] [PyTorch] DeepSeekV3Layer example README: remeasure all variants on the current code All rows re-taken after the router-gradient fix in the naive baseline and the per-forward EP buffer; kernel breakdowns refreshed. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 60 ++++++++++++-------------- 1 file changed, 28 insertions(+), 32 deletions(-) diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md index 793cec8465..e3d3520e72 100644 --- a/examples/pytorch/deepseek_v3/README.md +++ b/examples/pytorch/deepseek_v3/README.md @@ -13,26 +13,25 @@ See [benchmark configuration](#c-benchmark-configuration) for the full setup. Same layer and dimensions, BF16 unless noted. -The naive measurements predate the router-backward fix and omit probability-gradient -communication; updated timings are pending. They do not yet establish training speedups. +All variants run the router backward and clear parameter gradients before each step. | MoE implementation | 4 GPUs (32 experts) | 8 GPUs (64 experts) | |---|---:|---:| -| `naive`: all_to_all + loop over experts | 26.83 ms | 27.34 ms | -| `naive_grouped`: all_to_all + TE grouped GEMM | 16.84 ms | 17.20 ms | -| `te`: NCCL EP + grouped GEMM | 13.09 ms | 15.06 ms | -| `te`, mxfp8 (unfused grouped GEMM) | 11.02 ms | 12.88 ms | -| `te`, mxfp8 fused | 10.19 ms | 11.32 ms | +| `naive`: all_to_all + loop over experts | 27.27 ms | 27.55 ms | +| `naive_grouped`: all_to_all + TE grouped GEMM | 16.48 ms | 16.91 ms | +| `te`: NCCL EP + grouped GEMM | 12.39 ms | 14.03 ms | +| `te`, mxfp8 (unfused grouped GEMM) | 10.33 ms | 12.33 ms | +| `te`, mxfp8 fused | 9.51 ms | 11.94 ms | -At the small default dims (4 GPUs): `naive` 10.08 ms, `naive_grouped` 6.80 ms, `te` 6.17 ms. +At the small default dims (4 GPUs): `naive` 11.09 ms, `naive_grouped` 7.15 ms, `te` 6.27 ms. ### TE precision and throughput (`--dsv3`) | Precision | 4 GPUs · ms/iter | 4 GPUs · Mtok/s | 8 GPUs · ms/iter | 8 GPUs · Mtok/s | |---|---:|---:|---:|---:| -| BF16 | 13.09 | 1.25 | 15.06 | 2.18 | -| MXFP8 | 11.02 | 1.49 | 12.88 | 2.54 | -| MXFP8 fused | 10.19 | 1.61 | 11.32 | 2.89 | +| BF16 | 12.39 | 1.32 | 14.03 | 2.34 | +| MXFP8 | 10.33 | 1.59 | 12.33 | 2.66 | +| MXFP8 fused | 9.51 | 1.72 | 11.94 | 2.74 | 4 GPUs = 1 node / 32 experts; 8 GPUs = 2 nodes / 64 experts. Both nodes share one NVLink domain (MNNVL). @@ -43,8 +42,9 @@ Both nodes share one NVLink domain (MNNVL). | Precision | ms/iter | Mtok/s | |---|---:|---:| -| BF16 | 6.17 | 2.65 | -| MXFP8 fused | 8.4 | 1.95 | +| BF16 | 6.27 | 2.61 | +| MXFP8 | 8.12 | 2.02 | +| MXFP8 fused | 8.48 | 1.93 | “Fused” enables `NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`. At the small dimensions, MXFP8 fused is slower than BF16. @@ -158,17 +158,17 @@ about 1.5 ms per iteration to these numbers. ### E. TE kernel breakdown 8 GPUs, `--dsv3`, MXFP8 fused. Per GPU and iteration, from `nsys stats --report cuda_gpu_kern_sum` on one node -(kernel time 11.7 ms of an 11.3 ms iteration: the GPU is busy back to back): +(kernel time 12.9 ms of a 13.6 ms iteration under nsys, 11.9 ms without it): | Group | ms | Kernels | |---|---:|---| -| NCCL EP all-to-all | 2.3 | `nccl_ep_jit_ht_dispatch_kernel` (0.99), `nccl_ep_jit_ht_combine_kernel` (1.31), each twice per iteration (fwd + bwd) | -| NCCL EP local permute + routing all-gather | 1.3 | `local_permute_dup/reduce` (0.62), `ncclDevKernel_AllGather_RING_LL` (0.65, includes waiting for slower ranks) | -| fused grouped MLP (cuDNN, MXFP8) | 2.7 | fc1+SwiGLU fwd (0.58), fc2 fwd (0.9), dGLU bwd (0.35), wgrad (0.85) | +| NCCL EP all-to-all | 2.7 | `nccl_ep_jit_ht_dispatch_kernel` (1.36), `nccl_ep_jit_ht_combine_kernel` (1.32), each twice per iteration (fwd + bwd) | +| NCCL EP local permute + routing all-gather | 2.8 | `local_permute_dup/reduce` (0.62), `ncclDevKernel_AllGather_RING_LL` (2.16, mostly waiting for slower ranks) | +| fused grouped MLP (cuDNN, MXFP8) | 2.5 | fc1+SwiGLU fwd (0.55), fc2 fwd (0.84), dGLU bwd (0.32), wgrad (0.80) | | 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.3 | `nvjet_sm103_qqtst_*` | -| attention (cuDNN SDPA) | 0.7 | flash fprop (0.19) + bprop (0.51) | -| elementwise | 1.1 | residual/shared-expert adds (0.73), RMSNorm fwd+bwd (0.39) | +| 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) | Both nodes sit in one NVLink domain, so dispatch and combine move roughly 470 MB per GPU per call over NVLink at close to link bandwidth; on an InfiniBand-connected pair of nodes the @@ -176,25 +176,21 @@ all-to-all share would be much larger. ### F. Naive kernel profiles -The naive timings and profiles below predate the routing-probability autograd fix: -they omit router backward and probability-gradient communication. They are historical -measurements and must be rerun before drawing training-speedup conclusions. The current -benchmark also clears parameter gradients before each step. - -Per GPU and iteration, the `naive` MoE spends (8 GPUs, kernel time 25.9 ms of a 28.5 ms -iteration): +Per GPU and iteration, the `naive` MoE spends (8 GPUs, kernel time 24.0 ms of a 27.5 ms +iteration; the rest is host syncs and launch gaps): | Group | ms | Details | |---|---:|---| -| `ncclDevKernel_SendRecv` | 5.2 | 7 all_to_all launches per iteration (tokens fwd/bwd, results fwd/bwd, counts, indices, probs) | -| expert and dense GEMMs (`nvjet_*`) | 6.4 | 8 separate GEMM pairs per rank instead of one grouped GEMM, plus the MLA projections | -| elementwise adds | 4.1 | `index_add` and its backward, residuals | +| expert and dense GEMMs (`nvjet_*`) | 6.5 | 8 separate GEMM pairs per rank instead of one grouped GEMM, plus the MLA projections | +| elementwise adds | 3.6 | `index_add` and its backward, residuals | +| `ncclDevKernel_SendRecv` | 3.4 | 8 all_to_all launches per iteration (tokens and probs fwd/bwd, results fwd/bwd, counts, indices) | +| indexing kernels | 3.3 | `x[tok]`, `nonzero` masks, `index_copy`, `index_fill`, `indexing_backward`, sorts | | `FillFunctor` (zeros) | 2.5 | `zeros_like` for the per-expert output buffer and `index_add` targets | | device-to-device copies | 2.3 | `index_copy` and gathers materialising per-expert slices | -| indexing kernels | 3.0 | `x[tok]`, `nonzero` masks, `index_copy`, `indexing_backward` | | attention, norms | 1.1 | same as in the TE variant | -`naive_grouped` (8 GPUs, kernel time 17.1 ms of a 17.2 ms iteration): +`naive_grouped` (8 GPUs, kernel time 17.1 ms of a 17.2 ms iteration; profile taken before the +router-gradient fix, timings since then moved by less than 0.3 ms): | Group | ms | Details | |---|---:|---| From 4ffa639e4aa4a11dc517dbb87412a72b1e549a0c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 19:49:59 +0200 Subject: [PATCH 40/43] [PyTorch] DeepSeekV3Layer example README: split NCCL EP permute from rank-skew wait in the breakdown Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md index e3d3520e72..758592b1b0 100644 --- a/examples/pytorch/deepseek_v3/README.md +++ b/examples/pytorch/deepseek_v3/README.md @@ -163,7 +163,8 @@ about 1.5 ms per iteration to these numbers. | Group | ms | Kernels | |---|---:|---| | NCCL EP all-to-all | 2.7 | `nccl_ep_jit_ht_dispatch_kernel` (1.36), `nccl_ep_jit_ht_combine_kernel` (1.32), each twice per iteration (fwd + bwd) | -| NCCL EP local permute + routing all-gather | 2.8 | `local_permute_dup/reduce` (0.62), `ncclDevKernel_AllGather_RING_LL` (2.16, mostly waiting for slower ranks) | +| NCCL EP local permute | 0.6 | `local_permute_dup/reduce`: staging buffer to expert-major layout, zero-filled padding | +| rank desync wait | 2.2 | `ncclDevKernel_AllGather_RING_LL` (routing-map all-gather in prepare, ~0.06 ms of transfer); the first collective of the layer absorbs the skew between ranks | | fused grouped MLP (cuDNN, MXFP8) | 2.5 | fc1+SwiGLU fwd (0.55), fc2 fwd (0.84), dGLU bwd (0.32), wgrad (0.80) | | 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_*` | From 0367781444692fbbce86b12780af3ab40791e9db Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 20:14:34 +0200 Subject: [PATCH 41/43] [PyTorch] DeepSeekV3MoE: optional caller-owned EpBuffer DeepSeekV3MoE.make_ep_buffer() builds a buffer for the module's routing config; forward(ep_buffer=...) on the MoE and the layer reuses it instead of creating one per call. The example reuses one buffer by default (--ep-buffer-per-call restores per-call buffers). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- .../deepseek_v3/deepseek_v3_layer_ep.py | 13 ++++++++-- .../pytorch/models/deepseek_v3/moe.py | 24 +++++++++++++++---- .../models/deepseek_v3/transformer_layer.py | 6 ++++- 3 files changed, 35 insertions(+), 8 deletions(-) diff --git a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py index ca7144b332..dea62e313c 100644 --- a/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py +++ b/examples/pytorch/deepseek_v3/deepseek_v3_layer_ep.py @@ -53,6 +53,11 @@ def _parse_args(argv=None): ) p.add_argument("--impl", choices=["te", "naive", "naive_grouped"], default="te") p.add_argument("--recipe", choices=["none", "mxfp8"], default="none") + p.add_argument( + "--ep-buffer-per-call", + action="store_true", + help="Let the layer create a new EpBuffer per forward instead of reusing one.", + ) p.add_argument("--warmup", type=int, default=5) p.add_argument("--iters", type=int, default=10) args = p.parse_args(argv) @@ -244,11 +249,15 @@ def main(): seq = args.tokens_per_rank // 4 x = torch.randn(seq, 4, args.hidden, dtype=torch.bfloat16, device="cuda", requires_grad=True) + ep_buffer = None + if args.impl == "te" and not args.ep_buffer_per_call: + ep_buffer = layer.mlp.make_ep_buffer() + def step(): layer.zero_grad(set_to_none=True) x.grad = None with _autocast(args.recipe): - out = layer(x) + out = layer(x, ep_buffer=ep_buffer) out.backward(torch.ones_like(out)) return out @@ -276,7 +285,7 @@ def step(): tok_s = args.tokens_per_rank * world_size / (ms / 1e3) print( f"DeepSeekV3Layer impl={args.impl}:" - f" ranks={world_size} experts={num_experts} topk={args.topk} tokens/rank={args.tokens_per_rank} hidden={args.hidden} recipe={args.recipe} fused_mlp={os.environ.get('NVTE_CUTEDSL_FUSED_GROUPED_MLP', '0')} fwd+bwd" + f" ranks={world_size} experts={num_experts} topk={args.topk} tokens/rank={args.tokens_per_rank} hidden={args.hidden} recipe={args.recipe} fused_mlp={os.environ.get('NVTE_CUTEDSL_FUSED_GROUPED_MLP', '0')} ep_buffer={'per_call' if ep_buffer is None else 'reused'} fwd+bwd" f" {ms:.3f} ms/iter ({tok_s / 1e6:.2f} Mtok/s) finite={finite}", flush=True, ) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index 22be051749..fa606df2d7 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -220,11 +220,22 @@ def _forward_local(self, tokens: torch.Tensor) -> torch.Tensor: out[:num_rows], row_id_map, restore_shape=tokens.shape, pad_offsets=pad_offsets ) - def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: - from transformer_engine.pytorch.ep import EpBuffer, ep_dispatch, ep_combine + def make_ep_buffer(self, device: Optional[torch.device] = None): + """EP buffer for this module's routing config. ``forward`` creates one per call unless + given one; a buffer holds one call's routing state until its backward runs, so reuse + it only across calls whose backward has completed.""" + from transformer_engine.pytorch.ep import EpBuffer + + assert self._ep_buffer_kwargs is not None, "make_ep_buffer requires ep_group." + if device is None: + device = torch.device("cuda", torch.cuda.current_device()) + return EpBuffer(**self._ep_buffer_kwargs, device=device) + + def _forward_ep(self, tokens: torch.Tensor, ep_buffer=None) -> torch.Tensor: + from transformer_engine.pytorch.ep import ep_dispatch, ep_combine assert tokens.dtype == torch.bfloat16, "The EP path requires bfloat16 inputs." - buffer = EpBuffer(**self._ep_buffer_kwargs, device=tokens.device) + buffer = ep_buffer if ep_buffer is not None else self.make_ep_buffer(tokens.device) topk_idx = torch.empty( (tokens.shape[0], self.topk), dtype=torch.int64, device=tokens.device ) @@ -262,16 +273,19 @@ def _forward_ep(self, tokens: torch.Tensor) -> torch.Tensor: ) return ep_combine(buffer, expert_out, num_local_tokens=tokens.shape[0], grad_out=grad_out) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, ep_buffer=None) -> torch.Tensor: """ Parameters ---------- hidden_states : torch.Tensor input of shape ``[..., hidden_size]``. + ep_buffer : EpBuffer, optional + buffer from :meth:`make_ep_buffer` to reuse instead of creating one + per call (EP only). """ tokens = hidden_states.reshape(-1, self.hidden_size) if self.ep_group is not None: - out = self._forward_ep(tokens) + out = self._forward_ep(tokens, ep_buffer) else: out = self._forward_local(tokens) if self.shared_expert is not None: diff --git a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py index 741b235bd9..f491137dd2 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py +++ b/transformer_engine/pytorch/models/deepseek_v3/transformer_layer.py @@ -155,6 +155,7 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, checkpoint_core_attention: bool = False, + ep_buffer=None, ) -> torch.Tensor: """ Parameters @@ -165,6 +166,8 @@ def forward( boolean attention mask. checkpoint_core_attention : bool, default = False checkpoint the core attention computation. + ep_buffer : EpBuffer, optional + forwarded to :meth:`DeepSeekV3MoE.forward` (MoE layers with EP). """ attention_out = self.self_attention( self.input_layernorm(hidden_states), @@ -173,8 +176,9 @@ def forward( ) hidden_states = self._residual_add(attention_out, hidden_states) + mlp_kwargs = {"ep_buffer": ep_buffer} if ep_buffer is not None else {} if self.pre_mlp_layernorm is not None: - mlp_out = self.mlp(self.pre_mlp_layernorm(hidden_states)) + mlp_out = self.mlp(self.pre_mlp_layernorm(hidden_states), **mlp_kwargs) else: mlp_out = self.mlp(hidden_states) return self._residual_add(mlp_out, hidden_states) From da81d1d9d6b91c47db7773e495802888bbde8ee1 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 9 Sep 2026 20:43:26 +0200 Subject: [PATCH 42/43] [PyTorch] DeepSeekV3Layer example README: medians of three runs, breakdown from an unskewed node All timings are now medians of three runs (spread within 0.3 ms). The TE kernel breakdown is taken from the node not slowed down by nsys, so the rank-skew wait reflects load imbalance rather than profiler skew. Also folds in the pending README restructuring (throughput table, historical naive_grouped profile, results interpretation). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Pawel Gadzinski --- examples/pytorch/deepseek_v3/README.md | 70 ++++++++++++++------------ 1 file changed, 39 insertions(+), 31 deletions(-) diff --git a/examples/pytorch/deepseek_v3/README.md b/examples/pytorch/deepseek_v3/README.md index 758592b1b0..84bfca3ba5 100644 --- a/examples/pytorch/deepseek_v3/README.md +++ b/examples/pytorch/deepseek_v3/README.md @@ -7,6 +7,7 @@ and MoE with a shared expert. Routed experts are sharded across GPUs using NCCL GB300 GPUs, 4096 tokens per rank, top-k 8, 8 local experts per GPU. Times cover one layer's forward + backward; throughput is global, in millions of tokens/s. +Every number is the median of three runs (spread within 0.3 ms). See [benchmark configuration](#c-benchmark-configuration) for the full setup. ### TE vs. plain PyTorch (`--dsv3`) @@ -17,21 +18,21 @@ All variants run the router backward and clear parameter gradients before each s | MoE implementation | 4 GPUs (32 experts) | 8 GPUs (64 experts) | |---|---:|---:| -| `naive`: all_to_all + loop over experts | 27.27 ms | 27.55 ms | -| `naive_grouped`: all_to_all + TE grouped GEMM | 16.48 ms | 16.91 ms | -| `te`: NCCL EP + grouped GEMM | 12.39 ms | 14.03 ms | -| `te`, mxfp8 (unfused grouped GEMM) | 10.33 ms | 12.33 ms | -| `te`, mxfp8 fused | 9.51 ms | 11.94 ms | +| `naive`: all_to_all + loop over experts | 26.97 ms | 27.44 ms | +| `naive_grouped`: all_to_all + TE grouped GEMM | 16.48 ms | 16.89 ms | +| `te`: NCCL EP + grouped GEMM | 12.36 ms | 13.99 ms | +| `te`, mxfp8 (unfused grouped GEMM) | 10.28 ms | 12.10 ms | +| `te`, mxfp8 fused | 9.52 ms | 10.47 ms | -At the small default dims (4 GPUs): `naive` 11.09 ms, `naive_grouped` 7.15 ms, `te` 6.27 ms. +At the small default dims (4 GPUs): `naive` 11.13 ms, `naive_grouped` 7.04 ms, `te` 6.29 ms. -### TE precision and throughput (`--dsv3`) +### TE throughput (`--dsv3`) -| Precision | 4 GPUs · ms/iter | 4 GPUs · Mtok/s | 8 GPUs · ms/iter | 8 GPUs · Mtok/s | -|---|---:|---:|---:|---:| -| BF16 | 12.39 | 1.32 | 14.03 | 2.34 | -| MXFP8 | 10.33 | 1.59 | 12.33 | 2.66 | -| MXFP8 fused | 9.51 | 1.72 | 11.94 | 2.74 | +| Precision | 4 GPUs · Mtok/s | 8 GPUs · Mtok/s | +|---|---:|---:| +| BF16 | 1.33 | 2.34 | +| MXFP8 | 1.59 | 2.71 | +| MXFP8 fused | 1.72 | 3.13 | 4 GPUs = 1 node / 32 experts; 8 GPUs = 2 nodes / 64 experts. Both nodes share one NVLink domain (MNNVL). @@ -42,9 +43,9 @@ Both nodes share one NVLink domain (MNNVL). | Precision | ms/iter | Mtok/s | |---|---:|---:| -| BF16 | 6.27 | 2.61 | -| MXFP8 | 8.12 | 2.02 | -| MXFP8 fused | 8.48 | 1.93 | +| BF16 | 6.29 | 2.61 | +| MXFP8 | 8.13 | 2.02 | +| MXFP8 fused | 8.62 | 1.90 | “Fused” enables `NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`. At the small dimensions, MXFP8 fused is slower than BF16. @@ -157,15 +158,16 @@ about 1.5 ms per iteration to these numbers. ### E. TE kernel breakdown -8 GPUs, `--dsv3`, MXFP8 fused. Per GPU and iteration, from `nsys stats --report cuda_gpu_kern_sum` on one node -(kernel time 12.9 ms of a 13.6 ms iteration under nsys, 11.9 ms without it): +8 GPUs, `--dsv3`, MXFP8 fused. Per GPU and iteration, from `nsys stats --report cuda_gpu_kern_sum` +on one node (kernel time 10.8 ms; the iteration takes 10.5 ms without the profiler). Profiling adds +skew between nodes, so the wait row is taken from the node that was not slowed down by nsys. | Group | ms | Kernels | |---|---:|---| -| NCCL EP all-to-all | 2.7 | `nccl_ep_jit_ht_dispatch_kernel` (1.36), `nccl_ep_jit_ht_combine_kernel` (1.32), each twice per iteration (fwd + bwd) | +| 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 desync wait | 2.2 | `ncclDevKernel_AllGather_RING_LL` (routing-map all-gather in prepare, ~0.06 ms of transfer); the first collective of the layer absorbs the skew between ranks | -| fused grouped MLP (cuDNN, MXFP8) | 2.5 | fc1+SwiGLU fwd (0.55), fc2 fwd (0.84), dGLU bwd (0.32), wgrad (0.80) | +| 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 | @@ -190,26 +192,32 @@ iteration; the rest is host syncs and launch gaps): | device-to-device copies | 2.3 | `index_copy` and gathers materialising per-expert slices | | attention, norms | 1.1 | same as in the TE variant | -`naive_grouped` (8 GPUs, kernel time 17.1 ms of a 17.2 ms iteration; profile taken before the -router-gradient fix, timings since then moved by less than 0.3 ms): +#### Historical `naive_grouped` profile + +8 GPUs, kernel time 17.1 ms of a 17.2 ms iteration. This profile predates the +router-gradient fix and omits probability-gradient communication. It is retained for +reference, not for direct comparison with the current profiles or headline timings. | Group | ms | Details | |---|---:|---| -| `ncclDevKernel_SendRecv` | 3.8 | the same 7 all_to_all launches, less time because the GPU is no longer stalled between them | +| `ncclDevKernel_SendRecv` | 3.8 | 7 all_to_all launches in this historical run; the current implementation has 8 | | grouped GEMMs (`nvjet_*_ptrGroup_*`) | 4.1 | fc1 / fc2 forward, dgrad, wgrad as grouped GEMMs, same as in `te` | | sorting rows by expert and back | 2.6 | `argsort`, gathers (`x[tok]`, `x_recv[by_expert]`), `index_copy`, `indexing_backward` | | dense GEMMs (MLA projections, shared expert) | 2.4 | same as in `te` | | elementwise adds | 0.7 | `index_add`, residuals | | attention, norms | 1.1 | same as in `te` | -Reading the three rows of the table together: the Python loop over experts costs about -10 ms per iteration (`naive` -> `naive_grouped`, 8 separate GEMM pairs, `nonzero` masks, -zero-filled buffers, copies); replacing torch `all_to_all` plus the surrounding sort / gather / -scatter with NCCL EP dispatch and combine, which write straight into the expert-major layout -and zero-fill the padding, saves another 2 ms (`naive_grouped` -> `te`). Both `naive` -variants also synchronise with the host twice per layer to learn the all_to_all split sizes; -`te` does not. `--recipe mxfp8` is only supported by `te`: the naive variants would need the -per-expert row counts padded to the MXFP8 block size. +#### Interpreting the current results + +In the headline BF16 timings, `naive` -> `naive_grouped` reduces iteration time by +10.5 ms on 4 GPUs and 10.6 ms on 8 GPUs. `naive_grouped` -> `te` saves another +4.1 ms and 2.9 ms, respectively. These are end-to-end differences between +implementations, not isolated measurements of Python-loop or communication overhead. + +NCCL EP dispatch and combine write directly into the expert-major layout and zero-fill +the padding. Both naive variants synchronise with the host to obtain the all_to_all +split sizes; TE avoids those synchronizations. MXFP8 is only supported by `te`: the naive +variants would need per-expert row counts padded to the MXFP8 block size. ### G. EP implementation notes From 2f7859aecddcbc709c5b0ec81364229e19410f81 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 10 Sep 2026 12:26:51 +0200 Subject: [PATCH 43/43] [PyTorch] Validate MoE grouping and check synchronized expert bias updates Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_models.py | 13 ++++++++-- tests/pytorch/test_models.py | 25 +++++++++++++++++++ .../pytorch/models/deepseek_v3/moe.py | 20 +++++++++++++-- 3 files changed, 54 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/distributed/run_models.py b/tests/pytorch/distributed/run_models.py index 302da451c5..2462fe28aa 100644 --- a/tests/pytorch/distributed/run_models.py +++ b/tests/pytorch/distributed/run_models.py @@ -130,12 +130,21 @@ def test_layer_ep_matches_local( 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) + dist.all_reduce(counts, group=ep_group) last_num_tokens = microbatches[-1][0].numel() // HIDDEN assert counts.sum().item() == ep_size * last_num_tokens * TOP_K + expected_bias = ep_layer.mlp.expert_bias.clone() + expected_bias += ep_layer.mlp.expert_bias_update_rate * torch.sign( + counts.float().mean() - counts.float() + ) + ep_layer.mlp._last_tokens_per_expert = counts ep_layer.mlp.update_expert_bias() - assert torch.isfinite(ep_layer.mlp.expert_bias).all() + torch.testing.assert_close(ep_layer.mlp.expert_bias, expected_bias, rtol=0, atol=0) + biases = [torch.empty_like(expected_bias) for _ in range(ep_size)] + dist.all_gather(biases, ep_layer.mlp.expert_bias, group=ep_group) + for bias in biases: + torch.testing.assert_close(bias, expected_bias, rtol=0, atol=0) def test_mla_tp_matches_local(rank: int, tp_size: int, tp_group) -> None: diff --git a/tests/pytorch/test_models.py b/tests/pytorch/test_models.py index 5d3a2d5531..f0cf4be1af 100644 --- a/tests/pytorch/test_models.py +++ b/tests/pytorch/test_models.py @@ -129,6 +129,31 @@ def test_mla_yarn_softmax_scale(mscale_all_dim): assert mla.softmax_scale == pytest.approx(m * m / math.sqrt(qk_head_dim)) +@pytest.mark.parametrize( + "kwargs,match", + [ + ({"num_experts": 0}, "num_experts must be positive"), + ({"topk": 0}, "topk must be in"), + ({"topk": 9}, "topk must be in"), + ({"num_groups": 2}, "must be provided together"), + ({"group_topk": 1}, "must be provided together"), + ({"num_groups": 0, "group_topk": 1}, "num_groups must be positive"), + ({"num_groups": -2, "group_topk": 1}, "num_groups must be positive"), + ({"num_groups": 3, "group_topk": 1}, "divide num_experts"), + ({"num_groups": 2, "group_topk": 0}, "group_topk must be in"), + ({"num_groups": 2, "group_topk": -1}, "group_topk must be in"), + ({"num_groups": 2, "group_topk": 3}, "group_topk must be in"), + ({"num_groups": 2, "group_topk": 2, "topk": 3}, "topk must be divisible"), + ({"num_groups": 4, "group_topk": 1, "topk": 4}, "topk per group must not exceed"), + ], +) +def test_moe_rejects_invalid_routing_config(kwargs, match): + config = dict(num_experts=8, topk=2) + config.update(kwargs) + with pytest.raises(ValueError, match=match): + DeepSeekV3MoE(HIDDEN, moe_ffn_hidden_size=128, device="cpu", **config) + + @pytest.mark.parametrize("shared", [False, True], ids=["no_shared", "shared"]) @pytest.mark.parametrize("grouped", [False, True], ids=["ungrouped", "grouped"]) @pytest.mark.parametrize("topk", [2, 4]) diff --git a/transformer_engine/pytorch/models/deepseek_v3/moe.py b/transformer_engine/pytorch/models/deepseek_v3/moe.py index fa606df2d7..31b9294774 100644 --- a/transformer_engine/pytorch/models/deepseek_v3/moe.py +++ b/transformer_engine/pytorch/models/deepseek_v3/moe.py @@ -79,9 +79,9 @@ class DeepSeekV3MoE(torch.nn.Module): topk : int, default = 8 number of experts per token. num_groups : int, optional - number of expert groups for node-limited routing. + number of expert groups for node-limited routing; requires ``group_topk``. group_topk : int, optional - number of groups each token is limited to. + number of groups each token is limited to; requires ``num_groups``. routed_scaling_factor : float, default = 2.5 scaling applied to the routing probabilities. shared_expert_ffn_hidden_size : int, optional @@ -116,6 +116,22 @@ def __init__( ) -> None: super().__init__() + if num_experts <= 0: + raise ValueError("num_experts must be positive.") + if not 1 <= topk <= num_experts: + raise ValueError("topk must be in [1, num_experts].") + if (num_groups is None) != (group_topk is None): + raise ValueError("num_groups and group_topk must be provided together.") + if num_groups is not None: + if num_groups <= 0 or num_experts % num_groups != 0: + raise ValueError("num_groups must be positive and divide num_experts.") + if not 1 <= group_topk <= num_groups: + raise ValueError("group_topk must be in [1, num_groups].") + if topk % group_topk != 0: + raise ValueError("topk must be divisible by group_topk.") + if topk // group_topk > num_experts // num_groups: + raise ValueError("topk per group must not exceed the number of experts per group.") + dtype = params_dtype if params_dtype is not None else torch.get_default_dtype() self.hidden_size = hidden_size self.num_experts = num_experts