mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-17 00:58:44 +00:00
[rocm-libraries] ROCm/rocm-libraries#8333 (commit 69b3fc1)
Revert "[CK_TILE] Implement RTC API for a subset of FMHA functionality for MGX" (#8333) Reverts ROCm/rocm-libraries#6086 Need to revert as the codegen test for fmha is failing due to including std header: 2026-06-11T22:36:03.673Z] In file included from /tmp/comgr-953928-0-473822/include/ck/host/device_fmha_fwd/fmha_fwd_wrapper.hpp:8: [2026-06-11T22:36:03.673Z] In file included from /bin/../lib/gcc/x86_64-linux-gnu/13/../../../../include/c++/13/cmath:49: [2026-06-11T22:36:03.673Z] In file included from /bin/../lib/gcc/x86_64-linux-gnu/13/../../../../include/c++/13/bits/std_abs.h:38: [2026-06-11T22:36:03.673Z] /usr/include/stdlib.h:32:10: fatal error: 'stddef.h' file not found [2026-06-11T22:36:03.673Z] 32 | #include <stddef.h> [2026-06-11T22:36:03.673Z] | ^~~~~~~~~~ The ck_tile headers were never prepped for hiprtc compilation.
This commit is contained in:
Illia Silin
committed by
assistant-librarian[bot]
assistant-librarian[bot]
parent
c2601f38b7
commit
789ef38093
@@ -1,23 +0,0 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
|
||||
namespace ck {
|
||||
namespace host {
|
||||
|
||||
// Preprocesses a ck_tile header for HIPRTC compilation by replacing
|
||||
// the bodies of CK_TILE_HOST-only functions with stubs. This prevents
|
||||
// host-only code (which references APIs unavailable in HIPRTC) from
|
||||
// being type-checked by the device compiler.
|
||||
//
|
||||
// For non-constexpr functions: { __builtin_unreachable(); }
|
||||
// For constexpr functions: { return {}; }
|
||||
// For constexpr auto functions: { return 0; }
|
||||
std::string strip_host_bodies(std::string_view content);
|
||||
|
||||
} // namespace host
|
||||
} // namespace ck
|
||||
@@ -1,273 +0,0 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
// This header is designed to be embedded and used at RTC compilation time.
|
||||
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cassert>
|
||||
#include "ck_tile/core.hpp"
|
||||
#include "ck_tile/ops/fmha.hpp"
|
||||
#include "ck_tile/ops/epilogue/default_2d_epilogue.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
enum class FmhaPipelineTag
|
||||
{
|
||||
QR, // BlockFmhaPipelineQRKSVS
|
||||
QR_ASYNC, // BlockFmhaPipelineQRKSVSAsync
|
||||
QR_ASYNC_TRLOAD // BlockFmhaPipelineQRKSVSAsyncTrload
|
||||
};
|
||||
|
||||
template <typename DataType_,
|
||||
// Block tile
|
||||
index_t kBM0,
|
||||
index_t kBN0,
|
||||
index_t kBK0,
|
||||
index_t kBN1,
|
||||
index_t kBK1,
|
||||
index_t kBK0Max,
|
||||
// Gemm0 block warps
|
||||
index_t kRM0,
|
||||
index_t kRN0,
|
||||
index_t kRK0,
|
||||
// Gemm1 block warps
|
||||
index_t kRM1,
|
||||
index_t kRN1,
|
||||
index_t kRK1,
|
||||
// Gemm0 warp tile
|
||||
index_t kWM0,
|
||||
index_t kWN0,
|
||||
index_t kWK0,
|
||||
// Gemm1 warp tile
|
||||
index_t kWM1,
|
||||
index_t kWN1,
|
||||
index_t kWK1,
|
||||
//
|
||||
bool kIsCausal,
|
||||
bool kIsVRowMajor,
|
||||
bool kHasBias,
|
||||
//
|
||||
bool kPadM,
|
||||
bool kPadN,
|
||||
bool kPadK,
|
||||
bool kPadO,
|
||||
//
|
||||
FmhaPipelineTag kPipelineTag>
|
||||
struct FmhaFwdWrapper
|
||||
{
|
||||
using BlockTile = sequence<kBM0, kBN0, kBK0, kBN1, kBK1, kBK0Max>;
|
||||
|
||||
using Gemm0BlockWarps = sequence<kRM0, kRN0, kRK0>;
|
||||
using Gemm0WarpTile = sequence<kWM0, kWN0, kWK0>;
|
||||
using Gemm1BlockWarps = sequence<kRM1, kRN1, kRK1>;
|
||||
using Gemm1WarpTile = sequence<kWM1, kWN1, kWK1>;
|
||||
|
||||
using FmhaShape = TileFmhaShape<BlockTile,
|
||||
Gemm0BlockWarps,
|
||||
Gemm0WarpTile,
|
||||
Gemm1BlockWarps,
|
||||
Gemm1WarpTile,
|
||||
kIsVRowMajor>;
|
||||
|
||||
static constexpr auto BiasEnum =
|
||||
kHasBias ? BlockAttentionBiasEnum::ELEMENTWISE_BIAS : BlockAttentionBiasEnum::NO_BIAS;
|
||||
|
||||
using FmhaTraits = TileFmhaTraits<kPadM, // kPadSeqLenQ
|
||||
kPadN, // kPadSeqLenK
|
||||
kPadK, // kPadHeadDimQ
|
||||
kPadO, // kPadHeadDimV
|
||||
false, // kHasLogitsSoftCap
|
||||
BiasEnum,
|
||||
false, // kHasBiasGrad
|
||||
false, // kStoreLSE
|
||||
false, // kHasDropout
|
||||
BlockAttentionQuantScaleEnum::NO_SCALE,
|
||||
-1, // kBlockPerCu
|
||||
false, // kSkipMinSeqlenQ
|
||||
false>; // kHasSink
|
||||
|
||||
using FmhaMask = std::conditional_t<kIsCausal,
|
||||
SimplifiedGenericAttentionMask<true>,
|
||||
SimplifiedGenericAttentionMask<false>>;
|
||||
|
||||
static constexpr bool kUseTrLoad = (kPipelineTag == FmhaPipelineTag::QR_ASYNC_TRLOAD);
|
||||
|
||||
using PipelineProblem =
|
||||
BlockFmhaPipelineProblem<DataType_, // Q type
|
||||
DataType_, // K type
|
||||
DataType_, // V type
|
||||
float, // Sacc type
|
||||
float, // SMPLCompute type
|
||||
DataType_, // Bias type
|
||||
unsigned short, // RandVal type (unused)
|
||||
float, // LSE type (unused)
|
||||
DataType_, // P type
|
||||
float, // Oacc type
|
||||
DataType_, // O type
|
||||
FmhaShape,
|
||||
false, // mode (false = batch mode)
|
||||
ComposedAttention<false, CK_TILE_FMHA_FWD_FAST_EXP2>,
|
||||
FmhaMask,
|
||||
kUseTrLoad,
|
||||
FmhaTraits>;
|
||||
|
||||
using Pipeline =
|
||||
std::conditional_t<kPipelineTag == FmhaPipelineTag::QR_ASYNC_TRLOAD,
|
||||
BlockFmhaPipelineQRKSVSAsyncTrload<PipelineProblem>,
|
||||
std::conditional_t<kPipelineTag == FmhaPipelineTag::QR_ASYNC,
|
||||
BlockFmhaPipelineQRKSVSAsync<PipelineProblem>,
|
||||
BlockFmhaPipelineQRKSVS<PipelineProblem>>>;
|
||||
|
||||
using Epilogue = Default2DEpilogue<Default2DEpilogueProblem<float, DataType_, kPadM, kPadO>>;
|
||||
|
||||
using Kernel = FmhaFwdKernel<Pipeline, Epilogue>;
|
||||
|
||||
// Innermost dimension is always contiguous (stride=1):
|
||||
//
|
||||
// K is stored as [batch, nhead, N, K] (not transposed).
|
||||
// The kernel internally handles the transpose for Q @ K^T.
|
||||
//
|
||||
// Q: [batch, nhead, M, K]
|
||||
// K: [batch, nhead, N, K]
|
||||
// V: [batch, nhead, N, O] (rowmajor) or [batch, nhead, O, N] (colmajor)
|
||||
// O: [batch, nhead, M, O]
|
||||
// Bias: [batch, nhead, M, N]
|
||||
struct Descriptor
|
||||
{
|
||||
index_t batch, nhead, M, K;
|
||||
index_t q_stride_batch, q_stride_nhead, q_stride_m;
|
||||
|
||||
index_t N;
|
||||
index_t k_stride_batch, k_stride_nhead, k_stride_n;
|
||||
|
||||
index_t O;
|
||||
index_t v_stride_batch, v_stride_nhead, v_stride_n;
|
||||
|
||||
index_t o_stride_batch, o_stride_nhead, o_stride_m;
|
||||
|
||||
index_t bias_stride_batch, bias_stride_nhead, bias_stride_m;
|
||||
|
||||
// Only reflects compile time arch availability,
|
||||
// does not perform runtime descriptor validation.
|
||||
CK_TILE_HOST_DEVICE constexpr bool IsValid() const { return Kernel::kIsAvailable; }
|
||||
};
|
||||
|
||||
// Each tensor is specified as (batch, nhead, dim0, dim1) and (stride0, stride1, stride2)
|
||||
// Innermost stride is always 1 and not passed.
|
||||
template <typename QDims,
|
||||
typename QStrides,
|
||||
typename KDims,
|
||||
typename KStrides,
|
||||
typename VDims,
|
||||
typename VStrides,
|
||||
typename ODims,
|
||||
typename OStrides,
|
||||
typename BiasDims,
|
||||
typename BiasStrides>
|
||||
CK_TILE_HOST_DEVICE static constexpr auto make_descriptor(QDims q_dims,
|
||||
QStrides q_strides,
|
||||
KDims k_dims,
|
||||
KStrides k_strides,
|
||||
VDims v_dims,
|
||||
VStrides v_strides,
|
||||
ODims o_dims,
|
||||
OStrides o_strides,
|
||||
BiasDims bias_dims,
|
||||
BiasStrides bias_strides)
|
||||
{
|
||||
return Descriptor{q_dims[number<0>{}],
|
||||
q_dims[number<1>{}],
|
||||
q_dims[number<2>{}],
|
||||
q_dims[number<3>{}],
|
||||
q_strides[number<0>{}],
|
||||
q_strides[number<1>{}],
|
||||
q_strides[number<2>{}],
|
||||
//
|
||||
k_dims[number<2>{}],
|
||||
k_strides[number<0>{}],
|
||||
k_strides[number<1>{}],
|
||||
k_strides[number<2>{}],
|
||||
//
|
||||
v_dims[number<3>{}],
|
||||
v_strides[number<0>{}],
|
||||
v_strides[number<1>{}],
|
||||
v_strides[number<2>{}],
|
||||
//
|
||||
o_strides[number<0>{}],
|
||||
o_strides[number<1>{}],
|
||||
o_strides[number<2>{}],
|
||||
//
|
||||
bias_strides[number<0>{}],
|
||||
bias_strides[number<1>{}],
|
||||
bias_strides[number<2>{}]};
|
||||
}
|
||||
|
||||
CK_TILE_DEVICE static void Run(const Descriptor& desc,
|
||||
float scale_s,
|
||||
const DataType_* q_ptr,
|
||||
const DataType_* k_ptr,
|
||||
const DataType_* v_ptr,
|
||||
const DataType_* bias_ptr,
|
||||
DataType_* o_ptr)
|
||||
{
|
||||
using Kargs = typename Kernel::Kargs;
|
||||
Kargs kargs{};
|
||||
|
||||
kargs.q_ptr = q_ptr;
|
||||
kargs.k_ptr = k_ptr;
|
||||
kargs.v_ptr = v_ptr;
|
||||
kargs.o_ptr = o_ptr;
|
||||
kargs.sink_ptr = nullptr;
|
||||
|
||||
kargs.seqlen_q = desc.M;
|
||||
kargs.seqlen_k = desc.N;
|
||||
kargs.hdim_q = desc.K;
|
||||
kargs.hdim_v = desc.O;
|
||||
|
||||
kargs.num_head_q = desc.nhead;
|
||||
kargs.nhead_ratio_qk = 1; // nhead_q == nhead_k
|
||||
|
||||
kargs.scale_s = scale_s;
|
||||
|
||||
kargs.stride_q = desc.q_stride_m;
|
||||
kargs.stride_k = desc.k_stride_n;
|
||||
kargs.stride_v = desc.v_stride_n;
|
||||
kargs.stride_o = desc.o_stride_m;
|
||||
|
||||
kargs.nhead_stride_q = desc.q_stride_nhead;
|
||||
kargs.nhead_stride_k = desc.k_stride_nhead;
|
||||
kargs.nhead_stride_v = desc.v_stride_nhead;
|
||||
kargs.nhead_stride_o = desc.o_stride_nhead;
|
||||
|
||||
if constexpr(kHasBias)
|
||||
{
|
||||
kargs.bias_ptr = bias_ptr;
|
||||
kargs.stride_bias = desc.bias_stride_m;
|
||||
kargs.nhead_stride_bias = desc.bias_stride_nhead;
|
||||
kargs.batch_stride_bias = desc.bias_stride_batch;
|
||||
}
|
||||
|
||||
if constexpr(kIsCausal)
|
||||
{
|
||||
kargs.window_size_left = -1;
|
||||
kargs.window_size_right = 0;
|
||||
kargs.sink_size = 0;
|
||||
kargs.mask_type = GenericAttentionMaskEnum::MASK_FROM_BOTTOM_RIGHT;
|
||||
}
|
||||
|
||||
kargs.batch_stride_q = desc.q_stride_batch;
|
||||
kargs.batch_stride_k = desc.k_stride_batch;
|
||||
kargs.batch_stride_v = desc.v_stride_batch;
|
||||
kargs.batch_stride_o = desc.o_stride_batch;
|
||||
|
||||
kargs.cu_seqlen_q_ptr = nullptr;
|
||||
kargs.cu_seqlen_k_ptr = nullptr;
|
||||
|
||||
Kernel{}(kargs);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -1,92 +0,0 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include "ck/host/types.hpp"
|
||||
#include "ck/host/device_fmha_fwd/problem.hpp"
|
||||
|
||||
namespace ck {
|
||||
namespace host {
|
||||
namespace device_fmha_fwd {
|
||||
|
||||
// Derived from fmha_fwd.py FmhaFwdTileSize.
|
||||
struct TileConfig
|
||||
{
|
||||
// Block tile
|
||||
std::size_t bm0;
|
||||
std::size_t bn0;
|
||||
std::size_t bk0;
|
||||
std::size_t bn1;
|
||||
std::size_t bk1;
|
||||
std::size_t bk0max;
|
||||
|
||||
// Gemm0 block warps
|
||||
std::size_t rm0;
|
||||
std::size_t rn0;
|
||||
std::size_t rk0;
|
||||
|
||||
// Gemm1 block warps
|
||||
std::size_t rm1;
|
||||
std::size_t rn1;
|
||||
std::size_t rk1;
|
||||
|
||||
// Gemm0 warp tile
|
||||
std::size_t wm0;
|
||||
std::size_t wn0;
|
||||
std::size_t wk0;
|
||||
|
||||
// Gemm1 warp tile
|
||||
std::size_t wm1;
|
||||
std::size_t wn1;
|
||||
std::size_t wk1;
|
||||
};
|
||||
|
||||
struct PipelineConfig
|
||||
{
|
||||
std::string name;
|
||||
bool pad_m;
|
||||
bool pad_n;
|
||||
bool pad_k;
|
||||
bool pad_o;
|
||||
};
|
||||
|
||||
struct Operation
|
||||
{
|
||||
TileConfig tile = {};
|
||||
|
||||
std::string pipeline = "qr_async";
|
||||
|
||||
bool is_causal = false;
|
||||
bool is_v_rowmajor = true;
|
||||
bool has_bias = false;
|
||||
DataType dtype = DataType::Half;
|
||||
|
||||
bool pad_m = true; // pad seqlen_q
|
||||
bool pad_n = true; // pad seqlen_k
|
||||
bool pad_k = true; // pad hdim_q
|
||||
bool pad_o = true; // pad hdim_v
|
||||
|
||||
static std::vector<Operation> CreateOperations(const Problem& prob, const std::string& arch);
|
||||
|
||||
Solution ToSolution() const;
|
||||
};
|
||||
|
||||
struct HdimBucketResult
|
||||
{
|
||||
std::size_t bucket_hdim = 0;
|
||||
std::size_t bucket_hdim_v = 0;
|
||||
std::vector<TileConfig> tiles;
|
||||
};
|
||||
|
||||
HdimBucketResult
|
||||
GetTileConfigsForHdim(const std::string& arch, DataType dtype, std::size_t K, std::size_t O);
|
||||
|
||||
bool IsSupportedArch(const std::string& arch);
|
||||
|
||||
} // namespace device_fmha_fwd
|
||||
} // namespace host
|
||||
} // namespace ck
|
||||
@@ -1,38 +0,0 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include "ck/host/types.hpp"
|
||||
|
||||
namespace ck {
|
||||
namespace host {
|
||||
namespace device_fmha_fwd {
|
||||
|
||||
struct Problem
|
||||
{
|
||||
std::size_t M = 0; // seqlen_q
|
||||
std::size_t N = 0; // seqlen_k
|
||||
std::size_t K = 0; // hdim_q
|
||||
std::size_t O = 0; // hdim_v
|
||||
|
||||
std::size_t batch = 0;
|
||||
std::size_t nhead = 0; // nhead_q == nhead_k
|
||||
|
||||
DataType dtype = DataType::Half;
|
||||
|
||||
bool is_v_rowmajor = true; // true=[N,O], false=[O,N]
|
||||
bool is_causal = false;
|
||||
bool has_bias = false;
|
||||
|
||||
std::string GetIncludeHeader() const;
|
||||
|
||||
std::vector<Solution> GetSolutions(const std::string& arch) const;
|
||||
};
|
||||
|
||||
} // namespace device_fmha_fwd
|
||||
} // namespace host
|
||||
} // namespace ck
|
||||
@@ -13,7 +13,5 @@ namespace host {
|
||||
|
||||
std::unordered_map<std::string_view, std::string_view> GetHeaders();
|
||||
|
||||
std::unordered_map<std::string, std::string> GetTileHeaders();
|
||||
|
||||
} // namespace host
|
||||
} // namespace ck
|
||||
|
||||
Reference in New Issue
Block a user