leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 25 Packages Projects Releases Wiki Activity

Revert "[ex77] fix mla split; add fwd lse; add bwd varlen (#2366)" (#2370)

Browse Source 
This reverts commit f12b1d75c9.
This commit is contained in:
Manish Gupta
2025-06-05 23:14:57 -04:00
committed by GitHub
parent f12b1d75c9
commit 2e2af190bd
19 changed files with 326 additions and 846 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/77_blackwell_fmha/reference/fmha_bwd_reference.hpp
+33 -72
View File
@@ -44,10 +44,10 @@ template<
class Fusion
>
void __global__ fmha_bwd_reference_dQ_kernel(
ProblemShape problem_shape_in,
TensorQ mQ_in, TensorK mK_in, TensorV mV_in,
TensorO mO_in, TensorLSE mLSE_in, TensorDO mDO_in,
TensorDQ mDQ_in, /* TensorDK mDK, TensorDV mDV, */
ProblemShape problem_shape,
TensorQ mQ, TensorK mK, TensorV mV,
TensorO mO, TensorLSE mLSE, TensorDO mDO,
TensorDQ mDQ, /* TensorDK mDK, TensorDV mDV, */
Fusion fusion) {
using namespace cute;
@@ -58,28 +58,15 @@ void __global__ fmha_bwd_reference_dQ_kernel(
extern __shared__ char mS_mem[];
Element* mS = reinterpret_cast<Element*>(mS_mem);
Element softmax_scale = static_cast<Element>(1.0 / sqrt(1.0 * size<2>(problem_shape_in)));
Element softmax_scale = static_cast<Element>(1.0 / sqrt(1.0 * size<1>(mO)));
for (int idx_L = blockIdx.y; idx_L < size<3>(problem_shape_in); idx_L += gridDim.y) {
auto [problem_shape, offset] = apply_variable_length_offset(
problem_shape_in,
make_coord(_0{}, _0{}, _0{}, idx2crd(idx_L, get<3>(problem_shape_in)))
);
// problem_shape = problem_shape_in;
// offset = repeat_like(problem_shape_in, _0{});
auto mQ = domain_offset(select<0,2,3>(offset), mQ_in);
auto mK = domain_offset(select<1,2,3>(offset), mK_in);
auto mV = domain_offset(select<1,2,3>(offset), mV_in);
auto mO = domain_offset(select<0,2,3>(offset), mO_in);
auto mLSE = domain_offset(select<0,3>(offset), mLSE_in);
auto mDO = domain_offset(select<0,2,3>(offset), mDO_in);
auto mDQ = domain_offset(select<0,2,3>(offset), mDQ_in);
for (int idx_Q = blockIdx.x; idx_Q < size<0>(problem_shape); idx_Q += gridDim.x) {
for (int idx_K = threadIdx.x; idx_K < size<1>(problem_shape); idx_K += blockDim.x) {
for (int idx_L = blockIdx.y; idx_L < size<2>(mDQ); idx_L += gridDim.y) {
for (int idx_Q = blockIdx.x; idx_Q < size<0>(mDQ); idx_Q += gridDim.x) {
for (int idx_K = threadIdx.x; idx_K < size<0>(mK); idx_K += blockDim.x) {
ElementAccumulator acc_qk = 0;
ElementAccumulator acc_dov = 0;
ElementAccumulator acc_doo = 0;
for (int idx_D0 = 0; idx_D0 < size<2>(problem_shape); idx_D0++) {
for (int idx_D0 = 0; idx_D0 < size<1>(mK); idx_D0++) {
acc_qk += mQ(idx_Q, idx_D0, idx_L) * mK(idx_K, idx_D0, idx_L);
acc_dov += mDO(idx_Q, idx_D0, idx_L) * mV(idx_K, idx_D0, idx_L);
acc_doo += mDO(idx_Q, idx_D0, idx_L) * mO(idx_Q, idx_D0, idx_L);
@@ -96,9 +83,9 @@ void __global__ fmha_bwd_reference_dQ_kernel(
__syncthreads();
for (int idx_D = threadIdx.x; idx_D < size<2>(problem_shape); idx_D += blockDim.x) {
for (int idx_D = threadIdx.x; idx_D < size<1>(mDQ); idx_D += blockDim.x) {
ElementAccumulator acc = 0;
for (int idx_K = 0; idx_K < size<1>(problem_shape); idx_K++) {
for (int idx_K = 0; idx_K < size<0>(mK); idx_K++) {
acc += mS[idx_K] * mK(idx_K, idx_D, idx_L);
}
mDQ(idx_Q, idx_D, idx_L) = static_cast<typename TensorDQ::value_type>(acc);
@@ -117,10 +104,10 @@ template<
class Fusion
>
void __global__ fmha_bwd_reference_dK_kernel(
ProblemShape problem_shape_in,
TensorQ mQ_in, TensorK mK_in, TensorV mV_in,
TensorO mO_in, TensorLSE mLSE_in, TensorDO mDO_in,
/* TensorDQ mDQ_in, */ TensorDK mDK_in, /* TensorDV mDV_in, */
ProblemShape problem_shape,
TensorQ mQ, TensorK mK, TensorV mV,
TensorO mO, TensorLSE mLSE, TensorDO mDO,
/* TensorDQ mDQ, */ TensorDK mDK, /* TensorDV mDV, */
Fusion fusion) {
using namespace cute;
@@ -131,28 +118,15 @@ void __global__ fmha_bwd_reference_dK_kernel(
extern __shared__ char mS_mem[];
Element* mS = reinterpret_cast<Element*>(mS_mem);
Element softmax_scale = static_cast<Element>(1.0 / sqrt(1.0 * size<2>(problem_shape_in)));
Element softmax_scale = static_cast<Element>(1.0 / sqrt(1.0 * size<1>(mO)));
for (int idx_L = blockIdx.y; idx_L < size<3>(problem_shape_in); idx_L += gridDim.y) {
auto [problem_shape, offset] = apply_variable_length_offset(
problem_shape_in,
make_coord(_0{}, _0{}, _0{}, idx2crd(idx_L, get<3>(problem_shape_in)))
);
// problem_shape = problem_shape_in;
// offset = repeat_like(problem_shape_in, _0{});
auto mQ = domain_offset(select<0,2,3>(offset), mQ_in);
auto mK = domain_offset(select<1,2,3>(offset), mK_in);
auto mV = domain_offset(select<1,2,3>(offset), mV_in);
auto mO = domain_offset(select<0,2,3>(offset), mO_in);
auto mLSE = domain_offset(select<0,3>(offset), mLSE_in);
auto mDO = domain_offset(select<0,2,3>(offset), mDO_in);
auto mDK = domain_offset(select<1,2,3>(offset), mDK_in);
for (int idx_K = blockIdx.x; idx_K < size<1>(problem_shape); idx_K += gridDim.x) {
for (int idx_Q = threadIdx.x; idx_Q < size<0>(problem_shape); idx_Q += blockDim.x) {
for (int idx_L = blockIdx.y; idx_L < size<2>(mDK); idx_L += gridDim.y) {
for (int idx_K = blockIdx.x; idx_K < size<0>(mDK); idx_K += gridDim.x) {
for (int idx_Q = threadIdx.x; idx_Q < size<0>(mDO); idx_Q += blockDim.x) {
ElementAccumulator acc_qk = 0;
ElementAccumulator acc_dov = 0;
ElementAccumulator acc_doo = 0;
for (int idx_D0 = 0; idx_D0 < size<2>(problem_shape); idx_D0++) {
for (int idx_D0 = 0; idx_D0 < size<1>(mK); idx_D0++) {
acc_qk += mQ(idx_Q, idx_D0, idx_L) * mK(idx_K, idx_D0, idx_L);
acc_dov += mDO(idx_Q, idx_D0, idx_L) * mV(idx_K, idx_D0, idx_L);
acc_doo += mDO(idx_Q, idx_D0, idx_L) * mO(idx_Q, idx_D0, idx_L);
@@ -169,9 +143,9 @@ void __global__ fmha_bwd_reference_dK_kernel(
__syncthreads();
for (int idx_D = threadIdx.x; idx_D < size<2>(problem_shape); idx_D += blockDim.x) {
for (int idx_D = threadIdx.x; idx_D < size<1>(mDK); idx_D += blockDim.x) {
ElementAccumulator acc = 0;
for (int idx_Q = 0; idx_Q < size<0>(problem_shape); idx_Q++) {
for (int idx_Q = 0; idx_Q < size<0>(mDO); idx_Q++) {
acc += mS[idx_Q] * mQ(idx_Q, idx_D, idx_L);
}
mDK(idx_K, idx_D, idx_L) = static_cast<typename TensorDK::value_type>(acc);
@@ -190,10 +164,10 @@ template<
class Fusion
>
void __global__ fmha_bwd_reference_dV_kernel(
ProblemShape problem_shape_in,
TensorQ mQ_in, TensorK mK_in, TensorV mV_in,
TensorO mO_in, TensorLSE mLSE_in, TensorDO mDO_in,
/* TensorDQ mDQ_in, TensorDK mDK_in, */ TensorDV mDV_in,
ProblemShape problem_shape,
TensorQ mQ, TensorK mK, TensorV mV,
TensorO mO, TensorLSE mLSE, TensorDO mDO,
/* TensorDQ mDQ, TensorDK mDK, */ TensorDV mDV,
Fusion fusion) {
using namespace cute;
@@ -204,27 +178,14 @@ void __global__ fmha_bwd_reference_dV_kernel(
extern __shared__ char mS_mem[];
Element* mS = reinterpret_cast<Element*>(mS_mem);
Element softmax_scale = static_cast<Element>(1.0 / sqrt(1.0 * size<2>(problem_shape_in)));
ElementAcc softmax_scale = static_cast<ElementAcc>(1.0 / sqrt(1.0 * size<1>(mO)));
for (int idx_L = blockIdx.y; idx_L < size<3>(problem_shape_in); idx_L += gridDim.y) {
auto [problem_shape, offset] = apply_variable_length_offset(
problem_shape_in,
make_coord(_0{}, _0{}, _0{}, idx2crd(idx_L, get<3>(problem_shape_in)))
);
// problem_shape = problem_shape_in;
// offset = repeat_like(problem_shape_in, _0{});
auto mQ = domain_offset(select<0,2,3>(offset), mQ_in);
auto mK = domain_offset(select<1,2,3>(offset), mK_in);
auto mV = domain_offset(select<1,2,3>(offset), mV_in);
auto mO = domain_offset(select<0,2,3>(offset), mO_in);
auto mLSE = domain_offset(select<0,3>(offset), mLSE_in);
auto mDO = domain_offset(select<0,2,3>(offset), mDO_in);
auto mDV = domain_offset(select<1,2,3>(offset), mDV_in);
for (int idx_K = blockIdx.x; idx_K < size<1>(problem_shape); idx_K += gridDim.x) {
for (int idx_Q = threadIdx.x; idx_Q < size<0>(problem_shape); idx_Q += blockDim.x) {
for (int idx_L = blockIdx.y; idx_L < size<2>(mDV); idx_L += gridDim.y) {
for (int idx_K = blockIdx.x; idx_K < size<0>(mDV); idx_K += gridDim.x) {
for (int idx_Q = threadIdx.x; idx_Q < size<0>(mDO); idx_Q += blockDim.x) {
ElementAcc acc_qk = 0;
for (int idx_D0 = 0; idx_D0 < size<2>(problem_shape); idx_D0++) {
for (int idx_D0 = 0; idx_D0 < size<1>(mK); idx_D0++) {
ElementAcc rQ = mQ(idx_Q, idx_D0, idx_L);
ElementAcc rK = mK(idx_K, idx_D0, idx_L);
acc_qk += rQ * rK;
@@ -241,9 +202,9 @@ void __global__ fmha_bwd_reference_dV_kernel(
__syncthreads();
for (int idx_D = threadIdx.x; idx_D < size<2>(problem_shape); idx_D += blockDim.x) {
for (int idx_D = threadIdx.x; idx_D < size<1>(mDV); idx_D += blockDim.x) {
ElementAcc acc = 0;
for (int idx_Q = 0; idx_Q < size<0>(problem_shape); idx_Q++) {
for (int idx_Q = 0; idx_Q < size<0>(mDO); idx_Q++) {
ElementAcc rS = mS[idx_Q];
ElementAcc rDO = mDO(idx_Q, idx_D, idx_L);
acc += rS * rDO;
examples/77_blackwell_fmha/reference/fmha_fwd_reference.hpp
+1 -1
View File
@@ -127,7 +127,7 @@ void __global__ fmha_reference_kernel(
mO(idx_Q + offset_Q, idx_D, idx_L) = static_cast<typename TensorO::value_type>(acc * scale);
}
if (threadIdx.x == 0 && mLSE.data() != nullptr) {
if (threadIdx.x == 0) {
mLSE(idx_Q + offset_Q, idx_L) = log(sum) + softmax_scale * maxS;
}
examples/77_blackwell_fmha/reference/reference_abs_error.hpp
-2
View File
@@ -75,8 +75,6 @@ struct DeviceAllocation {
size_t size() const { return size_; }
size_t get_storage_size() const { return (size_ + offset_) * sizeof(T); }
void copy_from_host(const T* ptr, size_t sz) {
auto ret = cudaMemcpy(ptr_, ptr, sz * sizeof(T), cudaMemcpyDefault);
assert(ret == cudaSuccess);