perf(qwen3): Qwen3-30B-A3B FP8 on MI355X to 40k tokens/s/GPU with Primus-Turbo tensorwise

examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml

@@ -3,6 +3,16 @@ user_name: ${PRIMUS_USER:root}
 exp_name: ${PRIMUS_EXP_NAME:qwen3_30B_A3B-pretrain}
 workspace: ${PRIMUS_WORKSPACE:./output}
 
+# Fused packed-QKV split + Q/K RMSNorm + RoPE. head_dim 128 needs a Primus-Turbo
+# with QK_RMSNORM_ROPE_HEAD_DIMS; older Turbo keeps the unfused path.
+#
+# Full DeepEP fence at turbo_deepep_num_cu>=160: the cheap fence (#539) was only
+# row-validated at <=80 CUs; the full opt stack hit intermittent NaN forward loss
+# at 160/192 (~1/13). DISABLE=1 selects the release/acquire fence (~2-3% step).
+env:
+  PRIMUS_FUSED_QK_RMSNORM_ROPE: "1"
+  PRIMUS_TURBO_DEEPEP_DISABLE_CHEAP_FENCE: "1"
+
 modules:
   pre_trainer:
     framework: megatron
@@ -83,16 +93,21 @@ modules:
       # Turbo
       enable_primus_turbo: true
       use_turbo_attention: false
-      use_turbo_grouped_gemm: false
+      use_turbo_gemm: true
+      use_turbo_grouped_gemm: true
       use_turbo_rms_norm: false # bug
+      turbo_fused_grouped_gemm: true
+      use_turbo_fused_act_with_probs: true
+      # needs a Primus-Turbo whose moe_permute takes quantize_dtype
+      turbo_fp8_permute: true
 
       # deepep
       use_turbo_deepep: true
       moe_shared_expert_overlap: false
       moe_router_dtype: fp32
 
-      # 64 or 80 for ep8, 32 for ep16-64 is best practice
-      turbo_deepep_num_cu: 80
+      # 160 for ep8 on MI355X (192 gives a NaN loss, 256 is slower), 32 for ep16-64
+      turbo_deepep_num_cu: 160
       turbo_deepep_use_comm_stream: false
 
       # sync-free moe support stage 1-2, 0 means not use sync-free moe
@@ -105,5 +120,6 @@ modules:
 
       # enable fp8 training
       fp8: hybrid
-      moe_use_legacy_grouped_gemm: true
+      fp8_recipe: tensorwise
+      moe_use_legacy_grouped_gemm: false
       gradient_accumulation_fusion: true

primus/backends/megatron/core/extensions/primus_turbo.py

