diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_mxfp4.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_mxfp4.yaml
new file mode 100644
index 000000000..d495d7866
--- /dev/null
+++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_mxfp4.yaml
@@ -0,0 +1,204 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+# Flux 12B Schnell Training — DDP + TE Spec + MXFP4 (native-TE autocast) (MI355X)
+#
+# Combines Megatron DDP + distributed optimizer with the TransformerEngine spec,
+# routed through TE's native FP4 autocast (MXFP4BlockScaling) for MXFP4 training.
+#
+# Key configuration:
+#   - Megatron DDP with overlap_grad_reduce + overlap_param_gather
+#   - TEColumnParallelLinear / TERowParallelLinear / TEDotProductAttention / TENorm
+#   - Fused CK attention (attention_backend: fused), ~2.5x faster bwd than FlashAttn
+#   - MXFP4 (E2M1 + E8M0 block-of-32 scales) via TE-native fp4_use_native_te_autocast
+#   - Primus Turbo disabled (native TE linears consume the model-level FP4 autocast)
+#   - torch.compile disabled (TE FP4 autocast global state breaks graph capture)
+#   - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode)
+#
+# Set these environment variables before launching:
+#   export NVTE_USE_CAST_TRANSPOSE_TRITON=0
+#   export NVTE_FUSED_ATTN=1
+#   export NVTE_FUSED_ATTN_CK=1
+#   export NVTE_USE_HIPBLASLT=1
+#   export USE_HIPBLASLT=1
+#   export TORCH_BLAS_PREFER_HIPBLASLT=1
+#   export NVTE_FP4_USE_TUNED_GEMM=1
+#   export NVTE_FP4_LOG_GEMM_SHAPES=0
+#   unset NVTE_FLASH_ATTN 
+#   unset NVTE_UNFUSED_ATTN
+
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_mxfp4}
+
+workspace: ${PRIMUS_WORKSPACE:./output}
+
+modules:
+  pre_trainer:
+    framework: megatron
+    config: pre_trainer.yaml
+    model: diffusion/flux_12b.yaml
+
+    trainer_class: FluxPretrainTrainer
+
+    overrides:
+      model_type: flux_schnell
+
+      # ==========================================
+      # Flux Schnell training recipe
+      # ==========================================
+      timestep_sampling_strategy: "direct_uniform"
+      activation_func: "openai_gelu"
+      cfg_dropout_prob: 0.1
+
+      vae_scale: 0.3611
+      vae_shift: 0.1159
+      vae_latent_mode: resample
+
+      # ==========================================
+      # RoPE Fusion
+      # ==========================================
+      rotary_interleaved: true
+      apply_rope_fusion: true
+      position_embedding_type: rope
+
+      # ==========================================
+      # TransformerEngine Spec
+      # ==========================================
+      transformer_impl: "transformer_engine"
+
+      # ==========================================
+      # Attention backend — FUSED (TE CK fused attention)
+      # ==========================================
+      attention_backend: fused
+
+      # ==========================================
+      # MXFP4 — via TE-native MXFP4BlockScaling autocast (no Primus-Turbo)
+      # ==========================================
+      fp4: "mxfp4"
+      fp4_recipe: "mxfp4"
+      fp4_use_native_te_autocast: true
+      mxfp4_backward_precision: "mxfp4"
+
+      # ==========================================
+      # Energon Dataset
+      # ==========================================
+      data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset}
+
+      # Training iterations (example value; increase for a full training run)
+      train_iters: 1000
+      eval_interval: 10000
+      eval_iters: 0
+      log_interval: 1
+      save_interval: 10000
+
+      # ==========================================
+      # Batch Configuration
+      # ==========================================
+      micro_batch_size: 64
+      global_batch_size: 512
+      seq_length: 512  # 256 img tokens + 256 text tokens (schnell)
+
+      # ==========================================
+      # BF16 + MXFP4 Precision
+      # ==========================================
+      bf16: true
+      fp16: false
+      params_dtype: bfloat16
+      grad_reduce_in_bf16: true
+
+      # ==========================================
+      # Optimizer
+      # ==========================================
+      optimizer: adam
+      lr: 2.0e-4
+      min_lr: 2.0e-4
+      weight_decay: 0.1
+      adam_beta1: 0.9
+      adam_beta2: 0.95
+      adam_eps: 1.0e-8
+      clip_grad: 1.0
+      
+      # Learning rate scheduler (warmup-hold, no decay)
+      lr_warmup_iters: 1600
+      lr_decay_iters: 4000
+      lr_decay_style: constant
+
+      # DataLoader
+      num_workers: 8
+      prefetch_factor: 8
+      dataloader_type: external
+      max_samples_per_sequence: null
+
+      # ==========================================
+      # Megatron DDP + Distributed Optimizer
+      # ==========================================
+      tensor_model_parallel_size: 1
+      pipeline_model_parallel_size: 1
+      expert_model_parallel_size: 1
+      context_parallel_size: 1
+
+      use_torch_fsdp2: false
+      use_megatron_fsdp: false
+
+      use_distributed_optimizer: true
+      overlap_grad_reduce: true
+      overlap_param_gather: true
+      overlap_param_gather_with_optimizer_step: false
+      use_precision_aware_optimizer: false
+
+      ddp_bucket_size: 256000000
+      use_fsdp2_fp32_param_optimizer: false
+
+      ckpt_format: torch_dist
+
+      gradient_accumulation_fusion: false
+
+      check_for_nan_in_loss_and_grad: true
+
+      # ==========================================
+      # Memory / Misc
+      # ==========================================
+      empty_unused_memory_level: 0
+
+      # Manual GC — align GC timing across ranks to avoid stragglers
+      manual_gc: true
+      manual_gc_interval: 1000
+
+      distributed_timeout_minutes: 60
+      distributed_backend: nccl
+
+      # Checkpointing (disabled)
+      finetune: false
+      save: null
+      load: null
+      no_save_optim: true
+      no_save_rng: true
+      auto_continue_train: false
+      disable_last_saving: true
+
+      # Logging
+      log_throughput: true
+      wall_clock_step_timer: true
+
+      # Profiler — disabled
+      profile: false
+
+      # ==========================================
+      # NATIVE TE — Primus-Turbo disabled
+      # ==========================================
+      # The native-TE MXFP4 path uses TE's own linears + fp8_autocast(MXFP4BlockScaling).
+      # Primus-Turbo is disabled, so the model builds native TE linears that consume the
+      # model-level FP4 autocast returned by Flux.get_fp8_context().
+      enable_primus_turbo: false
+      use_turbo_attention: false
+
+      seed: 2025
+      te_rng_tracker: true
+
+      # torch.compile — DISABLED for native-TE MXFP4.
+      # TE's FP4 autocast toggles process-global quantization state, which breaks
+      # torch.compile graph capture (graph breaks / stale guards). Run eager, matching
+      # the validated native-TE MXFP4 runs.
+      torch_compile:
+        enable: false
diff --git a/primus/backends/megatron/core/fp4_utils.py b/primus/backends/megatron/core/fp4_utils.py
index b747d9b60..23ce8308b 100644
--- a/primus/backends/megatron/core/fp4_utils.py
+++ b/primus/backends/megatron/core/fp4_utils.py
@@ -144,7 +144,13 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool
                 fp4_context = nullcontext()
             else:
                 fp4_recipe, fp4_recipe_none_reason = get_fp4_recipe(config)
