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

fix multi batch and hack page idx core

Browse Source
This commit is contained in:
coderfeli
2025-04-15 01:21:23 +00:00
parent 234528e06c
commit ff281e135d
7 changed files with 51 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
include/ck_tile/core/tensor/tile_scatter_gather.hpp
View File

@@ -338,7 +338,6 @@ struct tile_scatter_gather
constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess);
constexpr auto idx_gather = idx_ys_start[number<YsGatherDim>{}];
const auto page_offset = page_idx_[idx_gather];
// read from bottom tensor
const vector_t vec_value =
get_bottom_tensor_view().template get_vectorized_elements<vector_t>(

50
include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp
View File

@@ -266,6 +266,7 @@ struct FmhaFwdKernel
const int32_t* seqstart_k_ptr;
const int32_t* seqlen_k_ptr;
const int32_t* page_idx;
int32_t page_num;
};
using Kargs = std::conditional_t<kIsGroupMode, FmhaFwdGroupModeKargs, FmhaFwdBatchModeKargs>;
@@ -598,6 +599,7 @@ struct FmhaFwdKernel
const void* seqstart_k_ptr,
const void* seqlen_k_ptr,
const void* page_idx_ptr,
ck_tile::index_t page_num,
ck_tile::index_t hdim_q,
ck_tile::index_t hdim_v,
ck_tile::index_t num_head_q,
@@ -657,7 +659,8 @@ struct FmhaFwdKernel
reinterpret_cast<const int32_t*>(seqstart_q_ptr),
reinterpret_cast<const int32_t*>(seqstart_k_ptr),
reinterpret_cast<const int32_t*>(seqlen_k_ptr),
reinterpret_cast<const int32_t*>(page_idx_ptr)};
reinterpret_cast<const int32_t*>(page_idx_ptr),
page_num};
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
{
@@ -724,6 +727,7 @@ struct FmhaFwdKernel
const void* seqstart_k_ptr,
const void* seqlen_k_ptr,
const void* page_idx_ptr,
ck_tile::index_t page_num,
ck_tile::index_t hdim_q,
ck_tile::index_t hdim_v,
ck_tile::index_t num_head_q,
@@ -763,6 +767,7 @@ struct FmhaFwdKernel
seqstart_k_ptr,
seqlen_k_ptr,
page_idx_ptr,
page_num,
hdim_q,
hdim_v,
num_head_q,
@@ -805,6 +810,7 @@ struct FmhaFwdKernel
const void* seqstart_k_ptr,
const void* seqlen_k_ptr,
const void* page_idx_ptr,
ck_tile::index_t page_num,
ck_tile::index_t hdim_q,
ck_tile::index_t hdim_v,
ck_tile::index_t num_head_q,
@@ -844,6 +850,7 @@ struct FmhaFwdKernel
seqstart_k_ptr,
seqlen_k_ptr,
page_idx_ptr,
page_num,
hdim_q,
hdim_v,
num_head_q,
@@ -966,7 +973,7 @@ struct FmhaFwdKernel
long_index_t batch_offset_q = 0;
// long_index_t batch_offset_k = 0;
long_index_t batch_offset_v = 0;
// long_index_t batch_offset_v = 0;
long_index_t batch_offset_bias = 0;
long_index_t batch_offset_randval = 0;
long_index_t batch_offset_lse = 0;
@@ -980,16 +987,19 @@ struct FmhaFwdKernel
batch_offset_q = query_start * kargs.stride_q;
// batch_offset_k = key_start * kargs.stride_k;
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
{
batch_offset_v = key_start * kargs.stride_v;
}
else
{
batch_offset_v = key_start;
}
// if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
// {
// batch_offset_v = key_start * kargs.stride_v;
// }
// else
// {
// batch_offset_v = key_start;
// }
kargs.page_idx += key_start;
// if(threadIdx.x==0){
// printf("\nbid %d %d page id %d pagev %d\n", blockIdx.z, i_batch, key_start, kargs.page_idx[0]);
// }
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
{
batch_offset_bias = query_start * kargs.stride_bias + key_start;
@@ -1050,17 +1060,17 @@ struct FmhaFwdKernel
const QDataType* q_ptr = reinterpret_cast<const QDataType*>(kargs.q_ptr) +
static_cast<long_index_t>(i_nhead) * kargs.nhead_stride_q +
batch_offset_q;
// const KDataType* k_ptr =
// reinterpret_cast<const KDataType*>(kargs.k_ptr) +
// static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) * kargs.nhead_stride_k +
// batch_offset_k;
const KDataType* k_ptr =
reinterpret_cast<const KDataType*>(kargs.k_ptr) +
static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) * kargs.nhead_stride_k;
static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) * kargs.nhead_stride_k;// +
// batch_offset_k;
// const KDataType* k_ptr =
// reinterpret_cast<const KDataType*>(kargs.k_ptr) +
// static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) * kargs.nhead_stride_k;
const VDataType* v_ptr =
reinterpret_cast<const VDataType*>(kargs.v_ptr) +
static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) * kargs.nhead_stride_v +
batch_offset_v;
static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) * kargs.nhead_stride_v;// +
// batch_offset_v;
ODataType* o_ptr = reinterpret_cast<ODataType*>(kargs.o_ptr) +
static_cast<long_index_t>(i_nhead) * kargs.nhead_stride_o +
batch_offset_o;
@@ -1091,7 +1101,7 @@ struct FmhaFwdKernel
const auto k_dram = [&]() {
const auto k_dram_naive = make_naive_tensor_view<address_space_enum::global>(
k_ptr,
make_tuple(kargs.seqlen_k, kargs.hdim_q),
make_tuple(kargs.page_num, kargs.hdim_q),
make_tuple(kargs.stride_k, 1),
number<FmhaPipeline::kAlignmentK>{},
number<1>{});
@@ -1107,7 +1117,7 @@ struct FmhaFwdKernel
{
const auto v_dram_naive = make_naive_tensor_view<address_space_enum::global>(
v_ptr,
make_tuple(kargs.seqlen_k, kargs.hdim_v),
make_tuple(kargs.page_num, kargs.hdim_v),
make_tuple(kargs.stride_v, 1),
number<FmhaPipeline::kAlignmentV>{},
number<1>{});
@@ -1115,7 +1125,7 @@ struct FmhaFwdKernel
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.seqlen_k)),
make_pass_through_transform(kargs.page_num)),
make_tuple(sequence<1>{}, sequence<0>{}),
make_tuple(sequence<0>{}, sequence<1>{}));

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