@@ -975,6 +975,7 @@ def __init__(
         self._num_heads_for_sinks = self.config.num_attention_heads
 
         self.offload = args.offload and "attn" in args.offload_ops
+        self.use_fp8_attention = args.enable_turbo_attention_float8
         if args.enable_turbo_attention_float8:
             self.attn = (
                 primus_turbo_torch.ops.flash_attn_fp8_usp_func
@@ -1129,6 +1130,15 @@ def forward(
             key = key.permute(0, 2, 1, 3)
             value = value.permute(0, 2, 1, 3)
 
+        if self.use_fp8_attention:
+            # flash_attn_fp8_func permutes non-bshd *storage* as if it were the
+            # logical layout, swapping b and s, so hand it bshd-contiguous q/k/v.
+            query = query.contiguous()
+            key = key.contiguous()
+            value = value.contiguous()
+
+        # flash_attn_fp8_func takes no ``sink`` argument.
+        sink_kwargs = {"sink": sink_tensor} if sink_tensor is not None else {}
         o = self.attn(
             query,
             key,
@@ -1142,7 +1152,7 @@ def forward(
             deterministic=self.deterministic_mode,
             return_lse=False,
             return_attn_probs=False,
-            sink=sink_tensor,  # PR 208: pass sink tensor to Primus-Turbo
+            **sink_kwargs,
             **self.attn_kwargs,
         )
 
@@ -2630,6 +2640,23 @@ def __init__(
         self.moe_router_force_load_balancing_type = getattr(
             args, "moe_router_force_load_balancing_type", "uniform"
         )
+        self.turbo_fp8_permute = getattr(args, "turbo_fp8_permute", False)
+
+    def _fp8_permute_dtypes(self):
+        """(input, grad) FP8 dtypes the fused grouped MLP expects, or (None, None).
+
+        Only tensorwise current scaling: that is the one recipe whose scale is a single scalar,
+        so quantizing the dispatched tokens before permute is exact.
+        """
+        if not self.turbo_fp8_permute or not PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled():
+            return None, None
+        quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
+        if not quant_config.current_scaling():
+            return None, None
+        from primus_turbo.pytorch.ops.utils import _get_fp8_dtype
+
+        fmt = quant_config.data().format
+        return _get_fp8_dtype(fmt, True), _get_fp8_dtype(fmt, False)
 
     def dispatch_preprocess(
         self, hidden_states: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor
@@ -2711,9 +2738,15 @@ def dispatch_postprocess(self, hidden_states: torch.Tensor, probs: torch.Tensor)
         Returns:
             A tuple of permuted tokens, token counts per expert, and permuted probabilities.
         """
-        permuted_input, tokens_per_expert, permuted_probs = self.deepep_dispatcher._post_dispatch(
-            hidden_states, probs
-        )
+        quantize_dtype, _ = self._fp8_permute_dtypes()
+        if quantize_dtype is None:
+            permuted_input, tokens_per_expert, permuted_probs = self.deepep_dispatcher._post_dispatch(
+                hidden_states, probs
+            )
+        else:
+            permuted_input, tokens_per_expert, permuted_probs = self.deepep_dispatcher._post_dispatch(
+                hidden_states, probs, quantize_dtype=quantize_dtype
+            )
         if self.config.moe_router_dtype == "fp64":
             permuted_probs = permuted_probs.to(torch.float64)
         return permuted_input, tokens_per_expert, permuted_probs
@@ -2724,7 +2757,13 @@ def combine_preprocess(self, hidden_states: torch.Tensor):
         This method restores the hidden states to their original ordering before expert processing
         by using the communication manager's restoration function.
         """
-        hidden_states = self.deepep_dispatcher._pre_combine(hidden_states)
+        _, grad_quantize_dtype = self._fp8_permute_dtypes()
+        if grad_quantize_dtype is None:
+            hidden_states = self.deepep_dispatcher._pre_combine(hidden_states)
+        else:
+            hidden_states = self.deepep_dispatcher._pre_combine(
+                hidden_states, grad_quantize_dtype=grad_quantize_dtype
+            )
         return hidden_states
 
     def token_combine(

primus/backends/megatron/patches/args/rocm_arg_validation.py

@@ -108,6 +108,19 @@ def validate_turbo_grouped_gemm_without_padding(args) -> None:
         raise ValueError(f"{option}=True requires moe_router_padding_for_quantization=False.")
 
 
+def validate_turbo_fp8_permute(args) -> None:
+    """The FP8 permute handoff only has one consumer: the fused grouped MLP behind Turbo DeepEP."""
+    option = "turbo_fp8_permute"
+    if not getattr(args, option, False):
+        return
+    for required in ("enable_primus_turbo", "use_turbo_deepep", "turbo_fused_grouped_gemm"):
+        if not getattr(args, required, False):
+            raise ValueError(f"{option}=True requires {required}=True.")
+    recompute_modules = getattr(args, "recompute_modules", None) or []
+    if getattr(args, "recompute_granularity", None) == "selective" and "moe_act" in recompute_modules:
+        raise ValueError(f"{option}=True does not support recompute of moe_act.")
+
+
 def validate_args_on_rocm(args):
     # Primus-Turbo auto-tuning
     use_turbo_autotune = getattr(args, "use_turbo_autotune", False)
@@ -216,6 +229,7 @@ def validate_args_on_rocm(args):
         )
 
     validate_turbo_grouped_gemm_without_padding(args)
+    validate_turbo_fp8_permute(args)
 
     # turbo deepep
     if args.use_turbo_deepep:

primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py

@@ -4,16 +4,16 @@
 # See LICENSE for license information.
 ###############################################################################
 
-"""Opt-in GPT-OSS packed-QKV RMSNorm + RoPE integration.
+"""Opt-in packed-QKV RMSNorm + RoPE integration (GPT-OSS head_dim 64, Qwen3 head_dim 128).
 
 ``PRIMUS_FUSED_QK_RMSNORM_ROPE=1`` replaces the training-only sequence
 
     split packed QKV -> Q/K RMSNorm -> Q/K RoPE
 
 with Primus-Turbo's packed FlyDSL operator.  The patch is intentionally narrow:
 TP=1, self attention, SBHD, non-interleaved full RoPE, no packed sequences,
-and PrimusTurboRMSNorm for both Q and K.  Unsupported calls continue through
-the original Megatron path.
+a head_dim the installed Turbo supports, and TE RMSNorm (or PrimusTurboRMSNorm)
+for both Q and K.  Unsupported calls continue through the original Megatron path.
 """
 
 from __future__ import annotations
@@ -79,7 +79,7 @@ def _enabled(_ctx: PatchContext) -> bool:
     "megatron.turbo.qk_rmsnorm_rope",
     backend="megatron",
     phase="before_train",
-    description="Fuse GPT-OSS packed QKV split, Q/K RMSNorm, and RoPE with FlyDSL",
+    description="Fuse packed QKV split, Q/K RMSNorm, and RoPE with FlyDSL",
     condition=_enabled,
     priority=70,
 )
@@ -92,8 +92,14 @@ def patch_qk_rmsnorm_rope(_ctx: PatchContext):
         qk_rmsnorm_rope_shape_error,
     )
     from primus_turbo.pytorch.ops.rope import fused_qkv_rmsnorm_rope
+    from transformer_engine.pytorch.module.rmsnorm import RMSNorm as TERMSNorm
 
-    from primus.backends.megatron.core.extensions.primus_turbo import PrimusTurboRMSNorm
+    try:
+        from primus_turbo.pytorch.kernels.rope.qk_rmsnorm_rope_impl import (
+            QK_RMSNORM_ROPE_HEAD_DIMS,
+        )
+    except ImportError:
+        QK_RMSNORM_ROPE_HEAD_DIMS = (64,)
 
     if getattr(attention_module, "_primus_qk_rmsnorm_rope_installed", False):
         return
@@ -121,15 +127,16 @@ def _eligible(self, rotary_pos_emb, inference_context, packed_seq_params) -> boo
             and rotary_pos_emb is not None
             and not no_rope
             and self.attention_type == "self"
-            and self.hidden_size_per_attention_head == 64
+            and self.hidden_size_per_attention_head in QK_RMSNORM_ROPE_HEAD_DIMS
             and not getattr(self.config, "rotary_interleaved", False)
             and getattr(self.config, "rotary_percent", 1.0) == 1.0
             and attention_module._yarn_get_concentration_factor_from_config(self.config) == 1.0
             and not getattr(self.config, "attention_output_gate", False)
             and getattr(self, "world_size", 1) == 1
             and not getattr(self, "offload_qkv_linear", False)
-            and isinstance(q_norm, PrimusTurboRMSNorm)
-            and isinstance(k_norm, PrimusTurboRMSNorm)
+            # Only weight / eps are consumed, so TE RMSNorm and its PrimusTurboRMSNorm subclass both qualify.
+            and isinstance(q_norm, TERMSNorm)
+            and isinstance(k_norm, TERMSNorm)
             and _supported_forward_pre_hooks(q_norm) is not None
             and _supported_forward_pre_hooks(k_norm) is not None
             and not getattr(q_norm, "zero_centered_gamma", False)
@@ -278,6 +285,5 @@ def _apply_rotary_pos_emb(t, *args, **kwargs):
     attention_module.apply_rotary_pos_emb = _apply_rotary_pos_emb
     attention_module._primus_qk_rmsnorm_rope_installed = True
     log_rank_0(
-        "[Patch:megatron.turbo.qk_rmsnorm_rope] Installed GPT-OSS FlyDSL "
-        "packed QKV + Q/K RMSNorm + RoPE fusion"
+        "[Patch:megatron.turbo.qk_rmsnorm_rope] Installed FlyDSL packed QKV + Q/K RMSNorm + RoPE fusion"
     )

primus/configs/modules/megatron/primus_turbo.yaml

@@ -61,6 +61,11 @@ use_turbo_grouped_gemm: false
 # Fuse the grouped MLP chain (fc1 + GLU + fc2, including WGRAD) into Primus-Turbo's
 # fused grouped MLP op.
 turbo_fused_grouped_gemm: false
+# Hand the fused grouped MLP its input and output gradient already in FP8 (tensorwise current
+# scaling only): the DeepEP dispatcher quantizes the dispatched tokens before permute, and the
+# combine gradient before un-permute's backward, instead of their topk-times larger permuted
+# copies. Bit-identical. Requires turbo_fused_grouped_gemm and use_turbo_deepep.
+turbo_fp8_permute: false
 # let Turbo grouped GEMMs consume non-aligned expert token counts directly,
 # without TE Fp8Padding/Fp8Unpadding
 turbo_grouped_gemm_without_padding: false

tests/unit_tests/backends/megatron/test_rocm_arg_validation.py

@@ -9,6 +9,7 @@
 import pytest
 
 from primus.backends.megatron.patches.args.rocm_arg_validation import (
+    validate_turbo_fp8_permute,
     validate_turbo_grouped_gemm_without_padding,
 )
 
@@ -62,3 +63,49 @@ def test_turbo_grouped_gemm_without_padding_disabled_is_noop():
 def test_turbo_grouped_gemm_without_padding_rejects_unsupported_config(overrides, message):
     with pytest.raises(ValueError, match=message):
         validate_turbo_grouped_gemm_without_padding(_no_padding_args(**overrides))
+
+
+def _fp8_permute_args(**overrides):
+    values = {
+        "turbo_fp8_permute": True,
+        "enable_primus_turbo": True,
+        "use_turbo_deepep": True,
+        "turbo_fused_grouped_gemm": True,
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+@pytest.mark.parametrize(
+    "overrides",
+    [
+        {},
+        {"recompute_granularity": "full", "recompute_modules": ["moe_act"]},
+        {"recompute_granularity": "selective", "recompute_modules": ["core_attn"]},
+        {"recompute_granularity": "selective", "recompute_modules": None},
+    ],
+)
+def test_turbo_fp8_permute_accepts_supported_config(overrides):
+    validate_turbo_fp8_permute(_fp8_permute_args(**overrides))
+
+
+def test_turbo_fp8_permute_disabled_is_noop():
+    validate_turbo_fp8_permute(SimpleNamespace())
+    validate_turbo_fp8_permute(SimpleNamespace(turbo_fp8_permute=False, use_turbo_deepep=False))
+
+
+@pytest.mark.parametrize(
+    ("overrides", "message"),
+    [
+        ({"enable_primus_turbo": False}, "requires enable_primus_turbo=True"),
+        ({"use_turbo_deepep": False}, "requires use_turbo_deepep=True"),
+        ({"turbo_fused_grouped_gemm": False}, "requires turbo_fused_grouped_gemm=True"),
+        (
+            {"recompute_granularity": "selective", "recompute_modules": ["core_attn", "moe_act"]},
+            "does not support recompute of moe_act",
+        ),
+    ],
+)
+def test_turbo_fp8_permute_rejects_unsupported_config(overrides, message):
+    with pytest.raises(ValueError, match=message):
+        validate_turbo_fp8_permute(_fp8_permute_args(**overrides))

tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py

@@ -0,0 +1,69 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""
+Unit tests for the FP8 dtypes the Turbo DeepEP dispatcher asks the permute for under
+``turbo_fp8_permute``. Quantizing before permute is only exact for tensorwise current
+scaling, so every other state must keep the bf16 permute.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tests.utils import skip_if_no_cuda
+
+skip_if_no_cuda()
+
+from primus_turbo.pytorch.core.low_precision import (  # noqa: E402  isort:skip
+    Format,
+    ScalingGranularity,
+    float8_e4m3,
+    float8_e5m2,
+)
+
+from primus.backends.megatron.core.extensions.primus_turbo import (  # noqa: E402  isort:skip
+    PrimusTurboDeepEPTokenDispatcher,
+    PrimusTurboLowPrecisionGlobalStateManager,
+    PrimusTurboQuantConfig,
+)
+
+
+def _dtypes(monkeypatch, *, flag=True, turbo_fp8=True, quant_config=None):
+    monkeypatch.setattr(PrimusTurboLowPrecisionGlobalStateManager, "PRIMUS_TURBO_FP8_ENABLED", turbo_fp8)
+    monkeypatch.setattr(PrimusTurboLowPrecisionGlobalStateManager, "PRIMUS_TURBO_QUANT_CONFIG", quant_config)
+    dispatcher = SimpleNamespace(turbo_fp8_permute=flag)
+    return PrimusTurboDeepEPTokenDispatcher._fp8_permute_dtypes(dispatcher)
+
+
+@pytest.mark.parametrize(
+    ("fmt", "expected"),
+    [
+        (Format.HYBRID, (float8_e4m3, float8_e5m2)),
+        (Format.E4M3, (float8_e4m3, float8_e4m3)),
+    ],
+)
+def test_tensorwise_returns_input_and_grad_dtypes(monkeypatch, fmt, expected):
+    quant_config = PrimusTurboQuantConfig(format=fmt, granularity=ScalingGranularity.TENSORWISE)
+    assert _dtypes(monkeypatch, quant_config=quant_config) == expected
+
+
+def test_flag_off_keeps_bf16_permute(monkeypatch):
+    quant_config = PrimusTurboQuantConfig(format=Format.HYBRID, granularity=ScalingGranularity.TENSORWISE)
+    assert _dtypes(monkeypatch, flag=False, quant_config=quant_config) == (None, None)
+
+
+def test_turbo_fp8_off_keeps_bf16_permute(monkeypatch):
+    """A layer outside Turbo FP8 (e.g. pinned to BF16) must not get FP8 tokens."""
+    quant_config = PrimusTurboQuantConfig(format=Format.HYBRID, granularity=ScalingGranularity.TENSORWISE)
+    assert _dtypes(monkeypatch, turbo_fp8=False, quant_config=quant_config) == (None, None)
+
+
+def test_blockwise_keeps_bf16_permute(monkeypatch):
+    quant_config = PrimusTurboQuantConfig(
+        format=Format.HYBRID, granularity=ScalingGranularity.BLOCKWISE, block_size=128
+    )
+    assert _dtypes(monkeypatch, quant_config=quant_config) == (None, None)