[CK_TILE] Add async workspace prepare to FMHA BWD launcher by DDEle · Pull Request #7331 · ROCm/rocm-libraries

DDEle · 2026-05-12T07:10:45Z

Motivation

aiter::mha_bwd in group mode currently issues two synchronous hipMemcpy D2H copies to read seqstart_q/k for launcher construction. These sync copies block the host (~10–30 µs each) and implicitly synchronize the device by draining the stream, breaking CPU/GPU overlap on hot training paths.

This PR adds a fully stream-async workspace preparation path on the FMHA BWD launcher so callers can pre-allocate the device workspace from upper-bound shapes and stage seqstart-dependent metadata via D2H/host-pack/H2D entirely on the user's stream.

Technical Details

FmhaBwdWorkspaceManager::GetWorkspaceDeviceSizeUpperBound (include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp): computes the worst-case device dq_acc size from (max_batch, hdim_q, nhead_q, max_seqlen_q, max_seqlen_k) without dereferencing any seqstart array. Mirrors PrepareWorkspaceHost's return value with worst-case bounds.
fmha_bwd_launcher::prepare_workspace_async (example/ck_tile/01_fmha/fmha_bwd.hpp): on the caller's stream, in order:
1. hipMemsetAsync of the dq_acc region (when NeedsZeroDqAcc())
2. group mode: hipMemcpyAsync D2H of seqstart_q/k into a pinned host staging buffer
3. hipLaunchHostFunc runs PrepareWorkspaceHost on the pinned buffer
4. hipMemcpyAsync H2D of the packed metadata into device_ws_ptr
The pinned staging buffer is held via std::shared_ptr<void> returned by a caller-provided pinned_host_alloc callback. Lifetime is extended past stream completion by a tail hipLaunchHostFunc scheduled in the launcher's destructor.
ck_tile::pinned_host_releaser (include/ck_tile/host/pinned_host_releaser.hpp): worker-thread utility for callers using bare hipHostMalloc. Defers hipHostFree off the HIP driver callback thread, which holds runtime locks and would deadlock against concurrent main-thread hipFree. PyTorch's CachingHostAllocator does not need this.
Example runner (example/ck_tile/01_fmha/fmha_bwd_runner.hpp): switched to the async path.

Test Plan

tile_example_fmha_bwd (gfx950, dev preset -Werror -Weverything):
- batch + nondet / batch + det / group + nondet / group + det
- group + det 4-batch varlen (-b=4 -h=8 -s=4096,3072,2048,1024 -d=128)
FA (flash-attention) integration on ROCm 7.1.1 + PyTorch 2.9.1:
- tests/test_flash_attn_ck.py::test_flash_attn_varlen_deterministic
- tests/test_flash_attn_ck.py::test_flash_attn_bwd_varlen_seqq_zero

Test Result

All CK runner cases valid:y.
FA pytest: 1952 passed in 44.82s.

Submission Checklist

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

Copilot

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds a stream-asynchronous workspace preparation path for FMHA backward (DQ/DK/DV) so device workspace can be allocated/prepared without requiring host-side seqstart availability at launcher construction time, and avoids HIP API calls from HIP callback context when freeing pinned memory.

Changes:

Add a “device workspace size upper bound” API for FMHA BWD kernels to pre-allocate worst-case workspace on device.
Introduce async workspace preparation in the FMHA BWD launcher (D2H seqstart staging + host packing + H2D metadata + dq_acc memset) with pinned host staging lifetime management.
Add a process-wide pinned host releaser worker to defer hipHostFree off HIP callback threads; update example runner to use the async prepare path.

Reviewed changes

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

Show a summary per file

File	Description
projects/composablekernel/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp	Adds kernel-level API to compute an upper bound on required device workspace size.
projects/composablekernel/include/ck_tile/host/pinned_host_releaser.hpp	New utility to defer `hipHostFree` to a worker thread (avoids HIP callback deadlocks).
projects/composablekernel/example/ck_tile/01_fmha/fmha_bwd_runner.hpp	Switches example to `prepare_workspace_async` and stages seqstart to device earlier.
projects/composablekernel/example/ck_tile/01_fmha/fmha_bwd.hpp	Redesigns launcher to support async prepare and pinned staging lifetime; adds device upper-bound sizing hook.
projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py	Codegen specialization to expose kernel upper-bound workspace sizing through the launcher APIs.

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

poyenc

Minor suggestions — take or leave.

GetWorkspaceDeviceSizeUpperBound was computing max_batch * nhead_q * max_seqlen_q * hdim_q in non-deterministic group mode, but PrepareWorkspaceHost actually returns nhead_q * seqstart_q[batch] * hdim_q i.e. it scales with the sum of *padded* per-batch seqlen_q, not max_batch times the *logical* max. When per-batch padding makes seqstart_q[batch] exceed max_batch * max_seqlen_q the launcher under-allocates dq_acc, the kernel writes past the buffer, and tests see either ~42% wrong QGrad values or a GPU page fault (e.g. test_ck_tile_fmha_bwd_bf16 QKVPadding/23,24,26 corrupt; /27 page-faults). Fix: replace the (max_batch, max_seqlen_q) pair with a single total_seqlen_q_padded parameter holding the true total padded q tokens. Launcher derives it from the trait (group: t.seqlen_q already is the padded total; batch: t.batch * t.seqlen_q). The four mode formulas collapse to one: size = nhead_q * nsplits_factor * total_seqlen_q_padded * hdim_q where nsplits_factor is 1 for non-deterministic, ceil(max_seqlen_k, kN0) for deterministic group, and the persistent worker computation for deterministic non-group (the only branch that still needs max_batch). No caller-side API change: FA, AITER and the CK runner already pass q.shape[0] (the padded total) as traits.seqlen_q in group mode. Verified on gfx1201: full test_ck_tile_fmha_bwd_{bf16,fp16} 672/672 PASS, 0 fail, 0 crash (was 27/28 QKVPadding fails + 1 GPU illegal access).

