ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 17:19:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

update the method of calculating coordinates

Browse Source
This commit is contained in:
shay-li77
2025-07-04 10:38:07 +08:00
parent c96cad7a11
commit a9840ea14c
1 changed files with 12 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
include/ck_tile/ops/fmha/block/block_masking.hpp
View File

@@ -283,11 +283,12 @@ struct SimplifiedGenericAttentionMask
}
CK_TILE_HOST_DEVICE
SimplifiedGenericAttentionMask(
index_t y_, index_t x_, index_t y_total_, index_t x_total_, mdiv y_ratio_mdiv_)
: y(y_),
index_t y_real_, index_t x_, index_t y_total_, index_t x_total_, mdiv y_ratio_mdiv_)
: y(y_real_ * static_cast<index_t>(y_ratio_mdiv_.get())),
x(x_),
y_total(y_total_),
x_total(x_total_),
y_real(y_real_),
y_ratio(static_cast<index_t>(y_ratio_mdiv_.get())),
y_ratio_mdiv(y_ratio_mdiv_)
{
@@ -335,10 +336,9 @@ struct SimplifiedGenericAttentionMask
{
// get the tile start/end range assum we loop over along X tile by tile
index_t x_start = [&]() {
index_t tmp_offset = -y + i_y + y_ratio;
index_t tmp_div =
static_cast<index_t>(y_ratio_mdiv.div(static_cast<uint32_t>(tmp_offset)));
index_t tmp = tmp_offset > 0 ? tmp_div : 0; // clamp by zero
index_t tmp =
-y_real +
static_cast<index_t>(y_ratio_mdiv.div(static_cast<uint32_t>(i_y))) + 1;
return (tmp / XTile) * XTile; // round to tile aligned
}();
@@ -443,13 +443,9 @@ struct SimplifiedGenericAttentionMask
}
else
{
index_t start_tmp = -y + i_y + y_ratio;
index_t start_tmp_div =
static_cast<index_t>(y_ratio_mdiv.div(static_cast<uint32_t>(start_tmp)));
index_t x_start = start_tmp > 0 ? start_tmp_div : 0; // clamp by zero
uint32_t end_tmp = static_cast<uint32_t>(i_y);
index_t x_end = min(static_cast<index_t>(y_ratio_mdiv.div(end_tmp)) + x,
index_t x_tmp = static_cast<index_t>(y_ratio_mdiv.div(static_cast<uint32_t>(i_y)));
index_t x_start = -y_real + x_tmp + 1;
index_t x_end = min(x_tmp + x,
x_total); // need min in case x is padded
return i_x < x_start || i_x >= x_end || i_y >= y_total;
}
@@ -504,7 +500,9 @@ struct SimplifiedGenericAttentionMask
private:
index_t y, x;
index_t y_total, x_total;
index_t y_ratio = 1;
// y_real is vertical axis before multiplying y_ratio. y_real * y_ratio = y
index_t y_real;
index_t y_ratio;
mdiv y_ratio_mdiv;
};