[CK_TILE] Use Unified Workspace for FMHA BWD - #182
Merged
Conversation
Bump composable_kernel submodule to mono-split/users/yiding12/fmha-bwd-workspace HEAD and adapt the FMHA BWD host wrappers to the new unified workspace API: - Replace dq_acc tensor argument with workspace_ptr in get_ck_fmha_bwd_args / get_ck_fmha_varlen_bwd_args - Drop dq_acc strides that have been removed from fmha_bwd_args - In mha_bwd / mha_varlen_bwd, allocate the device workspace based on fmha_bwd_launcher::workspace_size and call launcher.prepare_workspace() - Invoke launcher.run(args, stream_config) instead of fmha_bwd(...)
Collaborator
Author
2 tasks
rocking5566
added a commit
that referenced
this pull request
Jul 16, 2026
…ILab#2675) * Add sink_ptr/d_sink_ptr to fmha_bwd_args to match updated CK submodule Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * update submodule * [CK_TILE] Use Unified Workspace for FMHA BWD (#182) * [CK_TILE] Use Unified Workspace for FMHA BWD Bump composable_kernel submodule to mono-split/users/yiding12/fmha-bwd-workspace HEAD and adapt the FMHA BWD host wrappers to the new unified workspace API: - Replace dq_acc tensor argument with workspace_ptr in get_ck_fmha_bwd_args / get_ck_fmha_varlen_bwd_args - Drop dq_acc strides that have been removed from fmha_bwd_args - In mha_bwd / mha_varlen_bwd, allocate the device workspace based on fmha_bwd_launcher::workspace_size and call launcher.prepare_workspace() - Invoke launcher.run(args, stream_config) instead of fmha_bwd(...) * Update CK pin as ROCm/rocm-libraries#6152 merged * [CK_TILE] FMHA BWD: stream-async workspace prepare (#183) * [CK_TILE] FMHA BWD: stream-async workspace prepare Bump composable_kernel submodule to mono-split/users/yiding12/fmha-bwd- async-prepare HEAD and adapt the FMHA BWD host wrappers to the new async workspace prepare API (CK PR #7331): - Replace launcher.prepare_workspace() with prepare_workspace_async(), which enqueues the full workspace setup (dq_acc zero, group-mode D2H of seqstart, host-side metadata pack via hipLaunchHostFunc, H2D back to device) on the caller's stream. No host-blocking sync remains in the BWD launch path. - Pass a pinned_host_alloc lambda backed by PyTorch's CachingHostAllocator (torch::empty(..., pin_memory=true)). The launcher keeps the returned shared_ptr alive via a stream-tail hipLaunchHostFunc keepalive so the pinned buffer is not recycled while async copies are still in flight. - mha_varlen_bwd: drop the cu_seqlens_q.cpu() / cu_seqlens_k.cpu() host copies; the launcher now reads device seqstart directly via async D2H. get_ck_fmha_varlen_bwd_traits no longer takes seqstart_qs/ks. * [CK_TILE] FMHA BWD: bump CK submodule to develop tip (#7331 merged) ROCm/rocm-libraries#7331 (async workspace prepare for FMHA BWD launcher) landed on develop. Move csrc/composable_kernel from the pre-merge fork tip ce838e19e5 to ROCm/composable_kernel develop tip 83566edb0f, which is the split commit for #7331 (rocm-libraries 5692db0). * [CK_TILE] FMHA BWD: explicit at::kCPU on pinned host TensorOptions * Update CK and enable RDNA backward --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Yi DING <yi.ding@amd.com> Co-authored-by: Hosang Yoon <hosang.yoon@amd.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Bump composable_kernel submodule to mono-split/users/yiding12/fmha-bwd-workspace
HEAD and adapt the FMHA BWD host wrappers to the new unified workspace API:
/ get_ck_fmha_varlen_bwd_args
fmha_bwd_launcher::workspace_size and call launcher.prepare_workspace()