neural: Fix a sync issue (#10281)

Close #10272.

The root cause of this issue is the race condition. Add two syncs to
avoid this issue.
1. At beginning of the mma loop -> This prevent the fast warp start
writing the shared memory A in the i+1 iter when other slow warps are
still reading the shared memory A.

2. In the backward, after mma and before outerproduceAccumulate ->
becuase mma ends with warp-sync, so outerproduceAccumulate could start
executing for some fast warps while the slow warps are still running
mma.

This PR also fix a numerical issue in the activation functions.
Basically every function using exp() is not numerical stable, so update
them.

Author: kaizhangNV

source/standard-modules/neural/accelerate-vector-coopmat.slang

@@ -1044,6 +1044,11 @@ VISIBILITY_LEVEL struct MMAHelper<T, int InputSize, int OutputSize, int Subgroup
             }
             else
             {
+                // Ensure all warps finished reading shared memory from the previous
+                // tile iteration (or from the previous layer's output writeback in
+                // fused multi-layer kernels) before overwriting with new tile data.
+                GroupMemoryBarrierWithGroupSync();
+
                 loadShA<U, Address>(ptrAOffset[0], tileIndex, weightAddress);
                 loadVectorToShB(ptrBOffset[0], tileIndex, subgroupIndex, inputVector);
 
@@ -1271,6 +1276,13 @@ public struct WaveTangledVector<T, ShMemSize : ISharedMemorySize, int N, int Sub
         // outerProductAccumulate uses per-warp shared memory for both A (dOutput vectors)
         // and B (input vectors). The B region must start after ALL per-warp A regions to
         // avoid overlapping writes between warps (warp i's A would alias warp (i-1)'s B).
+        //
+        // Sync required: mma() above ends with warp-level sync only, but
+        // outerProductAccumulate reuses the same shared memory pool (starting at shA=0).
+        // Without a group sync, fast warps starting the outer product would corrupt
+        // slow warps still reading from mma's output writeback.
+        GroupMemoryBarrierWithGroupSync();
+
         static const int _outerRows = (MMA.M + MMA.CMShape.ROW_A - 1) / MMA.CMShape.ROW_A;
         static const int _outerPerWarpA = _outerRows * MMA.CMShape.ROW_A / MMA.CMShape.ElementCountPerVector * MMA.CMShape.COLUMN_A;
         uint shB_outer = shA + uint(getWaveCount() * _outerPerWarpA);
@@ -1284,6 +1296,8 @@ public struct WaveTangledVector<T, ShMemSize : ISharedMemorySize, int N, int Sub
                         >( doutput, shA, shB, shC, dWeightAddress.p, none);
         dthis = DifferentialPair<This>(dthis.p, dInput);
 
+        GroupMemoryBarrierWithGroupSync();
+
         static const int _outerRows2 = (MMA.M + MMA.CMShape.ROW_A - 1) / MMA.CMShape.ROW_A;
         static const int _outerPerWarpA2 = _outerRows2 * MMA.CMShape.ROW_A / MMA.CMShape.ElementCountPerVector * MMA.CMShape.COLUMN_A;
         uint shB_outer2 = shA + uint(getWaveCount() * _outerPerWarpA2);

source/standard-modules/neural/activations.slang

@@ -94,8 +94,10 @@ public struct Sigmoid<T> : IActivation<T>
         [ForceUnroll]
         for (int i = 0; i < Vector.Size; i++)
         {
+            // Numerically stable: exp() argument is always <= 0
             let x = input[i];
-            output[i] = T(1) / (T(1) + exp(-x));
+            let ex = exp(x >= T(0) ? -x : x);
+            output[i] = x >= T(0) ? T(1) / (T(1) + ex) : ex / (T(1) + ex);
         }
         return output;
     }
@@ -185,8 +187,10 @@ public struct SiLU<T> : IActivation<T>
         [ForceUnroll]
         for (int i = 0; i < Vector.Size; i++)
         {
+            // x * sigmoid(x), numerically stable: exp() argument is always <= 0
             let x = input[i];
-            output[i] = x / (T(1) + exp(-x));  // x * sigmoid(x)
+            let ex = exp(x >= T(0) ? -x : x);
+            output[i] = x >= T(0) ? x / (T(1) + ex) : x * ex / (T(1) + ex);
         }
         return output;
     }
@@ -212,8 +216,11 @@ public struct QuickGELU<T> : IActivation<T>
         [ForceUnroll]
         for (int i = 0; i < Vector.Size; i++)
         {
+            // x * sigmoid(1.702 * x), numerically stable: exp() argument is always <= 0
             let x = input[i];
-            output[i] = x / (T(1) + exp(T(-1.702) * x));  // x * sigmoid(1.702 * x)
+            let sx = T(1.702) * x;
+            let ex = exp(sx >= T(0) ? -sx : sx);
+            output[i] = sx >= T(0) ? x / (T(1) + ex) : x * ex / (T(1) + ex);
         }
         return output;
     }