perf(nemotron): enable CuteDSL fusion for GB MXFP8 (#5620)

Signed-off-by: Malay Nagda <malayn@nvidia.com>
Co-authored-by: Kamran Jafari <kjafarisadeg@nvidia.com>

Author: malay-nagda

src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py

@@ -143,6 +143,9 @@ def nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer:
     cfg.model.cuda_graph_scope = ["attn", "mamba", "moe_router", "moe_preprocess"]
 
     _apply_nemotron_3_super_perf_defaults(cfg)
+    cfg.model.use_transformer_engine_op_fuser = True
+    cfg.model.moe_mlp_glu_interleave_size = 32
+    cfg.mixed_precision.fp8_dot_product_attention = True
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -160,7 +163,9 @@ def nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer:
         "NVLINK_DOMAIN_SIZE": 72,
         "USE_MNNVL": 1,
         # Transformer Engine overlap settings for this model.
+        "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
     }
     return cfg
@@ -354,8 +359,7 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     return cfg
 
 
-def nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
-    """Nemotron 3 Nano pretrain: 8× GB200, MXFP8."""
+def _build_nemotron_3_nano_gb200_mxfp8() -> ConfigContainer:
     cfg = nemotron_3_nano_pretrain_config()
     _apply_nemotron_3_nano_perf_defaults(cfg)
     cfg.mixed_precision = _perf_precision("fp8_mx")
@@ -379,6 +383,15 @@ def nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
 
     _benchmark_common(cfg)
     cfg.model.moe_hybridep_num_sms = 32
+    return cfg
+
+
+def nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
+    """Nemotron 3 Nano pretrain: 8× GB200, MXFP8."""
+    cfg = _build_nemotron_3_nano_gb200_mxfp8()
+    cfg.model.use_transformer_engine_op_fuser = True
+    cfg.model.moe_mlp_glu_interleave_size = 32
+    cfg.mixed_precision.fp8_dot_product_attention = True
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -396,7 +409,9 @@ def nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
         "NVLINK_DOMAIN_SIZE": 72,
         "USE_MNNVL": 1,
         # Transformer Engine overlap settings for this model.
+        "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
         # Use cuDNN LayerNorm for this measured baseline.
         "NVTE_NORM_BWD_USE_CUDNN": 1,
@@ -489,9 +504,8 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     return cfg
 
 
-def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
-    """Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8."""
-    cfg = nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config()
+def _build_nemotron_3_5_lightning_gb200_mxfp8() -> ConfigContainer:
+    cfg = _build_nemotron_3_nano_gb200_mxfp8()
     cfg.model.moe_hybridep_num_sms = 16
     cfg.model.mtp_num_layers = 2
     cfg.model.mtp_hybrid_override_pattern = "*E"
@@ -502,6 +516,15 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer
     cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
     cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
     cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
+    return cfg
+
+
+def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
+    """Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8."""
+    cfg = _build_nemotron_3_5_lightning_gb200_mxfp8()
+    cfg.model.use_transformer_engine_op_fuser = True
+    cfg.model.moe_mlp_glu_interleave_size = 32
+    cfg.mixed_precision.fp8_dot_product_attention = True
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
         "CUDA_DEVICE_MAX_CONNECTIONS": 32,
@@ -513,7 +536,9 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer
         "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
         "NVLINK_DOMAIN_SIZE": 72,
         "USE_MNNVL": 1,
+        "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
         "NVTE_NORM_BWD_USE_CUDNN": 1,
         "NVTE_NORM_FWD_USE_CUDNN": 1,
@@ -523,7 +548,7 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer
 
 def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config() -> ConfigContainer:
     """Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8, Megatron FSDP."""
-    cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config()
+    cfg = _build_nemotron_3_5_lightning_gb200_mxfp8()
 
     # FSDP reduces the model-state footprint enough to use the larger measured
     # microbatch. Megatron FSDP registers module hooks that Transformer Engine

src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py

@@ -115,8 +115,7 @@ def nemotron_3_super_pretrain_64gpu_gb300_bf16_config() -> ConfigContainer:
     return cfg
 
 
-def nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config() -> ConfigContainer:
-    """Nemotron 3 Super pretrain: 64× GB300, MXFP8."""
+def _build_nemotron_3_super_gb300_mxfp8() -> ConfigContainer:
     cfg = nemotron_3_super_pretrain_config()
     cfg.mixed_precision = _perf_precision("fp8_mx")
 
