From c903ea2395cc4f8adb0c228661608ddb606b0909 Mon Sep 17 00:00:00 2001 From: xmpeng-dev Date: Sat, 10 Oct 2026 07:02:25 +0000 Subject: [PATCH 1/7] fix(turbo): make enable_turbo_attention_float8 runnable again Turbo FP8 attention failed on its first call for two reasons: - The core-attention call always passed ``sink=`` (added for sink attention), but ``flash_attn_fp8_func`` has no ``sink`` parameter, so every FP8 run raised a TypeError. Pass ``sink`` only when a sink tensor exists. - Megatron hands over sbhd storage viewed as bshd. ``flash_attn_fp8_func`` infers "sbhd" from the strides and then permutes the already-bshd logical tensor, swapping b and s (``shape '[4096, 32, 0, 64, 128]' is invalid``). Make q/k/v bshd-contiguous on the FP8 path until the Turbo fix lands. --- .../megatron/core/extensions/primus_turbo.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 6b8addc88..0f29cc87f 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/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, ) From f4c6e9110a51aeea14caae9e23e9095c9a1cd17e Mon Sep 17 00:00:00 2001 From: xmpeng-dev Date: Sat, 10 Oct 2026 07:02:25 +0000 Subject: [PATCH 2/7] perf(qwen3): tune Qwen3-30B-A3B FP8 MI355X config to Turbo tensorwise Switch the FP8 recipe to tensorwise with dense and grouped GEMMs on Primus-Turbo (FlyDSL grouped GEMM on gfx950) and the TE grouped-MLP path instead of the legacy grouped GEMM. The previous config hit a GPU memory fault (sum_and_scatter) at iteration 3 on the current image. 1 node x 8 MI355X, EP8, MBS 8, GBS 512, seq 4096, force-even routing: | config | ms/iter | TFLOP/s/GPU | tokens/s/GPU | | ------------------------------ | ------- | ----------- | ------------ | | BF16 config (legacy grouped) | 9587 | 622.1 | 27,344 | | FP8 config, before | memory fault at iter 3 | | FP8 config, this change | 8894 | 670.6 | 29,474 | --- .../megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml index a342485f4..4d710eab9 100644 --- a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml @@ -83,7 +83,8 @@ 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 # deepep @@ -105,5 +106,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 From f0a09523b9f19615c617bed0cba9bf8b4d855224 Mon Sep 17 00:00:00 2001 From: xmpeng-dev Date: Sat, 10 Oct 2026 10:00:17 +0000 Subject: [PATCH 3/7] perf(megatron): enable the fused packed-QKV RMSNorm + RoPE patch for head_dim 128 The PRIMUS_FUSED_QK_RMSNORM_ROPE=1 patch only fired for GPT-OSS: it required head_dim 64 and PrimusTurboRMSNorm Q/K norms. Qwen3-30B-A3B uses head_dim 128 and, with use_turbo_rms_norm false, TE RMSNorm, so it always fell back. - Accept any head_dim listed in Turbo's QK_RMSNORM_ROPE_HEAD_DIMS (64 and 128 with the companion Turbo change; falls back to (64,) on older Turbo). - Accept TE RMSNorm; PrimusTurboRMSNorm subclasses it and the fused path only reads weight / eps (zero_centered_gamma is still rejected). Qwen3-30B-A3B attention shape on MI355X (S=4096, B=8, 4 KV groups x 8 Q heads per group, D=128): split + Q/K RMSNorm + RoPE fwd+bwd 1243.6 -> 369.9 us (3.36x); 48 layers x 8 micro-batches is about -335 ms per iteration. End-to-end: Qwen3-30B-A3B FP8 tensorwise pretrain on 8x MI355X (EP8, MBS 8, GBS 512, seq 4096, even routing), mean of iterations 11-20 of 20-iteration runs; each row adds one change on top of the row above: ms/iter TFLOP/s tok/s/GPU mem loss@20 submitted FP8 tensorwise config (Turbo e2d9f1d7) 8894.0 670.6 29,474 277GB 11.33393 + use_turbo_fused_act_with_probs (Primus flag) 7915.6 753.5 33,117 223GB 11.33435 Turbo 1103b2df, base of these changes 7925.4 752.5 33,077 222GB 11.33427 + flat tensorwise FP8 quant kernel 7870.9 757.7 33,305 222GB 11.33415 + per-token get_dispatch_layout 7774.2 767.2 33,720 222GB 11.33431 + HIP permute in DeepEPTokenDispatcher 7643.5 780.3 34,296 222GB 11.33423 + turbo_deepep_num_cu 160 (Primus flag) 7353.4 811.1 35,649 223GB 11.33456 * + fused qk RMSNorm + RoPE, head_dim 128 6987.1 853.6 37,518 224GB 11.33373 This change (*): 7353.4 -> 6987.1 ms/iter (-4.98%), 35,649 -> 37,518 tokens/s/GPU (+5.24%). Loss differences are within run-to-run noise (about 4e-4 at iteration 20). Requires Primus-Turbo perf/flydsl/qk-rmsnorm-rope-hd128. --- .../patches/turbo/qk_rmsnorm_rope_patches.py | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py b/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py index a8193444c..a7aa6f151 100644 --- a/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py +++ b/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py @@ -4,7 +4,7 @@ # 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 @@ -12,8 +12,8 @@ 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" ) From 60ca6b386d77523dadd1e14bfaa846fba2276748 Mon Sep 17 00:00:00 2001 From: xmpeng-dev Date: Sat, 10 Oct 2026 12:49:58 +0000 Subject: [PATCH 4/7] perf(megatron): hand the Turbo fused grouped MLP FP8 tokens from the DeepEP permute New opt-in flag turbo_fp8_permute (default false). Under Primus-Turbo FP8 tensorwise current scaling with use_turbo_deepep and turbo_fused_grouped_gemm, PrimusTurboDeepEPTokenDispatcher asks Turbo's DeepEPTokenDispatcher to - quantize the dispatched tokens before permute and return the permuted tokens as a grouped tensorwise QuantizedTensor (_post_dispatch(quantize_dtype=...)), and - return the experts' output gradient as a tensorwise QuantizedTensor from un-permute's backward (_pre_combine(grad_quantize_dtype=...)), so the fused grouped MLP no longer quantizes the top-k times larger permuted copies itself. The dtypes follow the FP8 format (HYBRID: e4m3 input, e5m2 gradient); results are bit-identical to the default path. Both dtypes are derived in the call itself, so interleaved micro-batches need no dispatcher state. Any other recipe, or a layer with Turbo FP8 off, keeps the bf16 path. Argument validation requires enable_primus_turbo, use_turbo_deepep and turbo_fused_grouped_gemm, and rejects selective recompute of moe_act (that path checkpoints the fused MLP and is not validated with QuantizedTensor inputs and gradients). Requires Primus-Turbo perf/moe/fp8-permute-tensorwise (moe_permute quantize_dtype, moe_unpermute grad_quantize_dtype, QuantizedTensor grad_out in grouped_mlp_fp8). Qwen3-30B-A3B EP8 dispatcher call (32768 received tokens, hidden 2048, 16 local experts, top-8, pad_multiple 16), kernel time per call: forward 645 -> 167 us, backward 687 -> 174 us. End-to-end: Qwen3-30B-A3B FP8 tensorwise pretrain on 8x MI355X (EP8, MBS 8, GBS 512, seq 4096, even routing), mean of iterations 11-20 of 20-iteration runs; each row adds one change on top of the row above: ms/iter TFLOP/s tok/s/GPU mem loss@20 submitted FP8 tensorwise config (Turbo e2d9f1d7) 8894.0 670.6 29,474 277GB 11.33393 + use_turbo_fused_act_with_probs (Primus flag) 7915.6 753.5 33,117 223GB 11.33435 Turbo 1103b2df, base of these changes 7925.4 752.5 33,077 222GB 11.33427 + flat tensorwise FP8 quant kernel 7870.9 757.7 33,305 222GB 11.33415 + per-token get_dispatch_layout 7774.2 767.2 33,720 222GB 11.33431 + HIP permute in DeepEPTokenDispatcher 7643.5 780.3 34,296 222GB 11.33423 + turbo_deepep_num_cu 160 (Primus flag) 7353.4 811.1 35,649 223GB 11.33456 + fused qk RMSNorm + RoPE, head_dim 128 6987.1 853.6 37,518 224GB 11.33373 + turbo_fused_grouped_gemm (Primus flag) 6820.2 874.5 38,437 225GB 11.33412 * + turbo_fp8_permute 6442.4 925.8 40,690 223GB 11.33404 Repeat runs of the last two rows: turbo_fused_grouped_gemm alone 6820.2 and 6815.1 ms; with turbo_fp8_permute 6442.4, 6515.0, 6523.8 and 6549.1 ms (the later runs show more per-iteration jitter; the cause is not identified). This change (*), means of those runs: 6817.7 -> 6507.6 ms/iter (-4.55%), 38,451 -> 40,283 tokens/s/GPU (+4.76%). Loss differences are within run-to-run noise (about 4e-4 at iteration 20). Tests: test_rocm_arg_validation.py covers validate_turbo_fp8_permute (each required flag, selective recompute of moe_act, disabled no-op); test_turbo_fp8_permute_dtypes.py checks _fp8_permute_dtypes returns (e4m3, e5m2) for HYBRID and (e4m3, e4m3) for E4M3 tensorwise, and (None, None) with the flag off, Turbo FP8 off or blockwise scaling. --- .../megatron/core/extensions/primus_turbo.py | 37 ++++++++-- .../patches/args/rocm_arg_validation.py | 14 ++++ .../modules/megatron/primus_turbo.yaml | 5 ++ .../megatron/test_rocm_arg_validation.py | 47 +++++++++++++ .../megatron/test_turbo_fp8_permute_dtypes.py | 69 +++++++++++++++++++ 5 files changed, 168 insertions(+), 4 deletions(-) create mode 100644 tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 0f29cc87f..b5d5716ab 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -2640,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 @@ -2721,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 @@ -2734,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( diff --git a/primus/backends/megatron/patches/args/rocm_arg_validation.py b/primus/backends/megatron/patches/args/rocm_arg_validation.py index c00621a4e..c8c189e4c 100644 --- a/primus/backends/megatron/patches/args/rocm_arg_validation.py +++ b/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: diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index 437685e0e..52afd8f54 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/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 diff --git a/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py b/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py index dd52c5bd3..bcb868b12 100644 --- a/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py +++ b/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)) diff --git a/tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py b/tests/unit_tests/backends/megatron/test_turbo_fp8_permute_dtypes.py new file mode 100644 index 000000000..4fc7e3b12 --- /dev/null +++ b/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) From 16ae0307a6d35186a56518af35936a176113c232 Mon Sep 17 00:00:00 2001 From: sim Date: Sat, 10 Oct 2026 15:38:59 +0000 Subject: [PATCH 5/7] perf(qwen3): turn on the Turbo fusions in the Qwen3-30B-A3B FP8 MI355X config Turn on every switch behind the 40k tokens/s/GPU Qwen3-30B-A3B FP8 tensorwise pretrain on 8x MI355X: - env PRIMUS_FUSED_QK_RMSNORM_ROPE=1 (top-level env:): fused packed-QKV split + Q/K RMSNorm + RoPE for head_dim 128; - turbo_fused_grouped_gemm and use_turbo_fused_act_with_probs; - turbo_fp8_permute: the DeepEP permute hands the fused grouped MLP FP8 tokens and takes its FP8 output gradient; - turbo_deepep_num_cu 80 -> 160 (192 gives a NaN loss on one rank at iteration 5, 256 is slower). turbo_fp8_permute needs a Primus-Turbo whose moe_permute takes quantize_dtype (AMD-AGI/Primus-Turbo#555); the head_dim 128 fusion needs AMD-AGI/Primus-Turbo#556 and falls back to the unfused path without it. End-to-end: Qwen3-30B-A3B FP8 tensorwise pretrain on 8x MI355X (EP8, MBS 8, GBS 512, seq 4096, even routing), mean of iterations 11-20 of 20-iteration runs; each row adds one change on top of the row above: ms/iter TFLOP/s tok/s/GPU mem loss@20 FP8 tensorwise config (Turbo e2d9f1d7) 8894.0 670.6 29,474 277GB 11.33393 + use_turbo_fused_act_with_probs 7915.6 753.5 33,117 223GB 11.33435 Turbo 1103b2df 7925.4 752.5 33,077 222GB 11.33427 + flat tensorwise FP8 quant kernel (Turbo) 7870.9 757.7 33,305 222GB 11.33415 + per-token get_dispatch_layout (Turbo) 7774.2 767.2 33,720 222GB 11.33431 + HIP permute in DeepEPTokenDispatcher (Turbo) 7643.5 780.3 34,296 222GB 11.33423 + turbo_deepep_num_cu 160 7353.4 811.1 35,649 223GB 11.33456 + fused qk RMSNorm + RoPE, head_dim 128 6987.1 853.6 37,518 224GB 11.33373 + turbo_fused_grouped_gemm (mean of 2 runs) 6817.7 874.9 38,451 225GB 11.33412 + turbo_fp8_permute (mean of 4 runs) 6507.6 916.7 40,283 223GB 11.33404 This branch on main 0a68e5cd with the switches on: 6518.7 ms/iter with the switches on the command line, 6410.6 and 6459.6 ms/iter from this config alone (40,214 / 40,892 / 40,582 tokens/s/GPU). Total: 8894.0 -> 6507.6 ms/iter (-26.8%), 29,474 -> 40,283 tokens/s/GPU (+36.7%). Loss differences are within run-to-run noise (about 4e-4 at iteration 20). Known issue: one more config-only run hit a NaN forward loss on one rank at iteration 3; its parsed arguments match a passing run exactly. Across this work, 1 of 13 runs at 160 DeepEP CUs and the only run at 192 hit this; none of the 5 runs at 80 CUs did, but those predate the fused paths, so the cause is not isolated yet. --- .../configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml index 4d710eab9..0586e10d2 100644 --- a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml @@ -3,6 +3,11 @@ 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. +env: + PRIMUS_FUSED_QK_RMSNORM_ROPE: "1" + modules: pre_trainer: framework: megatron @@ -86,14 +91,18 @@ modules: 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 From b8d708f8d35af49ad5a1e58c27de527b10784c09 Mon Sep 17 00:00:00 2001 From: xiaoming-amd Date: Mon, 12 Oct 2026 02:40:06 -0500 Subject: [PATCH 6/7] refactor(megatron): gate the fused QK RMSNorm + RoPE patch on use_turbo_fused_qk_rmsnorm_rope Replace the PRIMUS_FUSED_QK_RMSNORM_ROPE environment switch with a use_turbo_fused_qk_rmsnorm_rope flag in primus_turbo.yaml (default false), like the other Turbo switches. The Qwen3-30B-A3B FP8 MI355X config sets the flag instead of a top-level env block. --- .../MI355X/qwen3_30B_A3B-FP8-pretrain.yaml | 7 ++---- .../patches/turbo/qk_rmsnorm_rope_patches.py | 13 +++++------ .../modules/megatron/primus_turbo.yaml | 4 ++++ .../megatron/test_qk_rmsnorm_rope_patches.py | 22 +++++++++++++++++++ 4 files changed, 34 insertions(+), 12 deletions(-) diff --git a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml index 0586e10d2..ecb5a9496 100644 --- a/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen3_30B_A3B-FP8-pretrain.yaml @@ -3,11 +3,6 @@ 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. -env: - PRIMUS_FUSED_QK_RMSNORM_ROPE: "1" - modules: pre_trainer: framework: megatron @@ -88,6 +83,8 @@ modules: # Turbo enable_primus_turbo: true use_turbo_attention: false + # head_dim 128 needs a Primus-Turbo with QK_RMSNORM_ROPE_HEAD_DIMS; older Turbo keeps the unfused path + use_turbo_fused_qk_rmsnorm_rope: true use_turbo_gemm: true use_turbo_grouped_gemm: true use_turbo_rms_norm: false # bug diff --git a/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py b/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py index a7aa6f151..c19dcf396 100644 --- a/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py +++ b/primus/backends/megatron/patches/turbo/qk_rmsnorm_rope_patches.py @@ -6,7 +6,7 @@ """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 +``use_turbo_fused_qk_rmsnorm_rope: true`` replaces the training-only sequence split packed QKV -> Q/K RMSNorm -> Q/K RoPE @@ -18,10 +18,8 @@ from __future__ import annotations -import os - from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch -from primus.core.patches import PatchContext, register_patch +from primus.core.patches import PatchContext, get_args, register_patch from primus.core.utils.module_utils import log_rank_0 _DDP_PARAM_GATHER_HOOK_ATTR = "_primus_ddp_param_gather_hook" @@ -70,9 +68,10 @@ def _run_forward_pre_hooks(module, hooks) -> None: raise RuntimeError("DDP parameter-gather pre-hook unexpectedly modified its inputs") -def _enabled(_ctx: PatchContext) -> bool: - value = os.environ.get("PRIMUS_FUSED_QK_RMSNORM_ROPE", "0").strip().lower() - return value in ("1", "true", "yes", "on") and is_primus_turbo_can_patch(_ctx) +def _enabled(ctx: PatchContext) -> bool: + args = get_args(ctx) + use_fused = bool(getattr(args, "use_turbo_fused_qk_rmsnorm_rope", False)) + return use_fused and is_primus_turbo_can_patch(ctx) @register_patch( diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index 52afd8f54..72598a59e 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/primus/configs/modules/megatron/primus_turbo.yaml @@ -45,6 +45,10 @@ sink_sliding_window: 0 sink_window_even_layers_only: true # attention float8 path enable_turbo_attention_float8: false +# Fuse the packed-QKV split, Q/K RMSNorm and RoPE into one Primus-Turbo FlyDSL op +# (training, TP=1, SBHD, head_dim supported by the installed Turbo). Layers that do +# not qualify keep the Megatron path. +use_turbo_fused_qk_rmsnorm_rope: false # ===== Parallel Linear ===== # operator switch diff --git a/tests/unit_tests/backends/megatron/test_qk_rmsnorm_rope_patches.py b/tests/unit_tests/backends/megatron/test_qk_rmsnorm_rope_patches.py index 6192f6e09..723b6f70f 100644 --- a/tests/unit_tests/backends/megatron/test_qk_rmsnorm_rope_patches.py +++ b/tests/unit_tests/backends/megatron/test_qk_rmsnorm_rope_patches.py @@ -5,15 +5,37 @@ ############################################################################### from collections import OrderedDict +from types import SimpleNamespace import pytest +from primus.backends.megatron.patches.turbo import qk_rmsnorm_rope_patches from primus.backends.megatron.patches.turbo.qk_rmsnorm_rope_patches import ( _DDP_PARAM_GATHER_HOOK_ATTR, + _enabled, _install_ddp_param_gather_hook_marker, _run_forward_pre_hooks, _supported_forward_pre_hooks, ) +from primus.core.patches import PatchContext + + +@pytest.mark.parametrize( + "use_flag, turbo_can_patch, expected", + [(True, True, True), (False, True, False), (True, False, False)], +) +def test_enabled_follows_use_turbo_fused_qk_rmsnorm_rope(monkeypatch, use_flag, turbo_can_patch, expected): + monkeypatch.setattr(qk_rmsnorm_rope_patches, "is_primus_turbo_can_patch", lambda _ctx: turbo_can_patch) + monkeypatch.setenv("PRIMUS_FUSED_QK_RMSNORM_ROPE", "1") + ctx = PatchContext( + backend="megatron", + phase="before_train", + extra={ + "module_config": SimpleNamespace(params=SimpleNamespace(use_turbo_fused_qk_rmsnorm_rope=use_flag)) + }, + ) + + assert _enabled(ctx) is expected class _FakeModule: From 9dc4523aa931812f7cd4544b52b1c33b82eed352 Mon Sep 17 00:00:00 2001 From: xiaoming-amd Date: Mon, 12 Oct 2026 02:58:49 -0500 Subject: [PATCH 7/7] ci: bump Primus-Turbo to 67b8e093 for the Qwen3 FP8 fusions turbo_fp8_permute needs moe_permute(quantize_dtype=...) (#555), and the head_dim 128 fused QK RMSNorm + RoPE needs #556. Turbo main 67b8e093 also carries the MXFP4 gather fusion (#546) that the previous pin 8f19706a provided, as the squash merge of the same branch. --- .github/workflows/ci.yaml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 95110681d..196226263 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -41,8 +41,9 @@ env: PRIMUS_CI_FULL: ${{ (github.event_name == 'schedule' || github.event.inputs.full_tests == 'true') && '1' || '0' }} # Includes fused expert wgrad output overwrite support (#534), MegaMoE-owned FlyDSL # kernels (#525), the repaired DeepEP intranode cheap fence, on by default (#539), - # and MegaMoE SwiGLU-OAI activations (#538). - PRIMUS_TURBO_COMMIT: 8f19706aa1c656b11a146eb00db32a6c2b46822d + # MegaMoE SwiGLU-OAI activations (#538), MXFP4 gather fusion (#546), tensorwise FP8 + # from the TURBO permute (#555) and head_dim 128 fused QK RMSNorm + RoPE (#556). + PRIMUS_TURBO_COMMIT: 67b8e0930e38f837d145ee89d44dedda034838fa PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py #ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355) UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710)