[main] moe(fix): Avoid TE cuda graph dummy attention masks

megatron/core/transformer/transformer_config.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import logging
 import math
@@ -957,6 +957,14 @@ class TransformerConfig(ModelParallelConfig):
     transformed to an empty list in __post_init__. The deprecated values "full_iteration" and
     "full_iteration_inference" are also accepted and migrated to the new API in __post_init__."""
 
+    create_attention_mask_in_dataloader: bool = True
+    """Whether training data loaders create and pass an attention_mask tensor.
+
+    This mirrors the training argument of the same name so Transformer Engine CUDA graph capture
+    can preserve the actual forward signature. When disabled for causal/no-mask attention,
+    the graph capture path should not allocate a synthetic global-sequence attention mask.
+    """
+
     inference_cuda_graph_scope: Optional[InferenceCudaGraphScope] = field(
         default=None,
         metadata={

megatron/core/transformer/transformer_layer.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 from __future__ import annotations
 
 import functools
@@ -19,7 +19,12 @@
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.cuda_graphs import is_graph_capturing
-from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope, LayerType
+from megatron.core.transformer.enums import (
+    AttnMaskType,
+    CudaGraphModule,
+    InferenceCudaGraphScope,
+    LayerType,
+)
 from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
 from megatron.core.transformer.mlp import MLP
 from megatron.core.transformer.module import GraphableMegatronModule
@@ -1049,13 +1054,29 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size):
             not self.config.cuda_graph_modules
             or CudaGraphModule.attn in self.config.cuda_graph_modules
         ):
-            slen_per_cp = seq_length // self.config.context_parallel_size
-            static_inputs["attention_mask"] = (
-                ~(torch.tril(torch.ones((slen_per_cp, seq_length))).bool())
-                .to(torch.cuda.current_device())
-                .reshape(1, 1, slen_per_cp, seq_length)
-                .tile(micro_batch_size, 1, 1, 1)
-            )
+            if not self.config.create_attention_mask_in_dataloader:
+                if self.self_attention.attn_mask_type not in (
+                    AttnMaskType.causal,
+                    AttnMaskType.no_mask,
+                    AttnMaskType.causal_bottom_right,
+                ):
+                    log_single_rank(
+                        logger,
+                        logging.WARNING,
+                        "TE CUDA graph capture is omitting attention_mask because "
+                        "create_attention_mask_in_dataloader is False, but "
+                        f"attn_mask_type={self.self_attention.attn_mask_type.name} may require "
+                        "an explicit mask. Ensure this is intended for the current workload.",
+                    )
+            else:
+                slen_per_cp = seq_length // self.config.context_parallel_size
+                static_inputs["attention_mask"] = (
+                    ~(torch.tril(torch.ones((slen_per_cp, seq_length))).bool())
+                    .to(torch.cuda.current_device())
+                    .reshape(1, 1, slen_per_cp, seq_length)
+                    .tile(micro_batch_size, 1, 1, 1)
+                )
+
         return static_inputs
 
     def _get_submodules_under_cudagraphs(self):
@@ -1300,11 +1321,17 @@ def get_zero_attention_mask(slen_per_tpcp, micro_batch_size):
                 'attention_mask' in cudagraph_kwargs and cudagraph_kwargs['attention_mask'] is None
             ):
                 # The attention_mask can be None when there is no padding to the input sequence.
-                # However, an attention_mask Tensor must be passed into cudagraph for replay, so
-                # we create an equivalent zero Tensor as the attention_mask.
-                cudagraph_kwargs["attention_mask"] = get_zero_attention_mask(
-                    hidden_states.size(0), hidden_states.size(1)
-                )
+                # If the dataloader never creates masks, the TE CUDA graph was captured without
+                # this kwarg. Preserve that signature instead of allocating a synthetic
+                # [local_seq, global_seq] zero mask.
+                if not self.config.create_attention_mask_in_dataloader:
+                    cudagraph_kwargs.pop("attention_mask")
+                else:
+                    # The graph was captured with an attention_mask tensor, so replay needs a
+                    # tensor input with the same semantic value as no padding.
+                    cudagraph_kwargs["attention_mask"] = get_zero_attention_mask(
+                        hidden_states.size(0), hidden_states.size(1)
+                    )
         except ImportError:
             raise RuntimeError("CUDAGraph requires TransformerEngine, but not installed")
         return tuple(cudagraph_args), cudagraph_kwargs

megatron/training/arguments.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 """Megatron arguments."""
 
@@ -276,12 +276,12 @@ def _normalize_cuda_graph_modules_args(args):
         args.cuda_graph_modules
     )
     validate_deprecated_cuda_graph_modules_migration_inputs(
-        deprecated_scopes,
-        args.cuda_graph_impl,
-        args.inference_cuda_graph_scope,
+        deprecated_scopes, args.cuda_graph_impl, args.inference_cuda_graph_scope
     )
     if used_full_scope:
-        warn_rank_0('full scope is deprecated. Use empty cuda_graph_modules to capture the whole layer.')
+        warn_rank_0(
+            'full scope is deprecated. Use empty cuda_graph_modules to capture the whole layer.'
+        )
 
     for scope, attr, value in deprecated_scopes:
         migration = get_deprecated_cuda_graph_modules_migration(
@@ -629,23 +629,18 @@ def validate_args(args, defaults={}):
             del args.external_cuda_graph
 
     if getattr(args, 'cuda_graph_scope_deprecated', None) is not None:
-        assert not args.cuda_graph_modules, (
-            "--cuda-graph-scope and --cuda-graph-modules cannot be used together."
-        )
-        warn_rank_0(
-            '--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.'
-        )
+        assert (
+            not args.cuda_graph_modules
+        ), "--cuda-graph-scope and --cuda-graph-modules cannot be used together."
+        warn_rank_0('--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.')
         args.cuda_graph_modules = args.cuda_graph_scope_deprecated
     del args.cuda_graph_scope_deprecated
 
     # Normalize cuda_graph_modules and inference_cuda_graph_scope early so that
     # all subsequent validation sees fully-typed enum values.
     _normalize_cuda_graph_modules_args(args)
     _normalize_inference_cuda_graph_scope_arg(args)
-    assert (
-        args.inference_cuda_graph_scope
-        in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl]
-    ), (
+    assert args.inference_cuda_graph_scope in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl], (
         "Invalid inference CUDA graph scope "
         f"{args.inference_cuda_graph_scope.name!r} for "
         f"--cuda-graph-impl={args.cuda_graph_impl!r}."
@@ -2095,6 +2090,8 @@ def _add_network_size_args(parser):
         "num_moe_experts",
         "fp8_param",
         "fp4_param",
+        # already generated by data args
+        "create_attention_mask_in_dataloader",
         # incompatible defaults in dataclass
         "gradient_accumulation_fusion",
         "overlap_p2p_comm",

tests/unit_tests/models/test_hybrid_moe_model.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import hashlib
 import inspect
@@ -68,6 +68,7 @@
     "cpu_offloading_weights": False,
     "cross_entropy_fusion_impl": "native",
     "cross_entropy_loss_fusion": True,
+    "create_attention_mask_in_dataloader": True,
     "cuda_graph_impl": "none",
     "cuda_graph_retain_backward_graph": False,
     "cuda_graph_modules": [],

tests/unit_tests/transformer/test_transformer_layer.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 
 import gc
@@ -20,7 +20,7 @@
     model_parallel_cuda_manual_seed,
 )
 from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord
-from megatron.core.transformer.enums import InferenceCudaGraphScope
+from megatron.core.transformer.enums import AttnMaskType, CudaGraphModule, InferenceCudaGraphScope
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.transformer_layer import (
     TransformerLayer,
@@ -73,6 +73,89 @@ def test_gpu_forward(self):
         assert hidden_states.shape[1] == micro_batch_size
         assert hidden_states.shape[2] == config.hidden_size
 
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+    @pytest.mark.skipif(
+        not (HAVE_TE and is_te_min_version("1.10.0")),
+        reason="TE CUDA graph kwargs require TransformerEngine >= 1.10",
+    )
+    def test_te_cuda_graph_omits_absent_attention_mask(self):
+        config = TransformerConfig(
+            num_layers=2,
+            hidden_size=12,
+            num_attention_heads=4,
+            use_cpu_initialization=True,
+            cuda_graph_impl="transformer_engine",
+            cuda_graph_modules=[CudaGraphModule.attn],
+            create_attention_mask_in_dataloader=False,
+        )
+        layer = TransformerLayer(config, get_gpt_layer_with_transformer_engine_submodules())
+
+        static_inputs = layer.get_layer_static_inputs(seq_length=32, micro_batch_size=1)
+        assert "attention_mask" not in static_inputs
+
+        hidden_states = torch.ones((32, 1, config.hidden_size), device="cuda")
+        _, cudagraph_kwargs = layer._get_te_cuda_graph_replay_args(
+            hidden_states, attention_mask=None
+        )
+        assert "attention_mask" not in cudagraph_kwargs
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+    @pytest.mark.skipif(
+        not (HAVE_TE and is_te_min_version("1.10.0")),
+        reason="TE CUDA graph kwargs require TransformerEngine >= 1.10",
+    )
+    def test_te_cuda_graph_keeps_configured_attention_mask(self):
+        config = TransformerConfig(
+            num_layers=2,
+            hidden_size=12,
+            num_attention_heads=4,
+            use_cpu_initialization=True,
+            cuda_graph_impl="transformer_engine",
+            cuda_graph_modules=[CudaGraphModule.attn],
+            create_attention_mask_in_dataloader=True,
+        )
+        layer = TransformerLayer(config, get_gpt_layer_with_transformer_engine_submodules())
+
+        static_inputs = layer.get_layer_static_inputs(seq_length=32, micro_batch_size=1)
+        assert static_inputs["attention_mask"].shape == (1, 1, 32, 32)
+
+        hidden_states = torch.ones((32, 1, config.hidden_size), device="cuda")
+        _, cudagraph_kwargs = layer._get_te_cuda_graph_replay_args(
+            hidden_states, attention_mask=None
+        )
+        assert cudagraph_kwargs["attention_mask"].shape == (1, 1, 32, 32)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+    @pytest.mark.skipif(
+        not (HAVE_TE and is_te_min_version("1.10.0")),
+        reason="TE CUDA graph kwargs require TransformerEngine >= 1.10",
+    )
+    def test_te_cuda_graph_warns_when_omitting_padding_attention_mask(self, monkeypatch):
+        config = TransformerConfig(
+            num_layers=2,
+            hidden_size=12,
+            num_attention_heads=4,
+            use_cpu_initialization=True,
+            cuda_graph_impl="transformer_engine",
+            cuda_graph_modules=[CudaGraphModule.attn],
+            create_attention_mask_in_dataloader=False,
+        )
+        layer = TransformerLayer(config, get_gpt_layer_with_transformer_engine_submodules())
+        layer.self_attention.attn_mask_type = AttnMaskType.padding
+
+        warning_messages = []
+        monkeypatch.setattr(
+            "megatron.core.transformer.transformer_layer.log_single_rank",
+            lambda _logger, _level, message: warning_messages.append(message),
+        )
+
+        static_inputs = layer.get_layer_static_inputs(seq_length=32, micro_batch_size=1)
+
+        assert "attention_mask" not in static_inputs
+        assert any(
+            "attn_mask_type=padding may require an explicit mask" in msg for msg in warning_messages
+        )
+
     def test_chunked_mlp(self):
         with torch.no_grad():
             num_layers = 2