Porting existing FMHA infra from users/shumway/ck/exp-kpack to develop

projects/composablekernel/rocm_ck/include/rocm_ck/args.hpp

@@ -12,8 +12,8 @@
 // an operation means adding a type, updating launch code, and changing the
 // kpack format. A generic buffer keeps the dispatcher open.
 //
-// Capacity limits (kMaxRank=6, kMaxTensors=16, kMaxScalars=16) are sized to
-// the most demanding current operation (FMHA backward: ~12 tensors, ~12
+// Capacity limits (kMaxRank=6, kMaxTensors=20, kMaxScalars=16) are sized to
+// the most demanding current operation (FMHA backward: ~18 tensors, ~12
 // scalars, rank-6 for grouped 3D conv). If a future operation exceeds these,
 // bump the constants - the layout is not versioned, and the 4KB HSA kernarg
 // budget has room. Don't over-provision speculatively.
@@ -38,7 +38,7 @@ namespace rocm_ck {
 
 // When changing these, update the byte-size comments on TensorArg and Args fields.
 constexpr int kMaxRank    = 6;  // grouped 3D conv = GNCDHW = rank 6
-constexpr int kMaxTensors = 16; // FMHA backward uses ~12
+constexpr int kMaxTensors = 20; // FMHA backward uses ~18
 constexpr int kMaxScalars = 16; // FMHA with masking+dropout needs ~12
 
 struct TensorArg
@@ -62,11 +62,11 @@ union ScalarValue
 // Slot ordering matches Signature: tensors[i] <-> Signature::tensors[i].
 struct Args
 {
-    std::array<TensorArg, kMaxTensors> tensors;   // 16 x 80 = 1280 bytes
+    std::array<TensorArg, kMaxTensors> tensors;   // 20 x 80 = 1600 bytes
     std::array<ScalarValue, kMaxScalars> scalars; // 16 x  8 =  128 bytes
 
     index_t batch_count                                 = 0;       //  4 bytes
-    std::array<long_index_t, kMaxTensors> batch_strides = {};      // 16 x 8 = 128 bytes
+    std::array<long_index_t, kMaxTensors> batch_strides = {};      // 20 x 8 = 160 bytes
     void* workspace_ptr                                 = nullptr; //  8 bytes
 };
 

projects/composablekernel/rocm_ck/include/rocm_ck/ck_type_map.hpp

@@ -0,0 +1,74 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+//
+// Role: device -- maps DataType to CK Tile types. Requires --cuda-device-only.
+//
+// Maps DataType enum values to CK Tile C++ numeric types.
+
+#pragma once
+
+#ifndef __HIP_DEVICE_COMPILE__
+#error "ck_type_map.hpp requires CK Tile headers (device compilation only)"
+#endif
+
+#include <rocm_ck/datatype.hpp>
+
+#include <ck_tile/core.hpp>
+
+namespace rocm_ck {
+
+/// Maps a DataType enum value to the corresponding CK Tile numeric type.
+/// Primary template is intentionally undefined -- only valid specializations compile.
+/// Add specializations as new DataType values are used in device kernels.
+template <DataType>
+struct CkTypeMap;
+
+template <>
+struct CkTypeMap<DataType::FP64>
+{
+    using type = double;
+};
+template <>
+struct CkTypeMap<DataType::FP32>
+{
+    using type = float;
+};
+template <>
+struct CkTypeMap<DataType::FP16>
+{
+    using type = ck_tile::half_t;
+};
+template <>
+struct CkTypeMap<DataType::BF16>
+{
+    using type = ck_tile::bf16_t;
+};
+template <>
+struct CkTypeMap<DataType::FP8_FNUZ>
+{
+    using type = ck_tile::fp8_t;
+};
+template <>
+struct CkTypeMap<DataType::BF8_FNUZ>
+{
+    using type = ck_tile::bf8_t;
+};
+// FP8_OCP/BF8_OCP: add when CK Tile exposes distinct OCP types.
+// Currently ck_tile::fp8_t/bf8_t are selected at compile time via CK_TILE_USE_OCP_FP8.
+template <>
+struct CkTypeMap<DataType::I8>
+{
+    using type = int8_t;
+};
+template <>
+struct CkTypeMap<DataType::I32>
+{
+    using type = int32_t;
+};
+template <>
+struct CkTypeMap<DataType::I4>
+{
+    using type = ck_tile::pk_int4_t;
+};
+
+} // namespace rocm_ck

projects/composablekernel/rocm_ck/include/rocm_ck/grid_dim.hpp

@@ -0,0 +1,19 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+//
+// HIP-free grid dimension type for host-only kernel launch calculations.
+
+#pragma once
+
+namespace rocm_ck {
+
+/// HIP-free replacement for dim3. Consumers convert to dim3 at the HIP
+/// call site: `dim3 grid(g.x, g.y, g.z)`.
+struct GridDim
+{
+    unsigned int x = 1;
+    unsigned int y = 1;
+    unsigned int z = 1;
+};
+
+} // namespace rocm_ck

projects/composablekernel/rocm_ck/include/rocm_ck/ops/fmha_bwd/common.hpp

@@ -0,0 +1,65 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+//
+// Shared types for all FMHA BWD kernel families (OGradDotO, DqDkDv, ConvertDq).
+//
+// This header has NO CK Tile dependency. It is included by both host code
+// (main.cpp) and device code (.hip files).
+
+#pragma once
+
+#include <rocm_ck/datatype.hpp>
+#include <rocm_ck/index_t.hpp>
+
+#include <cstdint>
+
+namespace rocm_ck {
+
+// Padding semantics vary per kernel family:
+//   OGradDotO / ConvertDQ: bool (pad or no-pad)
+//   DqDkDv: int {0=none, 1=small, 8=full vector-aligned}
+// The tri-valued int maps to CK Tile's TileFmhaBwdTraits::kPadHeadDimQ/V
+// which controls vector load widths. 0 = no padding, 1 = scalar fallback,
+// 8 = full 128-bit vector loads with padding. bool is sufficient for
+// OGradDotO/ConvertDQ which only need on/off.
+
+/// FMHA attention mode: fixed-length batches vs variable-length groups.
+enum class FmhaMode
+{
+    BATCH,
+    GROUP
+};
+
+/// Bias type for attention score modification.
+/// Values must match ck_tile::BlockAttentionBiasEnum.
+enum class FmhaBiasType
+{
+    NONE,
+    ELEMENTWISE,
+    ALIBI
+};
+
+/// Attention mask family.
+///
+/// Integer values must match ck_tile::GenericAttentionMaskEnum so the device
+/// bridge can forward the spec-time enum to the kernel via a static_cast
+/// without a translation table (see dqdkdv_dev.hpp's MASK_TYPE scalar wiring).
+///
+/// Both causal variants describe a sliding window with left=-1 (unbounded
+/// past) and right=0 (no lookahead). The two flavours differ in where the
+/// causal diagonal is anchored when seqlen_q != seqlen_k:
+///   * TOP_LEFT     -- diagonal at (0, 0); standard "predict next token".
+///   * BOTTOM_RIGHT -- diagonal at (seqlen_q-1, seqlen_k-1); used when the
+///                     query is the *tail* of a longer cached K/V (decode).
+/// GENERIC selects ck_tile's runtime (left, right, top-left/bottom-right)
+/// window descriptor and is intended for sliding-window / xformer-style
+/// attention. NO_MASK disables masking at compile time.
+enum class FmhaMaskType
+{
+    NO_MASK             = 0,
+    TOP_LEFT_CAUSAL     = 1,
+    BOTTOM_RIGHT_CAUSAL = 2,
+    GENERIC             = 3,
+};
+
+} // namespace rocm_ck

projects/composablekernel/rocm_ck/include/rocm_ck/ops/fmha_bwd/convert_dq_api.hpp

@@ -0,0 +1,119 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+//
+// Host-only helpers for the FMHA BWD ConvertDQ kernel family.
+//
+// HOST ONLY: this header must NOT be included from device code (.hip files).
+// Device code should include <rocm_ck/ops/fmha_bwd/convert_dq_dev.hpp>.
+//
+// Compilation boundary:
+//   _spec.hpp -- consteval factory + slot constants (both passes)
+//   _api.hpp (this) -- host-only helpers: grid_size (host pass only, #error on device)
+//   _dev.hpp -- CK Tile bridge + __device__ code (device pass only, #error on host)
+
+#pragma once
+
+#ifdef __HIP_DEVICE_COMPILE__
+#error "convert_dq_api.hpp is host-only." \
+       " Device code should include <rocm_ck/ops/fmha_bwd/convert_dq_dev.hpp>."
+#endif
+
+#include <rocm_ck/ops/fmha_bwd/convert_dq_spec.hpp>
+
+#include <rocm_ck/args.hpp>
+#include <rocm_ck/grid_dim.hpp>
+
+#ifndef NDEBUG
+#include <cstdio>
+#include <cstdlib>
+#endif
+
+namespace rocm_ck {
+
+// ---------------------------------------------------------------------------
+// Grid calculation
+// ---------------------------------------------------------------------------
+
+/// Compute the launch grid for ConvertDQ.
+/// Matches FmhaBwdConvertQGradKernel::GridSize():
+///   GridDim(ceil(seqlen_q / kM0), nhead, batch).
+/// kM0 = 64 (tile rows along seqlen_q for 1D kernels), NOT block_size.
+/// Precondition: tile_m0 > 0, seqlen_q >= 0, batch > 0, nhead > 0.
+inline GridDim convert_dq_grid_size(int batch, int nhead, int seqlen_q, int tile_m0 = 64)
+{
+#ifndef NDEBUG
+    if(tile_m0 <= 0)
+    {
+        std::fprintf(
+            stderr, "rocm_ck::convert_dq_grid_size: tile_m0 must be positive, got %d\n", tile_m0);
+        std::abort();
+    }
+    if(seqlen_q < 0)
+    {
+        std::fprintf(stderr,
+                     "rocm_ck::convert_dq_grid_size: seqlen_q must be non-negative, got %d\n",
+                     seqlen_q);
+        std::abort();
+    }
+    if(batch <= 0)
+    {
+        std::fprintf(
+            stderr, "rocm_ck::convert_dq_grid_size: batch must be positive, got %d\n", batch);
+        std::abort();
+    }
+    if(nhead <= 0)
+    {
+        std::fprintf(
+            stderr, "rocm_ck::convert_dq_grid_size: nhead must be positive, got %d\n", nhead);
+        std::abort();
+    }
+#endif
+    const auto uq = static_cast<unsigned>(seqlen_q);
+    const auto ut = static_cast<unsigned>(tile_m0);
+    return {(uq + ut - 1u) / ut, static_cast<unsigned>(nhead), static_cast<unsigned>(batch)};
+}
+
+// ---------------------------------------------------------------------------
+// Debug-only runtime Args validation
+// ---------------------------------------------------------------------------
+
+/// Validate that all required tensor slots for ConvertDQ are populated.
+/// Compiles to nothing in release builds.
+inline void validateArgs([[maybe_unused]] const Args& args, [[maybe_unused]] FmhaBwdConvertDQSpec k)
+{
+#ifndef NDEBUG
+    namespace S = fmha_bwd_convert_dq_slots;
+
+    // Slot 2 is mode-overloaded (BATCH: NSPLITS; GROUP: SEQSTART_Q), and GROUP
+    // mode packs the workspace slots after the seq* slots (see the slot-layout
+    // comment in convert_dq_spec.hpp). Use a per-mode name table sized to that
+    // mode's requiredTensors() so the lookup below never indexes out of bounds.
+    // clang-format off
+    static constexpr const char* batch_tensor_names[] = {
+        "DQ_ACC", "DQ", "NSPLITS"
+    };
+    static constexpr const char* group_tensor_names[] = {
+        "DQ_ACC", "DQ", "SEQSTART_Q", "SEQLEN_Q", "SEQSTART_K", "SEQLEN_K",
+        "NSPLITS", "DQ_ACC_BATCH_OFFSET"
+    };
+    // clang-format on
+
+    const bool is_group      = (k.mode == FmhaMode::GROUP);
+    const char* const* names = is_group ? group_tensor_names : batch_tensor_names;
+    const int n              = S::requiredTensors(k); // 3 (batch) or 8 (group)
+    for(int i = 0; i < n; ++i)
+    {
+        if(args.tensors[i].ptr == nullptr)
+        {
+            std::fprintf(stderr,
+                         "rocm_ck::validateArgs(ConvertDQ): tensor \"%s\" (slot %d)"
+                         " has null pointer\n",
+                         names[i],
+                         i);
+            std::abort();
+        }
+    }
+#endif
+}
+
+} // namespace rocm_ck