mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-24 19:43:35 +00:00
509cb54d6f
### Description Implement softcap for gqa. ### Motivation and Context Fixes certain models like Gemma-2 which need softcap to work so they don't output nan's.
100 lines
3.5 KiB
Diff
100 lines
3.5 KiB
Diff
diff --git a/examples/41_fused_multi_head_attention/kernel_forward.h b/examples/41_fused_multi_head_attention/kernel_forward.h
|
|
index 4c80f549..5ad610c8 100644
|
|
--- a/examples/41_fused_multi_head_attention/kernel_forward.h
|
|
+++ b/examples/41_fused_multi_head_attention/kernel_forward.h
|
|
@@ -189,6 +189,7 @@ struct AttentionKernel {
|
|
|
|
// Scale
|
|
accum_t scale = 0.0;
|
|
+ accum_t softcap = 0.0;
|
|
|
|
// Dimensions/strides
|
|
int32_t head_dim = 0;
|
|
@@ -221,6 +222,8 @@ struct AttentionKernel {
|
|
int32_t num_batches = 0;
|
|
int32_t num_heads = 0;
|
|
|
|
+ bool use_smooth_softmax = false;
|
|
+
|
|
// dropout
|
|
bool use_dropout = false;
|
|
unsigned long long dropout_batch_head_rng_offset = 0;
|
|
@@ -818,6 +821,15 @@ struct AttentionKernel {
|
|
accum =
|
|
cutlass::multiplies<typename MM0::Mma::FragmentC>()(p.scale, accum);
|
|
}
|
|
+
|
|
+ // apply softcap if applicable
|
|
+ if (p.softcap > 0.0) {
|
|
+ accum = cutlass::multiplies<typename MM0::Mma::FragmentC>()(1.0 / p.softcap, accum);
|
|
+ for (int i = 0; i < accum.size(); ++i) {
|
|
+ accum[i] = cutlass::fast_tanh(accum[i]);
|
|
+ }
|
|
+ accum = cutlass::multiplies<typename MM0::Mma::FragmentC>()(p.softcap, accum);
|
|
+ }
|
|
|
|
// apply attention bias if applicable
|
|
if (kSupportsBias && p.attn_bias_ptr != nullptr) {
|
|
@@ -897,7 +909,8 @@ struct AttentionKernel {
|
|
p.num_keys - iter_key_start,
|
|
iter_key_start == 0,
|
|
iteratorC_tile_offset,
|
|
- kSupportsBias ? 1.0f : p.scale);
|
|
+ kSupportsBias ? 1.0f : p.scale,
|
|
+ p.use_smooth_softmax);
|
|
|
|
// Output results to shared-memory
|
|
int warp_idx_mn_0 = my_warp_id %
|
|
@@ -1166,7 +1179,8 @@ struct AttentionKernel {
|
|
int max_col,
|
|
bool is_first,
|
|
typename WarpIteratorC::TensorCoord const& tile_offset,
|
|
- float scaling) {
|
|
+ float scaling,
|
|
+ bool use_smooth_softmax) {
|
|
/* Iterates on the accumulator and corresponding position on result matrix
|
|
|
|
(1) Update `mi[r]` to the max value of the row `r`
|
|
@@ -1257,7 +1271,7 @@ struct AttentionKernel {
|
|
accum_t mi_row, total_row;
|
|
LambdaIterator::iterateRows(
|
|
lane_offset,
|
|
- [&](int accum_m) { mi_row = mi[accum_m]; },
|
|
+ [&](int accum_m) { mi_row = mi[accum_m];},
|
|
[&](int accum_m, int accum_n, int idx) {
|
|
frag[idx] =
|
|
(accum_n < max_col) ? exp2f(frag[idx] - mi_row) : accum_t(0.0);
|
|
@@ -1294,7 +1308,7 @@ struct AttentionKernel {
|
|
for (int i = 0; i < MM0::MmaCore::WarpCount::kN; ++i) {
|
|
total_row += addition_storage[id + kQueriesPerBlock * i];
|
|
}
|
|
- s_prime[id] = total_row;
|
|
+ s_prime[id] = (use_smooth_softmax && (max_col <= kKeysPerBlock)) ? total_row + exp2f(-mi[id]) : total_row;
|
|
}
|
|
}
|
|
|
|
diff --git a/include/cutlass/functional.h b/include/cutlass/functional.h
|
|
index 964d2ff3..676ba768 100644
|
|
--- a/include/cutlass/functional.h
|
|
+++ b/include/cutlass/functional.h
|
|
@@ -39,6 +39,7 @@
|
|
#include "cutlass/numeric_types.h"
|
|
|
|
#include <cuda_runtime.h>
|
|
+#include <cuda_fp16.h>
|
|
|
|
#if defined(CUTLASS_ARCH_WMMA_ENABLED)
|
|
#include <mma.h>
|
|
@@ -230,8 +231,12 @@ struct inverse_square_root<half_t> {
|
|
CUTLASS_HOST_DEVICE
|
|
half_t operator()(half_t const &lhs) const {
|
|
#if defined(__CUDA_ARCH__)
|
|
+#if (__CUDA_ARCH__ >= 530)
|
|
auto result = hrsqrt(reinterpret_cast<__half const &>(lhs));
|
|
return reinterpret_cast<half_t const &>(result);
|
|
+#else
|
|
+ return half_t::convert((rsqrtf(half_t::convert(lhs))));
|
|
+#endif
|
|
#else
|
|
return half_t(1.f / std::sqrt(half_t::convert(lhs)));
|
|
#endif
|