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

remove padding for seq len Q

Browse Source
This commit is contained in:
ltqin
2025-07-08 10:25:38 -05:00
parent 60ec161e0c
commit db82ecaf78
2 changed files with 5 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
include/ck_tile/ops/fmha/kernel/fmha_fwd_pagedkv_kernel.hpp
View File

@@ -1016,14 +1016,14 @@ struct FmhaFwdPagedKVKernel
return pad_tensor_view(
q_dram_naive,
make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kSubQKHeaddim>{}),
sequence<kPadSeqLenQ, kPadHeadDimQ>{});
sequence<false, kPadHeadDimQ>{});
}
else
{
return pad_tensor_view(
q_dram_naive,
make_tuple(number<FmhaPipeline::kM0>{}, number<FmhaPipeline::kK0>{}),
sequence<kPadSeqLenQ, kPadHeadDimQ>{});
sequence<false, kPadHeadDimQ>{});
}
}();
@@ -1120,7 +1120,7 @@ struct FmhaFwdPagedKVKernel
const index_t num_blocks =
integer_divide_ceil(kv_l2p_offset + kargs.seqlen_k, kargs.page_block_size);
const long_index_t fixed_offset = i_nhead_ * kargs.nhead_stride_k;
const long_index_t fixed_offset = static_cast<long_index_t>(i_nhead_) * kargs.nhead_stride_k;
return make_page_block_navigator<const KDataType, 0>(
kargs.k_ptr,
@@ -1150,7 +1150,7 @@ struct FmhaFwdPagedKVKernel
const index_t num_blocks =
integer_divide_ceil(kv_l2p_offset + kargs.seqlen_k, kargs.page_block_size);
const long_index_t fixed_offset = i_nhead_ * kargs.nhead_stride_v;
const long_index_t fixed_offset = static_cast<long_index_t>(i_nhead_) * kargs.nhead_stride_v;
return make_page_block_navigator<const VDataType, 1>(
kargs.v_ptr,
@@ -1197,7 +1197,7 @@ struct FmhaFwdPagedKVKernel
return pad_tensor_view(bias_dram_naive,
bias_dram_window_lengths,
sequence<kPadSeqLenQ, kPadSeqLenK>{});
sequence<false, kPadSeqLenK>{});
}();
return make_tile_window(bias_dram, bias_dram_window_lengths, {i_m0, 0});

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

@@ -110,7 +110,6 @@ struct BlockFmhaFwdPagedKVPipelineProblem
static constexpr bool kStoreLSE = Traits::kStoreLSE;
static constexpr bool kDoFp8StaticQuant = Traits::kDoFp8StaticQuant;
static constexpr bool kIsPagedKV = Traits::kIsPagedKV;
static constexpr bool kMergeNumHeadGroupsSeqLenQ = Traits::kMergeNumHeadGroupsSeqLenQ;
static constexpr index_t kBlockPerCu = Traits::kBlockPerCu;
};