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

fix k origin

Browse Source
This commit is contained in:
coderfeli
2025-04-07 10:04:22 +00:00
parent 57c9d84eb1
commit 4e644a33ab
4 changed files with 23 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
include/ck_tile/core/tensor/tile_window_paged.hpp
View File

@@ -188,10 +188,10 @@ struct page_tile_with_static_distribution
array<index_t, NDimY>{0}));
#endif
BottomTensorIndex bottom_tensor_thread_origin_idx_tmp ={0,
window_origin[1] + window_adaptor_thread_coord_tmp.get_bottom_index()[1]};
// BottomTensorIndex bottom_tensor_thread_origin_idx_tmp =
// window_origin + window_adaptor_thread_coord_tmp.get_bottom_index();
BottomTensorIndex bottom_tensor_thread_origin_idx_tmp =
window_origin + tuple<index_t, index_t>(0, window_adaptor_thread_coord_tmp.get_bottom_index()[1]);
// tuple<index_t, index_t>(0, window_adaptor_thread_coord_tmp.get_bottom_index()[1]);
const auto bottom_tensor_thread_coord_tmp = make_tensor_coordinate(
bottom_tensor_view_.get_tensor_descriptor(), bottom_tensor_thread_origin_idx_tmp);

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

@@ -988,6 +988,8 @@ struct FmhaFwdKernel
{
batch_offset_v = key_start;
}
kargs.page_idx += key_start;
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
{
batch_offset_bias = query_start * kargs.stride_bias + key_start;

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

@@ -243,11 +243,10 @@ struct BlockFmhaPipelineQRKSVS
return o_acc;
}
}
auto k_dram_block_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});
{seqlen_k_start, 0}); //todo fixme felix
const auto bias_origin = bias_dram_block_window_tmp.get_window_origin();
auto bias_dram_window =
@@ -283,7 +282,7 @@ struct BlockFmhaPipelineQRKSVS
statically_indexed_array<index_t, NR> offsets;
static_for<0, NR, 1>{}([&](auto n0) {
offsets[n0] = page_idx[c_coord[0] + 64 * n0.value] * 128; // Problem::kN_;
offsets[n0] = page_idx[i_total_loops * kN0 + c_coord[0] + 64 * n0.value] * kQKHeaddim; // Problem::kN_;
});
auto k_dram_window = make_tile_window_paged(
k_dram_block_window.get_bottom_tensor_view(),