@@ -10,7 +10,6 @@
#include "ck_tile/ops/reduce/block/block_reduce.hpp"
#include "ck_tile/core/tensor/tile_distribution.hpp"
#include "ck_tile/core/tensor/tile_scatter_gather.hpp"
// #include "ck_tile/core/tensor/tile_scatter_gather_debug.hpp"
namespace ck_tile {
@@ -315,6 +314,11 @@ struct BlockFmhaPipelineQRKSVS
k_dram_block_window.get_window_origin(),
k_dist,
k_offsets); // K DRAM tile window for
// auto k_dram_window = make_tile_window_debug(
// k_dram_block_window.get_bottom_tensor_view(),
// k_dram_block_window.get_window_lengths(),
// k_dram_block_window.get_window_origin(),
// k_dist); // K DRAM tile window for
auto k_block_tile = load_tile(k_dram_window);
{

1
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp
View File

@@ -67,7 +67,6 @@ struct BlockFmhaPipelineQRKSVSAsync
static constexpr auto BiasEnum = Problem::BiasEnum;
static constexpr bool kStoreLSE = Problem::kStoreLSE;
static constexpr bool kHasDropout = Problem::kHasDropout;
// last dimension vector length used to create tensor view(and decide buffer_load vector length)
// ... together with tensor distribution. tensor dist should able to overwrite this
static constexpr index_t kAlignmentQ = Policy::template GetAlignmentQ<Problem>();