ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-13 02:27:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

merge

Browse Source
This commit is contained in:
Juuso Korhonen
2025-10-14 12:35:33 +00:00
parent 4d232d59cc b940a75328
commit 72fe8b311c
2 changed files with 28 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp
View File

@@ -97,7 +97,6 @@ struct FmhaFwdV3Kernel
ck_tile::index_t num_seqs; // number of batches for q
};
using Kargs = UnifiedAttentionVarlenKargs;
CK_TILE_HOST static constexpr Kargs MakeKargs(
@@ -333,6 +332,8 @@ struct FmhaFwdV3Kernel
index_t o_ptr_offset_0 = cur_batch_in_all_start_index * kargs.output_stride_0; // move the pointer to the batch start
index_t o_ptr_offset_1 = kv_head_idx * num_queries_per_kv * kargs.output_stride_1; // move the pointer to the correct head group start
index_t o_ptr_offset = o_ptr_offset_0 + o_ptr_offset_1;
index_t block_table_offset = seq_idx * kargs.block_table_stride;
const QDataType* q_ptr = reinterpret_cast<const QDataType*>(kargs.q_ptr) + q_ptr_offset;
const KDataType* k_ptr = reinterpret_cast<const KDataType*>(kargs.k_ptr) + kv_head_offset;
@@ -462,6 +463,8 @@ struct FmhaFwdV3Kernel
return FmhaPipeline{}(q_dram_window,
k_dram_window,
v_dram_window,
block_tables_ptr,
block_table_offset,
lse_dram_window,
mask,
kargs.scale_s,
@@ -502,7 +505,7 @@ struct FmhaFwdV3Kernel
auto o_dram_window =
make_tile_window(o_dram,
make_tuple(BLOCK_M, HEAD_SIZE_PADDED),
{q_block_global_idx * num_queries_per_kv * HEAD_SIZE_PADDED, 0});
{query_pos * num_queries_per_kv, 0});
EpiloguePipeline{}(o_dram_window, o_acc_tile, nullptr);
}

29
include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
View File

@@ -391,6 +391,8 @@ struct UnifiedAttentionPipeline
[[maybe_unused]] const KElementFunction& k_element_func,
const VDramBlockWindowTmp& v_dram_block_window_tmp, // N1*K1 tile
[[maybe_unused]] const VElementFunction& v_element_func,
const void* block_tables_ptr,
index_t block_table_offset,
LSEDramBlockWindowTmp& lse_dram_window_tmp, // M0*1 tile
const LSEElementFunction& lse_element_func,
[[maybe_unused]] const SAccElementFunction& s_acc_element_func,
@@ -402,6 +404,7 @@ struct UnifiedAttentionPipeline
{
using namespace ck_tile;
static_assert(
std::is_same_v<QDataType, remove_cvref_t<typename QDramBlockWindowTmp::DataType>> &&
std::is_same_v<KDataType, remove_cvref_t<typename KDramBlockWindowTmp::DataType>> &&
@@ -573,22 +576,26 @@ struct UnifiedAttentionPipeline
}
}
index_t i_total_loops = 0;
index_t kv_blk_idx = block_tables_ptr[block_table_offset + i_total_loops];
index_t kv_blk_idx_prev = 0;
auto k_dram_window =
make_tile_window(k_dram_block_window_tmp.get_bottom_tensor_view(),
k_dram_block_window_tmp.get_window_lengths(),
{seqlen_k_start, 0},
{(kv_blk_idx - kv_blk_idx_prev) * BLOCK_SIZE, 0},
Policy::template MakeKDramTileDistribution<Problem>());
k_dram_window.init_raw();
auto v_dram_window =
make_tile_window(v_dram_block_window_tmp.get_bottom_tensor_view(),
v_dram_block_window_tmp.get_window_lengths(),
{seqlen_k_start, 0}, // TODO: hdim split?
{(kv_blk_idx - kv_blk_idx_prev) * BLOCK_SIZE, 0}, // TODO: hdim split?
Policy::template MakeVDramTileDistribution<Problem>());
v_dram_window.init_raw();
// prefetch K tile
index_t i_total_loops = 0;
constexpr index_t k0_loops = kQKHeaddim / kK0;
constexpr index_t k1_loops = kN0 / kK1;
static_assert(1 == k0_loops);
@@ -678,10 +685,13 @@ struct UnifiedAttentionPipeline
auto K_mem_load = [&](auto k_lds_write_idx) {
async_load_tile_raw(k_lds_window_store(k_lds_write_idx), k_dram_window);
// TODO maybe needs i_total_loops as argument. Or maybe needs to use the k_lds_write_idx as the index
/// FIXME: use the future-predicting method to move the window
// move K tile windows
move_tile_window(k_dram_window, {kN0, 0});
auto k_dram_window = make_tile_window(k_dram_window.get_bottom_tensor_view(),
k_dram_window.get_window_lengths(),
{(block_tables_ptr[block_table_offset + i_total_loops]) * BLOCK_SIZE, 0},
Policy::template MakeVDramTileDistribution<Problem>());
};
auto K_lds_load = [&](auto k_lds_read_idx) {
@@ -692,7 +702,10 @@ struct UnifiedAttentionPipeline
async_load_tile_raw(v_lds_window_store(v_lds_write_idx), v_dram_window);
/// FIXME: use the future-predicting method to move the window
move_tile_window(v_dram_window, {kK1, 0});
auto v_dram_window = make_tile_window(v_dram_window.get_bottom_tensor_view(),
v_dram_window.get_window_lengths(),
{(block_tables_ptr[block_table_offset + i_total_loops]) * BLOCK_SIZE, 0},
Policy::template MakeVDramTileDistribution<Problem>());
};
auto V_lds_load = [&](auto v_lds_read_idx) {
@@ -1231,6 +1244,8 @@ struct UnifiedAttentionPipeline
CK_TILE_DEVICE auto operator()(const QDramBlockWindowTmp& q_dram_block_window_tmp, // M0*K0 tile
const KDramBlockWindowTmp& k_dram_block_window_tmp, // N0*K0 tile
const VDramBlockWindowTmp& v_dram_block_window_tmp, // N1*K1 tile
const void* block_tables_ptr,
index_t block_table_offset,
LSEDramBlockWindowTmp& lse_dram_block_window_tmp, // M0*1 tile
FmhaMask mask,
float scale_s,
@@ -1244,6 +1259,8 @@ struct UnifiedAttentionPipeline
identity{},
v_dram_block_window_tmp,
identity{},
block_tables_ptr,
block_table_offset,
lse_dram_block_window_tmp,
identity{},
identity{},