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 f4f4ebe726..fc067a40ce 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 @@ -97,7 +97,6 @@ struct FmhaFwdV3Kernel ck_tile::index_t num_seqs; // number of batches for q }; - using Kargs = UnifiedAttentionVarlenKargs; CK_TILE_HOST static constexpr Kargs MakeKargs( @@ -333,6 +332,8 @@ struct FmhaFwdV3Kernel index_t o_ptr_offset_0 = cur_batch_in_all_start_index * kargs.output_stride_0; // move the pointer to the batch start index_t o_ptr_offset_1 = kv_head_idx * num_queries_per_kv * kargs.output_stride_1; // move the pointer to the correct head group start index_t o_ptr_offset = o_ptr_offset_0 + o_ptr_offset_1; + index_t block_table_offset = seq_idx * kargs.block_table_stride; + const QDataType* q_ptr = reinterpret_cast(kargs.q_ptr) + q_ptr_offset; const KDataType* k_ptr = reinterpret_cast(kargs.k_ptr) + kv_head_offset; @@ -462,6 +463,8 @@ struct FmhaFwdV3Kernel return FmhaPipeline{}(q_dram_window, k_dram_window, v_dram_window, + block_tables_ptr, + block_table_offset, lse_dram_window, mask, kargs.scale_s, @@ -502,7 +505,7 @@ struct FmhaFwdV3Kernel auto o_dram_window = make_tile_window(o_dram, make_tuple(BLOCK_M, HEAD_SIZE_PADDED), - {q_block_global_idx * num_queries_per_kv * HEAD_SIZE_PADDED, 0}); + {query_pos * num_queries_per_kv, 0}); EpiloguePipeline{}(o_dram_window, o_acc_tile, nullptr); } 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 b151b61028..8d37d42291 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 @@ -391,6 +391,8 @@ struct UnifiedAttentionPipeline [[maybe_unused]] const KElementFunction& k_element_func, const VDramBlockWindowTmp& v_dram_block_window_tmp, // N1*K1 tile [[maybe_unused]] const VElementFunction& v_element_func, + 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, @@ -402,6 +404,7 @@ struct UnifiedAttentionPipeline { using namespace ck_tile; + static_assert( std::is_same_v> && std::is_same_v> && @@ -573,22 +576,26 @@ struct UnifiedAttentionPipeline } } + index_t i_total_loops = 0; + index_t kv_blk_idx = block_tables_ptr[block_table_offset + i_total_loops]; + index_t kv_blk_idx_prev = 0; + + auto k_dram_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}, + {(kv_blk_idx - kv_blk_idx_prev) * BLOCK_SIZE, 0}, Policy::template MakeKDramTileDistribution()); k_dram_window.init_raw(); auto v_dram_window = make_tile_window(v_dram_block_window_tmp.get_bottom_tensor_view(), v_dram_block_window_tmp.get_window_lengths(), - {seqlen_k_start, 0}, // TODO: hdim split? + {(kv_blk_idx - kv_blk_idx_prev) * BLOCK_SIZE, 0}, // TODO: hdim split? Policy::template MakeVDramTileDistribution()); v_dram_window.init_raw(); // prefetch K tile - index_t i_total_loops = 0; constexpr index_t k0_loops = kQKHeaddim / kK0; constexpr index_t k1_loops = kN0 / kK1; static_assert(1 == k0_loops); @@ -678,10 +685,13 @@ struct UnifiedAttentionPipeline auto K_mem_load = [&](auto k_lds_write_idx) { async_load_tile_raw(k_lds_window_store(k_lds_write_idx), k_dram_window); - + // TODO maybe needs i_total_loops as argument. Or maybe needs to use the k_lds_write_idx as the index /// FIXME: use the future-predicting method to move the window // move K tile windows - move_tile_window(k_dram_window, {kN0, 0}); + auto k_dram_window = make_tile_window(k_dram_window.get_bottom_tensor_view(), + k_dram_window.get_window_lengths(), + {(block_tables_ptr[block_table_offset + i_total_loops]) * BLOCK_SIZE, 0}, + Policy::template MakeVDramTileDistribution()); }; auto K_lds_load = [&](auto k_lds_read_idx) { @@ -692,7 +702,10 @@ struct UnifiedAttentionPipeline async_load_tile_raw(v_lds_window_store(v_lds_write_idx), v_dram_window); /// FIXME: use the future-predicting method to move the window - move_tile_window(v_dram_window, {kK1, 0}); + auto v_dram_window = make_tile_window(v_dram_window.get_bottom_tensor_view(), + v_dram_window.get_window_lengths(), + {(block_tables_ptr[block_table_offset + i_total_loops]) * BLOCK_SIZE, 0}, + Policy::template MakeVDramTileDistribution()); }; auto V_lds_load = [&](auto v_lds_read_idx) { @@ -1231,6 +1244,8 @@ struct UnifiedAttentionPipeline CK_TILE_DEVICE auto operator()(const QDramBlockWindowTmp& q_dram_block_window_tmp, // M0*K0 tile const KDramBlockWindowTmp& k_dram_block_window_tmp, // N0*K0 tile const VDramBlockWindowTmp& v_dram_block_window_tmp, // N1*K1 tile + const void* block_tables_ptr, + index_t block_table_offset, LSEDramBlockWindowTmp& lse_dram_block_window_tmp, // M0*1 tile FmhaMask mask, float scale_s, @@ -1244,6 +1259,8 @@ struct UnifiedAttentionPipeline identity{}, v_dram_block_window_tmp, identity{}, + block_tables_ptr, + block_table_offset, lse_dram_block_window_tmp, identity{}, identity{},