-                turbo_enabled = _primus_turbo_enabled()
+                # fp4_use_native_te_autocast forces the TE-native autocast branch
+                # (TE fp8_autocast + MXFP4BlockScaling -> AITER a4w4), bypassing
+                # Primus-Turbo entirely -- even if the Turbo autocast is otherwise
+                # enabled. This is the pure-TE MXFP4 path (enable_primus_turbo=false).
+                turbo_enabled = _primus_turbo_enabled() and not getattr(
+                    config, "fp4_use_native_te_autocast", False
+                )
 
                 global WARN_ONCE
                 if WARN_ONCE:
diff --git a/primus/backends/megatron/core/models/diffusion/flux/config.py b/primus/backends/megatron/core/models/diffusion/flux/config.py
index 83446615b..3a0bf1cb4 100644
--- a/primus/backends/megatron/core/models/diffusion/flux/config.py
+++ b/primus/backends/megatron/core/models/diffusion/flux/config.py
@@ -174,6 +174,15 @@ class FluxConfig(BaseDiffusionConfig):
     # quantization instead of the default inline implementation.
     use_cpp_fp8_quantize: bool = False
 
+    # FP4 autocast backend (TE-spec only). When True, the model-level FP4 autocast
+    # uses TE's native fp8_autocast with the MXFP4BlockScaling recipe (TE -> AITER
+    # a4w4) instead of Primus-Turbo's primus_turbo_fp4_autocast, forcing
+    # get_fp4_context onto the TE-native branch. Enables pure-TE MXFP4
+    # (enable_primus_turbo=false); requires use_turbo_gemm=false (native TE linears)
+    # because PrimusTurboLinear quantizes only under the Turbo autocast and would
+    # otherwise fall back to BF16. Enforced in validate_args_on_rocm.
+    fp4_use_native_te_autocast: bool = False
+
     # CUDA graph support
     enable_cuda_graph: bool = False
     cuda_graph_scope: Optional[str] = None  # Options: "full", "full_iteration"
