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

change kv layout to [num_total_pages, page_block_size, hdim]

Browse Source
This commit is contained in:
ltqin
2025-09-17 14:57:07 +00:00
parent bb7ff2978c
commit 7bcc4c5e7a
2 changed files with 41 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp
View File

@@ -816,14 +816,21 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
const auto k_dram = [&]() {
const auto k_dram_naive = make_naive_tensor_view<address_space_enum::global>(
k_ptr,
make_tuple(kargs.num_total_pages * kargs.page_block_size, kargs.hdim_q),
make_tuple(kargs.stride_k, 1),
make_tuple(kargs.num_total_pages, kargs.page_block_size, kargs.hdim_q),
make_tuple(kargs.batch_stride_k, kargs.stride_k, 1),
number<FmhaPipeline::kAlignmentK>{},
number<1>{});
auto k_dram_2d = transform_tensor_view(
k_dram_naive,
make_tuple(make_merge_transform(make_tuple(kargs.num_total_pages, kargs.page_block_size)),
make_pass_through_transform(kargs.hdim_q)),
make_tuple(sequence<0, 1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
return pad_tensor_view(
k_dram_naive,
k_dram_2d,
make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}),
sequence<kPadSeqLenK_, kPadHeadDimQ>{});
}();
@@ -832,13 +839,21 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
{
const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
v_ptr,
make_tuple(kargs.num_total_pages * kargs.page_block_size, kargs.hdim_v),
make_tuple(kargs.stride_v, 1),
make_tuple(kargs.num_total_pages, kargs.page_block_size, kargs.hdim_v),
make_tuple(kargs.batch_stride_v, kargs.stride_v, 1),
number<FmhaPipeline::kAlignmentV>{},
number<1>{});
const auto v_dram_transposed = transform_tensor_view(
auto v_dram_2d = transform_tensor_view(
v_dram_naive,
make_tuple(make_merge_transform(
make_tuple(kargs.num_total_pages, kargs.page_block_size)),
make_pass_through_transform(kargs.hdim_v)),
make_tuple(sequence<0, 1>{}, sequence<2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
const auto v_dram_transposed = transform_tensor_view(
v_dram_2d,
make_tuple(
make_pass_through_transform(kargs.hdim_v),
make_pass_through_transform(kargs.num_total_pages * kargs.page_block_size)),
@@ -1088,6 +1103,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
kargs.kv_page_indices,
kargs.stride_k,
kargs.stride_v,
kargs.batch_stride_k,
kargs.batch_stride_v,
dropout);
}
else
@@ -1108,6 +1125,8 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
kargs.kv_page_indices,
kargs.stride_k,
kargs.stride_v,
kargs.batch_stride_k,
kargs.batch_stride_v,
dropout);
}
}();

25
include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp
View File

@@ -24,6 +24,7 @@ template <typename OffsetVecType,
bool kIsKcache>
CK_TILE_HOST_DEVICE void kv_offset_array_transform(const index_t* page_vec,
const index_t& stride_kv,
const index_t& page_stride_kv,
const CoordVecType& coord_vec,
OffsetVecType& kv_offset_vec)
{
@@ -46,7 +47,7 @@ CK_TILE_HOST_DEVICE void kv_offset_array_transform(const index_t* page_vec,
const index_t page_offset =
(thread_coord_start + kLoopStride * k0.value) & kPageMask;
kv_offset_vec[k0] =
((page_vec[kPageId] << kPageShiftSize) + page_offset) * stride_kv;
page_vec[kPageId] * page_stride_kv + page_offset * stride_kv;
});
}
else
@@ -54,11 +55,11 @@ CK_TILE_HOST_DEVICE void kv_offset_array_transform(const index_t* page_vec,
// for v_offset_vec
const index_t lane0_start = __builtin_amdgcn_readfirstlane(thread_coord_start);
const index_t lane0_page_id = (lane0_start + kLoopStart) >> kPageShiftSize;
const index_t page_loc = page_vec[lane0_page_id] << kPageShiftSize;
const index_t page_loc = page_vec[lane0_page_id] * page_stride_kv;
static_for<0, kLoopCount, 1>{}([&](auto k0) {
const index_t page_offset =
(thread_coord_start + kLoopStart + k0.value) & kPageMask;
kv_offset_vec[k0] = (page_loc + page_offset) * stride_kv;
kv_offset_vec[k0] = page_loc + page_offset * stride_kv;
});
}
}
@@ -248,6 +249,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
const index_t* page_idx,
const index_t stride_k,
const index_t stride_v,
const index_t page_stride_k,
const index_t page_stride_v,
DropoutType& dropout) const
{
static_assert(
@@ -386,7 +389,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
NRepeat,
kN0 / NRepeat,
kIsSglangLayout,
true>(page_idx, stride_k, k_coord, k_offsets);
true>(page_idx, stride_k, page_stride_k, k_coord, k_offsets);
auto k_dram_window = make_tile_scatter_gather(k_dram_block_window.get_bottom_tensor_view(),
k_dram_block_window.get_window_lengths(),
@@ -429,7 +432,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
V_KRepeat,
1,
kIsSglangLayout,
false>(page_idx, stride_v, v_coord, v_offsets);
false>(page_idx, stride_v, page_stride_v, v_coord, v_offsets);
auto v_dram_window =
make_tile_scatter_gather(v_dram_block_window_tmp.get_bottom_tensor_view(),
@@ -502,7 +505,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
V_KRepeat,
1,
kIsSglangLayout,
false>(page_idx, stride_v, v_coord, v_offsets);
false>(page_idx, stride_v, page_stride_v, v_coord, v_offsets);
v_dram_window.update_page_idx(v_offsets);
__builtin_amdgcn_sched_barrier(0);
@@ -678,7 +681,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
V_KRepeat,
1,
kIsSglangLayout,
false>(page_idx, stride_v, v_coord, v_offsets);
false>(page_idx, stride_v, page_stride_v, v_coord, v_offsets);
v_dram_window.update_page_idx(v_offsets);
}
__builtin_amdgcn_sched_barrier(0);
@@ -815,7 +818,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
V_KRepeat,
1,
kIsSglangLayout,
false>(page_idx, stride_v, v_coord, v_offsets);
false>(page_idx, stride_v, page_stride_v, v_coord, v_offsets);
v_dram_window.update_page_idx(v_offsets);
}
block_sync_lds();
@@ -877,7 +880,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
NRepeat,
kN0 / NRepeat,
kIsSglangLayout,
true>(page_idx, stride_k, k_coord, k_offsets);
true>(page_idx, stride_k, page_stride_k, k_coord, k_offsets);
k_dram_window.update_page_idx(k_offsets);
if constexpr(k1_loops >= 2 &&
LdsSeq.at(number<0>{}) == LdsSeq.at(number<k0_loops + k1_loops - 2>{}))
@@ -985,6 +988,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
const index_t* page_idx,
const index_t stride_k,
const index_t stride_v,
const index_t page_stride_k,
const index_t page_stride_v,
DropoutType& dropout) const
{
return operator()(q_dram_block_window_tmp,
@@ -1011,6 +1016,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
page_idx,
stride_k,
stride_v,
page_stride_k,
page_stride_v,
dropout);
}
};