Fix ROCm float multiplication in sparse Marlin MMA

Update AMD GPU implementation to use __builtin_amdgcn_fmul_f32 instead of __builtin_amdgcn_fmul_legacy for more accurate float multiplication in the scale_floats function.

Author: petrex

torchao/csrc/cuda/sparse_marlin/mma.h

@@ -272,16 +272,16 @@ __device__ inline void scale_floats(float* c0, float* c1, float* c2, float* c3,
                                     FragS& s0, float* c4, float* c5, float* c6,
                                     float* c7, FragS& s1) {
   #ifdef USE_ROCM
-  // AMD implementation
-  *c0 = __builtin_amdgcn_fmul_legacy(*c0, __half2float(s0[0].x));
-  *c1 = __builtin_amdgcn_fmul_legacy(*c1, __half2float(s0[0].y));
-  *c2 = __builtin_amdgcn_fmul_legacy(*c2, __half2float(s0[1].x));
-  *c3 = __builtin_amdgcn_fmul_legacy(*c3, __half2float(s0[1].y));
+  // AMD implementation - fixed
+  *c0 = __builtin_amdgcn_fmul_f32(*c0, __half2float(s0[0].x));
+  *c1 = __builtin_amdgcn_fmul_f32(*c1, __half2float(s0[0].y));
+  *c2 = __builtin_amdgcn_fmul_f32(*c2, __half2float(s0[1].x));
+  *c3 = __builtin_amdgcn_fmul_f32(*c3, __half2float(s0[1].y));
 
-  *c4 = __builtin_amdgcn_fmul_legacy(*c4, __half2float(s1[0].x));
-  *c5 = __builtin_amdgcn_fmul_legacy(*c5, __half2float(s1[0].y));
-  *c6 = __builtin_amdgcn_fmul_legacy(*c6, __half2float(s1[1].x));
-  *c7 = __builtin_amdgcn_fmul_legacy(*c7, __half2float(s1[1].y));
+  *c4 = __builtin_amdgcn_fmul_f32(*c4, __half2float(s1[0].x));
+  *c5 = __builtin_amdgcn_fmul_f32(*c5, __half2float(s1[0].y));
+  *c6 = __builtin_amdgcn_fmul_f32(*c6, __half2float(s1[1].x));
+  *c7 = __builtin_amdgcn_fmul_f32(*c7, __half2float(s1[1].y));
   #else
   // NVIDIA implementation
   *c0 = __fmul_rn(*c0, __half2float(s0[0].x));