From 97e7527eb1721c5d243797b17da45d1d72b67f53 Mon Sep 17 00:00:00 2001 From: Juuso Korhonen <40278371+juuso-oskari@users.noreply.github.com> Date: Mon, 20 Oct 2025 14:03:15 +0000 Subject: [PATCH] fixing compile errors... --- .../unified_attention/block/block_masking.hpp | 12 +++---- .../kernel/unified_attention_kernel.hpp | 36 ++++++++++--------- .../pipeline/unified_attention_pipeline.hpp | 30 +++------------- 3 files changed, 30 insertions(+), 48 deletions(-) diff --git a/include/ck_tile/ops/unified_attention/block/block_masking.hpp b/include/ck_tile/ops/unified_attention/block/block_masking.hpp index 958f036edf..87868a56a1 100644 --- a/include/ck_tile/ops/unified_attention/block/block_masking.hpp +++ b/include/ck_tile/ops/unified_attention/block/block_masking.hpp @@ -86,24 +86,24 @@ struct GenericAttentionMask static constexpr const char* name = impl::MaskName::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 - 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_) { } diff --git a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp index 42e0dd25ff..b5d46c754f 100644 --- a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp +++ b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp @@ -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 @@ -30,9 +30,12 @@ struct UnifiedAttentionKernel using ODataType = ck_tile::remove_cvref_t; using SaccDataType = ck_tile::remove_cvref_t; using FmhaMask = ck_tile::remove_cvref_t; - + 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{} + sequence{} ); // pads to (seq_len_padded, num_head_q, HEAD_SIZE_PADDED) const auto o_dram_merged = transform_tensor_view( diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp index 892a5a3db0..5c1a91fb22 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp @@ -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(m.get_tile_distribution()); - - set_tile(lse, -numeric::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(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();