diff --git a/cpp/kernels/xqa/mha_sm90.cu b/cpp/kernels/xqa/mha_sm90.cu index 5b14f37aea0..038daef393d 100644 --- a/cpp/kernels/xqa/mha_sm90.cu +++ b/cpp/kernels/xqa/mha_sm90.cu @@ -2078,9 +2078,13 @@ __device__ inline RegColWiseVec loadGmemColWiseVecWithDup(ShmQWiseVec const& gme for (uint32_t i = 0; i < exactDiv(ShmQWiseVec::size, gmma::instNBase); i++) { static_assert(nbThrdsPerInstNBase * RegColWiseVec::size == exactDiv(ShmQWiseVec::size, GmmaAccCoreMat::cols)); - ret[i] = reinterpret_cast< - Vec, exactDiv(ShmQWiseVec::size, GmmaAccCoreMat::cols)> const&>( - gmemVec)[mha::min(i * nbThrdsPerInstNBase + idx, bound)]; + uint32_t const clampedIdx = mha::min(i * nbThrdsPerInstNBase + idx, bound); + uint32_t const baseOffset = clampedIdx * GmmaAccCoreMat::cols; +#pragma unroll + for (uint32_t j = 0; j < GmmaAccCoreMat::cols; j++) + { + ret[i][j] = gmemVec[baseOffset + j]; + } } return ret; }