ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 16:51:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

hack for cap logits

Browse Source
This commit is contained in:
Bernard
2025-04-18 08:38:25 +00:00
parent b20173a494
commit 838ff7034d
2 changed files with 56 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
include/ck_tile/core/numeric/math.hpp
View File

@@ -1380,6 +1380,43 @@ CK_TILE_DEVICE double exp<double>(double x)
return exp(x);
};
template <typename T>
CK_TILE_DEVICE T tanh_fast(T x)
{
return type_convert<T>((exp<T>(2.0 * type_convert<float>(x)) - 1.0) / (exp<T>(2.0 * type_convert<float>(x)) + 1.0));
};
template <>
CK_TILE_DEVICE float tanh_fast<float>(float x)
{
// float a = __builtin_amdgcn_sinh(x);
// float b = __builtin_amdgcn_cosh(x);
// float e = a * __builtin_amdgcn_rcpf(b);
// return e;
float a = 2.0f * log2e_v<float> * x;
a = __builtin_amdgcn_exp2f(a);
a = __builtin_amdgcn_rcpf(a + 1.0f);
a = 2 * a;
a = 1 - a;
return a;
// float e, r, s, t, d;
// float a = x;
// s = abs(a);
// t = -log2e_v<float> * 2.0f * s;
// e = __builtin_amdgcn_exp2f(t);
// d = e + 1.0f;
// r = __builtin_amdgcn_rcpf(d);
// r = e * (-r) + r;
// if (s < 4.997253418e-3f) r = a;
// union fipnr {float f; unsigned int i;};
// fipnr r_; r_.f = r;
// fipnr a_; a_.f = a;
// { r_.i = (r_.i|(a_.i&0x80000000)); r = r_.f; }
// return r;
};
template <typename T>
CK_TILE_DEVICE T log(T x)
{

27
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp
View File

@@ -192,6 +192,9 @@ struct BlockFmhaPipelineQRKSVSAsync
constexpr auto LdsSeq = Policy::template GetLdsBufferSequence<Problem>();
const float logits_cap = 30.0f;
// const float logits_cap_scale = scale_s * rcp<float>(logits_cap * log2e_v<float>);
// K tile in LDS
auto k_lds_ptr = reinterpret_cast<KDataType*>(smem_ptr);
auto k_lds_store = generate_tuple(
@@ -435,9 +438,15 @@ struct BlockFmhaPipelineQRKSVSAsync
else
{
s_acc = tile_elementwise_in(s_acc_element_func, s_acc);
#if !CK_TILE_FMHA_FWD_FAST_EXP2
// #if !CK_TILE_FMHA_FWD_FAST_EXP2
tile_elementwise_inout([&scale_s](auto& x) { x = x * scale_s; }, s_acc);
#endif
tile_elementwise_inout(
[&logits_cap](auto& x) {
x = log2e_v<float> * logits_cap * tanh_fast<float>(x * rcp<float>(logits_cap));
},
s_acc
);
// #endif
}
move_tile_window(bias_dram_window, {0, kN0});
if constexpr(kPadSeqLenK || FmhaMask::IsMasking)
@@ -530,9 +539,9 @@ struct BlockFmhaPipelineQRKSVSAsync
constexpr auto p_spans = decltype(p_compute)::get_distributed_spans();
sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) {
constexpr auto i_idx = make_tuple(idx0);
#if CK_TILE_FMHA_FWD_FAST_EXP2
auto row_max = scale_s * get_validated_m(m[i_idx]);
#endif
// #if CK_TILE_FMHA_FWD_FAST_EXP2
// auto row_max = scale_s * get_validated_m(m[i_idx]);
// #endif
sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) {
constexpr auto i_j_idx = make_tuple(idx0, idx1);
#if CK_TILE_FMHA_FWD_FAST_EXP2
@@ -543,7 +552,8 @@ struct BlockFmhaPipelineQRKSVSAsync
}
else
{
p_compute(i_j_idx) = exp2(scale_s * s[i_j_idx] - row_max);
// p_compute(i_j_idx) = exp2(scale_s * s[i_j_idx] - row_max);
p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m[i_idx]));
}
#else
p_compute(i_j_idx) = exp(s[i_j_idx] - get_validated_m(m[i_idx]));
@@ -568,8 +578,9 @@ struct BlockFmhaPipelineQRKSVSAsync
}
else
{
auto row_max = scale_s * get_validated_m(m[i_idx]);
return exp2(scale_s * m_old[i_idx] - row_max);
// auto row_max = scale_s * get_validated_m(m[i_idx]);
// return exp2(scale_s * m_old[i_idx] - row_max);
return exp2(m_old[i_idx] - get_validated_m(m[i_idx]));
}
}();
#else