@@ -139,6 +138,15 @@ def nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config() -> ConfigContainer:
     cfg.model.cuda_graph_scope = ["attn", "mamba", "moe_router", "moe_preprocess"]
 
     _apply_nemotron_3_super_perf_defaults(cfg)
+    return cfg
+
+
+def nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config() -> ConfigContainer:
+    """Nemotron 3 Super pretrain: 64× GB300, MXFP8."""
+    cfg = _build_nemotron_3_super_gb300_mxfp8()
+    cfg.model.use_transformer_engine_op_fuser = True
+    cfg.model.moe_mlp_glu_interleave_size = 32
+    cfg.mixed_precision.fp8_dot_product_attention = True
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -156,7 +164,9 @@ def nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config() -> ConfigContainer:
         "NVLINK_DOMAIN_SIZE": 72,
         "USE_MNNVL": 1,
         # Transformer Engine overlap settings for this model.
+        "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
     }
     return cfg
@@ -359,8 +369,7 @@ def nemotron_3_nano_pretrain_8gpu_gb300_bf16_config() -> ConfigContainer:
     return cfg
 
 
-def nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config() -> ConfigContainer:
-    """Nemotron 3 Nano pretrain: 8× GB300, MXFP8."""
+def _build_nemotron_3_nano_gb300_mxfp8() -> ConfigContainer:
     cfg = nemotron_3_nano_pretrain_config()
     _apply_nemotron_3_nano_perf_defaults(cfg)
     cfg.mixed_precision = _perf_precision("fp8_mx")
@@ -384,6 +393,15 @@ def nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config() -> ConfigContainer:
 
     _benchmark_common(cfg)
     cfg.model.moe_hybridep_num_sms = 16
+    return cfg
+
+
+def nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config() -> ConfigContainer:
+    """Nemotron 3 Nano pretrain: 8× GB300, MXFP8."""
+    cfg = _build_nemotron_3_nano_gb300_mxfp8()
+    cfg.model.use_transformer_engine_op_fuser = True
+    cfg.model.moe_mlp_glu_interleave_size = 32
+    cfg.mixed_precision.fp8_dot_product_attention = True
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -401,7 +419,9 @@ def nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config() -> ConfigContainer:
         "NVLINK_DOMAIN_SIZE": 72,
         "USE_MNNVL": 1,
         # Transformer Engine overlap settings for this model.
+        "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
         # Use cuDNN LayerNorm for this measured baseline.
         "NVTE_NORM_BWD_USE_CUDNN": 1,

src/megatron/bridge/perf_recipes/nemotronh/vr200/nemotronh.py

@@ -18,11 +18,11 @@
     ConfigContainer,
 )
 from megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh import (
+    _build_nemotron_3_nano_gb300_mxfp8,
+    _build_nemotron_3_super_gb300_mxfp8,
     nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
-    nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
     nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,
     nemotron_3_super_pretrain_64gpu_gb300_bf16_config,
-    nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config,
     nemotron_3_super_pretrain_64gpu_gb300_nvfp4_config,
 )
 
@@ -58,7 +58,7 @@ def nemotron_3_nano_pretrain_8gpu_vr200_bf16_config() -> ConfigContainer:
 
 def nemotron_3_nano_pretrain_8gpu_vr200_fp8mx_config() -> ConfigContainer:
     """Nemotron 3 Nano pretrain: 8× VR200, FP8-MX (alias of GB300)."""
-    cfg = nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config()
+    cfg = _build_nemotron_3_nano_gb300_mxfp8()
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -144,7 +144,7 @@ def nemotron_3_super_pretrain_64gpu_vr200_bf16_config() -> ConfigContainer:
 
 def nemotron_3_super_pretrain_64gpu_vr200_fp8mx_config() -> ConfigContainer:
     """Nemotron 3 Super pretrain: 64× VR200, FP8-MX (alias of GB300)."""
-    cfg = nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config()
+    cfg = _build_nemotron_3_super_gb300_mxfp8()
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,

tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py

