Skip to content

[CK_TILE] FMHA BWD: stream-async workspace prepare #183

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
DDEle merged 3 commits into from
May 19, 2026
Merged

[CK_TILE] FMHA BWD: stream-async workspace prepare #183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e38b717
[CK_TILE] FMHA BWD: stream-async workspace prepare
DDEle May 13, 2026
958a108
[CK_TILE] FMHA BWD: bump CK submodule to develop tip (#7331 merged)
DDEle May 14, 2026
605fb88
[CK_TILE] FMHA BWD: explicit at::kCPU on pinned host TensorOptions
DDEle May 19, 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 csrc/composable_kernel
Open in desktop
Submodule composable_kernel updated 116 files
18 changes: 17 additions & 1 deletion csrc/flash_attn_ck/mha_bwd.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,23 @@ mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num
workspace = torch::empty({static_cast<int64_t>(launcher.workspace_size)},
opts.dtype(at::kByte));
workspace_ptr = workspace.data_ptr();
launcher.prepare_workspace(workspace_ptr);
// Pinned host buffer allocator backed by PyTorch's CachingHostAllocator.
// The returned shared_ptr owns the at::Tensor; the launcher keeps it
// alive via a stream-tail hipLaunchHostFunc keepalive. Required when
// the launcher needs host-side workspace metadata (deterministic mode
// and/or non-trivial worker state); harmless when it doesn't.
auto pinned_host_alloc = [](size_t bytes) -> std::shared_ptr<void> {
auto t = std::make_shared<at::Tensor>(torch::empty(
{static_cast<int64_t>(bytes)},
torch::TensorOptions().dtype(at::kByte).device(at::kCPU).pinned_memory(true)));
return std::shared_ptr<void>(t, t->data_ptr());
};
ck_tile::stream_config prep_cfg{stream};
launcher.prepare_workspace_async(workspace_ptr,
/*seqstart_q_dev=*/nullptr,
/*seqstart_k_dev=*/nullptr,
prep_cfg,
pinned_host_alloc);
}

at::Tensor dk_expanded, dv_expanded;
Expand Down
35 changes: 20 additions & 15 deletions csrc/flash_attn_ck/mha_varlen_bwd.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,7 @@ fmha_bwd_traits get_ck_fmha_varlen_bwd_traits(const mask_info &mask,
int nhead_k,
bool has_dropout,
bool enable_alibi,
bool deterministic,
const int* seqstart_qs,
const int* seqstart_ks)
bool deterministic)
{
return fmha_bwd_traits{seqlen_q,
seqlen_k,
Expand All @@ -39,9 +37,7 @@ fmha_bwd_traits get_ck_fmha_varlen_bwd_traits(const mask_info &mask,
false, // has_dbias
has_dropout,
false, // s_randval
deterministic,
seqstart_qs,
seqstart_ks};
deterministic};
}
fmha_bwd_args get_ck_fmha_varlen_bwd_args(const mask_info &mask,
// sizes
Expand Down Expand Up @@ -339,11 +335,6 @@ mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads
dv = torch::empty_like(v);
}

// seqstart_qs/ks are dereferenced on the HOST inside fmha_bwd_launcher constructor
// (to compute workspace sizes), so we must use host copies.
at::Tensor cu_seqlens_q_host = cu_seqlens_q.cpu();
at::Tensor cu_seqlens_k_host = cu_seqlens_k.cpu();

const auto traits = get_ck_fmha_varlen_bwd_traits(
mask,
q_dtype_str,
Expand All @@ -357,9 +348,7 @@ mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads
num_heads_k,
is_dropout,
alibi_slopes_.has_value(),
deterministic,
reinterpret_cast<const int*>(cu_seqlens_q_host.data_ptr()),
reinterpret_cast<const int*>(cu_seqlens_k_host.data_ptr()));
deterministic);
fmha_bwd_launcher launcher(traits);

at::cuda::CUDAGuard device_guard{q.device()};
Expand All @@ -377,7 +366,23 @@ mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads
workspace = torch::empty({static_cast<int64_t>(launcher.workspace_size)},
opts.dtype(at::kByte));
workspace_ptr = workspace.data_ptr();
launcher.prepare_workspace(workspace_ptr);
// Pinned host buffer allocator backed by PyTorch's CachingHostAllocator.
// The returned shared_ptr owns the at::Tensor; the launcher keeps it
// alive via a stream-tail hipLaunchHostFunc keepalive so the buffer
// is not recycled while async D2H/H2D copies are still in flight.
auto pinned_host_alloc = [](size_t bytes) -> std::shared_ptr<void> {
auto t = std::make_shared<at::Tensor>(torch::empty(
{static_cast<int64_t>(bytes)},
torch::TensorOptions().dtype(at::kByte).device(at::kCPU).pinned_memory(true)));
return std::shared_ptr<void>(t, t->data_ptr());
};
ck_tile::stream_config prep_cfg{stream};
launcher.prepare_workspace_async(
workspace_ptr,
reinterpret_cast<const int*>(cu_seqlens_q.data_ptr()),
reinterpret_cast<const int*>(cu_seqlens_k.data_ptr()),
prep_cfg,
pinned_host_alloc);
}

at::Tensor dk_expanded, dv_expanded;
Expand Down