ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-30 19:57:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'tianxing/unified-attention' of https://github.com/ROCm/composable_kernel into tianxing/unified-attention

Browse Source
This commit is contained in:
Tianxing Wu
2025-10-17 09:05:12 +00:00
parent af9167abad 9940bd07f6
commit 995c6701d3
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -449,10 +449,10 @@ struct UnifiedAttentionKernel
FmhaMask mask = [&]() {
if constexpr(kHasMask)
return ck_tile::make_generic_attention_mask_from_lr_window<FmhaMask>(
cur_batch_query_len, // x (i.e. extend)
seq_len - cur_batch_query_len, // y (i.e. context)
cur_batch_query_len, // x_total
cur_batch_query_len, // x (i.e. extend)
seq_len, // y_total (x + y)
cur_batch_query_len, // x_total
num_queries_per_kv, // the same sequence index is repeated num_queries_per_kv times along x dim of the tile
kargs.mask_type == GenericAttentionMaskEnum::MASK_FROM_TOP_LEFT);
else