Skip to content

CK mha bwd: add sink attention score gradient support #2321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
poyenc merged 16 commits into from
Apr 4, 2026
Merged

CK mha bwd: add sink attention score gradient support #2321

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
d3fd124
CK mha bwd: add sink attention score gradient support
LJ-underdog Mar 18, 2026
d974b3e
test: add varlen sink bwd tests to test_mha_sink_bwd
LJ-underdog Mar 18, 2026
b065c60
Update op_tests/test_mha_sink_bwd.py
LJ-underdog Mar 18, 2026
a7bf4ae
style: apply black formatting to test_mha_sink_bwd
LJ-underdog Mar 18, 2026
4ac64db
test: move sink bwd tests into test_mha.py and test_mha_varlen.py
LJ-underdog Mar 18, 2026
3690ad1
style: apply black formatting to sink bwd tests in test_mha and test_…
LJ-underdog Mar 18, 2026
b559c5a
Merge branch 'main' into lj_ck_sink_bwd_v2
LJ-underdog Apr 1, 2026
35200bf
fix: adapt mha bwd to updated CK fmha_bwd API and zero dq_accum
LJ-underdog Apr 1, 2026
7481fd6
update ck to ROCm/rocm-libraries#5504
LJ-underdog Apr 2, 2026
852a009
Revert "update ck to ROCm/rocm-libraries#5504"
LJ-underdog Apr 2, 2026
aaf8519
update ck commit
LJ-underdog Apr 2, 2026
81baf75
Merge remote-tracking branch 'origin/main' into lj_ck_sink_bwd_v2
LJ-underdog Apr 2, 2026
d93c1b2
update bwd args
LJ-underdog Apr 2, 2026
8dd0643
[CK] update mha bwd traits args and fix sink_ptr comments
LJ-underdog Apr 3, 2026
7120f1c
[CK] fix mha_bwd_args initializer in benchmark_mha_bwd.cpp for sink_p…
LJ-underdog Apr 3, 2026
8d2d932
Merge branch 'main' into lj_ck_sink_bwd_v2
poyenc Apr 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion 3rdparty/composable_kernel
Open in desktop
Submodule composable_kernel updated 234 files
18 changes: 12 additions & 6 deletions aiter/ops/mha.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -624,7 +624,8 @@ def cmdGenFunc_mha_bwd(
alibi_slopes: Optional[Tensor] = None,
rng_state: Optional[Tensor] = None,
gen: Optional[Generator] = None,
sink_ptr: Optional[Tensor] = None,
sink: Optional[Tensor] = None,
d_sink: Optional[Tensor] = None,
):
md_name = "mha_bwd"
filter1 = "*" # get_bwd_dot_do_o_blobs()
Expand Down Expand Up @@ -775,7 +776,8 @@ def gen_mha_bwd_fake_tensors(
alibi_slopes: Optional[Tensor] = None,
rng_state: Optional[Tensor] = None,
gen: Optional[Generator] = None,
sink_ptr: Optional[Tensor] = None,
sink: Optional[Tensor] = None,
d_sink: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
return common_mha_bwd_fake_tensors(q, k, v, dq, dk, dv)

Expand Down Expand Up @@ -807,7 +809,8 @@ def mha_bwd(
alibi_slopes: Optional[Tensor] = None,
rng_state: Optional[Tensor] = None,
gen: Optional[Generator] = None,
sink_ptr: Optional[Tensor] = None,
sink: Optional[Tensor] = None,
d_sink: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor, Tensor, Tensor]: ...


Expand Down Expand Up @@ -889,7 +892,8 @@ def cmdGenFunc_mha_varlen_bwd(
gen: Optional[Generator] = None,
cu_seqlens_q_padded: Optional[Tensor] = None,
cu_seqlens_k_padded: Optional[Tensor] = None,
sink_ptr: Optional[Tensor] = None,
sink: Optional[Tensor] = None,
d_sink: Optional[Tensor] = None,
) -> dict[str, Any]:
md_name = "mha_varlen_bwd"
filter1 = "*" # get_bwd_dot_do_o_blobs()
Expand Down Expand Up @@ -1117,7 +1121,8 @@ def gen_mha_varlen_bwd_fake_tensors(
alibi_slopes: Optional[Tensor] = None,
rng_state: Optional[Tensor] = None,
gen: Optional[Generator] = None,
sink_ptr: Optional[Tensor] = None,
sink: Optional[Tensor] = None,
d_sink: Optional[Tensor] = None,
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
return gen_mha_varlen_bwd_fake_tensors_common(
q, k, v, cu_seqlens_q, max_seqlen_q, zero_tensors, dq, dk, dv
Expand Down Expand Up @@ -1157,7 +1162,8 @@ def mha_varlen_bwd(
gen: Optional[Generator] = None,
cu_seqlens_q_padded: Optional[Tensor] = None,
cu_seqlens_k_padded: Optional[Tensor] = None,
sink_ptr: Optional[Tensor] = None,
sink: Optional[Tensor] = None,
d_sink: Optional[Tensor] = None,
Comment thread
LJ-underdog marked this conversation as resolved.
Resolved
Comment thread
LJ-underdog marked this conversation as resolved.
Resolved
) -> Tuple[Tensor, Tensor, Tensor, Tensor]: ...


Expand Down
13 changes: 4 additions & 9 deletions csrc/cpp_itfs/mha_bwd.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,15 +134,8 @@ float mha_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
return asm_ret;
#else // !ONLY_FAV3
const fmha_bwd_traits traits{
a.seqlen_q,
a.seqlen_k,
a.batch,
a.max_seqlen_q,
a.max_seqlen_k,
a.hdim_q,
a.hdim_v,
a.nhead_q,
a.nhead_k,
a.data_type,
a.is_group_mode,
static_cast<mask_enum>(a.mask_type),
Expand All @@ -168,6 +161,8 @@ float mha_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
/* dv_ptr */ a.dv_ptr,
/* dbias_ptr */ a.dbias_ptr,
/* dq_acc_ptr */ a.dq_acc_ptr,
/* sink_ptr */ a.sink_ptr,
/* d_sink_ptr */ a.d_sink_ptr,

/* seqstart_q_ptr */ a.seqstart_q_ptr,
/* seqstart_k_ptr */ a.seqstart_k_ptr,
Expand Down Expand Up @@ -208,7 +203,7 @@ float mha_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
/* nhead_stride_randval*/ a.nhead_stride_randval,
/* nhead_stride_do */ a.nhead_stride_do,
/* nhead_stride_lsed */ a.nhead_stride_lsed,
/* nhead_stride_dq_acc*/ a.nhead_stride_dq_acc,
/* nhead_stride_dq_acc*/ static_cast<ck_tile::index_t>(a.nhead_stride_dq_acc),
/* nhead_stride_dq */ a.nhead_stride_dq,
/* nhead_stride_dk */ a.nhead_stride_dk,
/* nhead_stride_dv */ a.nhead_stride_dv,
Expand All @@ -222,7 +217,7 @@ float mha_bwd(mha_bwd_args a, const ck_tile::stream_config& s)
/* batch_stride_randval*/ a.batch_stride_randval,
/* batch_stride_do */ a.batch_stride_do,
/* batch_stride_lsed */ a.batch_stride_lsed,
/* batch_stride_dq_acc*/ a.batch_stride_dq_acc,
/* batch_stride_dq_acc*/ static_cast<ck_tile::index_t>(a.batch_stride_dq_acc),
/* batch_stride_dq */ a.batch_stride_dq,
/* batch_stride_dk */ a.batch_stride_dk,
/* batch_stride_dv */ a.batch_stride_dv,
Expand Down
2 changes: 2 additions & 0 deletions csrc/include/mha_bwd.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@ struct mha_bwd_args
void* dv_ptr;
void* dbias_ptr;
void* dq_acc_ptr;
const void* sink_ptr = nullptr; // sink scores [batch, nhead] log-space (LSEDataType=float); nullptr disables sink
void* d_sink_ptr = nullptr; // sink gradient accumulator [nhead] (LSEDataType=float); nullptr disables sink grad
Comment thread
LJ-underdog marked this conversation as resolved.
Resolved
// Usage notes for sequence length pointer parameters:
//
// [Note: Define "Group mode" vs "Batch mode" here if possible, e.g., "Group mode handles
Expand Down
8 changes: 6 additions & 2 deletions csrc/include/rocm_ops.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -816,7 +816,9 @@ namespace py = pybind11;
py::arg("bias") = std::nullopt, \
py::arg("alibi_slopes") = std::nullopt, \
py::arg("rng_state") = std::nullopt, \
py::arg("gen") = std::nullopt);
py::arg("gen") = std::nullopt, \
py::arg("sink") = std::nullopt, \
py::arg("d_sink") = std::nullopt);

#define MHA_FWD_ASM_PYBIND \
m.def("fmha_v3_fwd", \
Expand Down Expand Up @@ -948,7 +950,9 @@ namespace py = pybind11;
py::arg("rng_state") = std::nullopt, \
py::arg("gen") = std::nullopt, \
py::arg("cu_seqlens_q_padded") = std::nullopt, \
py::arg("cu_seqlens_k_padded") = std::nullopt);
py::arg("cu_seqlens_k_padded") = std::nullopt, \
py::arg("sink") = std::nullopt, \
py::arg("d_sink") = std::nullopt);

#define MOE_CK_2STAGES_PYBIND \
m.def("ck_moe_stage1", \
Expand Down
4 changes: 3 additions & 1 deletion csrc/include/torch/mha_bwd.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ std::vector<at::Tensor> mha_bwd(const at::Tensor& dout, // [b, sq, hq, d]
std::optional<const at::Tensor> bias_, // [sq, sk]
std::optional<const at::Tensor> alibi_slopes, // [hq] or [b, hq]
std::optional<const at::Tensor> rng_state,
std::optional<at::Generator> gen);
std::optional<at::Generator> gen,
std::optional<const at::Tensor> sink, // [b, hq] log-space sink scores (float)
std::optional<at::Tensor> d_sink); // [hq] sink gradient output (float)
} // namespace torch_itfs
} // namespace aiter
4 changes: 3 additions & 1 deletion csrc/include/torch/mha_varlen_bwd.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,9 @@ mha_varlen_bwd(const at::Tensor& dout, // [total_q, hq, d]
std::optional<const at::Tensor> rng_state,
std::optional<at::Generator> gen,
std::optional<const at::Tensor> cu_seqlens_q_padded, // [b+1]
std::optional<const at::Tensor> cu_seqlens_k_padded // [b+1]
std::optional<const at::Tensor> cu_seqlens_k_padded, // [b+1]
std::optional<const at::Tensor> sink, // [b, hq] log-space sink scores (float)
std::optional<at::Tensor> d_sink // [hq] sink gradient output (float)
);
} // namespace torch_itfs
} // namespace aiter
31 changes: 16 additions & 15 deletions csrc/py_itfs_ck/mha_bwd_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@ mha_bwd(const at::Tensor &dout, // [b, sq, hq, d_v]
std::optional<const at::Tensor> bias_, // [sq, sk]
std::optional<const at::Tensor> alibi_slopes_, // [hq] or [b, hq]
std::optional<const at::Tensor> rng_state_,
std::optional<at::Generator> gen_)
std::optional<at::Generator> gen_,
std::optional<const at::Tensor> sink_, // [b, hq] log-space sink scores (float)
std::optional<at::Tensor> d_sink_) // [hq] sink gradient output (float)
Comment thread
LJ-underdog marked this conversation as resolved.
Resolved
{
if (is_causal) { window_size_right = 0; }

Expand Down Expand Up @@ -133,15 +135,8 @@ mha_bwd(const at::Tensor &dout, // [b, sq, hq, d_v]
bool has_dbias = dbias_.has_value();
auto opts = q.options();
const fmha_bwd_traits traits{
seqlen_q,
seqlen_k,
batch_size,
seqlen_q, // max_seqlen_q
seqlen_k, // max_seqlen_k
head_size_q,
head_size_v,
num_heads,
num_heads_k,
q_dtype_str,
false, // is_group_mode
mask.type,
Expand All @@ -155,13 +150,14 @@ mha_bwd(const at::Tensor &dout, // [b, sq, hq, d_v]
const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard{q.device()};

auto softmax_d = torch::empty({batch_size, num_heads, seqlen_q}, opts.dtype(at::kFloat));
const fmha_bwd_launcher launcher(traits);
const ck_tile::index_t nsplits = launcher.dq_acc_splits;
at::Tensor dq_accum;
if (launcher.needs_zero_dq_acc)
dq_accum = torch::zeros({batch_size, num_heads, nsplits, seqlen_q, head_size_q}, opts.dtype(at::kFloat));
else
dq_accum = torch::empty({batch_size, num_heads, nsplits, seqlen_q, head_size_q}, opts.dtype(at::kFloat));
// nsplits: deterministic mode splits dK into ceil(seqlen_k/16) pieces for atomic-free accumulation.
constexpr ck_tile::index_t kN0 = 16;
const ck_tile::index_t nsplits = deterministic
? ck_tile::integer_divide_ceil(seqlen_k, kN0)
: 1;
// Always zero dq_accum: the dq_dk_dv kernel writes via atomicAdd regardless of
// deterministic mode, so an uninitialized accumulator would corrupt dQ.
at::Tensor dq_accum = torch::zeros({batch_size, num_heads, nsplits, seqlen_q, head_size_q}, opts.dtype(at::kFloat));

at::Tensor dk_expanded, dv_expanded;
if (num_heads_k != num_heads) { // MQA / GQA
Expand Down Expand Up @@ -198,6 +194,9 @@ mha_bwd(const at::Tensor &dout, // [b, sq, hq, d_v]
hipLaunchKernelGGL(
aiter::ParsePhiloxCudaState, dim3(1), dim3(64), 0, 0,
philox_args, reinterpret_cast<uint64_t*>(rng_state.data_ptr()));
} else {
// No dropout: allocate a dummy tensor so data_ptr() is always valid.
rng_state = torch::empty({2}, opts.dtype(torch::kInt64));
}

if (seqlen_q > 0) {
Expand Down Expand Up @@ -329,6 +328,8 @@ mha_bwd(const at::Tensor &dout, // [b, sq, hq, d_v]
dv_expanded.data_ptr(),
dbias_ptr,
dq_accum.data_ptr(),
(sink_.has_value() && sink_.value().defined()) ? sink_.value().data_ptr() : nullptr, // sink_ptr [b, hq]
(d_sink_.has_value() && d_sink_.value().defined()) ? d_sink_.value().data_ptr() : nullptr, // d_sink_ptr [hq]
Comment thread
LJ-underdog marked this conversation as resolved.
Resolved
nullptr, // seqstart_q_ptr
nullptr, // seqstart_k_ptr
nullptr, // seqlen_q_ptr
Expand Down
29 changes: 13 additions & 16 deletions csrc/py_itfs_ck/mha_varlen_bwd_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,9 @@ mha_varlen_bwd(const at::Tensor &dout, // [total_q, hq, d_v]
std::optional<const at::Tensor> rng_state_,
std::optional<at::Generator> gen_,
std::optional<const at::Tensor> cu_seqlens_q_padded, // [b+1]
std::optional<const at::Tensor> cu_seqlens_k_padded // [b+1]
)
std::optional<const at::Tensor> cu_seqlens_k_padded, // [b+1]
std::optional<const at::Tensor> sink_, // [b, hq] log-space sink scores (float)
std::optional<at::Tensor> d_sink_) // [hq] sink gradient output (float)
Comment thread
LJ-underdog marked this conversation as resolved.
Resolved
{
if (is_causal) { window_size_right = 0; }

Expand Down Expand Up @@ -151,15 +152,8 @@ mha_varlen_bwd(const at::Tensor &dout, // [total_q, hq, d_v]
bias_enum bias_type = alibi_slopes_.has_value() ? bias_enum::alibi : bias_enum::no_bias;
auto opts = q.options();
const fmha_bwd_traits traits{
total_q,
total_k,
batch_size,
max_seqlen_q,
max_seqlen_k,
head_size_q,
head_size_v,
num_heads,
num_heads_k,
q_dtype_str,
true, // is_group_mode
mask.type,
Expand All @@ -169,17 +163,18 @@ mha_varlen_bwd(const at::Tensor &dout, // [total_q, hq, d_v]
false, // is_store_randval
deterministic,
};
const fmha_bwd_launcher launcher(traits);
const ck_tile::index_t nsplits = launcher.dq_acc_splits;
// nsplits: deterministic mode splits dK into ceil(max_seqlen_k/16) pieces for atomic-free accumulation.
constexpr ck_tile::index_t kN0 = 16;
const ck_tile::index_t nsplits = deterministic
? ck_tile::integer_divide_ceil(max_seqlen_k, kN0)
: 1;

const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard{q.device()};

auto softmax_d = torch::empty({batch_size, num_heads, total_q}, opts.dtype(at::kFloat));
at::Tensor dq_accum;
if (launcher.needs_zero_dq_acc)
dq_accum = torch::zeros({num_heads, nsplits, total_q, head_size_q}, opts.dtype(at::kFloat));
else
dq_accum = torch::empty({num_heads, nsplits, total_q, head_size_q}, opts.dtype(at::kFloat));
// Always zero dq_accum: the dq_dk_dv kernel writes via atomicAdd regardless of
// deterministic mode, so an uninitialized accumulator would corrupt dQ.
at::Tensor dq_accum = torch::zeros({num_heads, nsplits, total_q, head_size_q}, opts.dtype(at::kFloat));

at::Tensor dk_expanded, dv_expanded;
if (num_heads_k != num_heads) { // MQA / GQA
Expand Down Expand Up @@ -341,6 +336,8 @@ mha_varlen_bwd(const at::Tensor &dout, // [total_q, hq, d_v]
dv_expanded.data_ptr(),
nullptr, // dbias
dq_accum.data_ptr(), // dq_acc
(sink_.has_value() && sink_.value().defined()) ? sink_.value().data_ptr() : nullptr, // sink_ptr [b, hq]
(d_sink_.has_value() && d_sink_.value().defined()) ? d_sink_.value().data_ptr() : nullptr, // d_sink_ptr [hq]
Comment thread
LJ-underdog marked this conversation as resolved.
Resolved
seqstart_q_ptr, // seqstart_q_ptr (physical cumulative)
seqstart_k_ptr, // seqstart_k_ptr (physical cumulative)
nullptr, // seqlen_q_ptr (per-sequence logical)
Expand Down
1 change: 1 addition & 0 deletions csrc/py_itfs_cu/fmha_bwd_pre_post_kernel_generate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ def get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype: str) -> Optional[dict]
typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::ODataType,
typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::OGradDataType,
typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::DDataType,
typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::LSEDataType,
/* BlockSize = */ 64,
{F_hdim},
{F_mode},
Expand Down
Loading