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

fixing compile errors...

Browse Source
This commit is contained in:
Juuso Korhonen
2025-10-20 14:03:15 +00:00
parent 9fda954253
commit 97e7527eb1
3 changed files with 30 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
include/ck_tile/ops/unified_attention/block/block_masking.hpp
View File

@@ -86,24 +86,24 @@ struct GenericAttentionMask
static constexpr const char* name = impl::MaskName<IsMasking, IsLocal>::name;
// New constructor accepting repeat_idx with default value 1
CK_TILE_HOST_DEVICE GenericAttentionMask(index_t y_total_, index_t x_total_, index_t repeat_idx = 1)
: GenericAttentionMask(0, 0, y_total_, x_total_, repeat_idx)
CK_TILE_HOST_DEVICE GenericAttentionMask(index_t y_total_, index_t x_total_, index_t repeat_idx_ = 1)
: GenericAttentionMask(0, 0, y_total_, x_total_, repeat_idx_)
{
}
CK_TILE_HOST_DEVICE
GenericAttentionMask(index_t y_, index_t x_, index_t y_total_, index_t x_total_, index_t repeat_idx = 1)
: y(y_), x(x_), y_total(y_total_), x_total(x_total_), repeat_idx(repeat_idx)
GenericAttentionMask(index_t y_, index_t x_, index_t y_total_, index_t x_total_, index_t repeat_idx_ = 1)
: y(y_), x(x_), y_total(y_total_), x_total(x_total_), repeat_idx(repeat_idx_)
{
}
template <typename MaskCoordinates>
CK_TILE_HOST_DEVICE GenericAttentionMask(const MaskCoordinates& mask_coord, index_t repeat_idx = 1)
CK_TILE_HOST_DEVICE GenericAttentionMask(const MaskCoordinates& mask_coord, index_t repeat_idx_ = 1)
: y(mask_coord.at(number<0>{})),
x(mask_coord.at(number<1>{})),
y_total(mask_coord.at(number<2>{})),
x_total(mask_coord.at(number<3>{})),
repeat_idx(repeat_idx)
repeat_idx(repeat_idx_)
{
}

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

