diff --git a/example/ck_tile/01_fmha/fmha_bwd.hpp b/example/ck_tile/01_fmha/fmha_bwd.hpp index 20e4a82241..93e770a0cd 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd.hpp @@ -18,6 +18,17 @@ #include #include +// Function pointer type for workspace packing (no captures, points to code segment) +using PrepareWorkspaceHostFunc = size_t (*)(void*, // host_ws + ck_tile::index_t, // batch + ck_tile::index_t, // hdim_q + ck_tile::index_t, // nhead_q + ck_tile::index_t, // seqlen_q + ck_tile::index_t, // seqlen_k + const ck_tile::index_t*, // seqstart_q + const ck_tile::index_t* // seqstart_k +); + struct FmhaBwdFp32 { }; @@ -588,6 +599,56 @@ float fmha_bwd(const fmha_bwd_traits&, fmha_bwd_args, const ck_tile::stream_conf struct fmha_bwd_launcher { + // POD closure for graph capture (trivially destructible, no heap allocation) + struct GraphClosure + { + PrepareWorkspaceHostFunc func_ptr; + void* pin_w_ptr; + const ck_tile::index_t* seqstart_q_ptr; + const ck_tile::index_t* seqstart_k_ptr; + ck_tile::index_t batch; + ck_tile::index_t hdim_q; + ck_tile::index_t nhead_q; + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_k; + + static void invoke(void* ud) + { + auto* closure = static_cast(ud); + if(closure->func_ptr) + { + // Callback runs on the HIP driver helper thread across a C ABI boundary; + // any exception escaping it would call std::terminate. + try + { + closure->func_ptr(closure->pin_w_ptr, + closure->batch, + closure->hdim_q, + closure->nhead_q, + closure->seqlen_q, + closure->seqlen_k, + closure->seqstart_q_ptr, + closure->seqstart_k_ptr); + } + catch(const std::exception& e) + { + // The H2D queued after this callback will copy indeterminate + // metadata to device and the kernel will produce wrong results; + // unlikely in practice since pack_workspace_host only throws on + // precondition violations. + std::cerr << "fmha_bwd_launcher: pack_workspace_host threw: " << e.what() + << '\n'; + } + catch(...) + { + std::cerr << "fmha_bwd_launcher: pack_workspace_host threw unknown\n"; + } + } + } + }; + static_assert(std::is_trivially_destructible_v, + "GraphClosure must be trivially destructible for placement-new without dtor"); + std::function run{ [](fmha_bwd_args, const ck_tile::stream_config&) { std::cerr << "fmha_bwd: no kernel found for given traits, skipping run\n"; @@ -603,8 +664,28 @@ struct fmha_bwd_launcher ~fmha_bwd_launcher() noexcept { schedule_pin_staging_release(); } // Stream-async: zero dq_acc, D2H seqstart, host-pack metadata, H2D into device_ws. - // `pinned_host_alloc` returns a shared_ptr to a pinned host buffer; its deleter - // is invoked on the stream tail after the H2D completes. + // + // `pinned_host_alloc` returns a shared_ptr to a pinned host buffer. + // + // **Deleter behavior differs by mode**: + // - Normal mode: deleter is invoked on the stream tail after H2D completes. + // - Graph capture mode: deleter is NOT invoked; caller must keep buffer alive + // until hipGraphDestroy (see precondition #2 below). + // + // REQUIRED PRECONDITIONS for `pinned_host_alloc`: + // + // 1. **Capture-safe allocation**: The allocator must NOT call synchronizing APIs + // (e.g., bare hipHostMalloc) during active stream capture, as these invalidate + // the capture and cause hipStreamEndCapture to fail. Use a caching allocator + // that serves from cache during capture (e.g., PyTorch CachingHostAllocator). + // + // 2. **Buffer lifetime in graph mode**: In graph capture mode (detected via + // hipStreamIsCapturing), the returned buffer is NOT freed automatically after + // the H2D completes. The caller MUST keep the buffer alive for the graph's + // entire lifetime (until hipGraphDestroy), otherwise graph replay will read + // from freed memory. In normal (non-graph) mode, the buffer is automatically + // freed on the stream tail after the H2D. + // void prepare_workspace_async( // void* device_ws_ptr, const int* seqstart_q_dev, @@ -634,8 +715,20 @@ struct fmha_bwd_launcher const size_t seqstart_stride = ck_tile::integer_least_multiple(seqstart_bytes, static_cast(16)); const size_t pin_w_offset = 2 * seqstart_stride; - const size_t total_bytes = pin_w_offset + host_ws_size_; - auto pin_base = pinned_host_alloc(total_bytes); + const size_t data_size = pin_w_offset + host_ws_size_; + + // Check if we're in graph capture mode + hipStreamCaptureStatus capture_status; + HIP_CHECK_ERROR(hipStreamIsCapturing(stream, &capture_status)); + const bool is_graph_capture = (capture_status == hipStreamCaptureStatusActive); + + // Allocate pinned buffer with extra aligned space for closure + // Both modes use placement new for closure (POD, trivially destructible) + constexpr size_t closure_align = alignof(GraphClosure); + constexpr size_t closure_size = sizeof(GraphClosure); + const size_t aligned_data = (data_size + closure_align - 1) & ~(closure_align - 1); + const size_t total_bytes = aligned_data + closure_size; + auto pin_base = pinned_host_alloc(total_bytes); if(needs_zero_dq_acc_ && workspace_size > host_ws_size_) HIP_CHECK_ERROR(hipMemsetAsync(static_cast(device_ws_ptr) + host_ws_size_, @@ -643,12 +736,14 @@ struct fmha_bwd_launcher workspace_size - host_ws_size_, stream)); - char* base = static_cast(pin_base.get()); - int* pin_q = reinterpret_cast(base); - int* pin_k = reinterpret_cast(base + seqstart_stride); - void* pin_w = base + pin_w_offset; - const int* seqstart_q_pinned = traits_.is_group_mode ? pin_q : nullptr; - const int* seqstart_k_pinned = traits_.is_group_mode ? pin_k : nullptr; + char* base = static_cast(pin_base.get()); + int* pin_q = reinterpret_cast(base); + int* pin_k = reinterpret_cast(base + seqstart_stride); + void* pin_w = base + pin_w_offset; + const ck_tile::index_t* seqstart_q_pinned = + traits_.is_group_mode ? reinterpret_cast(pin_q) : nullptr; + const ck_tile::index_t* seqstart_k_pinned = + traits_.is_group_mode ? reinterpret_cast(pin_k) : nullptr; if(traits_.is_group_mode) { @@ -662,55 +757,56 @@ struct fmha_bwd_launcher pin_k, seqstart_k_dev, seqstart_bytes, hipMemcpyDeviceToHost, stream)); } - auto pack_closure = std::make_unique>( - [=, fn = pack_workspace_host_]() { fn(pin_w, seqstart_q_pinned, seqstart_k_pinned); }); - // Callback runs on the HIP driver helper thread across a C ABI boundary; - // any exception escaping it would call std::terminate. - HIP_CHECK_ERROR(hipLaunchHostFunc( - stream, - [](void* ud) { - std::unique_ptr> c{static_cast*>(ud)}; - try - { - (*c)(); - } - catch(const std::exception& e) - { - // The H2D queued after this callback will copy indeterminate - // metadata to device and the kernel will produce wrong results; - // unlikely in practice since pack_workspace_host_ only throws on - // precondition violations. - std::cerr << "fmha_bwd_launcher: pack_workspace_host threw: " << e.what() - << '\n'; - } - catch(...) - { - std::cerr << "fmha_bwd_launcher: pack_workspace_host threw unknown\n"; - } - }, - pack_closure.get())); - // Ownership transferred to the callback only after a successful launch. - pack_closure.release(); + // === UNIFIED PATH: Both modes use placement new for POD closure === + // Construct closure in aligned location within pinned buffer (all POD, trivially + // destructible) + void* closure_addr = base + aligned_data; + auto* closure = new(closure_addr) GraphClosure{ + prepare_ws_func_, // Function pointer (points to code segment, always valid) + pin_w, + seqstart_q_pinned, + seqstart_k_pinned, + batch_, // Copy captured data (independent of launcher lifetime) + hdim_q_, + nhead_q_, + seqlen_q_, + seqlen_k_}; + + HIP_CHECK_ERROR(hipLaunchHostFunc(stream, GraphClosure::invoke, closure)); HIP_CHECK_ERROR( hipMemcpyAsync(device_ws_ptr, pin_w, host_ws_size_, hipMemcpyHostToDevice, stream)); - // Release any previous in-flight buffer before taking a new one. - schedule_pin_staging_release(); - pin_staging_ = std::move(pin_base); - release_stream_ = stream; + if(!is_graph_capture) + { + // Normal mode: transfer ownership of pinned buffer and schedule release. + // The H2D memcpy above copies host-prepared metadata to device workspace. + // Pinned buffer must stay alive until H2D completes, so we schedule its + // release on the stream tail (via hipLaunchHostFunc callback). + schedule_pin_staging_release(); // Release any previous in-flight buffer first + pin_staging_ = std::move(pin_base); + release_stream_ = stream; + } + // Graph mode: pin_base stays owned by caller (via shared_ptr ref-count). + // Caller must keep it alive for the graph's lifetime. Closure and data pointers + // within pin_base remain valid as long as caller holds pin_base. } private: fmha_bwd_traits traits_{}; size_t host_ws_size_ = 0; bool needs_zero_dq_acc_ = false; - // Pure CPU; safe to invoke from a hipLaunchHostFunc callback. - std::function - pack_workspace_host_{[](void*, const int*, const int*) { - std::cerr - << "fmha_bwd: no kernel found for given traits, skipping pack_workspace_host\n"; - }}; + + // Function pointer (points to code segment, survives launcher destruction) + PrepareWorkspaceHostFunc prepare_ws_func_ = nullptr; + + // Captured data (copied to closure in graph mode) + ck_tile::index_t batch_ = 0; + ck_tile::index_t hdim_q_ = 0; + ck_tile::index_t nhead_q_ = 0; + ck_tile::index_t seqlen_q_ = 0; + ck_tile::index_t seqlen_k_ = 0; + std::shared_ptr pin_staging_; hipStream_t release_stream_ = nullptr; @@ -756,15 +852,16 @@ struct fmha_bwd_launcher t.is_group_mode ? t.seqlen_q : t.batch * t.seqlen_q; device_ws_size = fmha_bwd_dq_dk_dv_dq_ws_device_upper_bound_( t.batch, t.hdim_q, t.nhead_q, total_seqlen_q_padded, t.max_seqlen_k); - pack_workspace_host_ = [batch = t.batch, - hdim_q = t.hdim_q, - nhead_q = t.nhead_q, - seqlen_q = t.seqlen_q, - seqlen_k = t.seqlen_k // - ](void* host_ws, const int* seqstart_q, const int* seqstart_k) { - fmha_bwd_dq_dk_dv_dq_prepare_ws_host_( - host_ws, batch, hdim_q, nhead_q, seqlen_q, seqlen_k, seqstart_q, seqstart_k); - }; + + // Store function pointer (directly assign template-instantiated function) + prepare_ws_func_ = &fmha_bwd_dq_dk_dv_dq_prepare_ws_host_; + + // Store captured data as member variables + batch_ = t.batch; + hdim_q_ = t.hdim_q; + nhead_q_ = t.nhead_q; + seqlen_q_ = t.seqlen_q; + seqlen_k_ = t.seqlen_k; } workspace_size = host_ws_size_ + device_ws_size; needs_zero_dq_acc_ = fmha_bwd_dq_dk_dv_needs_zero_dq_acc_();