- prepare_workspace_async: allocate pinned host staging before enqueuing the dq_acc memset. If pinned_host_alloc throws, no stream work has been issued yet, so the workspace is left cleanly un-prepared rather than half-initialized. - pack_workspace_host catch: note that the H2D queued after the callback will copy indeterminate metadata if the catch fires (kernel will produce wrong results); unlikely since pack only throws on precondition violations. - schedule_pin_staging_release: std::move pin_staging_ into the heap shared_ptr; the next line in prepare_workspace_async overwrites it, so the extra atomic inc/dec from a copy is wasted.

poyenc

LGTM. All comments addressed in a76471b.

[CK_TILE] Add async workspace prepare to FMHA BWD launcher (#7331) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation `aiter::mha_bwd` in group mode currently issues two synchronous `hipMemcpy` D2H copies to read `seqstart_q/k` for launcher construction. These sync copies block the host (~10–30 µs each) and implicitly synchronize the device by draining the stream, breaking CPU/GPU overlap on hot training paths. This PR adds a fully stream-async workspace preparation path on the FMHA BWD launcher so callers can pre-allocate the device workspace from upper-bound shapes and stage seqstart-dependent metadata via D2H/host-pack/H2D entirely on the user's stream. ## Technical Details - `FmhaBwdWorkspaceManager::GetWorkspaceDeviceSizeUpperBound` (`include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp`): computes the worst-case device dq_acc size from `(max_batch, hdim_q, nhead_q, max_seqlen_q, max_seqlen_k)` without dereferencing any seqstart array. Mirrors `PrepareWorkspaceHost`'s return value with worst-case bounds. - `fmha_bwd_launcher::prepare_workspace_async` (`example/ck_tile/01_fmha/fmha_bwd.hpp`): on the caller's stream, in order: 1. `hipMemsetAsync` of the dq_acc region (when `NeedsZeroDqAcc()`) 2. group mode: `hipMemcpyAsync` D2H of `seqstart_q/k` into a pinned host staging buffer 3. `hipLaunchHostFunc` runs `PrepareWorkspaceHost` on the pinned buffer 4. `hipMemcpyAsync` H2D of the packed metadata into `device_ws_ptr` The pinned staging buffer is held via `std::shared_ptr<void>` returned by a caller-provided `pinned_host_alloc` callback. Lifetime is extended past stream completion by a tail `hipLaunchHostFunc` scheduled in the launcher's destructor. - `ck_tile::pinned_host_releaser` (`include/ck_tile/host/pinned_host_releaser.hpp`): worker-thread utility for callers using bare `hipHostMalloc`. Defers `hipHostFree` off the HIP driver callback thread, which holds runtime locks and would deadlock against concurrent main-thread `hipFree`. PyTorch's `CachingHostAllocator` does not need this. - Example runner (`example/ck_tile/01_fmha/fmha_bwd_runner.hpp`): switched to the async path. ## Test Plan - `tile_example_fmha_bwd` (gfx950, dev preset `-Werror -Weverything`): - batch + nondet / batch + det / group + nondet / group + det - group + det 4-batch varlen (`-b=4 -h=8 -s=4096,3072,2048,1024 -d=128`) - FA (`flash-attention`) integration on ROCm 7.1.1 + PyTorch 2.9.1: - `tests/test_flash_attn_ck.py::test_flash_attn_varlen_deterministic` - `tests/test_flash_attn_ck.py::test_flash_attn_bwd_varlen_seqq_zero` ## Test Result - All CK runner cases `valid:y`. - FA pytest: **1952 passed in 44.82s**. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

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

…ies#7331 merged)

) * [CK_TILE] mha_bwd: use async workspace prepare pipeline * [CK_TILE] mha_bwd: address PR #3150 review comments - mha_bwd.h: pinned_host_alloc doc no longer claims a synchronous D2H fallback (CK launcher now throws when missing). Doc says it's required in group mode and unused in batch mode. - mha_bwd.cu: explicitly check pinned_host_alloc before entering the group-mode async path, matching the existing seqstart_*_ptr precondition check, so callers see a clean AITER_LOG_ERROR rather than a launcher exception. - benchmark_mha_bwd.cpp: remove the dead fmha_bwd_traits local that was left behind when launcher construction moved into aiter::mha_bwd. Triggered -Wunused-variable; the workspace_alloc design comment is kept. * [CK_TILE] mha_bwd: bump CK submodule to develop tip (ROCm/rocm-libraries#7331 merged) --------- Co-authored-by: Xin Huang <Xin.Huang@amd.com>

