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

modify fmha_batch_prefill_kernel layout into vllm

Browse Source
This commit is contained in:
zanzhang
2025-05-14 13:14:48 +08:00
parent b324bbe45f
commit cd294e76e4
1 changed files with 73 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -809,36 +809,72 @@ 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),
// number<FmhaPipeline::kAlignmentK>{},
// number<1>{});
// constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
// return pad_tensor_view(
// k_dram_naive,
// make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}),
// sequence<kPadSeqLenK_, kPadHeadDimQ>{});
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.seqlen_k / 16, kargs.hdim_q / 8, 16, 8),
make_tuple(kargs.hdim_q * 16, 16 * 8, 8, 1),
number<FmhaPipeline::kAlignmentK>{},
number<1>{});
constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
        const auto k_dram_transposed = transform_tensor_view(
            k_dram_naive,
            make_tuple(make_merge_transform(ck_tile::make_tuple(kargs.seqlen_k /16, 16)),
                       make_merge_transform(ck_tile::make_tuple(kargs.hdim_q / 8, 8))),
            ck_tile::make_tuple(ck_tile::sequence<0, 2>{}, ck_tile::sequence<1, 3>{}),
            ck_tile::make_tuple(ck_tile::sequence<0>{}, ck_tile::sequence<1>{}));
constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : false;
return pad_tensor_view(
k_dram_naive,
k_dram_transposed,
make_tuple(number<FmhaPipeline::kN0>{}, number<FmhaPipeline::kK0>{}),
sequence<kPadSeqLenK_, kPadHeadDimQ>{});
}();
const auto v_dram = [&]() {
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
{
// 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),
// number<FmhaPipeline::kAlignmentV>{},
// number<1>{});
//
// const auto v_dram_transposed = transform_tensor_view(
// v_dram_naive,
// make_tuple(
// make_pass_through_transform(kargs.hdim_v),
// make_pass_through_transform(kargs.num_total_pages * kargs.page_block_size)),
// make_tuple(sequence<1>{}, sequence<0>{}),
// make_tuple(sequence<0>{}, sequence<1>{}));
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.seqlen_k / 16, kargs.hdim_v, 16),
make_tuple(kargs.hdim_v * 16, 16, 1),
// make_tuple(kargs.seqlen_k, kargs.hdim_v),
// make_tuple(kargs.stride_v, 1),
number<FmhaPipeline::kAlignmentV>{},
number<1>{});
const auto v_dram_transposed = transform_tensor_view(
v_dram_naive,
make_tuple(
make_pass_through_transform(kargs.hdim_v),
make_pass_through_transform(kargs.num_total_pages * kargs.page_block_size)),
make_tuple(sequence<1>{}, sequence<0>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
const auto v_dram_transposed =
transform_tensor_view(v_dram_naive,
make_tuple(make_pass_through_transform(kargs.hdim_v),
                     make_merge_transform(ck_tile::make_tuple(kargs.seqlen_k /16, 16))),
make_tuple(sequence<1>{}, sequence<0, 2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
constexpr bool kPadSeqLenK_ = kUseAsyncCopy ? kPadSeqLenK : true;
return pad_tensor_view(
@@ -848,16 +884,37 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
}
else
{
// const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
// v_ptr,
// make_tuple(kargs.hdim_v, kargs.num_total_pages * kargs.page_block_size),
// make_tuple(kargs.stride_v, 1),
// number<FmhaPipeline::kAlignmentV>{},
// number<1>{});
//
// constexpr bool kPadHeadDimV_ = kUseAsyncCopy ? kPadHeadDimV : false;
// return pad_tensor_view(
// v_dram_naive,
// make_tuple(number<FmhaPipeline::kN1>{}, number<FmhaPipeline::kK1>{}),
// sequence<kPadHeadDimV_, kPadSeqLenK>{});
const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
v_ptr,
make_tuple(kargs.hdim_v, kargs.num_total_pages * kargs.page_block_size),
make_tuple(kargs.stride_v, 1),
// make_tuple(kargs.hdim_v, kargs.seqlen_k),
// make_tuple(kargs.stride_v, 1),
make_tuple(kargs.seqlen_k / 16, kargs.hdim_v, 16),
make_tuple(kargs.hdim_v * 16, 16, 1),
number<FmhaPipeline::kAlignmentV>{},
number<1>{});
const auto v_dram_transposed =
transform_tensor_view(v_dram_naive,
make_tuple(make_pass_through_transform(kargs.hdim_v),
                     make_merge_transform(ck_tile::make_tuple(kargs.seqlen_k /16, 16))),
make_tuple(sequence<1>{}, sequence<0, 2>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
constexpr bool kPadHeadDimV_ = kUseAsyncCopy ? kPadHeadDimV : false;
return pad_tensor_view(
v_dram_naive,
v_dram_transposed,
make_tuple(number<FmhaPipeline::kN1>{}, number<FmhaPipeline::kK1>{}),
sequence<kPadHeadDimV_, kPadSeqLenK>{});
}