[CK_TILE] Use Unified Workspace for FMHA BWD

[DDEle]

Motivation

dq_acc is the intermediate accumulation buffer used in FMHA backward pass for deterministic mode. The current implementation allocates it as a single rectangular tensor:

shape = [shape_batch, nhead, nsplits, shape_seqlen_q, hdim_q]

where nsplits = launcher.dq_acc_splits (a single scalar), computed from max_seqlen_k and shared across all batches.

Problems

1. Memory waste: In group mode, each batch may have a different seqlen_k, but nsplits is computed from max_seqlen_k, causing batches with shorter seqlen_k to over-allocate in the split dimension.

2. Interface coupling: fmha_bwd_args exposes internal layout details such as stride_dq_acc, nhead_stride_dq_acc, batch_stride_dq_acc, and split_stride_dq_acc. The caller is responsible for computing these strides, but this logic belongs inside the kernel.

Goals

1. Switch dq_acc buffer to a compact layout: batches are concatenated contiguously, with each batch occupying nhead * nsplits_i * seqq_i * hdim_q elements (nhead outermost).
2. Remove all *_stride_dq_acc fields from fmha_bwd_args, replacing them with a single workspace_ptr; the kernel splits this internally using a fixed layout.
3. fmha_bwd_launcher provides a workspace management interface: the caller only needs to allocate GPU memory and call prepare_workspace() — no layout computation required.
4. Isolate kernel internals from the caller API: the dq_acc layout (nsplits, strides, buffer size) is determined entirely inside the launcher/kernel. Future changes to block shape, pipeline type, or persistent kernel strategy require no modifications to the caller's fmha_bwd_args or workspace allocation logic.

Technical Details

Interface Design

New fields in fmha_bwd_traits

struct fmha_bwd_traits
{
    int seqlen_q;
    int seqlen_k;
    int batch;
    int max_seqlen_q;
    int max_seqlen_k;
    int hdim_q;
    int hdim_v;
    int nhead_q;
    int nhead_k;
    std::string data_type;
    bool is_group_mode;
    mask_enum mask_type;
    bias_enum bias_type;
    bool has_dbias;
    bool has_dropout;
    bool is_store_randval;
    bool is_deterministic;
    // New: cumulative physical seqlen pointers for group mode (pass nullptr for batch mode).
    // seqstart_qs[i+1] - seqstart_qs[i] = physical seqlen_q of batch i (including padding); length = batch+1
    // seqstart_ks[i+1] - seqstart_ks[i] = physical seqlen_k of batch i (including padding); length = batch+1
    const int* seqstart_qs = nullptr;
    const int* seqstart_ks = nullptr;
};

fmha_bwd_launcher actual structure

struct fmha_bwd_launcher
{
    std::function<float(fmha_bwd_args, const ck_tile::stream_config&)> run{};

    // Total workspace size in bytes (host_ws_size + device_ws_size), computed by init().
    // Zero for kUseQrQtrDorPipeline (writes dq directly, no acc buffer needed).
    size_t workspace_size = 0;

    fmha_bwd_launcher(const fmha_bwd_traits&);

    // Copies auxiliary data (nsplits[], offsets[]) via hipMemcpy to the head of the GPU workspace,
    // and zeros the dq_acc buffer portion (tail of workspace) if required.
    // The memory pointed to by device_ws must be >= workspace_size bytes.
    std::function<void(void* device_ws)> prepare_workspace{};

    template <typename... Args>
    float operator()(Args&&... args) const { return run(std::forward<Args>(args)...); }

private:
    size_t host_ws_size   = 0;  // CPU workspace size (nsplits[] + offsets[] arrays)
    size_t device_ws_size = 0;  // GPU-only data size (dq_acc buffer)
    std::unique_ptr<char[]> ws_host;  // host-side workspace buffer

public:
    template <typename T0, typename T1, typename T2, typename Arch>
    void init(const fmha_bwd_traits& traits);
};

The init<>() template method (invoked by codegen dispatch branches as this->init<...>(t)) is responsible for:

1. Setting the run lambda
2. Calling FmhaBwdDQDKDVKernel::GetWorkspaceHostSize(batch) to obtain host_ws_size
3. Allocating ws_host (host memory)
4. Calling FmhaBwdDQDKDVKernel::PrepareWorkspaceHost(ws_host.get(), ...) to fill nsplits/offsets; return value is device_ws_size
5. workspace_size = host_ws_size + device_ws_size
6. Setting the prepare_workspace lambda (captures this, calls PrepareWorkspaceDevice)