@@ -188,6 +188,36 @@ def test_nemotron_3_5_perf_recipes_inherit_nemotron_3_policy(
     assert cfg.tokenizer.tokenizer_model != base_cfg.tokenizer.tokenizer_model
 
 
+def test_gb200_mxfp8_enables_cutedsl_fusion() -> None:
+    """The non-FSDP Lightning recipe enables CutDSL without MoE A2A overlap."""
+    cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config()
+
+    assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1
+    assert cfg.env_vars["CUDNNFE_CLUSTER_OVERLAP_MARGIN"] == 8
+    assert cfg.model.use_transformer_engine_op_fuser is True
+    assert cfg.model.moe_mlp_glu_interleave_size == 32
+    assert cfg.model.high_priority_a2a_comm_stream is False
+    assert cfg.model.moe_hybridep_num_sms_preprocessing == 108
+    assert cfg.mixed_precision.fp8_dot_product_attention is True
+    assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is False
+    assert cfg.comm_overlap.delay_wgrad_compute is False
+
+
+def test_gb200_mxfp8_fsdp_skips_cutedsl_fusion() -> None:
+    """The Lightning MXFP8 FSDP variant remains outside the CutDSL tuning scope."""
+    cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config()
+
+    assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars
+    assert "CUDNNFE_CLUSTER_OVERLAP_MARGIN" not in cfg.env_vars
+    assert getattr(cfg.model, "use_transformer_engine_op_fuser", False) is False
+    assert getattr(cfg.model, "moe_mlp_glu_interleave_size", None) is None
+    assert getattr(cfg.model, "high_priority_a2a_comm_stream", False) is False
+    assert getattr(cfg.model, "moe_hybridep_num_sms_preprocessing", None) != 32
+    assert cfg.mixed_precision.fp8_dot_product_attention is False
+    assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is False
+    assert cfg.comm_overlap.delay_wgrad_compute is False
+
+
 @pytest.mark.parametrize("recipe_factory", _H100_RECIPES, ids=lambda recipe: recipe.__name__)
 def test_h100_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer]) -> None:
     """H100 Nemotron 3.5 Lightning variants retain the established performance topology."""

tests/unit_tests/recipes/test_nemotronh_recipes.py

@@ -54,6 +54,8 @@ class _FakeModelProvider:
 
     def __init__(self) -> None:
         self.vocab_size = 256
+        self.high_priority_a2a_comm_stream = False
+        self.moe_hybridep_num_sms_preprocessing = 108
 
     def finalize(self) -> None:
         return None
@@ -162,6 +164,59 @@ def test_nemotron_3_nano_gb200_defers_vocab_size_to_training_tokenizer():
     assert cfg.model.vocab_size is None
 
 
+@pytest.mark.parametrize(
+    ("module_name", "factory_name", "has_comm_overlap"),
+    [
+        pytest.param(
+            "megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh",
+            "nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config",
+            True,
+            id="nano-gb200",
+        ),
+        pytest.param(
+            "megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
+            "nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config",
+            True,
+            id="nano-gb300",
+        ),
+        pytest.param(
+            "megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh",
+            "nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config",
+            False,
+            id="super-gb200",
+        ),
+        pytest.param(
+            "megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
+            "nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config",
+            False,
+            id="super-gb300",
+        ),
+    ],
+)
+def test_nemotron_3_gb_mxfp8_perf_recipes_enable_cutedsl_fusion(
+    module_name: str,
+    factory_name: str,
+    has_comm_overlap: bool,
+) -> None:
+    """GB MXFP8 recipes match the measured CutDSL and MoE overlap settings."""
+    module = importlib.import_module(module_name)
+    cfg = getattr(module, factory_name)()
+
+    assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1
+    assert cfg.env_vars["CUDNNFE_CLUSTER_OVERLAP_MARGIN"] == 8
+    assert cfg.model.use_transformer_engine_op_fuser is True
+    assert cfg.model.moe_mlp_glu_interleave_size == 32
+    assert cfg.model.high_priority_a2a_comm_stream is False
+    assert cfg.model.moe_hybridep_num_sms_preprocessing == 108
+    assert cfg.mixed_precision.fp8_dot_product_attention is True
+    if has_comm_overlap:
+        assert cfg.comm_overlap is not None
+        assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is False
+        assert cfg.comm_overlap.delay_wgrad_compute is False
+    else:
+        assert cfg.comm_overlap is None
+
+
 def test_nemotron_3_super_64gpu_gb200_matches_benchmark_hardware_configuration():
     """The training recipe should share the tuned GB200 layout without benchmark-only behavior."""
     from megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh import (