Use relaxed memory ordering for Triton atomics on AMDGPU. (#3945)

Author: wenchenvincent

torchao/prototype/common/triton/matmul.py

@@ -259,7 +259,11 @@ def _kernel(
     if SPLIT_K == 1:
         tl.store(C, acc, mask=mask)
     else:
-        tl.atomic_add(C, acc, mask=mask)
+        # AMD GPUs need relaxed semantics for better performance
+        if tl.constexpr(torch.version.hip is not None):
+            tl.atomic_add(C, acc, mask=mask, sem="relaxed")
+        else:
+            tl.atomic_add(C, acc, mask=mask)
 
 
 class _matmul(torch.autograd.Function):

torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise.py

@@ -473,7 +473,11 @@ def _amax_atomic(
     block_mask = block_offs < num_elements
     vals = tl.load(input_ptr + block_offs, mask=block_mask).to(input_dtype)
     block_amax = tl.max(tl.abs(vals))
-    tl.atomic_max(amax_ptr, block_amax)
+    # AMD GPUs need relaxed semantics for better performance
+    if tl.constexpr(torch.version.hip is not None):
+        tl.atomic_max(amax_ptr, block_amax, sem="relaxed")
+    else:
+        tl.atomic_max(amax_ptr, block_amax)
 
 
 @triton.jit

torchao/prototype/hqq/kernels.py

@@ -3,6 +3,7 @@
 #
 # This source code is licensed under the BSD 3-Clause license found in the
 # LICENSE file in the root directory of this source tree.
+import torch
 import triton
 import triton.language as tl
 from triton import Config
@@ -389,7 +390,11 @@ def _mixed_mm_kernel(
     if SPLIT_K == 1:
         tl.store(C, acc, mask=mask)
     else:
-        tl.atomic_add(C, acc, mask=mask)
+        # AMD GPUs need relaxed semantics for better performance
+        if tl.constexpr(torch.version.hip is not None):
+            tl.atomic_add(C, acc, mask=mask, sem="relaxed")
+        else:
+            tl.atomic_add(C, acc, mask=mask)
 
 
 _mixed_mm = triton.heuristics(MIXED_MM_HEURISTICS)(_mixed_mm_kernel)

torchao/prototype/moe_training/kernels/float8_rowwise.py

@@ -207,7 +207,16 @@ def _triton_fp8_rowwise_3d_transpose_scales_rhs_kernel(
             + k_offs[None, :] * stride_scales_dim1
         )
         scales_mask = k_offs[None, :] < K
-        tl.atomic_min(scales_ptr + scales_offs, scales[None, :], mask=scales_mask)
+        # AMD GPUs need relaxed semantics for better performance
+        if tl.constexpr(torch.version.hip is not None):
+            tl.atomic_min(
+                scales_ptr + scales_offs,
+                scales[None, :],
+                mask=scales_mask,
+                sem="relaxed",
+            )
+        else:
+            tl.atomic_min(scales_ptr + scales_offs, scales[None, :], mask=scales_mask)
 
     @triton.autotune(configs=atomic_kernel_configs_2D, key=["num_elements"])
     @triton.jit