ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-06 15:54:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

Donot perform write again if already in last page-block

Browse Source
This commit is contained in:
PoYen, Chen
2024-08-07 12:00:15 +00:00
parent f265742b63
commit 1b96dc2592
2 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
include/ck_tile/ops/fmha/block/block_tile_window_navigator.hpp
View File

@@ -140,10 +140,12 @@ struct PagedTileWindowNavigator
}
template <typename TileWindow>
CK_TILE_HOST_DEVICE bool is_closs_block(const TileWindow& tile_window) const
CK_TILE_HOST_DEVICE bool is_cross_block(index_t block_index,
const TileWindow& tile_window) const
{
return page_block_size < (tile_window.get_window_origin().at(number<VirtualDim>{}) +
tile_window.get_window_lengths().at(number<VirtualDim>{}));
const index_t origin = tile_window.get_window_origin().at(number<VirtualDim>{});
const index_t length = tile_window.get_window_lengths().at(number<VirtualDim>{});
return (block_index < num_blocks - 1) && (page_block_size < origin + length);
}
template <typename TileWindow>

6
include/ck_tile/ops/fmha/pipeline/block_fmha_fwd_appendkv_pipeline.hpp
View File

@@ -151,9 +151,8 @@ struct BlockFmhaFwdAppendKVPipeline
if constexpr(kIsPagedKV)
{
store_tile(k_dram_block_window, knew_tile);
// write tile to another block if nesscary
if(k_tile_navigator.is_closs_block(k_dram_block_window))
if(k_tile_navigator.is_cross_block(i_block0, k_dram_block_window))
{
k_tile_navigator.move_to_block(i_block0, k_dram_block_window, i_block0 + 1);
store_tile(k_dram_block_window, knew_tile);
@@ -176,9 +175,8 @@ struct BlockFmhaFwdAppendKVPipeline
if constexpr(kIsPagedKV)
{
store_tile(v_dram_block_window, vnew_tile);
// write tile to another block if nesscary
if(v_tile_navigator.is_closs_block(v_dram_block_window))
if(v_tile_navigator.is_cross_block(i_block1, v_dram_block_window))
{
v_tile_navigator.move_to_block(i_block1, v_dram_block_window, i_block1 + 1);
store_tile(v_dram_block_window, vnew_tile);