perf(recipe): llama3 8B GB200 NVFP4 — run DPA in FP8 current scaling (#5394)

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

Author: nvegesna-netizen

src/megatron/bridge/perf_recipes/llama/gb200/llama3.py

@@ -160,6 +160,8 @@ def llama3_8b_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer:
     """Llama3 8B pretrain: 8× GB200, NVFP4."""
     cfg = llama3_8b_pretrain_config()
     cfg.mixed_precision = _perf_precision("nvfp4")
+    # NVFP4BlockScaling takes fp8_dpa from this field, so DPA runs in FP8 under the FP4 recipe.
+    cfg.mixed_precision.fp8_dot_product_attention = True
     cfg.tokenizer.vocab_size = 128256
     cfg.model.should_pad_vocab = True
 
@@ -193,6 +195,8 @@ def llama3_8b_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer:
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
         # NVFP4 fast-math path.
         "NVTE_USE_FAST_MATH": 1,
+        # FP8 recipe TE uses for the DPA enabled above.
+        "NVTE_DPA_FP8_RECIPE": "Float8CurrentScaling",
     }
     return cfg
 

tests/unit_tests/recipes/test_llama_perf_recipes.py

@@ -20,6 +20,7 @@
 import pytest
 
 from megatron.bridge.perf_recipes.llama.gb200.llama3 import (
+    llama3_8b_pretrain_8gpu_gb200_nvfp4_config,
     llama3_70b_pretrain_64gpu_gb200_nvfp4_config,
 )
 from megatron.bridge.perf_recipes.llama.h100.llama3 import (
@@ -139,3 +140,28 @@ def test_llama3_70b_gb200_nvfp4_captures_whole_transformer_layer(monkeypatch: py
     assert cfg.model.pipeline_model_parallel_size == 4
     assert cfg.model.virtual_pipeline_model_parallel_size == 5
     assert cfg.train.global_batch_size == 256
+
+
+@pytest.mark.unit
+def test_llama3_8b_gb200_nvfp4_runs_dpa_in_fp8_current_scaling(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The 8-GPU GB200 NVFP4 recipe runs dot-product attention in FP8 current scaling.
+
+    NVFP4BlockScaling takes ``fp8_dpa`` from ``fp8_dot_product_attention`` independently of
+    ``config.fp8``, so the flag is live under the FP4 recipe rather than inert.
+    ``NVTE_DPA_FP8_RECIPE`` selects the FP8 recipe TE uses for that attention path, and must be
+    part of the inline ``env_vars`` dict so the recipe-environment test sees it.
+    """
+    patch_recipe_construction_dependencies(monkeypatch)
+
+    cfg = llama3_8b_pretrain_8gpu_gb200_nvfp4_config()
+
+    assert cfg.mixed_precision.fp8_dot_product_attention is True
+    assert cfg.mixed_precision.fp4 == "e2m1"
+    assert cfg.mixed_precision.fp4_recipe == "nvfp4"
+    assert cfg.env_vars["NVTE_DPA_FP8_RECIPE"] == "Float8CurrentScaling"
+
+    # Parallelism and batch are unchanged by this setting.
+    assert cfg.model.tensor_model_parallel_size == 1
+    assert cfg.model.pipeline_model_parallel_size == 1
+    assert cfg.train.global_batch_size == 128
+    assert cfg.train.micro_batch_size == 4