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

[CK_TILE] fmha: Add query padding support to backward pass

Browse Source 
Introduces support for query sequence padding (q_padding) in the FMHA backward pass kernels.
- Passing `seqlen_q_ptr` to the backward kernels to distinguish logical from physical sequence lengths.
- Updating `OGradDotO`, `ConvertQGrad`, and `DQDKDV` kernels to respect logical lengths and handle zero-length sequences.
- Aligning LSE indexing in the forward kernel with the padded layout for consistency.
- Adding a new GTest suite (`test_fmha_bwd_kernel_padding.cpp`) with comprehensive tests for various padding scenarios, including zero-length
  sequences and deterministic mode.
This commit is contained in:
Jeff Huang
2025-10-08 15:35:02 +08:00
parent 86d542f663
commit 4e06eaa417
10 changed files with 839 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp
View File

@@ -313,6 +313,7 @@ struct FmhaBwdDQDKDVKernel
{
const int32_t* seqstart_q_ptr;
const int32_t* seqstart_k_ptr;
const int32_t* seqlen_q_ptr;
const int32_t* seqlen_k_ptr;
};
@@ -520,6 +521,7 @@ struct FmhaBwdDQDKDVKernel
void* dq_acc_ptr,
const void* seqstart_q_ptr,
const void* seqstart_k_ptr,
const void* seqlen_q_ptr,
const void* seqlen_k_ptr,
ck_tile::index_t hdim_q,
ck_tile::index_t hdim_v,
@@ -594,6 +596,7 @@ struct FmhaBwdDQDKDVKernel
{}, // placeholder for deterministic
reinterpret_cast<const int32_t*>(seqstart_q_ptr),
reinterpret_cast<const int32_t*>(seqstart_k_ptr),
reinterpret_cast<const int32_t*>(seqlen_q_ptr),
reinterpret_cast<const int32_t*>(seqlen_k_ptr)};
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
@@ -738,7 +741,10 @@ struct FmhaBwdDQDKDVKernel
// get real # queries & # keys under group mode
const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch;
kargs.seqlen_q = adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0];
const ck_tile::index_t physical_seqlen_q =
adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0];
kargs.seqlen_q = kargs.seqlen_q_ptr ? kargs.seqlen_q_ptr[i_batch] : physical_seqlen_q;
if(kargs.seqlen_k_ptr != nullptr)
{
kargs.seqlen_k = kargs.seqlen_k_ptr[i_batch];
@@ -749,6 +755,12 @@ struct FmhaBwdDQDKDVKernel
kargs.seqlen_k = adjusted_seqstart_k_ptr[1] - adjusted_seqstart_k_ptr[0];
}
// skip if logical lengths are zero
if(kargs.seqlen_q == 0 || kargs.seqlen_k == 0)
{
return;
}
// # of required blocks is different in each groups, terminate unnecessary blocks
// earlier
if constexpr(!kUseQrQtrDorPipeline)
@@ -1246,6 +1258,7 @@ struct FmhaBwdOGradDotOKernel
struct FmhaBwdOGradDotOGroupModeKargs : FmhaBwdOGradDotOCommonKargs
{
const int32_t* seqstart_q_ptr;
const int32_t* seqlen_q_ptr;
};
using Kargs = std::
@@ -1293,6 +1306,7 @@ struct FmhaBwdOGradDotOKernel
void* d_ptr,
float p_undrop,
const void* seqstart_q_ptr,
const void* seqlen_q_ptr,
ck_tile::index_t hdim_v,
ck_tile::index_t stride_do,
ck_tile::index_t stride_o,
@@ -1311,7 +1325,8 @@ struct FmhaBwdOGradDotOKernel
nhead_stride_do,
nhead_stride_o,
nhead_stride_d},
reinterpret_cast<const int32_t*>(seqstart_q_ptr)};
reinterpret_cast<const int32_t*>(seqstart_q_ptr),
reinterpret_cast<const int32_t*>(seqlen_q_ptr)};
return kargs;
}
@@ -1357,7 +1372,12 @@ struct FmhaBwdOGradDotOKernel
// get real # queries & # keys under group mode
const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch;
kargs.seqlen_q = adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0];
const ck_tile::index_t physical_seqlen_q =
adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0];
const ck_tile::index_t logical_seqlen_q =
kargs.seqlen_q_ptr ? static_cast<ck_tile::index_t>(kargs.seqlen_q_ptr[i_batch])
: physical_seqlen_q;
kargs.seqlen_q = logical_seqlen_q;
// # of required blocks is different in each groups, terminate unnecessary blocks
// earlier
if(kargs.seqlen_q <= i_m0)
@@ -1521,6 +1541,8 @@ struct FmhaBwdConvertQGradKernel
{
const int32_t* seqstart_q_ptr;
const int32_t* seqstart_k_ptr;
const int32_t* seqlen_q_ptr;
const int32_t* seqlen_k_ptr;
};
using Kargs = std::conditional_t<kIsGroupMode,
@@ -1569,6 +1591,8 @@ struct FmhaBwdConvertQGradKernel
void* dq_ptr,
const void* seqstart_q_ptr,
const void* seqstart_k_ptr,
const void* seqlen_q_ptr,
const void* seqlen_k_ptr,
ck_tile::index_t hdim_q,
ck_tile::index_t stride_dq,
ck_tile::index_t stride_dq_acc,
@@ -1587,7 +1611,9 @@ struct FmhaBwdConvertQGradKernel
nhead_stride_dq_acc},
{},
reinterpret_cast<const int32_t*>(seqstart_q_ptr),
reinterpret_cast<const int32_t*>(seqstart_k_ptr)};
reinterpret_cast<const int32_t*>(seqstart_k_ptr),
reinterpret_cast<const int32_t*>(seqlen_q_ptr),
reinterpret_cast<const int32_t*>(seqlen_k_ptr)};
if constexpr(kIsDeterministic)
{
@@ -1634,11 +1660,20 @@ struct FmhaBwdConvertQGradKernel
// get real # queries & # keys under group mode
const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch;
kargs.seqlen_q = adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0];
const ck_tile::index_t physical_seqlen_q =
adjusted_seqstart_q_ptr[1] - adjusted_seqstart_q_ptr[0];
const ck_tile::index_t logical_seqlen_q =
kargs.seqlen_q_ptr ? static_cast<ck_tile::index_t>(kargs.seqlen_q_ptr[i_batch])
: physical_seqlen_q;
kargs.seqlen_q = logical_seqlen_q;
if constexpr(kIsDeterministic)
{
const auto adjusted_seqstart_k_ptr = kargs.seqstart_k_ptr + i_batch;
kargs.seqlen_k = adjusted_seqstart_k_ptr[1] - adjusted_seqstart_k_ptr[0];
const ck_tile::index_t physical_seqlen_k =
adjusted_seqstart_k_ptr[1] - adjusted_seqstart_k_ptr[0];
kargs.seqlen_k = kargs.seqlen_k_ptr
? static_cast<ck_tile::index_t>(kargs.seqlen_k_ptr[i_batch])
: physical_seqlen_k;
}
// # of required blocks is different in each groups, terminate unnecessary blocks
// earlier

8
include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp
View File

@@ -1137,8 +1137,8 @@ struct FmhaFwdKernel
}
if constexpr(kStoreLSE)
{
// LSE stays indexed by unpadded starts
batch_offset_lse = query_start_unpadded;
// LSE follows the padded layout to stay consistent with other tensors
batch_offset_lse = query_start_padded;
}
if constexpr(kHasDropout)
{
@@ -1630,8 +1630,8 @@ struct FmhaFwdKernel
batch_offset_bias = query_start_padded * kargs.stride_bias;
}
// LSE layout is [nhead, total_seqlen], index by unpadded start
batch_offset_lse = query_start_unpadded;
// LSE layout is [nhead, total_seqlen] following the padded layout for Q/O
batch_offset_lse = query_start_padded;
batch_offset_o = query_start_padded * kargs.stride_o;
// get real # queries & # keys under group mode