Fix race condition in mxfp8 CUDA kernels (#4278)

A race condition was present in two kernels due to a bad ordering between a syncthreads and an async-proxy fence. The fence is needed because it makes sure that the calling thread's writes to shmem are visible in the async proxy. However, the operation we're synchronizing with is the TMA write issues by thread 0, hence we need to establish a causality link between _all_ the fences performed by _all_ threads, and the issuing of the TMA load by thread 0. Thus the syncthreads must be inserted in between these two operations.

The CUDA programming guide is very explicit about this, in [section 10.29.1. Using TMA to transfer one-dimensional arrays](https://docs.nvidia.com/cuda/cuda-c-programming-guide/#using-tma-to-transfer-one-dimensional-arrays).

Author: lw

torchao/csrc/cuda/mx_kernels/mx_block_rearrange_2d_M_groups.cu

@@ -316,10 +316,9 @@ __global__ void mx_blocked_layout_2d_M_groups_kernel(
                 ) = data;
             }
         }
-        __syncthreads();
-
         // Fence to ensure SMEM writes are visible to TMA async proxy
         ptx::fence_proxy_async_shared_cta();
+        __syncthreads();
 
         // Compute output tile coordinates
         const int chunk_sf_row_tile = (superblock_idx_in_group * CHUNKS_PER_TB) + chunk_idx;
@@ -785,8 +784,8 @@ __global__ void mx_blocked_layout_2d_simple_kernel(
     }
 
     // Ensure threads finish their smem writes and use explicit fence to ensure visibility to async proxy for TMA
-    __syncthreads();
     ptx::fence_proxy_async_shared_cta();
+    __syncthreads();
 
     if (is_master_thread) {
         // Issue separate 1D TMA stores for each valid SF tile