[CK_TILE] FMHA BWD: stream-async workspace prepare

[DDEle]

• Depends on [CK_TILE] Add async workspace prepare to FMHA BWD launcher rocm-libraries#7331
• Stacked on [CK_TILE] Use Unified Workspace for FMHA BWD #182

Adapts the FMHA BWD host wrappers to the new stream-async workspace prepare
API in CK PR #7331, eliminating the host-blocking cu_seqlens.cpu() D2H
sync that PR #182 had to perform on every group-mode bwd call.

Changes

• csrc/composable_kernel: bump submodule to mono-split/users/yiding12/fmha-bwd-async-prepare HEAD.
• csrc/flash_attn_ck/mha_bwd.cpp / mha_varlen_bwd.cpp:
  • 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.
  • 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 (group mode): 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.

Equivalence with aiter

This PR is the FA-side parallel of aiter PR ROCm/aiter#3150, just as #182 is
the FA-side parallel of aiter PR ROCm/aiter#2948.

Validation

• pytest tests/test_flash_attn_ck.py on gfx950 with
  rocm/pytorch:rocm7.1.1_ubuntu24.04_py3.12_pytorch_release_2.9.1:
  260680 passed, 152076 skipped, 0 failed (22:43)

[DDEle]

jenkins / rocm-libraries-folder/Composable Kernel / PR-7331 / #13
==>
jenkins / rocm-libraries-folder/Composable Kernel / PR-7331 / #17

[Copilot]

Pull request overview

This PR updates the CK_TILE FlashAttention backward (FMHA BWD) host wrappers to use the new stream-asynchronous workspace preparation API, removing the per-call host-blocking D2H sync previously needed in varlen/group-mode backward.

Changes:

• Switch from launcher.prepare_workspace() to launcher.prepare_workspace_async() and enqueue workspace preparation work on the caller’s stream.
• Add a pinned-host allocator backed by PyTorch pinned memory and pass it into async workspace preparation.
• In varlen/group mode, remove cu_seqlens_*.cpu() host copies and pass device cu_seqlens pointers into async workspace preparation.

Reviewed changes

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

File	Description
csrc/flash_attn_ck/mha_varlen_bwd.cpp	Removes host cu_seqlens copies and uses stream-async workspace prep with a pinned-host allocator.
csrc/flash_attn_ck/mha_bwd.cpp	Uses stream-async workspace prep with a pinned-host allocator for non-varlen backward.

Comments suppressed due to low confidence (2)

csrc/flash_attn_ck/mha_varlen_bwd.cpp:385

• prepare_workspace_async uses stream obtained earlier via getCurrent*Stream() before the CUDAGuard is set to q.device(). If the caller’s current device differs from q.device(), this can enqueue the async workspace prep on a stream belonging to the wrong device. Consider moving at::cuda::CUDAGuard device_guard{q.device()}; before retrieving the current stream (and then fetch stream after the guard) so the prep and subsequent launcher.run(...) use the correct per-device stream.

        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);

csrc/flash_attn_ck/mha_bwd.cpp:368

• prepare_workspace_async is passed stream that is captured before at::cuda::CUDAGuard device_guard{q.device()}; is constructed. If the current device at entry differs from q.device(), the workspace prep can be enqueued on the wrong device’s stream. Move the device guard earlier and retrieve the current stream after the guard so async preparation and the kernel launch run on the intended device/stream.

        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);

---

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

[rocking5566]

LGTM