* [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

## Motivation `aiter::mha_bwd` in group mode currently issues two synchronous `hipMemcpy` D2H copies to read `seqstart_q/k` for launcher construction. These sync copies block the host (~10–30 µs each) and implicitly synchronize the device by draining the stream, breaking CPU/GPU overlap on hot training paths. This PR adds a fully stream-async workspace preparation path on the FMHA BWD launcher so callers can pre-allocate the device workspace from upper-bound shapes and stage seqstart-dependent metadata via D2H/host-pack/H2D entirely on the user's stream. ## Technical Details - `FmhaBwdWorkspaceManager::GetWorkspaceDeviceSizeUpperBound` (`include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp`): computes the worst-case device dq_acc size from `(max_batch, hdim_q, nhead_q, max_seqlen_q, max_seqlen_k)` without dereferencing any seqstart array. Mirrors `PrepareWorkspaceHost`'s return value with worst-case bounds. - `fmha_bwd_launcher::prepare_workspace_async` (`example/ck_tile/01_fmha/fmha_bwd.hpp`): on the caller's stream, in order: 1. `hipMemsetAsync` of the dq_acc region (when `NeedsZeroDqAcc()`) 2. group mode: `hipMemcpyAsync` D2H of `seqstart_q/k` into a pinned host staging buffer 3. `hipLaunchHostFunc` runs `PrepareWorkspaceHost` on the pinned buffer 4. `hipMemcpyAsync` H2D of the packed metadata into `device_ws_ptr` The pinned staging buffer is held via `std::shared_ptr<void>` returned by a caller-provided `pinned_host_alloc` callback. Lifetime is extended past stream completion by a tail `hipLaunchHostFunc` scheduled in the launcher's destructor. - `ck_tile::pinned_host_releaser` (`include/ck_tile/host/pinned_host_releaser.hpp`): worker-thread utility for callers using bare `hipHostMalloc`. Defers `hipHostFree` off the HIP driver callback thread, which holds runtime locks and would deadlock against concurrent main-thread `hipFree`. PyTorch's `CachingHostAllocator` does not need this. - Example runner (`example/ck_tile/01_fmha/fmha_bwd_runner.hpp`): switched to the async path. ## Test Plan - `tile_example_fmha_bwd` (gfx950, dev preset `-Werror -Weverything`): - batch + nondet / batch + det / group + nondet / group + det - group + det 4-batch varlen (`-b=4 -h=8 -s=4096,3072,2048,1024 -d=128`) - FA (`flash-attention`) integration on ROCm 7.1.1 + PyTorch 2.9.1: - `tests/test_flash_attn_ck.py::test_flash_attn_varlen_deterministic` - `tests/test_flash_attn_ck.py::test_flash_attn_bwd_varlen_seqq_zero` ## Test Result - All CK runner cases `valid:y`. - FA pytest: **1952 passed in 44.82s**. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

…7450) ## Motivation FMHA BWD group-mode deterministic currently uses a non-persistent scheduler: each `(batch, head, K-row)` work-item is launched as its own block, with no work-stealing across CUs. On uneven workloads (varlen, GQA, many heads with few K-rows) this leaves CUs idle and forces a larger dq_acc workspace than necessary. This PR ports the persistent + deterministic scheduling already used in batch mode to group mode: a fixed-grid kernel that pre-computes per-CU work ranges on the host and uses sparse dq_acc slot indexing so multiple K-rows handled by the same CU share one accumulator slot via intra-CU atomic adds. Stacked on #7331; merge that first. ## Technical Details Single file changed: `ops/fmha/kernel/fmha_bwd_kernel.hpp`. A new `kUsePersistent` path is added to the group-mode deterministic kernel, mirroring the batch-mode persistent scheduler. The host pre-computes a fixed per-CU partition of the total `(batch, head, K-row)` work and packs it into `cu_states[]` so the GPU consumes it in a single launch. Host preparation happens in four steps: 1. Build per-batch `seqstart` prefix sums. 2. Fill per-batch `(sq_w, nc)` with a placeholder `nsplits` (bumped in step 3). 3. Two-pointer scan over CUs to fill `cu_states[c]` (`isplit`, `head_start`, `c_start`, `w_lo`, `w_hi`), accumulating `nsplits[b]` as `max(cs->isplit + 1)`. 4. Compute compact per-batch dq_acc offsets from the finalized `nsplits`. `isplit` is the sparse dq_acc slot index — one CU's multi-K-row writes share slot `ceil(wc_start / denom)`, enabling intra-CU atomic accumulation instead of one slot per K-row. `denom = max(sq_w, target_w)`, splitting two regimes: - `target_w >= sq_w` (large work): `denom = target_w`, intra-CU atomic optimization engaged. - `target_w < sq_w` (sub-K-row sharding, multiple CUs sharing one K-row): `denom = sq_w` collapses to per-K-row indexing (`= c_start`), keeping `isplit ∈ [0, nc-1]` and matching the `nsplits_max = ceil(s_k/kN0) = nc` upper bound that #7331's `GetWorkspaceDeviceSizeUpperBound` assumes for group+det. `isplit` is additionally clamped to `nc-1` to absorb empty CUs (rounded-up `wc_start` past the last K-row); they don't write dq_acc on GPU so the slot value is harmless. `nsplits[b]` is accumulated dynamically in step 3 rather than via a closed form so it tightly matches the actual sparse slots used; step 4 (offsets) follows step 3 since offsets now depend on the dynamic `nsplits`. Group mode also allows batches with `seqlen_q == 0`. The persistent scheduler skips them on the dQ path (no work) but dK/dV are still zero-filled. ## Test Plan Built `tile_example_fmha_bwd` with receipt 5 (fp16, no-bias, no-dropout, `dpad == dvpad`, group + batch) on gfx950 (MI355X). - 8-case smoke (shapes that exercise the sub-K-row regime). - 44-case sweep covering: mask 0/1/2, GQA, var seqlen, `d != d_v`, extreme small seqlen / `nc=1`, CU >> work, huge batch, batch-mode regression. - 12-case perf comparison vs the non-persistent baseline (warmup=10, repeat=50). ## Test Result - All 8 + 44 cases `valid:y`. - Perf: ±5% noise, average -0.4% across the 12 cases — neutral. - Batch-mode deterministic / non-deterministic regression unchanged. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[CK_TILE] Add async workspace prepare to FMHA BWD launcher (#7331) ## Motivation `aiter::mha_bwd` in group mode currently issues two synchronous `hipMemcpy` D2H copies to read `seqstart_q/k` for launcher construction. These sync copies block the host (~10–30 µs each) and implicitly synchronize the device by draining the stream, breaking CPU/GPU overlap on hot training paths. This PR adds a fully stream-async workspace preparation path on the FMHA BWD launcher so callers can pre-allocate the device workspace from upper-bound shapes and stage seqstart-dependent metadata via D2H/host-pack/H2D entirely on the user's stream. ## Technical Details - `FmhaBwdWorkspaceManager::GetWorkspaceDeviceSizeUpperBound` (`include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp`): computes the worst-case device dq_acc size from `(max_batch, hdim_q, nhead_q, max_seqlen_q, max_seqlen_k)` without dereferencing any seqstart array. Mirrors `PrepareWorkspaceHost`'s return value with worst-case bounds. - `fmha_bwd_launcher::prepare_workspace_async` (`example/ck_tile/01_fmha/fmha_bwd.hpp`): on the caller's stream, in order: 1. `hipMemsetAsync` of the dq_acc region (when `NeedsZeroDqAcc()`) 2. group mode: `hipMemcpyAsync` D2H of `seqstart_q/k` into a pinned host staging buffer 3. `hipLaunchHostFunc` runs `PrepareWorkspaceHost` on the pinned buffer 4. `hipMemcpyAsync` H2D of the packed metadata into `device_ws_ptr` The pinned staging buffer is held via `std::shared_ptr<void>` returned by a caller-provided `pinned_host_alloc` callback. Lifetime is extended past stream completion by a tail `hipLaunchHostFunc` scheduled in the launcher's destructor. - `ck_tile::pinned_host_releaser` (`include/ck_tile/host/pinned_host_releaser.hpp`): worker-thread utility for callers using bare `hipHostMalloc`. Defers `hipHostFree` off the HIP driver callback thread, which holds runtime locks and would deadlock against concurrent main-thread `hipFree`. PyTorch's `CachingHostAllocator` does not need this. - Example runner (`example/ck_tile/01_fmha/fmha_bwd_runner.hpp`): switched to the async path. ## Test Plan - `tile_example_fmha_bwd` (gfx950, dev preset `-Werror -Weverything`): - batch + nondet / batch + det / group + nondet / group + det - group + det 4-batch varlen (`-b=4 -h=8 -s=4096,3072,2048,1024 -d=128`) - FA (`flash-attention`) integration on ROCm 7.1.1 + PyTorch 2.9.1: - `tests/test_flash_attn_ck.py::test_flash_attn_varlen_deterministic` - `tests/test_flash_attn_ck.py::test_flash_attn_bwd_varlen_seqq_zero` ## Test Result - All CK runner cases `valid:y`. - FA pytest: **1952 passed in 44.82s**. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