When no kernel matches the given traits, both run and prepare_workspace are initialized to default lambdas that print a warning to std::cerr and return gracefully (no exception).

Workspace overall layout

The workspace is managed by FmhaBwdWorkspaceManager and consists of two segments:

Offset 0 (CPU-prepared segment, host_ws_size bytes; also hipMemcpy'd to the head of GPU workspace):
  index_t nsplits[batch or 1]       — per-batch nsplits array
                                      group mode: batch elements
                                      batch mode / non-deterministic: 1 element
  [group mode only] long_index_t dq_acc_offsets[batch+1]
                                    — per-batch element offset (inclusive prefix sum)
                                      offsets[0]=0, offsets[i+1] = offsets[i] + nhead*nsplits_i*seqq_i*hdim_q

Offset host_ws_size (device data segment, device_ws_size bytes):
  AccDataType dq_acc[total_elements] — compact dq_acc buffer (zeroed if required)
                                       total_elements = sum_i(nhead * nsplits_i * seqq_i * hdim_q)
                                       layout within each batch: [nhead, nsplits_i, seqq_i, hdim_q]
                                       note: seqq_i uses the physical length (including padding)

Alignment constant (ALIGNMENT = 16):

nsplits_size  = align_up(sizeof(index_t) * N, 16)          // N = batch (group) or 1 (batch/non-det)
offsets_size  = align_up(sizeof(long_index_t) * (batch+1), 16)  // group mode only
host_ws_size  = nsplits_size + offsets_size
dq_acc_offset = host_ws_size  // GetDqAccDataOffset(batch)

Key benefits:

• The kernel reads nsplits/offsets directly from the workspace head — no device-side recomputation.
• FmhaBwdConvertQGradKernel is completely decoupled from the pipeline block shape (kN0): nsplits is read from nsplits_ptr, kN0 is no longer a template parameter, and multiple dq_dk_dv tiles with different F_bn0 values now share a single convert_dq kernel instance (under receipt 1/2, deterministic convert_dq kernel count drops from ~300 to 60).
• nsplits/offsets are computed on the host and transferred in one hipMemcpy; the dq_acc buffer follows immediately, at the offset given by GetDqAccDataOffset.

Workspace size by scenario

Scenario	workspace_size	Notes
kUseQrQtrDorPipeline (any mode)	0	Writes dq directly; no acc buffer; PrepareWorkspaceHost returns 0
Non-deterministic + batch mode	> 0	nsplits[1]=1; dq_acc used for atomic add; workspace_size = host_ws_size + batch*nhead*seqlen_q*hdim_q*ebytes
Non-deterministic + group mode	> 0	nsplits[1]=1; dq_acc contiguous layout; workspace_size = host_ws_size + nhead*seqstart_qs[batch]*hdim_q*ebytes
Deterministic + group mode	> 0	nsplits[batch], offsets[batch+1], compact dq_acc; nsplits_i computed independently per batch
Deterministic + batch mode persistent	> 0	nsplits[1] (uniform across batches); dq_acc batch*nhead*nsplits*seqlen_q*hdim_q

NeedsZeroDqAcc (determines whether PrepareWorkspaceDevice calls hipMemset):

• Persistent kernel (deterministic batch mode) or non-deterministic: must zero (atomic add requires zero initialization)
• Deterministic group mode + no mask: no zeroing needed (every tile writes its full region)
• Deterministic + with mask: must zero (some blocks are skipped, leaving uninitialized tiles that would contribute to the reduction)

Caller usage

// 1. Create launcher (traits include seqstart_qs/ks pointers; workspace_size is computed during construction)
fmha_bwd_launcher launcher(fmha_traits);

// 2. Read launcher.workspace_size directly
const auto ws_size = launcher.workspace_size;

// 3. Allocate a single GPU workspace
ck_tile::DeviceMem ws_buf(ws_size);

// 4. Copy nsplits/offsets to GPU head and zero dq_acc if required
launcher.prepare_workspace(ws_buf.GetDeviceBuffer());

// 5. Build args with a single workspace pointer; the kernel splits it internally
fmha_bwd_args args{
    ...,
    ws_size > 0 ? ws_buf.GetDeviceBuffer() : nullptr,  // workspace_ptr
};
launcher(args, stream_config);

---

Key Code Structure

FmhaBwdWorkspaceManager (fmha_bwd_kernel.hpp, new class)

template <typename AccDataType, bool kIsGroupMode, bool kIsDeterministic>
struct FmhaBwdWorkspaceManager
{
    static constexpr size_t ALIGNMENT = 16;

    // CPU workspace (nsplits + offsets) sizes
    static size_t GetDqAccSplitsSize(int batch);   // align_up(sizeof(index_t)*N, 16)
    static size_t GetDqAccOffsetsSize(int batch);  // group mode only: align_up(sizeof(long_index_t)*(batch+1), 16)
    static size_t GetWorkspaceHostSize(int batch);  // = SplitsSize + OffsetsSize

    // Starting offset of dq_acc data within the full workspace (= host_ws_size)
    static size_t GetDqAccDataOffset(int batch);   // = GetWorkspaceHostSize(batch)

    // Fills nsplits/offsets in the CPU workspace; returns device_ws_size (dq_acc buffer bytes)
    template <bool kUseQrQtrDorPipeline, index_t kN0>
    static size_t PrepareWorkspaceHost(void* cpu_ws, index_t batch_size, index_t hdim_q,
                                       index_t nhead_q, index_t seqlen_q, index_t seqlen_k,
                                       const index_t* seqstart_qs, const index_t* seqstart_ks);

    // hipMemcpy's cpu_ws to device_ws head; hipMemset's the dq_acc portion to 0 if required
    template <bool kUseQrQtrDorPipeline, bool kHasMask>
    static void PrepareWorkspaceDevice(void* device_ws, const void* host_ws,
                                       size_t device_ws_size, size_t host_ws_size);
};

workspace_ptr parsing (inside the kernel)

The kernel parses three address regions from kargs.workspace_ptr:

Group mode (FmhaBwdDQDKDVKernel::MakeKargs):

const uint8_t* ws = reinterpret_cast<uint8_t*>(workspace_ptr);
// dq_acc_ptr (stored in FmhaBwdCommonKargs)
ws + WorkspaceManager::GetDqAccDataOffset(batch)
// dq_acc_batch_offset_ptr (FmhaBwdGroupModeKargs field)
reinterpret_cast<const long_index_t*>(ws + WorkspaceManager::GetDqAccOffsetsOffset(batch))

Batch mode:

ws + WorkspaceManager::GetDqAccDataOffset(batch)  // dq_acc_ptr
// No offsets pointer; batch offset is computed inside run_() from nsplits

FmhaBwdConvertQGradKernel follows the same pattern:

• Group mode: extracts dq_acc_ptr, dq_acc_batch_offset_ptr, and nsplits_ptr (GetDqAccSplitsOffset(batch)) from workspace
• Batch mode: reads nsplits from nsplits_ptr[0]; batch offset computed internally

Addressing in run_() (group mode)

// Per-batch processing:
const long_index_t batch_offset_dq_acc = kargs.dq_acc_batch_offset_ptr[i_batch];
// seqq_i (physical length) derived from seqstart_q_ptr
const index_t seqq_i = kargs.seqstart_q_ptr[i_batch+1] - kargs.seqstart_q_ptr[i_batch];
// nsplits_i read from nsplits_ptr (convert_dq kernel) or from GetDqAccSplits
const long_index_t split_stride_i = static_cast<long_index_t>(seqq_i) * kargs.hdim_q;
const long_index_t nhead_stride_i = static_cast<long_index_t>(nsplits_i) * split_stride_i;
// Final address:
dq_acc_base + batch_offset_dq_acc + i_nhead * nhead_stride_i + i_split * split_stride_i

nsplits computation (PrepareWorkspaceHost)

PrepareWorkspaceHost is a template method of FmhaBwdWorkspaceManager that still takes kN0 as a template parameter (from BlockFmhaShape::kN0 of the dq_dk_dv pipeline). However, this parameter is only used inside this host-side function to compute nsplits — it is no longer passed into the convert_dq kernel.

Mode	nsplits computation
kUseQrQtrDorPipeline	Writes dq directly; nsplits[0]=0; returns device_ws_size=0
Non-deterministic	nsplits[0]=1; dq_acc used for atomic add
Deterministic + group mode	ceil((seqstart_ks[i+1]-seqstart_ks[i]) / kN0) computed per batch
Deterministic + batch mode persistent	Same logic as the original GetDqAccSplits (dqdqkdv_workers based)

Removing kN0 dependency from FmhaBwdConvertQGradKernel

FmhaBwdConvertQGradKernel previously required kN0 as a template parameter (via BlockFmhaBwdConvertQGradPipelineProblem) for two purposes:

1. In batch mode operator(): self-computing nsplits = ceil(seqlen_k / kN0)
2. The b{kM0}x{kN0} component of the kernel name string

Both have been removed in this refactor:

• Batch mode: now reads kargs.nsplits_ptr[0] directly (guarded by if constexpr(kIsDeterministic) to avoid accessing a non-existent field in non-deterministic instances)
• Kernel name: simplified to b{kM0}, no longer includes kN0
• Template parameters: BlockFmhaBwdConvertQGradPipelineProblem drops the kN0_ parameter; fmha_bwd_convert_dq_traits_ drops the kN0 parameter; F_bn0/convert_dq_bn0 fields removed from codegen

Effect: all dq_dk_dv tiles sharing the same (hdim, dtype, mode, pad, deterministic) combination — regardless of F_bn0 value (16/64/128/192/256) — now share a single convert_dq kernel instance.

---

Test Plan

Test Result

Submission Checklist

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

[poyenc]

Code Review

Nice refactoring — the unified workspace approach is a clear win for API cleanliness, memory efficiency in group mode, and binary size reduction (convert_dq kernel deduplication). A few issues to address:

---

Bug

Batch-mode deterministic (persistent) offset mismatch in FmhaBwdDQDKDVKernel::run_()

Two places in run_() use integer_divide_ceil(seqlen_k, kN0) (i.e. jobs_per_head) as nsplits to compute offsets into dq_acc:

// batch offset
batch_offset_dq_acc = i_batch * nhead_q *
    integer_divide_ceil(seqlen_k, kN0) * seqlen_q * hdim_q;

// nhead+split offset
const auto nsplits = integer_divide_ceil(kargs.seqlen_k, FmhaPipeline::kN0);
return batch_offset_dq_acc + (i_nhead_ * nsplits + i_split) * split_stride;

But PrepareWorkspaceHost sizes the workspace using the persistent-formula nsplits (nsplits[0]), which is typically smaller than jobs_per_head. For example with batch=4, nhead=32, seqlen_k=2048, kN0=64, num_cus=120: jobs_per_head=32 but nsplits[0]=1. The kernel would use a stride 32x larger than what was allocated.

The convert_dq kernel handles this correctly by reading kargs.nsplits_ptr[0] from the workspace. The DQ-DK-DV kernel should do the same.

---

Latent UB

1. PrepareWorkspaceHost QrQtrDor path writes to unallocated memory. When kUseQrQtrDorPipeline=true, GetDqAccSplitsSize returns 0 so ws_host is never allocated, but the function unconditionally writes nsplits[0] = 0. Today unreachable (guarded by if(host_ws_size > 0) in init()), but the function itself is unsafe if called directly. Consider an early return before touching nsplits.

2. Null pointer arithmetic in group-mode QrQtrDor MakeKargs. dq_acc_batch_offset_ptr is always initialized from ws + GetDqAccOffsetsOffset(batch). When QrQtrDor, workspace_ptr is nullptr, so ws = nullptr. The arithmetic nullptr + 0 is UB in C++ even though the pointer is never dereferenced. Gate with if constexpr(!kUseQrQtrDorPipeline).

---

Design

3. prepare_workspace lambda captures this. The move/copy deletions prevent dangling today, but the design is fragile — future changes allowing move/copy would silently break. Capturing the needed state by value (ws_host.get(), device_ws_size, host_ws_size) would be self-contained and let the class become movable if needed.

4. init() is public. It's only called by codegen dispatch inside the constructor. Being public, a caller could accidentally reinitialize the launcher. Consider making it private.

---

Cleanup

5. Dead kN0 in block_fmha_bwd_convert_dq.hpp:18. static constexpr index_t kN0 = Problem::kN0; — Problem no longer has kN0. Compiles because the class template member is never ODR-used, but should be removed.

6. Unnecessary const_cast in PrepareWorkspaceDevice. hipMemcpy source is const void* — no cast needed.

7. GetDqAccOffsetsSize lacks kUseQrQtrDorPipeline template guard. GetDqAccSplitsSize has it but GetDqAccOffsetsSize does not. Results happen to be consistent, but the asymmetry is confusing.

---

Minor

8. fmha_bwd_traits lifetime requirement. The new seqstart_qs/seqstart_ks raw pointers must remain valid through launcher construction. Worth a comment documenting this.

9. Lost diagnostic info. Old runner output: workspace:42MiB|8splits. New: workspace:42MiB. The nsplits info was useful for debugging deterministic mode.

[DDEle]

@poyenc,

Comments 1/2/3/5/6/7/8/9: Fixed, thanks. For comment 8: GetDqAccOffsetsSize is kept without a template parameter — GetWorkspaceHostSize<true> now returns 0 early, so GetDqAccOffsetsSize is never reached on the QrQtrDor path.

Comment 4: Capturing by raw pointer creates the same lifetime dependency, just less explicitly. Since move/copy are deleted the lambda cannot dangle, so I'd prefer to keep this capture as the clearer expression of ownership.

Comment 10: In group mode each batch now has its own nsplits_i, so there's no single scalar to print. Left it out intentionally.

[poyenc]

All review comments addressed — the critical bug (persistent offset mismatch) and UB issues are fixed, and the declined items (this-capture, nsplits diagnostic) are reasonable. LGTM once marked ready for review.

[shumway]

This PR may further break the device/host code separation required for integration in rocm-libraries. We need to review this carefully. @afagaj @vidyasagar-amd @dpeabody-amd @aosewski

[shumway]

This PR may further break the device/host code separation required for integration in rocm-libraries. We need to review this carefully. @afagaj @vidyasagar-amd @dpeabody-amd @aosewski

I really like this design. We’ll need to think carefully about how to include it rock-ck, the dispatcher, and kpack/rfc0008/host-device separation. It’s a great example of the kinds of generalization for kernel launches we’ll need to support.

[aosewski]

Thanks @DDEle — the direction here looks right for the host/device separation we need on the rocm-libraries integration side. A few observations and one concrete ask:

What works well for us

• Collapsing dq_acc_ptr + the four *_stride_dq_acc fields into a single opaque workspace_ptr is exactly the ABI shape a host-only driver wants — no layout knowledge leaks past the boundary.
• Having the kernel read its own nsplits[] / per-batch offsets from the head of the workspace makes the device side self-describing given a correctly laid-out buffer. That removes a real coupling from kargs.
• Per-batch nsplits_i in group mode is a meaningful memory + correctness win independent of integration concerns.
• Dropping kN0 from FmhaBwdConvertQGradKernel collapses many precompiled instances into one — strictly good for a binary-distribution path.

Where we hit friction
The remaining coupling is in FmhaBwdWorkspaceManager<Pipeline, ConvertDqPipeline>::PrepareWorkspaceHost(...) and fmha_bwd_launcher::init<T0,T1,T2,Arch>(traits). Both are pure host code by body (CPU arithmetic, plus hipMemcpy/hipMemset driven from the host), but they are templated on device-side pipeline types just to read Pipeline::kN0, the persistent/deterministic flavor, and the accumulator element size. For a multi-arch host TU that does not (and should not) instantiate CK Tile pipelines, that's the awkward part.

Concrete ask
Would you be open to splitting FmhaBwdWorkspaceManager into:

1. a small POD FmhaBwdWorkspaceSpec { kN0; is_persistent; is_deterministic; is_group_mode; acc_element_size; needs_zero_dq_acc; }
2. a non-template host function prepare_workspace_host(const FmhaBwdWorkspaceSpec&, …) containing today's body
3. a template factory make_fmha_bwd_workspace_spec<Pipeline, ConvertDqPipeline>() that lives in the device-side template header and produces the POD

The in-tree example launcher uses the factory at compile time and is unchanged in behavior. A host-only driver fills the POD from operator metadata. The kernel, the workspace memory layout, and the codegen branches stay as you have them.

The same POD is also the natural extension point on our side for expressing the workspace contract under a Signature/Algorithm split — which we'd otherwise have to bolt on adapter-side anyway.

Minor: worth a one-line comment in fmha_bwd_traits documenting the seqstart_qs/ks lifetime requirement (must remain valid through launcher construction). poyenc already noted this as a "minor"; it'd help downstream callers.

[DDEle]

Thanks @aosewski — glad the direction works for the integration side.

On the FmhaBwdWorkspaceSpec proposal

Agree with the direction. The POD + factory + non-template host function split is the right shape, both for an out-of-tree host-only driver and as the extension point for the Signature/Algorithm split.

I'd like to defer it to a follow-up PR though. I prototyped locally and the surface is larger than the workspace API alone — kargs assembly wants the same treatment (it currently pulls in pipeline types for stride / padding / group-mode offset computation), and once that's pulled out the change spreads across kernel template, codegen, launcher, and the new host TU. One thing I'd want to be careful about in the layout: host-side workspace prep and kernel-side workspace consumption are tightly coupled by contract, so from a kernel developer's perspective keeping them at nearby paths is more intuitive than splitting purely along host-only / device-only TU boundaries — worth thinking through how to satisfy both the integration requirement and code locality. I'd rather keep this PR scoped to the unified-workspace ABI and land the host/device split cleanly in a follow-up where we can agree on the spec boundary and file layout up front.

Re minor (seqstart_qs/ks lifetime)

Already documented at fmha_bwd.hpp:542-544 since bf364d7 — likely an earlier diff snapshot.