import aicc-master-dev

csrc/flash_attn_hg/include/fwd/gfx938/pv_gemm_utils_mls_ds.h

@@ -34,9 +34,7 @@ __forceinline__ __device__ void prefetch_v_to_lds_mls_ds(
     int lds_offset = (lds_stage_id * WARP_K * kHeadDim_OPT + warp_id * 32 * 32) * ELEMENT_BYTES;
 
     flash::wait_all_warp_arrived(); // 防止写 v lds 和读 k lds 冲突, qk 可能有的 warp 没结束
-    union union_vec4_uint v_rsrc_bits;
-    v_rsrc_bits.v32 = v_srsrc;
-    size_t lds_addr_warp = reinterpret_cast<size_t>(v_lds) + lds_offset;
-    matrix_load_b16_lds_builtin<32, 32, 1, 0>(lds_addr_warp, v_rsrc_bits.i32, 0);
+    inline_matrix_load_32x32_b16_lds<0, 1>(v_lds, v_srsrc, lds_offset, 0);
     __builtin_amdgcn_sched_barrier(0);
 }
+