* 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>

* [Fwd,Sm100] fix: decode↔prefill exp2 emulation consistency (Dao-AILab#2595) apply_exp2_convert selected the exp2 implementation based on mask_fn presence: hardware ex2.approx.ftz for causal-masked tiles, polynomial emulation for unmasked tiles. Different q_stage values (1 for decode, 2 for prefill) compute different m_block for the same logical Q row, shifting which tiles are processed with vs without mask_fn. The same K tile could receive different exp2 methods across variants. Fix: always pass self.ex2_emu_freq regardless of mask_fn presence. Add regression test for decode↔prefill bitwise consistency on MLA (192,128) shapes. * replace deprecated apis (Dao-AILab#2602) * Bump nvidia-cutlass-dsl to >=4.5.2 and quack-kernels to >=0.5.0 (Dao-AILab#2605) cutlass 4.5.2 is safe to update, and quack 0.5.0 has been published, so bump the FA4 (flash_attn/cute) requirement floors to match. Updates the dependencies and the cu13 extra in pyproject.toml, and the documented versions in CLAUDE.md. Verified on NVIDIA GB300 (SM100, CUDA 13.2): deps resolve cleanly (nvidia-cutlass-dsl 4.5.2 base+cu13, quack-kernels 0.5.0), imports OK, and a representative GPU sample of tests/cute/test_flash_attn.py passes (6 passed / 6 skipped / 0 failed across hd 64/96/128/192, causal, mha/gqa/mqa, fwd+bwd). * [CuTe,Fwd,Sm100] refactor mla sm100 forward and add page table (Dao-AILab#2558) * refactor mla sm100 forward * add benchmark; address deprecation warnings; tweak ptx gemm dispatch * update interface and tests * ci: bump Jimver/cuda-toolkit to v0.2.35 for CUDA 13.2 support (Dao-AILab#2617) v0.2.30 only ships URLs up to CUDA 13.1.0; bumping to v0.2.35 adds 13.1.1, 13.2.0, and the matching aarch64 SBSA installers. Signed-off-by: oliver könig <okoenig@nvidia.com> * [ROCm] Bump Triton to >=3.6.0 and aiter submodule (Dao-AILab#2614) * [Triton] Fix graph capture issues and env var (Dao-AILab#2620) * graph capture fix * rm env flag * [CuTe,Bwd,Sm100] allow 2cta with score mod and mask mod in bwd (Dao-AILab#2557) * [CuTe] Fix lint failures (Dao-AILab#2625) stack-info: PR: Dao-AILab#2625, branch: drisspg/stack/42 * [CuTe] Fix lint failure in flash_bwd_sm100.py (Dao-AILab#2627) ruff format flagged flash_attn/cute/flash_bwd_sm100.py (trailing whitespace in a comment and an over-split call). It was missed by the lint sweep in Dao-AILab#2625. * fix: add weights_only=True to all torch.load call sites (Dao-AILab#2622) Passing weights_only=False (the pre-2.4 default) to torch.load allows arbitrary Python object deserialization from the checkpoint file. A malicious .pt/.pth file can execute arbitrary code on the machine loading it — a well-known PyTorch deserialization vector (CWE-502). Four call sites updated: training/src/utils/checkpoint.py load_checkpoint() training/src/eval.py eval checkpoint loader flash_attn/utils/pretrained.py partial(torch.load, ...) loader flash_attn/models/llama.py state_dicts_from_checkpoint() weights_only=True restricts deserialization to tensors, dicts, lists, tuples, and other primitive types — no arbitrary Python objects. Requires PyTorch >= 1.13; FA4's CuTeDSL dependency already requires a modern PyTorch 2.x build, so no compatibility regression. Fixes Dao-AILab#2583 * use correction warps if not tma store; remove outdated packgqa guard (Dao-AILab#2629) * Add aux-scalars to interface to enable dynamic ints and floats in expressions (Dao-AILab#2616) stack-info: PR: Dao-AILab#2616, branch: drisspg/stack/41 * fix: build and select cu13.2 prebuilt wheels (Dao-AILab#2618) * ci: use 1 ninja job for cu13.2 Signed-off-by: oliver könig <okoenig@nvidia.com> * fix(setup): request cu13 prebuilt wheels for CUDA 13 torch get_wheel_url() binned every CUDA >= 12 to major '12', so under a CUDA 13 torch it requested cu12 wheels and never matched the published cu13 artifacts, falling back to a multi-hour source build. Add a CUDA 13 branch so the guessed wheel name uses cu13, matching WHEEL_CUDA_VERSION in _build.yml. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: oliver könig <okoenig@nvidia.com> --------- Signed-off-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * ci(fa4): enforce cutlass-dsl/quack dep floors and rebake cu130 image (Dao-AILab#2636) * ci(fa4): assert cute dep floors in CI; fail loudly on a stale SIF run_fa4_ci.py installs FA4 with --no-deps (to keep the SIF's baked torch/cudnn), so the nvidia-cutlass-dsl>=4.5.2 / quack-kernels>=0.5.0 floors in flash_attn/cute/pyproject.toml are not enforced at install time. A SIF baked before a floor bump keeps a stale dep — e.g. cutlass-dsl 4.4.2, which can't convert the AuxData JIT arg and dies with a cryptic DSLRuntimeError deep in SM100 kernel launch (reproduced on B200). Upgrading the dep in-place is not viable: the --writable-tmpfs overlay is RAM-backed and too small for a cutlass-dsl reinstall (ENOSPC, and a partial removal corrupts the baked torch). So instead of installing, add assert_dsl_floor.py — it reads the floors from pyproject (no hardcoded version to drift) and fails with an actionable "rebake the image" message when the installed cutlass-dsl/quack are below them. Wired into run_step right after the editable install. The durable fix is to rebake the image at the current floors and bump the digest in .github/workflows/ci.yml; this guard makes future drift fail fast instead of silently. * ci(fa4): bump cu130 image to 26.06.10 (cutlass-dsl 4.5.2 / quack 0.5.0) * ci(fa4): fall back to tomli when tomllib is unavailable (Python 3.10) * Fix SM100 FP8 fwd with cutlass-dsl >=4.5.2 (MmaF8F6F4Op) (Dao-AILab#2640) cutlass-dsl >=4.5.2 changed make_trivial_tiled_mma to build plain FP8 MMAs as MmaF8F6F4Op (its _F8F6F4_TYPES branch) instead of the now-legacy MmaFP8Op. The two are siblings under MmaOp, so _tcgen05_mma_kind's isinstance(op, MmaFP8Op) check missed the new type and raised "Unsupported tcgen05 MMA op kind: MmaF8F6F4Op", breaking the FP8 forward path on Blackwell. Worked on 4.4.2. Accept both ops in the f8f6f4 branch (both map to kind::f8f6f4). mma_op_to_idesc only reads generic op attrs and is unaffected. Validated on B200: FP8 fwd passes for all configs in the issue (incl. hd=64) plus hd=128, causal and non-causal; mean abs err vs bf16 ~0.002-0.01. Fixes Dao-AILab#2639 * [cute] Fix int32 overflow in SM100 LPT tile scheduler for long context (Dao-AILab#2662) The LPT tile scheduler sizes its L2 swizzle from seqlen_k * (headdim + headdim_v) * element_size in int32. For long context this overflows once it exceeds 2**31 (seqlen_k > ~4M for hdim-128 bf16), making size_one_head negative. That corrupts the swizzle and the L2 divmods, so get_current_work decodes an out-of-bounds batch_idx and the kernel performs an illegal memory access (cudaErrorIllegalAddress) on SM100. Compute the byte size in int64. swizzle stays small and is cast back to int32 for the device-side divmods, so there is no behavior or perf change for non-overflowing shapes. Fixes both SingleTileLPTScheduler (forward; selected for causal/local) and SingleTileLPTBwdScheduler (backward; its extra seqlen_k * headdim * 4 term overflows even sooner). Repro on SM100 (e.g. GB200), causal forward at seqlen_k = 2**22: import torch from flash_attn.cute.interface import flash_attn_func sq, sk = 2048, 4_194_304 # seqlen_k = 2**22 -> int32 overflow q = torch.randn(1, sq, 8, 128, dtype=torch.bfloat16, device="cuda") k = torch.randn(1, sk, 1, 128, dtype=torch.bfloat16, device="cuda") v = torch.randn(1, sk, 1, 128, dtype=torch.bfloat16, device="cuda") out = flash_attn_func(q, k, v, causal=True) torch.cuda.synchronize() # cudaErrorIllegalAddress here before the fix Crashes before this change, runs clean after; seqlen_k = 2**22 - 128 is clean both ways (the int32 boundary). Verified clean under compute-sanitizer memcheck. * [Fwd,Sm100] Tune FP8 causal hd128 ex2_emu_freq (8 vs inherited 16) (Dao-AILab#2642) FP8 fwd is MUFU/ex2-bound on Blackwell, so the optimal exp2-emulation frequency differs from bf16. The causal hd128 key (False,True,128,False) had no FP8 entry and inherited bf16's freq=16; freq=8 offloads more exp from the MUFU unit. Thermally-matched back-to-back A/B on B200 (locked-ish clock, hot GPU, median of 300 iters, nheads=16 = benchmark default) across the official benchmark's causal hd128 shapes: b s f16 TFLOP f8 TFLOP delta 32 512 500.6 516.3 +3.1% 16 1024 796.5 832.5 +4.5% 8 2048 1124.6 1175.1 +4.5% 4 4096 1407.9 1481.3 +5.2% 2 8192 1604.1 1661.7 +3.6% 1 16384 1683.7 1726.0 +2.5% Accuracy-neutral (FP8-vs-bf16 mean-abs-err unchanged; benchmark --check passes 24/24). Keyed on is_causal=True only: freq=8 would regress non-causal hd128 (0.94x), which keeps its existing freq=10. * Make q_subtile_factor default to identity (Dao-AILab#2660) * fix(hd256/sm100): make q/k/v contiguous before dedicated hd256 kernel (Dao-AILab#2666) The BlackwellFusedMultiHeadAttentionForward kernel builds tensor layouts with hardcoded contiguous strides computed from shape dimensions, so non-contiguous inputs (e.g. from .transpose()) cause wrong memory accesses and silently corrupt outputs on B200 (SM100) with head_dim=256. maybe_contiguous() only guarantees stride(-1)==1; add explicit full contiguity checks in both the forward and backward paths when the hd256 dedicated kernel is selected. Fixes: Dao-AILab#2665 * [Cute,Bwd,Sm100] add sparse MLA (Deepseek v4) backward kernels (Dao-AILab#2621) * add backward sparse mla kernels * add dk gemm * fix errors * fix dq errors * rename bwd kernels * refactor interface * fix predicate error in dq kernel * update tests * mla fwd fixes * improve varlen fwd perf * use cluster idx scheduling in fwd * use packed scheduler for mqa 128 * fix int32 overflow in swizzle * simplify bwd preprocess * refactor bwd * simplify preprocess * update benchmark script * add safety check * remove test code * ruff format * ensure scale is 0 for masked out rows * fix: sync callers with new _flash_attn_fwd 4-tuple return signature (Dao-AILab#2674) * Fix compatibility issues with CuTe DSL 4.6.0+ (Dao-AILab#2648) * Prepare for 4.6 release * Bump version * Update pyproject.toml * Update nvidia-cutlass-dsl version in pyproject.toml * Pass tmem scalar fields as .ptr to TmemAllocator on SM100 (Dao-AILab#2679) The DSL now warns when a struct scalar is used directly as a pointer ("Use explicit struct.scalar.ptr for pointer instead"), so these fire on every tmem_holding_buf / dealloc mbar access. Just pass .ptr like the other SM100 kernels already do. * Add FLASHATTENTION_DISABLE_SPLIT_ALIGNMENT (Dao-AILab#2680) * ci: rebake cu130 image for cutlass-dsl 4.6.0.dev0 floor (Dao-AILab#2684) PR Dao-AILab#2648 bumped the flash_attn/cute/pyproject.toml floor to nvidia-cutlass-dsl==4.6.0.dev0, but the CI image (26.06.10) still ships 4.5.2. assert_dsl_floor.py correctly fails every push to main with "installed 4.5.2 does not satisfy floor ==4.6.0.dev0", so FA4 CI has been red since Dao-AILab#2648 landed. - Dockerfile: add --prerelease=allow to the FA4 install. The dev-build floor pulls transitive pre-releases (nvidia-cutlass-dsl-libs-base== 4.6.0.dev0 ...) that uv refuses without it; the old stable 4.5.2 floor didn't need it. - ci.yml: bump fa4_image_cu130 to the rebaked 26.06.27 image (cutlass-dsl 4.6.0.dev0, quack-kernels 0.5.3, torch 2.12.1). E2e verified on B200: assert_dsl_floor passes, compile + run + benchmark all green (run_fa4_ci.py, exit 0). * Update FA4 cute quack compatibility (Dao-AILab#2676) * Update FA4 cute quack compatibility * Use quack 0.5.3 make_smem_layout instead of vendored copy Tri re-added the major_mode_size arg to quack.sm90_utils.make_smem_layout in quack 0.5.3 (commit 68888e2), so FA4 no longer needs the local sm90_layout helper. Revert the 4 backward call sites to quack's helper and bump the floor to >=0.5.3 (0.5.2 lacks the arg). --------- Co-authored-by: Johnsonms <lizhaofu@gmail.com> * ci: install cutlass-dsl/quack at runtime to decouple from the baked image (Dao-AILab#2685) * [Cute,Bwd,Sm100] Assume 16B stride divisibility for LSE/dPsum bulk-copy inputs (Dao-AILab#2686) The SM100 backward stats (LSE, dPsum) are loaded via cp.async.bulk (CopyBulkG2SOp), which - unlike cp.async.bulk.tensor - needs the source pointer alignment provable at compile time. After slicing, the newer cute-dsl can't deduce 16B alignment unless the input strides carry the divisibility assumption, so the bulk copy fails to compile on real tensors (the FakeTensor path masks it). - flash_bwd_mla_sm100.py: add mdPsum to the new_stride divisibility list (it already covered ScaleP and the other stats; mdPsum was omitted). - flash_bwd_sm100.py: the ordinary backward had no divisibility assumption at all; add it for both mLSE and mdPsum. Only these two SM100 kernels use CopyBulkG2SOp; the SM90/SM80/SM120 and MLA dK/dQ backward kernels use other copy paths and are unaffected. Addresses the dPsum stride-divisibility finding (Finding 1) in Dao-AILab#2677. * fix(hd256/sm100): forward reads actual input strides, drop .contiguous() patch (Dao-AILab#2670) * follow up to Dao-AILab#2666: fixing the layouts in the sm100 hd256 kernels and removing the temporary fix of calling .contiguous everywhere * respond to PR comments * respond to PR comments-2: move to utils file * Add tests --------- Co-authored-by: drisspg <drisspguessous@gmail.com> * ci: run MLA backward cases so CI exercises flash_bwd_mla_sm100.py (Dao-AILab#2690) FA4_TEST_FILTER selected no MLA test, so the MLA backward kernels (flash_bwd_mla_sm100.py + dq_dqv + dk) had zero CI coverage. Add four small test_flash_attn_mla_absorbed cases covering the distinct backward paths: sparse (kv_sparsity=True) non-causal and causal, dense (kv_sparsity=False), and shared_kv=True. The ordinary SM100 backward is already covered by the existing test_flash_attn_output cases. Cold-cache cost on B200 (full 8-case filter): pass-1 compile ~4:54, GPU run ~1:03 — well under the 60-min job timeout. Stacked on Dao-AILab#2685 (runtime cutlass-dsl/quack install). * Parallelize splitkv alignment templated kernels, remove flag (Dao-AILab#2683) * [FA3] uv installation support (Dao-AILab#2458) * Expose flash_attn_3 as package so imports work correctly. * Add flash_attn_config package shim and fix uv packaging details Builds on the flash_attn_3 package exposure so both import styles work for downstream frameworks and uv/pyproject.toml installs: - Add flash_attn_3/flash_attn_config.py re-export so `from flash_attn_3 import flash_attn_config` works (previously only the top-level module was importable), matching the interface shim. - Un-ignore the committed shim in .gitignore; the bare `flash_attn_config.py` pattern (for the build-time generated top-level file) also matched the package shim and would have silently dropped it from the commit. - Read flash_attn_3.__version__ from installed package metadata with a fallback, avoiding drift from setup.py's version source. - README: move `dependencies` under `[project]` so the uv snippet is valid PEP 621. Verified on H100 (SM90): editable `uv pip install -e .` now succeeds (fails on main), both `import flash_attn_interface` and `from flash_attn_3 import flash_attn_interface` resolve, `flash_attn_config` imports both ways, and fp16 hdim128 forward matches a torch reference (max_abs_err <= 2e-3). ruff check passes. --------- Co-authored-by: Johnsonms <lizhaofu@gmail.com> * [AMD ROCm] Enable RDNA backward and adopt CK unified workspace (Dao-AILab#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 (Dao-AILab#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 (Dao-AILab#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> * Fix CuTe SM120 compile-time argument handling (Dao-AILab#2671) * Fix CuTe SM120 compile-time argument handling * clean up * guard empty SM120 local backward tiles --------- Co-authored-by: Kevin-Li-2025 <2242139@qq.com> Co-authored-by: drisspg <drisspguessous@gmail.com> * [NVIDIA][CuTe,Fwd,sm120] Implement Pack-GQA on SM120 (+ graceful SplitKV fallback) (Dao-AILab#2656) * [CuTe,Fwd,sm120] Fix use_tma_O crash on SM120 (issue Dao-AILab#2649) On SM120 (Blackwell GeForce / RTX PRO 6000 / DGX Spark) the forward kernel set `use_tma_O = self.arch >= Arch.sm_90`, enabling the TMA-based O-store epilogue. But SM120 does not build the TMA store atom (tma_atom_O is None), so any forward call crashes in cpasync.tma_partition with: AttributeError: 'NoneType' object has no attribute '_trait' This makes the CuTe-DSL forward unusable on every SM120 GPU. Restrict the TMA O-store to sm_90..sm_119, which is where the WGMMA-era epilogue path is actually available: self.use_tma_O = Arch.sm_90 <= self.arch < Arch.sm_120 SM120 falls back to the non-TMA register->gmem O store (already used for the SM80 path), which is correct and what the CpAsync SM120 kernel expects. Verified on RTX PRO 6000 Blackwell (sm_120, cc 12.0), torch 2.12.0+cu130, nvidia-cutlass-dsl 4.5.2: forward now runs and matches PyTorch SDPA reference for hdim 64/96/128, causal and non-causal (max abs err <= 8e-3 in bf16). Before this fix every SM120 forward call raised the AttributeError above. * [CuTe,Fwd,sm120] Implement Pack-GQA on SM120; graceful SplitKV fallback Pack-GQA was only half-wired in the SM80/SM120 CpAsync forward: the epilogue referenced PackGQA.store_O/store_LSE, but the Q-load and head-indexing used the plain (unpacked) path. So pack_gqa=True crashed in pack_gqa.store_O (crd2idx on a packed (h_idx, m_idx) coordinate against an unpacked mO layout). This implements Pack-GQA end to end on SM120 (and SM80), mirroring the SM90 path: - Reshape mQ/mO (head_idx=2) and mLSE (head_idx=1) via pack_gqa_layout so qhead_per_kvhead folds into the seqlen mode ((qhead, seqlen)). - Scheduler args use cute.size(mQ.shape[0]) (packed total rows) and seqlen_q_static = mQ.shape[0][1] (logical seqlen), so causal/mask q_idx stay correct. - Kernel head-indexing: when pack_gqa, num_head from the scheduler already indexes the KV head (mQ/mK share nheads_kv); no division. - Q-load: gather rows via PackGQA.load_Q (per-row (h_idx, m_idx) gmem pointers) instead of the contiguous local_tile path. SplitKV (num_splits>1) is an SM100-only feature (SM80/SM90 also assert it unsupported); SM120 has no forward+combine path. Fall back to num_splits=1, which is numerically correct, instead of crashing in _check_type on the fp32 partials. Verified on RTX PRO 6000 Blackwell (sm_120): pack_gqa=True matches PyTorch SDPA GQA/MQA reference (err <= 8.4e-3 bf16) AND is bit-identical to the unpacked path (max |packed - unpacked| = 0.0) across MHA/GQA/MQA, causal/non-causal, hd 64/128, seqlen 512-2048. num_splits=3 falls back and matches reference (err 6.8e-4). Stacked on the SM120 use_tma_O fix (Dao-AILab#2649). * re-enable SM120 pack-gqa after rebase * clean up SM120 pack-gqa split handling * fix SM120 varlen pack-gqa offset --------- Co-authored-by: drisspg <drisspguessous@gmail.com> * Fix pre-commit Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> --------- Signed-off-by: oliver könig <okoenig@nvidia.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: 鐘天楽 <tianle.zhong@bytedance.com> Co-authored-by: brandonsun <brandons@nvidia.com> Co-authored-by: Johnsonms <lizhaofu@gmail.com> Co-authored-by: jayhshah <jayhshah@gmail.com> Co-authored-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Michael Melesse <micmelesse@gmail.com> Co-authored-by: Reuben Stern <107093092+reubenconducts@users.noreply.github.com> Co-authored-by: Driss Guessous <32754868+drisspg@users.noreply.github.com> Co-authored-by: aryan <aryansputta@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: sryap <17482891+sryap@users.noreply.github.com> Co-authored-by: Yunwei Li <yunweili372423@gmail.com> Co-authored-by: Zihao Wang <rekind133@outlook.com> Co-authored-by: Anakin(Yancheng) Zheng <103552181+anakinxc@users.noreply.github.com> Co-authored-by: Prashant Kumar <prashant.kumar@cohere.com> Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com> Co-authored-by: Omar Attia <oy.attia@gmail.com> Co-authored-by: drisspg <drisspguessous@gmail.com> Co-authored-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com> Co-authored-by: rocking <ChunYu.Lai@amd.com> Co-authored-by: Yi DING <yi.ding@amd.com> Co-authored-by: Hosang Yoon <hosang.yoon@amd.com> Co-authored-by: Yin Li <kxl474@student.bham.ac.uk> Co-authored-by: Kevin-Li-2025 <2242139@qq.com> Co-authored-by: Johnny <johnnynuca14@gmail.com>

[CK_TILE] Add async workspace prepare to FMHA BWD launcher

53e3fb5

DDEle requested a review from a team as a code owner May 12, 2026 07:10

github-actions Bot added the project: composablekernel label May 12, 2026

DDEle requested a review from Copilot May 12, 2026 07:11

Copilot started reviewing on behalf of DDEle May 12, 2026 07:17 View session

Copilot AI reviewed May 12, 2026

View reviewed changes

assistant-librarian Bot added the organization: ROCm label May 12, 2026

[CK_TILE] FMHA BWD launcher: address PR #7331 review comments

1be1b45

DDEle mentioned this pull request May 12, 2026

[CK_TILE] FMHA BWD: stream-async workspace prepare for group mode ROCm/aiter#3150

Merged

1 task

poyenc reviewed May 13, 2026

View reviewed changes

DDEle added 2 commits May 13, 2026 02:20

DDEle requested a review from poyenc May 13, 2026 07:02

DDEle mentioned this pull request May 13, 2026

[CK_TILE] FMHA BWD: stream-async workspace prepare ROCm/flash-attention#183

Merged

2 tasks

poyenc approved these changes May 14, 2026

View reviewed changes

DDEle mentioned this pull request May 14, 2026

[CK_TILE] Use Persistent Scheduling for FMHA BWD Group Deterministic #7450

Merged

1 task

DDEle merged commit 5692db0 into develop May 14, 2026
34 checks passed

DDEle deleted the users/yiding12/fmha-bwd-async-prepare branch May 14, 2026 13:33

DDEle added a commit to ROCm/aiter that referenced this pull request May 15, 2026

[CK_TILE] mha_bwd: bump CK submodule to develop tip (ROCm/rocm-librar…

270fa1e

…ies#7331 merged)

qianfengz mentioned this pull request May 16, 2026

Update to random uniform kernel and adapt to latest ck_tile fmha_bwd_kernel ROCm/xformers#88

Merged

Copilot AI mentioned this pull request May 18, 2026

[Don't Merge] Isolate MLA failure: CK pre-#7331 (do not merge) ROCm/aiter#3246

Closed

Navigation Menu

Uh oh!

[CK_TILE] Add async workspace prepare to FMHA BWD launcher - #7331

[CK_TILE] Add async workspace prepare to FMHA BWD launcher#7331
DDEle merged 4 commits into
developfrom
users/yiding12/fmha-bwd-async-prepare

DDEle commented May 12, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

poyenc left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

poyenc left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

DDEle commented May 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

poyenc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

poyenc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

DDEle commented May 12, 2026 •

edited

Loading