diff --git a/scripts/performance/perf_plugins.py b/scripts/performance/perf_plugins.py
index 89f2133b88..a919561568 100644
--- a/scripts/performance/perf_plugins.py
+++ b/scripts/performance/perf_plugins.py
@@ -287,7 +287,7 @@ def _set_model_specific_environment_variables(
             model_family_name in ["llama"]
             and model_recipe_name in ["llama31_405b"]
             and train_task == "pretrain"
-            and gpu in ["gb200", "gb300"]
+            and gpu in ["gb200", "gb300", "b200", "b300"]
         ):
             if compute_dtype in ["fp8_cs", "fp8_mx", "nvfp4"]:
                 executor.env_vars["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"