diff --git a/primus/backends/megatron/core/models/diffusion/flux/model.py b/primus/backends/megatron/core/models/diffusion/flux/model.py
index b451045fa..8f05a4aa0 100644
--- a/primus/backends/megatron/core/models/diffusion/flux/model.py
+++ b/primus/backends/megatron/core/models/diffusion/flux/model.py
@@ -634,25 +634,46 @@ def sharded_state_dict(
 
     def get_fp8_context(self):
         """
-        Return FP8 context manager for FP8 training.
+        Return the active low-precision (FP8 or FP4) quantization context manager.
 
-        When transformer_impl="local", FP8 is handled per-module at init time
-        (inside Float8ColumnParallelLinear / Float8RowParallelLinear), so no
-        global FP8 context manager is needed. This avoids mutable global state
-        that would cause torch.compile graph breaks.
+        For the TransformerEngine spec, a TE linear only issues an FP8/FP4 GEMM
+        while the corresponding autocast context is active. The Flux forward wraps
+        every compile strategy in this single method, so it must return whichever
+        quantization context the config requests:
+
+        - FP8 (config.fp8): TE FP8 autocast via megatron.core.fp8_utils.
+        - FP4 (config.fp4, e.g. MXFP4): TE FP4 autocast via megatron.core.fp4_utils.
+          Without this branch a pure-FP4 run would get nullcontext and silently
+          fall back to BF16 GEMMs (no AITER FP4 routing).
+
+        FP8 takes precedence over FP4 when both are set, mirroring stock Megatron's
+        transformer_block (`if config.fp8 ... elif config.fp4 ...`).
+
+        When transformer_impl="local", FP8/FP4 is handled per-module at init time
+        (inside Float8/MXFP4 ColumnParallelLinear / RowParallelLinear), so no global
+        context manager is needed. This avoids mutable global state that would cause
+        torch.compile graph breaks.
 
         Returns:
-            Context manager for FP8 operations (or nullcontext if FP8 disabled
-            or using local spec)
+            Context manager for the active quantization mode (or nullcontext if
+            quantization is disabled or using local spec).
         """
         from contextlib import nullcontext
 
         if self.config.transformer_impl == "local":
             return nullcontext()
 
-        from megatron.core.fp8_utils import get_fp8_context as get_fp8_context_util
+        if getattr(self.config, "fp8", None):
+            from megatron.core.fp8_utils import get_fp8_context as get_fp8_context_util
+
+            return get_fp8_context_util(self.config)
+
+        if getattr(self.config, "fp4", None):
+            from megatron.core.fp4_utils import get_fp4_context as get_fp4_context_util
+
+            return get_fp4_context_util(self.config)
 
-        return get_fp8_context_util(self.config)
+        return nullcontext()
 
     def forward(
         self,
diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py
index 36da05ce3..602a21d85 100644
--- a/primus/backends/megatron/flux_pretrain_trainer.py
+++ b/primus/backends/megatron/flux_pretrain_trainer.py
@@ -466,6 +466,30 @@ def _build_flux_config_from_yaml(self):
         transformer_impl = getattr(params, "transformer_impl", "transformer_engine")
         config_params["transformer_impl"] = transformer_impl
 
+        # Convert attention_backend (string) to the AttnBackend enum. A YAML value
+        # like "fused" arrives as a plain string, but _set_attention_backend()
+        # compares it against AttnBackend enum members -- a raw string matches no
+        # branch, so the backend is silently left unset and falls back to
+        # FlashAttention. AttnBackend["fused"] returns AttnBackend.fused.
+        from megatron.core.transformer.enums import AttnBackend
+
+        # The getattr default only applies when the key is absent; an explicit
+        # `attention_backend: null` in YAML yields None, so coerce that (and an
+        # empty string) to auto rather than crashing on attn_backend.name below.
+        attn_backend = getattr(params, "attention_backend", AttnBackend.auto)
+        if attn_backend is None or (isinstance(attn_backend, str) and not attn_backend.strip()):
+            attn_backend = AttnBackend.auto
+        if isinstance(attn_backend, str):
+            try:
+                attn_backend = AttnBackend[attn_backend.strip()]
+            except KeyError as exc:
+                raise ValueError(
+                    f"Unknown attention_backend '{attn_backend}'. Choose from: "
+                    f"{[b.name for b in AttnBackend]}."
+                ) from exc
+        config_params["attention_backend"] = attn_backend
+        log_rank_0(f"Attention backend: {attn_backend.name}")
+
         # Precision settings from params (not in model YAML)
         config_params.update(
             {
@@ -505,6 +529,7 @@ def _build_flux_config_from_yaml(self):
                 "fp4": fp4_enabled,
                 "fp4_recipe": fp4_recipe,
                 "mxfp4_backward_precision": getattr(params, "mxfp4_backward_precision", "mxfp4"),
+                "fp4_use_native_te_autocast": getattr(params, "fp4_use_native_te_autocast", False),
             }
         )
 
diff --git a/primus/backends/megatron/patches/args/rocm_arg_validation.py b/primus/backends/megatron/patches/args/rocm_arg_validation.py
index 95aff1b15..ca8ef3911 100644
--- a/primus/backends/megatron/patches/args/rocm_arg_validation.py
+++ b/primus/backends/megatron/patches/args/rocm_arg_validation.py
@@ -143,6 +143,23 @@ def validate_args_on_rocm(args):
             args.fp4_recipe in support_fp4_recipe
         ), f"{args.fp4_recipe} recipe is not support when enable `use_turbo_gemm`."
 
+    # Turbo FP4 autocast check: PrimusTurboLinear quantizes only under the Turbo
+    # autocast, which the TE-native path never enters. Local specs quantize
+    # in-module rather than off that state, so they are exempt.
+    if (
+        args.fp4
+        and use_turbo_gemm
+        and getattr(args, "fp4_use_native_te_autocast", False)
+        and getattr(args, "transformer_impl", "transformer_engine") != "local"
+    ):
+        raise ValueError(
+            "fp4_use_native_te_autocast=True is incompatible with use_turbo_gemm=True. "
+            "Native TE FP4 autocast does not set the Primus-Turbo FP4 state that "
+            "PrimusTurboLinear reads, so Turbo GEMM linears would silently run BF16 "
+            "instead of FP4. Set use_turbo_gemm=false to use native TE MXFP4 linears, "
+            "or set fp4_use_native_te_autocast=false to use the Primus-Turbo FP4 path."
+        )
+
     # NOTE: mxfp8 environment variable must be set to 1 to enable mxfp8 recipe on ROCm.
     if args.fp8_recipe == "mxfp8":
         assert (
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py
index 2bd83cef0..395acd8b0 100644
--- a/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py
+++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py
@@ -4,8 +4,10 @@
 """
 Unit tests for FluxModel.get_fp8_context() behavior.
 
-Verifies that the local spec FP8 path returns nullcontext (no global state),
-while the non-local path delegates to Megatron's get_fp8_context utility.
+Verifies that the local spec path returns nullcontext (no global state), while
+the non-local (TransformerEngine) path delegates to the appropriate Megatron
+quantization context: FP8 via megatron.core.fp8_utils, FP4 via
+megatron.core.fp4_utils, with FP8 taking precedence when both are set.
 """
 
 import contextlib
@@ -49,3 +51,66 @@ def test_get_fp8_context_non_local_delegates(self):
             ctx = Flux.get_fp8_context(mock_model)
             mock_get.assert_called_once_with(mock_model.config)
             assert ctx is sentinel
+
+    def test_get_fp8_context_local_with_fp4_returns_nullcontext(self):
+        """Local spec must stay nullcontext even when fp4 is set (per-module FP4)."""
+        mock_model = MagicMock(spec=Flux)
+        mock_model.config = SimpleNamespace(
+            transformer_impl="local",
+            fp8=None,
+            fp4="mxfp4",
+        )
+
+        ctx = Flux.get_fp8_context(mock_model)
+        assert isinstance(ctx, contextlib.nullcontext)
+
+    def test_get_fp8_context_non_local_fp4_delegates(self):
+        """Non-local fp4-only run should delegate to megatron.core.fp4_utils.get_fp4_context."""
+        mock_model = MagicMock(spec=Flux)
+        mock_model.config = SimpleNamespace(
+            transformer_impl="transformer_engine",
+            fp8=None,
+            fp4="mxfp4",
+        )
+
+        sentinel = contextlib.nullcontext()
+        with patch(
+            "megatron.core.fp4_utils.get_fp4_context",
+            return_value=sentinel,
+        ) as mock_get:
+            ctx = Flux.get_fp8_context(mock_model)
+            mock_get.assert_called_once_with(mock_model.config)
+            assert ctx is sentinel
+
+    def test_get_fp8_context_fp8_takes_precedence_over_fp4(self):
+        """When both fp8 and fp4 are set, FP8 wins and FP4 context is not built."""
+        mock_model = MagicMock(spec=Flux)
+        mock_model.config = SimpleNamespace(
+            transformer_impl="transformer_engine",
+            fp8="e4m3",
+            fp4="mxfp4",
+        )
+
+        fp8_sentinel = contextlib.nullcontext()
+        with patch(
+            "megatron.core.fp8_utils.get_fp8_context",
+            return_value=fp8_sentinel,
+        ) as mock_fp8, patch(
+            "megatron.core.fp4_utils.get_fp4_context",
+        ) as mock_fp4:
+            ctx = Flux.get_fp8_context(mock_model)
+            mock_fp8.assert_called_once_with(mock_model.config)
+            mock_fp4.assert_not_called()
+            assert ctx is fp8_sentinel
+
+    def test_get_fp8_context_no_quantization_returns_nullcontext(self):
+        """Non-local with neither fp8 nor fp4 set should return nullcontext."""
+        mock_model = MagicMock(spec=Flux)
+        mock_model.config = SimpleNamespace(
+            transformer_impl="transformer_engine",
+            fp8=None,
+            fp4=None,
+        )
+
+        ctx = Flux.get_fp8_context(mock_model)
+        assert isinstance(ctx, contextlib.nullcontext)
diff --git a/tests/unit_tests/backends/megatron/diffusion/test_fp4_context.py b/tests/unit_tests/backends/megatron/diffusion/test_fp4_context.py
new file mode 100644
index 000000000..d6fa238af
--- /dev/null
+++ b/tests/unit_tests/backends/megatron/diffusion/test_fp4_context.py
@@ -0,0 +1,101 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# Licensed under the Apache License, Version 2.0.
+
+"""
+Unit tests for get_fp4_context() backend selection in
+primus.backends.megatron.core.fp4_utils.
+
+Focus: the fp4_use_native_te_autocast override. When set, get_fp4_context must
+take TransformerEngine's native fp8_autocast branch (TE MXFP4BlockScaling ->
+AITER a4w4) even with Primus-Turbo enabled; when unset (with Primus-Turbo
+enabled) it must take the primus_turbo_fp4_autocast branch.
+"""
+
+import contextlib
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus.backends.megatron.core import fp4_utils
+
+# The two-branch get_fp4_context (Primus-Turbo vs TE-native) only exists when
+# both TransformerEngine and Primus-Turbo are importable; otherwise the fallback
+# definitions have no branch to exercise.
+pytestmark = pytest.mark.skipif(
+    not (fp4_utils.HAVE_TE and fp4_utils.HAVE_TURBO),
+    reason="requires both TransformerEngine and Primus-Turbo installed",
+)
+
+
+def _fp4_config(**overrides):
+    """Minimal config that reaches the Turbo-vs-TE-native branch of get_fp4_context."""
+    config = SimpleNamespace(
+        fp4="mxfp4",
+        fp4_recipe="mxfp4",
+        transformer_impl="transformer_engine",
+        first_last_layers_bf16=False,
+        num_layers=10,
+        tp_only_amax_red=False,
+        fp4_use_native_te_autocast=False,
+    )
+    for key, value in overrides.items():
+        setattr(config, key, value)
+    return config
+
+
+class TestGetFp4ContextBackendSelection:
+    """get_fp4_context() Turbo-vs-TE-native branch selection."""
+
+    def test_native_te_autocast_forces_te_branch_over_turbo(self):
+        """fp4_use_native_te_autocast=True -> TE fp8_autocast, even with Turbo enabled."""
+        config = _fp4_config(fp4_use_native_te_autocast=True)
+        te_sentinel = contextlib.nullcontext()
+
+        with patch.object(
+            fp4_utils, "_primus_turbo_enabled", return_value=True
+        ), patch.object(
+            fp4_utils, "get_fp4_recipe", return_value=(MagicMock(name="recipe"), "")
+        ), patch(
+            "megatron.core.parallel_state.model_parallel_is_initialized",
+            return_value=False,
+        ), patch(
+            "transformer_engine.pytorch.fp8_autocast", return_value=te_sentinel
+        ) as mock_te, patch(
+            "primus.backends.megatron.core.extensions.primus_turbo.primus_turbo_fp4_autocast"
+        ) as mock_turbo:
+            ctx = fp4_utils.get_fp4_context(config)
+
+        mock_te.assert_called_once()
+        mock_turbo.assert_not_called()
+        assert ctx is te_sentinel
+
+    def test_turbo_branch_taken_when_native_te_autocast_disabled(self):
+        """fp4_use_native_te_autocast=False + Turbo enabled -> primus_turbo_fp4_autocast."""
+        config = _fp4_config(fp4_use_native_te_autocast=False)
+        turbo_sentinel = contextlib.nullcontext()
+
+        with patch.object(
+            fp4_utils, "_primus_turbo_enabled", return_value=True
+        ), patch.object(
+            fp4_utils, "get_fp4_recipe", return_value=(MagicMock(name="recipe"), "")
+        ), patch.object(
+            fp4_utils, "get_fp4_quant_config", return_value=(MagicMock(name="quant"), "")
+        ), patch(
+            "megatron.core.parallel_state.model_parallel_is_initialized",
+            return_value=False,
+        ), patch(
+            "transformer_engine.pytorch.fp8_autocast"
+        ) as mock_te, patch(
+            "primus.backends.megatron.core.extensions.primus_turbo.primus_turbo_fp4_autocast",
+            return_value=turbo_sentinel,
+        ) as mock_turbo:
+            ctx = fp4_utils.get_fp4_context(config)
+
+        mock_turbo.assert_called_once()
+        mock_te.assert_not_called()
+        assert ctx is turbo_sentinel
diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py
index 9a965ccfc..9b0b67e5f 100644
--- a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py
+++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py
@@ -286,3 +286,93 @@ def test_create_model_does_not_overwrite_existing_torch_compile_attrs(
         # New attributes should still be set
         assert mock_args.torch_compile_mode == "reduce-overhead"
         assert mock_args.torch_compile_fullgraph is True
+
+    def test_build_flux_config_from_yaml_attention_backend_string(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """A YAML string like 'fused' is converted to the AttnBackend enum."""
+        from megatron.core.transformer.enums import AttnBackend
+
+        backend_args = SimpleNamespace(mock_data=True, attention_backend="fused")
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        config = trainer._build_flux_config_from_yaml()
+
+        assert config.attention_backend is AttnBackend.fused
+
+    def test_build_flux_config_from_yaml_attention_backend_default(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """An omitted attention_backend falls back to AttnBackend.auto."""
+        from megatron.core.transformer.enums import AttnBackend
+
+        backend_args = SimpleNamespace(mock_data=True)
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        config = trainer._build_flux_config_from_yaml()
+
+        assert config.attention_backend is AttnBackend.auto
+
+    def test_build_flux_config_from_yaml_attention_backend_null(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """An explicit `attention_backend: null` (None) is coerced to auto.
+
+        Guards against the AttributeError regression where None flowed past the
+        string branch and crashed on attn_backend.name.
+        """
+        from megatron.core.transformer.enums import AttnBackend
+
+        backend_args = SimpleNamespace(mock_data=True, attention_backend=None)
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        config = trainer._build_flux_config_from_yaml()
+
+        assert config.attention_backend is AttnBackend.auto
+
+    @pytest.mark.parametrize("value", ["", "   "])
+    def test_build_flux_config_from_yaml_attention_backend_blank(
+        self, monkeypatch: pytest.MonkeyPatch, value: str
+    ):
+        """Empty/whitespace-only strings are treated as auto, not errors."""
+        from megatron.core.transformer.enums import AttnBackend
+
+        backend_args = SimpleNamespace(mock_data=True, attention_backend=value)
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        config = trainer._build_flux_config_from_yaml()
+
+        assert config.attention_backend is AttnBackend.auto
+
+    def test_build_flux_config_from_yaml_attention_backend_strips_whitespace(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """Surrounding whitespace is stripped before the enum lookup."""
+        from megatron.core.transformer.enums import AttnBackend
+
+        backend_args = SimpleNamespace(mock_data=True, attention_backend="  fused  ")
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        config = trainer._build_flux_config_from_yaml()
+
+        assert config.attention_backend is AttnBackend.fused
+
+    def test_build_flux_config_from_yaml_attention_backend_invalid(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """An unknown attention_backend raises a clear ValueError listing choices."""
+        backend_args = SimpleNamespace(mock_data=True, attention_backend="bogus")
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        with pytest.raises(ValueError, match=r"Unknown attention_backend 'bogus'"):
+            trainer._build_flux_config_from_yaml()
+
+    def test_build_flux_config_from_yaml_attention_backend_case_sensitive(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        """Enum lookup is case-sensitive: 'FUSED' is rejected, not matched to 'fused'."""
+        backend_args = SimpleNamespace(mock_data=True, attention_backend="FUSED")
+        trainer = _build_flux_trainer(monkeypatch, backend_args)
+
+        with pytest.raises(ValueError, match=r"Unknown attention_backend 'FUSED'"):
+            trainer._build_flux_config_from_yaml()
diff --git a/tests/unit_tests/backends/megatron/test_validate_args_patches.py b/tests/unit_tests/backends/megatron/test_validate_args_patches.py
index bd58796cf..047886362 100644
--- a/tests/unit_tests/backends/megatron/test_validate_args_patches.py
+++ b/tests/unit_tests/backends/megatron/test_validate_args_patches.py
@@ -120,6 +120,8 @@ def _make_args(**overrides):
         fp4=False,
         fp8_recipe=None,
         fp4_recipe=None,
+        fp4_use_native_te_autocast=False,
+        transformer_impl="transformer_engine",
         use_turbo_gemm=False,
         dump_pp_data=False,
         pipeline_model_parallel_size=1,
@@ -336,3 +338,71 @@ def test_pipeline_split_and_rocm_validation(self, monkeypatch):
         assert args._validate_args_called is True
         assert len(rocm_calls) == 1
         assert result is args
+
+
+class TestFp4NativeTeAutocastTurboGemmGuard:
+    """``validate_args_on_rocm`` rejects TE-native FP4 autocast + Turbo GEMM linears.
+
+    The TE-native autocast never sets ``PRIMUS_TURBO_FP4_ENABLED``, which is the
+    only signal ``PrimusTurboLinear`` consults, so the combination would silently
+    run BF16 GEMMs instead of FP4.
+    """
+
+    def _validate(self, monkeypatch, **overrides):
+        _install_fake_megatron(monkeypatch)
+
+        from primus.backends.megatron.patches.args.rocm_arg_validation import (
+            validate_args_on_rocm,
+        )
+
+        validate_args_on_rocm(_make_args(**overrides))
+
+    def test_rejects_native_te_autocast_with_turbo_gemm(self, monkeypatch):
+        with pytest.raises(ValueError, match="fp4_use_native_te_autocast"):
+            self._validate(
+                monkeypatch,
+                fp4=True,
+                fp4_recipe="mxfp4",
+                use_turbo_gemm=True,
+                fp4_use_native_te_autocast=True,
+            )
+
+    def test_allows_native_te_autocast_without_turbo_gemm(self, monkeypatch):
+        """Pure-TE MXFP4: native TE linears honour the TE autocast."""
+        self._validate(
+            monkeypatch,
+            fp4=True,
+            fp4_recipe="mxfp4",
+            use_turbo_gemm=False,
+            fp4_use_native_te_autocast=True,
+        )
+
+    def test_allows_turbo_gemm_without_native_te_autocast(self, monkeypatch):
+        """Turbo FP4 autocast drives the Turbo linears."""
+        self._validate(
+            monkeypatch,
+            fp4=True,
+            fp4_recipe="mxfp4",
+            use_turbo_gemm=True,
+            fp4_use_native_te_autocast=False,
+        )
+
+    def test_allows_both_flags_when_fp4_disabled(self, monkeypatch):
+        """Neither autocast is built when FP4 is off, so there is nothing to mismatch."""
+        self._validate(
+            monkeypatch,
+            fp4=False,
+            use_turbo_gemm=True,
+            fp4_use_native_te_autocast=True,
+        )
+
+    def test_allows_local_transformer_impl(self, monkeypatch):
+        """Local spec providers quantize inside the module, ignoring the global flag."""
+        self._validate(
+            monkeypatch,
+            fp4=True,
+            fp4_recipe="mxfp4",
+            use_turbo_gemm=True,
+            fp4_use_native_te_autocast=True,
+            transformer_impl="local",
+        )