fix: work around Triton buffer-store miscompile and ROCm compiled-training NaNs

docs/03-configuration-reference/environment-variables.md

@@ -121,6 +121,7 @@ Primus seeds many of these in `runner/helpers/envs/base_env.sh`. RCCL honors NCC
 | `HSA_NO_SCRATCH_RECLAIM` | `1` | `base_env.sh`, container passthrough | ROCm runtime; documented for MoE stability | `1` keeps scratch allocated (often used for MoE stability). See [ROCR environment](https://rocm.docs.amd.com/projects/ROCR-Runtime/en/docs-7.1.1/environment_variables.html). |
 | `HIP_VISIBLE_DEVICES` | `0..GPUS_PER_NODE-1` | `base_env.sh` | ROCm device visibility | Restricts which GPU indices ROCm exposes. |
 | `ROCBLAS_DEFAULT_ATOMICS_MODE` | (unset) | User | `primus/backends/megatron/patches/args/rocm_arg_validation.py` | Read for deterministic / accuracy-sensitive GEMM behavior. |
+| `TORCH_COMPILE_CACHE_KEY_TAG` | (unset) | `primus/core/patches/triton_bufops_war_patches.py`, on ROCm only | torch.compile caching | Gets `primus-bufops-war-v<N>` appended so graphs compiled before (or under a different version of) the workaround are not reused. An existing value is appended to, not replaced. |
 
 ---
 

docs/04-technical-guides/performance-tuning.md

@@ -233,6 +233,35 @@ From `trainer_base.yaml` and model settings:
 
 ---
 
+## 8. Known ROCm correctness workarounds
+
+The AMD Triton backend can emit a `buffer_store_dwordx4` whose data registers are redefined by a
+later instruction with no `s_waitcnt vmcnt` in between, so the store writes whatever the
+clobbering instruction left behind. Nothing faults and nothing warns; a few percent of the
+kernel's output elements simply hold garbage. Inductor's SwiGLU backward fusion trips it and
+poisons the MLP weight gradients, which diverges TorchTitan training. Megatron is exposed to the
+same kernels, since its fused activations are inductor-generated as well (`megatron/core/jit.py`
+sets `jit_fuser = torch.compile` on torch 2.2+).
+
+`primus/core/patches/triton_bufops_war_patches.py` handles this automatically on ROCm: every
+kernel compiles normally, and only those whose emitted AMDGCN contains the hazard are recompiled
+with buffer ops disabled. It has no switch: the recompile is a correctness fix, and a kernel that
+does not need one is left alone anyway.
+
+Two things about the design are deliberate. It selects per kernel rather than setting
+`AMDGCN_USE_BUFFER_OPS=0` globally, because the hazardous kernels lose nothing without buffer
+addressing while the Primus-Turbo grouped-GEMM kernels spill without it and carry no hazard — the
+global switch costs ~18% on MoE recipes for kernels that never needed fixing. And it decides from
+the emitted machine code rather than from a list of affected architectures, so a newly shipped
+architecture cannot silently fall out of coverage.
+
+It also appends to `TORCH_COMPILE_CACHE_KEY_TAG`, because on an FX graph cache hit inductor never
+calls `triton.compile`; without that, a cache filled before the patch existed would keep serving
+hazardous binaries. Note that every Triton build tried so far still emits the pattern, so this is
+not something an image upgrade removes.
+
+---
+
 ## Quick checklist
 
 1. **GEMMs:** run HipBLASLt stages 1–3 or use `offline_tune_gemm.py` for custom workflows.

primus/backends/torchtitan/patches/__init__.py

@@ -31,6 +31,7 @@
     embedding_amp_patches,
     flex_attention_patches,
     fsdp_weight_tying_patches,
+    inductor_precision_casts_patches,
     logger_patches,
     metrics_output_format,
     mock_dataset_patches,

primus/backends/torchtitan/patches/dsv3_v022_perf_patches.py

@@ -10,14 +10,22 @@
 Reduce DSv3-16B HBM on MI355X with torchtitan v0.2.2 without modifying the
 upstream submodule. The balanced-routing field migration
 (``training.debug_moe_force_load_balance`` -> ``debug.moe_force_load_balance``)
-is already handled in the DeepSeek configs on main, so this module only carries
-the two memory fixes:
-
-1. Replace v0.2.2 ``apply_compile`` (capture_scalar_outputs + per-submodule
-   compile) with whole-block compile matching v0.1.0 behavior (drops the
-   fragmented reserved HBM re-introduced by per-submodule compilation).
-2. Replace MoE combine fp32 ``bmm`` with a bf16 weighted sum (drops the fp32
-   activation copy retained across the MoE layers).
+is already handled in the DeepSeek configs on main.
+
+Whole-block compilation is still needed to avoid the fragmented reserved HBM
+caused by v0.2.2's per-submodule compilation. Compiling the entire MoE forward
+is unsafe, however: it pulls expert-parallel token dispatch/combine and
+GroupedExperts FSDP hooks into one inductor graph. On MI300X that silently
+produces a NaN forward loss from the first step.
+
+Use a safe dense-only boundary instead: dense blocks compile as a full graph,
+while MoE TransformerBlocks remain eager. Compiling even the outer MoE block
+with a graph break around ``MoE.forward`` still produces the NaN, whereas
+v0.2.2's per-submodule policy reintroduces the fragmented reserved HBM this
+patch was created to avoid. Keeping MoE blocks eager avoids both failure modes;
+the expensive attention and grouped-GEMM work still uses Primus-Turbo kernels.
+The MoE forward replacement also changes its fp32 ``bmm`` combine into a bf16
+weighted sum, dropping the fp32 activation copy retained across MoE layers.
 """
 
 from __future__ import annotations
@@ -52,36 +60,38 @@ def _compile_enabled(ctx: PatchContext) -> bool:
     return bool(get_param(ctx, "compile.enable", False))
 
 
+def _apply_dense_only_compile(model: nn.Module, compile_config: Any, ep_enabled: bool) -> None:
+    """Whole-block compile dense layers and leave MoE TransformerBlocks eager."""
+    from torchtitan.tools.logging import logger
+
+    del ep_enabled  # MoE blocks stay eager, including EP dispatch.
+    for layer_id, transformer_block in model.layers.named_children():
+        if not transformer_block.moe_enabled:
+            transformer_block = torch.compile(
+                transformer_block,
+                backend=compile_config.backend,
+                fullgraph=True,
+            )
+        model.layers.register_module(layer_id, transformer_block)
+
+    logger.info("Compiling dense TransformerBlocks; leaving MoE TransformerBlocks eager (Primus patch)")
+
+
 @register_patch(
     "torchtitan.dsv3.whole_block_compile",
     backend="torchtitan",
     phase="setup",
-    description="Whole-block torch.compile; leave capture_scalar_outputs disabled",
+    description="Whole-block compile dense layers; leave MoE TransformerBlocks eager",
     condition=lambda ctx: _is_deepseek_model(ctx) and _compile_enabled(ctx),
 )
 def patch_whole_block_compile(ctx: PatchContext) -> None:
-    """Replace v0.2.2 apply_compile with v0.1.0-style whole TransformerBlock compile."""
-    from torchtitan.config.job_config import Compile as CompileConfig
+    """Install low-fragment dense-only compilation."""
     from torchtitan.models.llama4.infra import parallelize as parallelize_module
-    from torchtitan.tools.logging import logger
-
-    def apply_compile_patched(model: nn.Module, compile_config: CompileConfig, ep_enabled: bool) -> None:
-        # Match v0.1.0: do NOT set capture_scalar_outputs=True (avoids fragmentation).
-        for layer_id, transformer_block in model.layers.named_children():
-            fullgraph = not transformer_block.moe_enabled
-            transformer_block = torch.compile(
-                transformer_block,
-                backend=compile_config.backend,
-                fullgraph=fullgraph,
-            )
-            model.layers.register_module(layer_id, transformer_block)
-
-        logger.info("Compiling each TransformerBlock with torch.compile (Primus whole-block patch)")
 
-    parallelize_module.apply_compile = apply_compile_patched
+    parallelize_module.apply_compile = _apply_dense_only_compile
     log_rank_0(
         "[Patch:torchtitan.dsv3.whole_block_compile] "
-        "Patched torchtitan.models.llama4.infra.parallelize.apply_compile",
+        "Patched apply_compile with eager MoE TransformerBlocks",
     )
 
 
@@ -143,8 +153,10 @@ def patch_moe_bf16_combine(ctx: PatchContext) -> None:
     """Replace MoE.forward to drop the fp32 bmm copy in the combine step."""
     import torchtitan.models.moe.moe as moe_module
 
-    moe_module.MoE.forward = _moe_forward_bf16_combine
+    # Keep this explicit eager boundary in case a caller wraps a larger parent
+    # module in torch.compile outside the DeepSeek parallelization path.
+    moe_module.MoE.forward = torch.compiler.disable(_moe_forward_bf16_combine)
     log_rank_0(
         "[Patch:torchtitan.dsv3.moe_bf16_combine] "
-        "Patched torchtitan.models.moe.moe.MoE.forward (bf16 weighted combine)",
+        "Patched MoE.forward (eager boundary, bf16 weighted combine)",
     )

primus/backends/torchtitan/patches/inductor_precision_casts_patches.py

@@ -0,0 +1,53 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Preserve eager low-precision casts in ROCm Inductor graphs.
+
+On the pinned ROCm torch 2.12 build, compiled FSDP BF16 llama/qwen recipes
+produce a finite first loss but NaN gradients unless
+``emulate_precision_casts`` is enabled. Eager and ``aot_eager`` remain finite.
+
+This is not attributed to pytorch#150859: that issue covered float8 rowwise
+training with selective activation checkpointing and is no longer reproducible
+on upstream 2.12 nightly. The plain-BF16 ROCm reproducer here is different even
+though the same compatibility flag fixes it.
+
+The flag is part of Inductor's cache key, so no manual cache invalidation is
+required.
+"""
+
+from primus.core.patches import PatchContext, get_param, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_PREFIX = "[Patch:torchtitan.torch.inductor_precision_casts]"
+
+
+def _compile_enabled(ctx: PatchContext) -> bool:
+    if not bool(get_param(ctx, "compile.enable", False)):
+        return False
+
+    import torch
+
+    return torch.version.hip is not None
+
+
+@register_patch(
+    "torchtitan.torch.inductor_precision_casts",
+    backend="torchtitan",
+    phase="setup",  # before parallelize_fn reaches apply_compile
+    description="Emulate eager precision casts in ROCm Inductor BF16/FSDP graphs",
+    condition=_compile_enabled,
+)
+def patch_inductor_precision_casts(ctx: PatchContext) -> None:
+    """Make inductor round intermediates the way eager does."""
+    import torch._inductor.config as inductor_config
+
+    if inductor_config.emulate_precision_casts:
+        log_rank_0(f"{_PREFIX} already enabled; leaving it alone")
+        return
+
+    inductor_config.emulate_precision_casts = True
+    log_rank_0(f"{_PREFIX} torch._inductor.config.emulate_precision_casts = True")

primus/backends/torchtitan/patches/turbo/__init__.py

@@ -13,6 +13,7 @@
 """
 
 from primus.backends.torchtitan.patches.turbo import (  # noqa: F401
+    aiter_mla_capability_patches,
     async_tp_patches,
     attention_patches,
     deepseek_v3_classic_attention_patches,

primus/backends/torchtitan/patches/turbo/aiter_mla_capability_patches.py

@@ -0,0 +1,117 @@
+###############################################################################
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Route the failing gfx942 DeepSeek MLA backward directly to aiter CK."""
+
+from typing import Any
+
+from primus.backends.torchtitan.patches.turbo.attention_safety import (
+    requires_ck_mla_backward,
+)
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+_LOG_PREFIX = "[Patch:torchtitan.primus_turbo.aiter_mla_capability]"
+_WRAPPED_ATTR = "_primus_gfx942_mla_ck_backward"
+
+
+def _can_patch(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    return requires_ck_mla_backward(ctx.model_name, args.primus_turbo)
+
+
+def _is_gfx942_tensor(tensor) -> bool:
+    try:
+        import torch
+
+        arch = torch.cuda.get_device_properties(tensor.device).gcnArchName or ""
+        return "gfx942" in arch
+    except Exception:  # pragma: no cover - defensive
+        return False
+
+
+def _should_use_ck(kwargs: dict[str, Any]) -> bool:
+    q = kwargs["q"]
+    v = kwargs["v"]
+    return bool(
+        kwargs.get("sink") is None
+        and kwargs.get("qkv_format", "bshd") == "bshd"
+        and q.ndim == 4
+        and v.ndim == 4
+        and q.shape[-1] in (128, 192)
+        and v.shape[-1] == 128
+        and kwargs["dropout_p"] == 0.0
+        and kwargs.get("bias") is None
+        and kwargs.get("alibi_slopes") is None
+        and kwargs.get("dbias") is None
+        and _is_gfx942_tensor(q)
+    )
+
+
+def _run_ck_backward(mha, kwargs: dict[str, Any]):
+    softmax_d = mha.mha_bwd(
+        kwargs["dout"],
+        kwargs["q"],
+        kwargs["k"],
+        kwargs["v"],
+        kwargs["out"],
+        kwargs["softmax_lse"],
+        kwargs["dropout_p"],
+        kwargs["softmax_scale"],
+        kwargs["causal"],
+        kwargs["window_size_left"],
+        kwargs["window_size_right"],
+        kwargs["deterministic"],
+        kwargs["dq"],
+        kwargs["dk"],
+        kwargs["dv"],
+        kwargs["dbias"],
+        kwargs["bias"],
+        kwargs["alibi_slopes"],
+        kwargs["rng_state"],
+        None,
+        kwargs["sink"],
+        kwargs["dsink"],
+    )
+    return (
+        softmax_d,
+        kwargs["dq"],
+        kwargs["dk"],
+        kwargs["dv"],
+        kwargs["dbias"],
+        kwargs["dsink"],
+    )
+
+
+def _make_execute_wrapper(original, mha):
+    def execute(*args, **kwargs):
+        if not args and _should_use_ck(kwargs):
+            return _run_ck_backward(mha, kwargs)
+        return original(*args, **kwargs)
+
+    setattr(execute, _WRAPPED_ATTR, True)
+    return execute
+
+
+@register_patch(
+    "torchtitan.primus_turbo.aiter_mla_capability",
+    backend="torchtitan",
+    phase="setup",
+    description="Route failing gfx942 DeepSeek MLA backward from v3 to aiter CK",
+    condition=_can_patch,
+    priority=51,
+)
+def patch_aiter_mla_capability(ctx: PatchContext) -> None:
+    import aiter.ops.mha as mha
+    from primus_turbo.pytorch.kernels.attention import attention_aiter_impl
+
+    backend = attention_aiter_impl.AttnBwdAiterBackend
+    original = backend.execute
+    if getattr(original, _WRAPPED_ATTR, False):
+        return
+
+    backend.execute = staticmethod(_make_execute_wrapper(original, mha))
+    log_rank_0(f"{_LOG_PREFIX} routing gfx942 DeepSeek MLA backward directly to aiter CK")

primus/backends/torchtitan/patches/turbo/attention_patches.py

@@ -15,19 +15,30 @@
 expressed as a backend patch so it can be managed via the Primus patch system.
 """
 
-from primus.core.patches import PatchContext, get_param, register_patch
+from primus.backends.torchtitan.patches.turbo.attention_safety import (
+    requires_ck_mla_backward,
+    should_use_turbo_attention,
+)
+from primus.core.patches import PatchContext, get_args, register_patch
 from primus.core.utils.module_utils import log_rank_0
 
 
+def _can_patch_turbo_attention(ctx: PatchContext) -> bool:
+    args = get_args(ctx)
+    if requires_ck_mla_backward(ctx.model_name, args.primus_turbo):
+        log_rank_0(
+            "[Patch:torchtitan.primus_turbo.turbo_attention] "
+            "Using Turbo attention with CK MLA backward on gfx942."
+        )
+    return should_use_turbo_attention(ctx.model_name, args.primus_turbo)
+
+
 @register_patch(
     "torchtitan.primus_turbo.turbo_attention",
     backend="torchtitan",
     phase="setup",
     description="Use Primus-Turbo Attention kernels for supported models",
-    condition=lambda ctx: (
-        get_param(ctx, "primus_turbo.enable_primus_turbo", False)
-        and get_param(ctx, "primus_turbo.use_turbo_attention", False)
-    ),
+    condition=_can_patch_turbo_attention,
 )
 def patch_turbo_attention(ctx: PatchContext) -> None:
     """