@@ -5,7 +5,7 @@
#include "ck_tile/core.hpp"
#include "ck_tile/ops/common.hpp"
#include "ck_tile/ops/fmha/block/block_masking.hpp"
#include "ck_tile/ops/unified_attention/block/block_masking.hpp"
#include "ck_tile/core/numeric/math.hpp"
#include <string>
@@ -30,9 +30,12 @@ struct UnifiedAttentionKernel
using ODataType = ck_tile::remove_cvref_t<typename UnifiedAttentionPipeline::ODataType>;
using SaccDataType = ck_tile::remove_cvref_t<typename UnifiedAttentionPipeline::SaccDataType>;
using FmhaMask = ck_tile::remove_cvref_t<typename UnifiedAttentionPipeline::FmhaMask>;
static constexpr bool kHasMask = FmhaMask::IsMasking;
static constexpr bool kPadSeqLenK = UnifiedAttentionPipeline::kPadSeqLenK;
static constexpr bool kPadSeqLenQ = UnifiedAttentionPipeline::kPadSeqLenQ;
static constexpr bool kPadHeadDim = UnifiedAttentionPipeline::kPadHeadDim;
static constexpr bool kPadHeadDimQ = UnifiedAttentionPipeline::kPadHeadDimQ;
static constexpr bool kPadHeadDimV = UnifiedAttentionPipeline::kPadHeadDimV;
// TODO add yjese
static constexpr index_t HEAD_SIZE = UnifiedAttentionPipeline::HEAD_SIZE;
@@ -181,7 +184,7 @@ struct UnifiedAttentionKernel
find_seq_idx(const int32_t* query_start_len_ptr,
ck_tile::index_t target_idx,
ck_tile::index_t num_seqs,
ck_tile::index_t BLOCK_Q,
ck_tile::index_t block_q,
bool use_q_block_mode)
{
ck_tile::index_t left = 0;
@@ -191,7 +194,7 @@ struct UnifiedAttentionKernel
{
ck_tile::index_t mid = (left + right) / 2;
ck_tile::index_t val = query_start_len_ptr[mid];
ck_tile::index_t mid_val = use_q_block_mode ? (val / BLOCK_Q + mid) : val;
ck_tile::index_t mid_val = use_q_block_mode ? (val / block_q + mid) : val;
if (mid_val <= target_idx)
{
@@ -276,9 +279,9 @@ struct UnifiedAttentionKernel
const index_t BLOCK_M = BLOCK_Q * kargs.num_queries_per_kv;
// for simplicity, batch stride we just modify the pointer
const index_t num_head_q = kargs.num_head_q;
// const index_t num_head_q = kargs.num_head_q;
const index_t num_queries_per_kv = kargs.num_queries_per_kv;
const index_t num_head_k = num_head_q / num_queries_per_kv;
// const index_t num_head_k = num_head_q / num_queries_per_kv;
pid = RemapTileIndices(pid, kargs);
@@ -311,14 +314,14 @@ struct UnifiedAttentionKernel
const index_t query_pos = q_block_local_idx * BLOCK_Q;
const index_t seq_len = kargs.seq_lens_ptr[seq_idx];
const index_t context_len = seq_len - cur_batch_query_len;
// const index_t context_len = seq_len - cur_batch_query_len;
const index_t max_seq_prefix_len = (
context_len
+ q_block_local_idx * BLOCK_Q
+ (BLOCK_M - 1) // num_queries_per_kv
+ 1
);
// const index_t max_seq_prefix_len = (
// context_len
// + q_block_local_idx * BLOCK_Q
// + (BLOCK_M - 1) // num_queries_per_kv
// + 1
// );
index_t kv_head_offset = kv_head_idx * kargs.stride_k_cache_2;
@@ -463,9 +466,8 @@ struct UnifiedAttentionKernel
return UnifiedAttentionPipeline{}(q_dram_window,
k_dram_window,
v_dram_window,
block_tables_ptr,
kargs.block_tables_ptr,
block_table_offset,
lse_dram_window,
mask,
kargs.scale_s,
smem_ptr);
@@ -484,7 +486,7 @@ struct UnifiedAttentionKernel
o_dram_base,
// block sizes
make_tuple(BLOCK_Q, 1, HEAD_SIZE_PADDED),
sequence<is_seq_len_aligned, false, kPadHeadDimQ>{}
sequence<is_query_len_padded, false, kPadHeadDimQ>{}
); // pads to (seq_len_padded, num_head_q, HEAD_SIZE_PADDED)
const auto o_dram_merged = transform_tensor_view(

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

@@ -277,8 +277,10 @@ struct UnifiedAttentionPipeline
static_assert(HEAD_SIZE_PADDED <= 256, "hdim bigger than 256 is not suitable for this pipeline!");
static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ;
static constexpr bool kPadHeadDim = Problem::kPadHeadDim;
static constexpr bool kStoreLSE = Problem::kStoreLSE;
static constexpr bool kPadSeqLenK = Problem::kPadSeqLenK;
static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV;
// static constexpr bool kStoreLSE = Problem::kStoreLSE;
// 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
@@ -387,7 +389,6 @@ struct UnifiedAttentionPipeline
index_t num_queries_per_kv,
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,
const PComputeElementFunction& p_compute_element_func,
@@ -554,15 +555,7 @@ struct UnifiedAttentionPipeline
{
if(num_total_loop <= 0)
{
if constexpr(kStoreLSE)
{
auto lse =
make_static_distributed_tensor<LSEDataType>(m.get_tile_distribution());
set_tile(lse, -numeric<SMPLComputeDataType>::infinity());
store_tile(lse_dram_window_tmp, tile_elementwise_in(lse_element_func, lse));
}
// Note: here occ are all cleard, return it
// Note: q loaded but no fence, ignore it.
@@ -1193,19 +1186,6 @@ struct UnifiedAttentionPipeline
fmha_post_process(number<0>{});
}
// store lse
if constexpr(kStoreLSE)
{
auto lse = make_static_distributed_tensor<LSEDataType>(m.get_tile_distribution());
constexpr auto lse_spans = decltype(lse)::get_distributed_spans();
sweep_tile_span(lse_spans[number<0>{}], [&](auto idx0) {
constexpr auto i_idx = make_tuple(idx0);
lse(i_idx) = m[i_idx] / C_LOG2E + log(l[i_idx]);
});
store_tile(lse_dram_window_tmp, tile_elementwise_in(lse_element_func, lse));
}
// finally, O
constexpr auto o_spans = decltype(o_acc)::get_distributed_spans();