[CK_TILE] FMHA BWD: stream-async workspace prepare for group mode

[DDEle]

Motivation

Adapts AITER to ROCm/rocm-libraries#7331, which exposes a fully stream-async workspace preparation API on fmha_bwd_launcher.

PR #2948 left two synchronous hipMemcpy D2H copies in aiter::mha_bwd's group-mode path (reading seqstart_q/k_ptr for launcher construction). These sync copies block the host (~10–30 µs each) and implicitly synchronize the device by draining the user's stream, breaking CPU/GPU overlap on hot training paths.

This PR drops the two sync copies and stages the seqstart-dependent workspace metadata via a fully async D2H → host-pack → H2D pipeline on the user's stream.

Technical Details

• csrc/include/mha_bwd.h: adds a pinned_host_alloc callback to mha_bwd_args that returns a std::shared_ptr<void> over a pinned host buffer. The deleter is invoked once the launcher releases the buffer on a stream-tail hipLaunchHostFunc.

• csrc/cpp_itfs/mha_bwd.cu: drops the two sync hipMemcpy D2H copies in the group-mode path. Calls launcher.prepare_workspace_async(...) on the user's stream, which schedules hipMemsetAsync (zero dq_acc) → hipMemcpyAsync D2H → hipLaunchHostFunc (host-pack) → hipMemcpyAsync H2D, all on the same stream.

• PyTorch wrappers (csrc/py_itfs_ck/mha_bwd_kernels.cu, mha_varlen_bwd_kernels.cu): pinned_host_alloc lambda backed by at::empty(..., pin_memory=true) (CachingHostAllocator). The returned shared_ptr keepalive is held until the launcher's stream-tail callback fires, so the buffer cannot be recycled while still in flight.

• op_tests/cpp/mha/benchmark_mha_bwd.cpp: pinned_host_alloc backed by bare hipHostMalloc; release routed through ck_tile::pinned_host_releaser to defer hipHostFree off the HIP driver callback thread (which holds runtime locks and would deadlock against concurrent main-thread hipFree).

Bumps composable_kernel submodule to 1f4cc34e6 (PR #7331 tip).

Dependencies

• CK side: ROCm/rocm-libraries#7331. Submodule pinned to a pre-merge commit; will rebase to develop tip once #7331 lands.
• Stacks on [CK_TILE] Use Unified Workspace for FMHA BWD #2948 (base branch is ck-fmha-bwd-workspace).

Test Plan

On gfx950 inside rocm/pytorch:latest:

• python op_tests/test_mha.py -b 2 -n 16 -q 512 -k 512 -d bf16 -d_qk_v 128,128 -gr 1
• python op_tests/test_mha_varlen.py -b 4 -nh 16 -s 512,512 -dt bf16 -d_qk_v 128,128 -gr 1
• op_tests/cpp/mha/bwd.exe -prec=bf16 -b=2 -h=4 -h_k=2 -d=128 -s=512 -mode=1 -v=1

Test Result

• test_mha.py: pass; padding diff 0.0020 < pytorch tol 0.0078
• test_mha_varlen.py: pass; 4 Group Mode Test variants (mixed, q_only, k_only, no_padding) all within tolerance; dQ/dK/dV CK max diff ≤ pytorch ref
• C++ bwd.exe group mode: valid:y

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

This PR updates AITER’s CK_TILE FMHA backward (mha_bwd) “group mode” path to avoid host-blocking, stream-synchronizing D2H copies by switching workspace preparation to a fully stream-ordered async pipeline, enabled by newer Composable Kernel launcher APIs.

Changes:

• Extends aiter::mha_bwd_args with a pinned_host_alloc callback to provide pinned host staging buffers for the async workspace pipeline.
• Switches CK workspace preparation in aiter::mha_bwd to prepare_workspace_async(...), removing synchronous seqstart D2H copies in group mode.
• Updates PyTorch CK wrappers and the C++ benchmark to provide pinned host allocators and to accommodate launcher/workspace handling changes.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
csrc/include/mha_bwd.h	Adds pinned_host_alloc API and documents intended fallback behavior.
csrc/cpp_itfs/mha_bwd.cu	Uses CK launcher async workspace preparation and removes sync seqstart host staging.
csrc/py_itfs_ck/mha_bwd_kernels.cu	Supplies a pinned host allocator via PyTorch pinned-memory tensors.
csrc/py_itfs_ck/mha_varlen_bwd_kernels.cu	Same pinned host allocator approach for varlen backward wrapper.
op_tests/cpp/mha/benchmark_mha_bwd.cpp	Adds pinned host allocation/free strategy and updates workspace allocation approach.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.