mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
refactor some code (#15566)
### Description <!-- Describe your changes. --> 1. moved onnxruntime/contrib_ops/cuda/decoder to onnxruntime/contrib_ops/cuda/bert 2. create utils.cuh under /bert for shared implementations in decoder_masked_multihead_attention_impl_utils.h and rotary_embedding_util.h 3. refactored relative_attn_bias_impl.cu by reusing the template specializations in utils.cuh ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> --------- Co-authored-by: Ubuntu <wy@v100-2.0cdb2e52twzevn1i4fi45bylyg.jx.internal.cloudapp.net>
This commit is contained in:
Ye Wang
GitHub
parent
b5a1941835
commit
633dec0b17
15 changed files with 485 additions and 1006 deletions
|
|
@ -9,6 +9,11 @@ set(contrib_ops_excluded_files
|
|||
"bert/attention.h"
|
||||
"bert/attention_impl.cu"
|
||||
"bert/attention_softmax.h"
|
||||
"bert/decoder_masked_multihead_attention.h"
|
||||
"bert/decoder_masked_multihead_attention.cc"
|
||||
"bert/decoder_masked_self_attention.h"
|
||||
"bert/decoder_masked_self_attention.cc"
|
||||
"bert/fastertransformer_decoder_attention/*"
|
||||
"bert/multihead_attention.cc"
|
||||
"bert/multihead_attention.h"
|
||||
"bert/fast_gelu_impl.cu"
|
||||
|
|
@ -92,16 +97,6 @@ set(contrib_ops_excluded_files
|
|||
"cuda_contrib_kernels.h"
|
||||
"inverse.cc"
|
||||
"fused_conv.cc"
|
||||
"decoder/decoder_masked_multihead_attention.h"
|
||||
"decoder/decoder_masked_multihead_attention.cc"
|
||||
"decoder/decoder_masked_self_attention.h"
|
||||
"decoder/decoder_masked_self_attention.cc"
|
||||
"decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h"
|
||||
"decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention.h"
|
||||
"decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu"
|
||||
"decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_32.cu"
|
||||
"decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_64.cu"
|
||||
"decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_128.cu"
|
||||
)
|
||||
|
||||
if (NOT onnxruntime_ENABLE_ATEN)
|
||||
|
|
|
|||
|
|
@ -276,9 +276,9 @@ __global__ void AddBiasTransposeQKV(int M, const T* input, const T* biases, T* o
|
|||
v_bias = *reinterpret_cast<const Vec_t*>(&biases[2 * NH + n * H + head_idx]);
|
||||
}
|
||||
|
||||
q = add(q, q_bias);
|
||||
k = add(k, k_bias);
|
||||
v = add(v, v_bias);
|
||||
q = add_vec(q, q_bias);
|
||||
k = add_vec(k, k_bias);
|
||||
v = add_vec(v, v_bias);
|
||||
|
||||
const bool do_rotary = !is_masked && vec_size * tidx < rotary_embedding_dim;
|
||||
|
||||
|
|
|
|||
|
|
@ -5,8 +5,8 @@
|
|||
#include "core/providers/cuda/shared_inc/fpgeneric.h"
|
||||
#include "core/platform/env_var_utils.h"
|
||||
#include "contrib_ops/cpu/bert/multihead_attention_helper.h"
|
||||
#include "contrib_ops/cuda/decoder/decoder_masked_multihead_attention.h"
|
||||
#include "contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h"
|
||||
#include "contrib_ops/cuda/bert/decoder_masked_multihead_attention.h"
|
||||
#include "contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h"
|
||||
|
||||
using namespace onnxruntime::cuda;
|
||||
using namespace ::onnxruntime::common;
|
||||
|
|
@ -4,8 +4,8 @@
|
|||
#include "core/providers/cuda/cuda_common.h"
|
||||
#include "core/providers/cuda/shared_inc/fpgeneric.h"
|
||||
#include "core/platform/env_var_utils.h"
|
||||
#include "contrib_ops/cuda/decoder/decoder_masked_self_attention.h"
|
||||
#include "contrib_ops/cuda/decoder/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h"
|
||||
#include "contrib_ops/cuda/bert/decoder_masked_self_attention.h"
|
||||
#include "contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.h"
|
||||
|
||||
using namespace onnxruntime::cuda;
|
||||
using namespace ::onnxruntime::common;
|
||||
|
|
@ -0,0 +1,173 @@
|
|||
/*
|
||||
* The implementation of this file is based on code provided by https://github.com/NVIDIA/FasterTransformer
|
||||
*
|
||||
* Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
// Modifications Copyright (c) Microsoft.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
// Modifications:
|
||||
// (1) Minor routine name changes for integration into the ORT code-base
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "contrib_ops/cuda/bert/utils.cuh"
|
||||
|
||||
using namespace onnxruntime::cuda;
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
namespace cuda {
|
||||
namespace decoder_masked_self_attention_details {
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Block_sum
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <int WARPS_PER_BLOCK, int WARP_SIZE = 32>
|
||||
inline __device__ float block_sum(float* red_smem, float sum) {
|
||||
// Decompose the thread index into warp / lane.
|
||||
int warp = threadIdx.x / WARP_SIZE;
|
||||
int lane = threadIdx.x % WARP_SIZE;
|
||||
|
||||
// Compute the sum per warp.
|
||||
#pragma unroll
|
||||
for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
|
||||
sum += __shfl_xor_sync(uint32_t(-1), sum, mask);
|
||||
}
|
||||
|
||||
// Warp leaders store the data to shared memory.
|
||||
if (lane == 0) {
|
||||
red_smem[warp] = sum;
|
||||
}
|
||||
|
||||
// Make sure the data is in shared memory.
|
||||
__syncthreads();
|
||||
|
||||
// The warps compute the final sums.
|
||||
if (lane < WARPS_PER_BLOCK) {
|
||||
sum = red_smem[lane];
|
||||
}
|
||||
|
||||
// Parallel reduction inside the warp.
|
||||
#pragma unroll
|
||||
for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) {
|
||||
sum += __shfl_xor_sync(uint32_t(-1), sum, mask);
|
||||
}
|
||||
|
||||
// Broadcast to other threads.
|
||||
return __shfl_sync(uint32_t(-1), sum, 0);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Shfl_Mask
|
||||
//------------------------------------------------------------
|
||||
|
||||
inline __device__ constexpr uint32_t shfl_mask(int threads) {
|
||||
return threads == 32 ? uint32_t(-1) : (1u << threads) - 1u;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Dot
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <typename T>
|
||||
inline __device__ float dot(T a, T b) {
|
||||
return sum(mul<T, T, T>(a, b));
|
||||
}
|
||||
|
||||
template <typename A, typename T>
|
||||
inline __device__ float dot(T a, T b) {
|
||||
return sum(mul<A, T, T>(a, b));
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Qk_Dot
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <int THREADS_PER_KEY, typename K_vec, int N>
|
||||
inline __device__ float qk_dot_(const K_vec (&q)[N], const K_vec (&k)[N]) {
|
||||
using K_vec_acum = K_vec;
|
||||
|
||||
// Compute the parallel products for Q*K^T (treat vector lanes separately).
|
||||
K_vec_acum qk_vec = mul<K_vec_acum, K_vec, K_vec>(q[0], k[0]);
|
||||
#pragma unroll
|
||||
for (int ii = 1; ii < N; ++ii) {
|
||||
qk_vec = onnxruntime::cuda::fma(q[ii], k[ii], qk_vec);
|
||||
}
|
||||
|
||||
// Finalize the reduction across lanes.
|
||||
float qk = sum(qk_vec);
|
||||
#pragma unroll
|
||||
for (int mask = THREADS_PER_KEY / 2; mask >= 1; mask /= 2) {
|
||||
qk += __shfl_xor_sync(uint32_t(-1), qk, mask);
|
||||
}
|
||||
return qk;
|
||||
}
|
||||
|
||||
template <typename T, int THREADS_PER_KEY>
|
||||
struct Qk_dot {
|
||||
template <typename K_vec, int N>
|
||||
static inline __device__ float dot(const K_vec (&q)[N], const K_vec (&k)[N]) {
|
||||
return qk_dot_<THREADS_PER_KEY>(q, k);
|
||||
}
|
||||
};
|
||||
|
||||
//------------------------------------------------------------
|
||||
// ThreadsPerValue
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <typename T, int head_size>
|
||||
struct ThreadsPerValue {
|
||||
static const int value = head_size * sizeof(T) / 16;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------
|
||||
// CalcDynamicBlockMemory
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <typename T>
|
||||
inline size_t CalcDynamicBlockMemory(const DecoderMaskedMultiHeadAttentionParams& params,
|
||||
int threads_per_value, int threads_per_block) {
|
||||
// The amount of shared memory needed to store the Q*K^T values in float.
|
||||
|
||||
const int total_sequence_length = params.total_sequence_length;
|
||||
size_t qk_sz = (((total_sequence_length + 3) / 4) * 16);
|
||||
|
||||
// The extra memory needed if we are not using floats for the final logits.
|
||||
size_t logits_sz = 0;
|
||||
|
||||
if (sizeof(T) != 4) {
|
||||
logits_sz = (((total_sequence_length + 3) / 4) * 4 * sizeof(T));
|
||||
}
|
||||
|
||||
// The total size needed during softmax.
|
||||
size_t softmax_sz = qk_sz + logits_sz;
|
||||
|
||||
// The number of partial rows to reduce in the final reduction.
|
||||
int rows_per_red = threads_per_block / threads_per_value;
|
||||
|
||||
// The amount of storage needed to finalize the outputs.
|
||||
size_t red_sz = rows_per_red * params.head_size * sizeof(T) / 2;
|
||||
|
||||
// The max.
|
||||
return std::max(softmax_sz, red_sz);
|
||||
}
|
||||
|
||||
} // namespace decoder_masked_self_attention_details
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -21,7 +21,7 @@ Licensed under the MIT License.
|
|||
#include <cuda_fp16.h>
|
||||
#include "core/providers/cuda/cu_inc/common.cuh"
|
||||
#include "contrib_ops/cuda/bert/relative_attn_bias_impl.h"
|
||||
#include "contrib_ops/cuda/bert/rotary_embedding_util.h"
|
||||
#include "contrib_ops/cuda/bert/utils.cuh"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
|
|
@ -149,29 +149,11 @@ template Status LaunchRelPosAttnBiasKernel<half>(cudaStream_t stream,
|
|||
const int max_threads_per_block);
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename T, size_t size>
|
||||
struct TypeMapper;
|
||||
|
||||
template <>
|
||||
struct TypeMapper<float, 2> {
|
||||
using type = float2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct TypeMapper<float, 4> {
|
||||
using type = float4;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct TypeMapper<half, 2> {
|
||||
using type = half2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct TypeMapper<half, 4> {
|
||||
using type = Half4;
|
||||
};
|
||||
struct TypeMapper : public V_vec_m_<T, size> {};
|
||||
|
||||
// The following operator overriding is not common so we put it in anonymous namespace
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ > 530
|
||||
inline __device__ half2 operator*(const float a, const half2 b) {
|
||||
return __hmul2_rn(__float2half2_rn(a), b);
|
||||
|
|
@ -292,13 +274,13 @@ Status LaunchGatedRelativePositionBiasKernel(
|
|||
dim3 grid(seq_len, num_heads, batch_size);
|
||||
|
||||
if (seq_len % 4 == 0) {
|
||||
using vec_type = typename TypeMapper<T, 4>::type;
|
||||
using vec_type = typename TypeMapper<T, 4>::Type;
|
||||
GatedRelativePositionBiasKernelSmallD<<<grid, block, sizeof(float), stream>>>(
|
||||
reinterpret_cast<vec_type*>(output),
|
||||
reinterpret_cast<const vec_type*>(rel_pos),
|
||||
qw, bias, eco_a, D, ldqw, equiv_seq_len);
|
||||
} else if (seq_len & 1 == 0) {
|
||||
using vec_type = typename TypeMapper<T, 2>::type;
|
||||
using vec_type = typename TypeMapper<T, 2>::Type;
|
||||
GatedRelativePositionBiasKernelSmallD<<<grid, block, sizeof(float), stream>>>(
|
||||
reinterpret_cast<vec_type*>(output),
|
||||
reinterpret_cast<const vec_type*>(rel_pos),
|
||||
|
|
|
|||
|
|
@ -20,821 +20,14 @@
|
|||
|
||||
#include "core/providers/cuda/cuda_common.h"
|
||||
#include "core/providers/cuda/cu_inc/common.cuh"
|
||||
#include "contrib_ops/cuda/bert/utils.cuh"
|
||||
|
||||
using namespace onnxruntime::cuda;
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace cuda {
|
||||
|
||||
struct __align__(8) Half4 {
|
||||
half2 x;
|
||||
half2 y;
|
||||
};
|
||||
|
||||
__device__ __forceinline__ Half4 operator+(const Half4& a, const Half4& b) {
|
||||
Half4 r;
|
||||
r.x = a.x + b.x;
|
||||
r.y = a.y + b.y;
|
||||
return r;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float2 operator+(const float2& a, const float2& b) {
|
||||
return make_float2(a.x + b.x, a.y + b.y);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float4 operator+(const float4& a, const float4& b) {
|
||||
return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Float8_ {
|
||||
float2 x;
|
||||
float2 y;
|
||||
float2 z;
|
||||
float2 w;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct Float4_ {
|
||||
float2 x;
|
||||
float2 y;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
struct num_elems;
|
||||
template <>
|
||||
struct num_elems<float> {
|
||||
static constexpr int value = 1;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<float2> {
|
||||
static constexpr int value = 2;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<float4> {
|
||||
static constexpr int value = 4;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<Float4_> {
|
||||
static constexpr int value = 4;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<Float8_> {
|
||||
static constexpr int value = 8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct num_elems<uint32_t> {
|
||||
static constexpr int value = 2;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<uint2> {
|
||||
static constexpr int value = 4;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<uint4> {
|
||||
static constexpr int value = 8;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
struct Vec_t {
|
||||
static constexpr int size = 0;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<float> {
|
||||
using Type = float2;
|
||||
static constexpr int size = 2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<float2> {
|
||||
using Type = float4;
|
||||
static constexpr int size = 4;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<float4> {
|
||||
using Type = Float8_;
|
||||
static constexpr int size = 8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<half> {
|
||||
using Type = uint32_t;
|
||||
static constexpr int size = 2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<half2> {
|
||||
using Type = uint2;
|
||||
static constexpr int size = 4;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<Half4> {
|
||||
using Type = uint4;
|
||||
static constexpr int size = 8;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float add(float a, float b) {
|
||||
return a + b;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float2 add(float2 a, float2 b) {
|
||||
float2 c;
|
||||
c.x = add(a.x, b.x);
|
||||
c.y = add(a.y, b.y);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float4 add(float4 a, float4 b) {
|
||||
float4 c;
|
||||
c.x = add(a.x, b.x);
|
||||
c.y = add(a.y, b.y);
|
||||
c.z = add(a.z, b.z);
|
||||
c.w = add(a.w, b.w);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float8_ add(Float8_ a, Float8_ b) {
|
||||
Float8_ c;
|
||||
c.x = add(a.x, b.x);
|
||||
c.y = add(a.y, b.y);
|
||||
c.z = add(a.z, b.z);
|
||||
c.w = add(a.w, b.w);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint16_t add(uint16_t a, uint16_t b) {
|
||||
uint16_t c;
|
||||
asm volatile("add.f16 %0, %1, %2;\n"
|
||||
: "=h"(c)
|
||||
: "h"(a), "h"(b));
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint32_t add(uint32_t a, uint32_t b) {
|
||||
uint32_t c;
|
||||
asm volatile("add.f16x2 %0, %1, %2;\n"
|
||||
: "=r"(c)
|
||||
: "r"(a), "r"(b));
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint2 add(uint2 a, uint2 b) {
|
||||
uint2 c;
|
||||
c.x = add(a.x, b.x);
|
||||
c.y = add(a.y, b.y);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint4 add(uint4 a, uint4 b) {
|
||||
uint4 c;
|
||||
c.x = add(a.x, b.x);
|
||||
c.y = add(a.y, b.y);
|
||||
c.z = add(a.z, b.z);
|
||||
c.w = add(a.w, b.w);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint16_t float_to_half(float f) {
|
||||
union {
|
||||
uint32_t u32;
|
||||
uint16_t u16[2];
|
||||
} tmp;
|
||||
#if 0 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 // Is it better?
|
||||
float zero = 0.f;
|
||||
asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(tmp.u32) : "f"(zero), "f"(f));
|
||||
#else
|
||||
asm volatile("cvt.rn.f16.f32 %0, %1;\n"
|
||||
: "=h"(tmp.u16[0])
|
||||
: "f"(f));
|
||||
#endif
|
||||
return tmp.u16[0];
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint32_t float2_to_half2(float2 f) {
|
||||
union {
|
||||
uint32_t u32;
|
||||
uint16_t u16[2];
|
||||
} tmp;
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n"
|
||||
: "=r"(tmp.u32)
|
||||
: "f"(f.y), "f"(f.x));
|
||||
#else
|
||||
asm volatile("cvt.rn.f16.f32 %0, %1;\n"
|
||||
: "=h"(tmp.u16[0])
|
||||
: "f"(f.x));
|
||||
asm volatile("cvt.rn.f16.f32 %0, %1;\n"
|
||||
: "=h"(tmp.u16[1])
|
||||
: "f"(f.y));
|
||||
#endif
|
||||
return tmp.u32;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float half_to_float(uint16_t h) {
|
||||
float f;
|
||||
asm volatile("cvt.f32.f16 %0, %1;\n"
|
||||
: "=f"(f)
|
||||
: "h"(h));
|
||||
return f;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float2 half2_to_float2(uint32_t v) {
|
||||
uint16_t lo, hi;
|
||||
asm volatile("mov.b32 {%0, %1}, %2;\n"
|
||||
: "=h"(lo), "=h"(hi)
|
||||
: "r"(v));
|
||||
return make_float2(half_to_float(lo), half_to_float(hi));
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float add(float a, uint16_t b) {
|
||||
return a + half_to_float(b);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float2 add(uint32_t a, float2 fb) {
|
||||
float2 fa = half2_to_float2(a);
|
||||
return add(fa, fb);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float4_ add(uint2 a, Float4_ fb) {
|
||||
Float4_ fc;
|
||||
fc.x = add(a.x, fb.x);
|
||||
fc.y = add(a.y, fb.y);
|
||||
return fc;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float8_ add(uint4 a, Float8_ fb) {
|
||||
Float8_ fc;
|
||||
fc.x = add(a.x, fb.x);
|
||||
fc.y = add(a.y, fb.y);
|
||||
fc.z = add(a.z, fb.z);
|
||||
fc.w = add(a.w, fb.w);
|
||||
return fc;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint32_t h0_h0(uint16_t a) {
|
||||
uint32_t b;
|
||||
asm volatile("mov.b32 %0, {%1, %1};"
|
||||
: "=r"(b)
|
||||
: "h"(a));
|
||||
return b;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float fma(float a, float b, float c) {
|
||||
return a * b + c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float2 fma(float2 a, float2 b, float2 c) {
|
||||
float2 d;
|
||||
d.x = fma(a.x, b.x, c.x);
|
||||
d.y = fma(a.y, b.y, c.y);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float2 fma(float a, float2 b, float2 c) {
|
||||
float2 d;
|
||||
d.x = fma(a, b.x, c.x);
|
||||
d.y = fma(a, b.y, c.y);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float4 fma(float4 a, float4 b, float4 c) {
|
||||
float4 d;
|
||||
d.x = fma(a.x, b.x, c.x);
|
||||
d.y = fma(a.y, b.y, c.y);
|
||||
d.z = fma(a.z, b.z, c.z);
|
||||
d.w = fma(a.w, b.w, c.w);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float4 fma(float a, float4 b, float4 c) {
|
||||
float4 d;
|
||||
d.x = fma(a, b.x, c.x);
|
||||
d.y = fma(a, b.y, c.y);
|
||||
d.z = fma(a, b.z, c.z);
|
||||
d.w = fma(a, b.w, c.w);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float4_ fma(float a, Float4_ b, Float4_ c) {
|
||||
Float4_ d;
|
||||
d.x = fma(a, b.x, c.x);
|
||||
d.y = fma(a, b.y, c.y);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float8_ fma(float a, Float8_ b, Float8_ c) {
|
||||
Float8_ d;
|
||||
d.x = fma(a, b.x, c.x);
|
||||
d.y = fma(a, b.y, c.y);
|
||||
d.z = fma(a, b.z, c.z);
|
||||
d.w = fma(a, b.w, c.w);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) {
|
||||
uint32_t d;
|
||||
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
|
||||
: "=r"(d)
|
||||
: "r"(a), "r"(b), "r"(c));
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint32_t fma(uint16_t a, uint32_t b, uint32_t c) {
|
||||
return fma(h0_h0(a), b, c);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint2 fma(uint2 a, uint2 b, uint2 c) {
|
||||
uint2 d;
|
||||
d.x = fma(a.x, b.x, c.x);
|
||||
d.y = fma(a.y, b.y, c.y);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint2 fma(uint16_t a, uint2 b, uint2 c) {
|
||||
uint32_t s = h0_h0(a);
|
||||
uint2 d;
|
||||
d.x = fma(s, b.x, c.x);
|
||||
d.y = fma(s, b.y, c.y);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint4 fma(uint4 a, uint4 b, uint4 c) {
|
||||
uint4 d;
|
||||
d.x = fma(a.x, b.x, c.x);
|
||||
d.y = fma(a.y, b.y, c.y);
|
||||
d.z = fma(a.z, b.z, c.z);
|
||||
d.w = fma(a.w, b.w, c.w);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ uint4 fma(uint16_t a, uint4 b, uint4 c) {
|
||||
uint32_t s = h0_h0(a);
|
||||
uint4 d;
|
||||
d.x = fma(s, b.x, c.x);
|
||||
d.y = fma(s, b.y, c.y);
|
||||
d.z = fma(s, b.z, c.z);
|
||||
d.w = fma(s, b.w, c.w);
|
||||
return d;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float fma(uint16_t a, uint16_t b, float fc) {
|
||||
float fa = half_to_float(a);
|
||||
float fb = half_to_float(b);
|
||||
return fa * fb + fc;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float2 fma(uint32_t a, uint32_t b, float2 fc) {
|
||||
float2 fa = half2_to_float2(a);
|
||||
float2 fb = half2_to_float2(b);
|
||||
return fma(fa, fb, fc);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc) {
|
||||
return fma(h0_h0(a), b, fc);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc) {
|
||||
Float4_ fd;
|
||||
fd.x = fma(a.x, b.x, fc.x);
|
||||
fd.y = fma(a.y, b.y, fc.y);
|
||||
return fd;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) {
|
||||
uint32_t s = h0_h0(a);
|
||||
Float4_ fd;
|
||||
fd.x = fma(s, b.x, fc.x);
|
||||
fd.y = fma(s, b.y, fc.y);
|
||||
return fd;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc) {
|
||||
Float8_ fd;
|
||||
fd.x = fma(a.x, b.x, fc.x);
|
||||
fd.y = fma(a.y, b.y, fc.y);
|
||||
fd.z = fma(a.z, b.z, fc.z);
|
||||
fd.w = fma(a.w, b.w, fc.w);
|
||||
return fd;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) {
|
||||
uint32_t s = h0_h0(a);
|
||||
Float8_ fd;
|
||||
fd.x = fma(s, b.x, fc.x);
|
||||
fd.y = fma(s, b.y, fc.y);
|
||||
fd.z = fma(s, b.z, fc.z);
|
||||
fd.w = fma(s, b.w, fc.w);
|
||||
return fd;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename Acc, typename A, typename B>
|
||||
inline __device__ Acc mul(A a, B b) {
|
||||
return a * b;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float mul<float, float>(float a, float b) {
|
||||
return a * b;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float2 mul(float2 a, float2 b) {
|
||||
float2 c;
|
||||
c.x = a.x * b.x;
|
||||
c.y = a.y * b.y;
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float2 mul(float a, float2 b) {
|
||||
float2 c;
|
||||
c.x = a * b.x;
|
||||
c.y = a * b.y;
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float4 mul(float4 a, float4 b) {
|
||||
float4 c;
|
||||
c.x = a.x * b.x;
|
||||
c.y = a.y * b.y;
|
||||
c.z = a.z * b.z;
|
||||
c.w = a.w * b.w;
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float4 mul(float a, float4 b) {
|
||||
float4 c;
|
||||
c.x = a * b.x;
|
||||
c.y = a * b.y;
|
||||
c.z = a * b.z;
|
||||
c.w = a * b.w;
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ Float8_ mul(float a, Float8_ b) {
|
||||
Float8_ c;
|
||||
c.x = make_float2(a * b.x.x, a * b.x.y);
|
||||
c.y = make_float2(a * b.y.x, a * b.y.y);
|
||||
c.z = make_float2(a * b.z.x, a * b.z.y);
|
||||
c.w = make_float2(a * b.w.x, a * b.w.y);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ uint16_t mul(uint16_t a, uint16_t b) {
|
||||
uint16_t c;
|
||||
asm volatile("mul.f16 %0, %1, %2;\n"
|
||||
: "=h"(c)
|
||||
: "h"(a), "h"(b));
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ uint32_t mul(uint32_t a, uint32_t b) {
|
||||
uint32_t c;
|
||||
asm volatile("mul.f16x2 %0, %1, %2;\n"
|
||||
: "=r"(c)
|
||||
: "r"(a), "r"(b));
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ uint32_t mul(uint16_t a, uint32_t b) {
|
||||
return mul<uint32_t, uint32_t, uint32_t>(h0_h0(a), b);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ uint2 mul(uint2 a, uint2 b) {
|
||||
uint2 c;
|
||||
c.x = mul<uint32_t, uint32_t, uint32_t>(a.x, b.x);
|
||||
c.y = mul<uint32_t, uint32_t, uint32_t>(a.y, b.y);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ uint2 mul(uint16_t a, uint2 b) {
|
||||
uint32_t s = h0_h0(a);
|
||||
uint2 c;
|
||||
c.x = mul<uint32_t, uint32_t, uint32_t>(s, b.x);
|
||||
c.y = mul<uint32_t, uint32_t, uint32_t>(s, b.y);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ uint4 mul(uint4 a, uint4 b) {
|
||||
uint4 c;
|
||||
c.x = mul<uint32_t, uint32_t, uint32_t>(a.x, b.x);
|
||||
c.y = mul<uint32_t, uint32_t, uint32_t>(a.y, b.y);
|
||||
c.z = mul<uint32_t, uint32_t, uint32_t>(a.z, b.z);
|
||||
c.w = mul<uint32_t, uint32_t, uint32_t>(a.w, b.w);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ uint4 mul(uint16_t a, uint4 b) {
|
||||
uint32_t s = h0_h0(a);
|
||||
uint4 c;
|
||||
c.x = mul<uint32_t, uint32_t, uint32_t>(s, b.x);
|
||||
c.y = mul<uint32_t, uint32_t, uint32_t>(s, b.y);
|
||||
c.z = mul<uint32_t, uint32_t, uint32_t>(s, b.z);
|
||||
c.w = mul<uint32_t, uint32_t, uint32_t>(s, b.w);
|
||||
return c;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float mul(uint16_t a, uint16_t b) {
|
||||
float fa = half_to_float(a);
|
||||
float fb = half_to_float(b);
|
||||
return fa * fb;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float mul(uint16_t a, float b) {
|
||||
return half_to_float(a) * b;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float2 mul(uint32_t a, uint32_t b) {
|
||||
float2 fa = half2_to_float2(a);
|
||||
float2 fb = half2_to_float2(b);
|
||||
return mul<float2, float2, float2>(fa, fb);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ float2 mul(uint16_t a, uint32_t b) {
|
||||
return mul<float2, uint32_t, uint32_t>(h0_h0(a), b);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ Float4_ mul(uint2 a, uint2 b) {
|
||||
Float4_ fc;
|
||||
fc.x = mul<float2, uint32_t, uint32_t>(a.x, b.x);
|
||||
fc.y = mul<float2, uint32_t, uint32_t>(a.y, b.y);
|
||||
return fc;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ Float4_ mul(uint16_t a, uint2 b) {
|
||||
uint32_t s = h0_h0(a);
|
||||
Float4_ fc;
|
||||
fc.x = mul<float2, uint32_t, uint32_t>(s, b.x);
|
||||
fc.y = mul<float2, uint32_t, uint32_t>(s, b.y);
|
||||
return fc;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ Float8_ mul(uint4 a, uint4 b) {
|
||||
Float8_ fc;
|
||||
fc.x = mul<float2, uint32_t, uint32_t>(a.x, b.x);
|
||||
fc.y = mul<float2, uint32_t, uint32_t>(a.y, b.y);
|
||||
fc.z = mul<float2, uint32_t, uint32_t>(a.z, b.z);
|
||||
fc.w = mul<float2, uint32_t, uint32_t>(a.w, b.w);
|
||||
return fc;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
inline __device__ Float8_ mul(uint16_t a, uint4 b) {
|
||||
uint32_t s = h0_h0(a);
|
||||
Float8_ fc;
|
||||
fc.x = mul<float2, uint32_t, uint32_t>(s, b.x);
|
||||
fc.y = mul<float2, uint32_t, uint32_t>(s, b.y);
|
||||
fc.z = mul<float2, uint32_t, uint32_t>(s, b.z);
|
||||
fc.w = mul<float2, uint32_t, uint32_t>(s, b.w);
|
||||
return fc;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(float v) {
|
||||
return v;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(float2 v) {
|
||||
return v.x + v.y;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(float4 v) {
|
||||
return v.x + v.y + v.z + v.w;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(uint16_t v) {
|
||||
return half_to_float(v);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(uint32_t v) {
|
||||
float2 tmp = half2_to_float2(v);
|
||||
return tmp.x + tmp.y;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(uint2 v) {
|
||||
uint32_t c = add(v.x, v.y);
|
||||
return sum(c);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(uint4 v) {
|
||||
#if 1
|
||||
uint32_t c = add(v.x, v.y);
|
||||
c = add(c, v.z);
|
||||
c = add(c, v.w);
|
||||
#else
|
||||
uint32_t c = add(v.x, v.y);
|
||||
uint32_t d = add(v.z, v.w);
|
||||
c = add(c, d);
|
||||
#endif
|
||||
return sum(c);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(Float4_ v) {
|
||||
return v.x.x + v.x.y + v.y.x + v.y.y;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float sum(Float8_ v) {
|
||||
return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
inline __device__ float dot(T a, T b) {
|
||||
return sum(mul<T, T, T>(a, b));
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename A, typename T>
|
||||
inline __device__ float dot(T a, T b) {
|
||||
return sum(mul<A, T, T>(a, b));
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ void zero(uint16_t& dst) {
|
||||
dst = uint16_t(0);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
inline __device__ void zero(T& dst) {
|
||||
constexpr int WORDS = sizeof(T) / 4;
|
||||
union {
|
||||
T raw;
|
||||
uint32_t words[WORDS];
|
||||
} tmp;
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < WORDS; ++ii) {
|
||||
tmp.words[ii] = 0u;
|
||||
}
|
||||
dst = tmp.raw;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline __device__ float2 rotary_embedding_coefficient(const int zid, const int rot_embed_dim, const float t_step) {
|
||||
const float inv_freq = t_step / pow(10000.0f, zid / (float)rot_embed_dim);
|
||||
|
|
@ -849,9 +42,9 @@ inline __device__ float2 rotary_embedding_transform(const float2 v, const float2
|
|||
}
|
||||
|
||||
inline __device__ uint32_t rotary_embedding_transform(const uint32_t v, const float2 coef) {
|
||||
float2 fv = half2_to_float2(v);
|
||||
float2 fv = Half2ToFloat2(v);
|
||||
float2 rot_fv = rotary_embedding_transform(fv, coef);
|
||||
return float2_to_half2(rot_fv);
|
||||
return Float2ToHalf2(rot_fv);
|
||||
}
|
||||
|
||||
inline __device__ void apply_rotary_embedding(float& q, int zid, int rot_embed_dim, int t_step) {
|
||||
|
|
|
|||
|
|
@ -25,12 +25,123 @@
|
|||
#pragma once
|
||||
|
||||
#include "core/providers/cuda/cuda_common.h"
|
||||
#include "decoder_masked_multihead_attention_impl_utils.h"
|
||||
#include "core/providers/cuda/cu_inc/common.cuh"
|
||||
|
||||
using namespace onnxruntime::cuda;
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace contrib {
|
||||
namespace cuda {
|
||||
namespace decoder_masked_self_attention_details {
|
||||
|
||||
struct __align__(8) Half4 {
|
||||
half2 x;
|
||||
half2 y;
|
||||
};
|
||||
|
||||
__device__ __forceinline__ Half4 operator+(const Half4& a, const Half4& b) {
|
||||
Half4 r;
|
||||
r.x = a.x + b.x;
|
||||
r.y = a.y + b.y;
|
||||
return r;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float2 operator+(const float2& a, const float2& b) {
|
||||
return make_float2(a.x + b.x, a.y + b.y);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float4 operator+(const float4& a, const float4& b) {
|
||||
return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
|
||||
}
|
||||
|
||||
struct Float8_ {
|
||||
float2 x;
|
||||
float2 y;
|
||||
float2 z;
|
||||
float2 w;
|
||||
};
|
||||
|
||||
struct Float4_ {
|
||||
float2 x;
|
||||
float2 y;
|
||||
};
|
||||
|
||||
#ifndef USE_ROCM
|
||||
|
||||
template <typename T>
|
||||
struct num_elems;
|
||||
template <>
|
||||
struct num_elems<float> {
|
||||
static constexpr int value = 1;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<float2> {
|
||||
static constexpr int value = 2;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<float4> {
|
||||
static constexpr int value = 4;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<Float4_> {
|
||||
static constexpr int value = 4;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<Float8_> {
|
||||
static constexpr int value = 8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct num_elems<uint32_t> {
|
||||
static constexpr int value = 2;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<uint2> {
|
||||
static constexpr int value = 4;
|
||||
};
|
||||
template <>
|
||||
struct num_elems<uint4> {
|
||||
static constexpr int value = 8;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Vec_t {
|
||||
static constexpr int size = 0;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<float> {
|
||||
using Type = float2;
|
||||
static constexpr int size = 2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<float2> {
|
||||
using Type = float4;
|
||||
static constexpr int size = 4;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<float4> {
|
||||
using Type = Float8_;
|
||||
static constexpr int size = 8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<half> {
|
||||
using Type = uint32_t;
|
||||
static constexpr int size = 2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<half2> {
|
||||
using Type = uint2;
|
||||
static constexpr int size = 4;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Vec_t<Half4> {
|
||||
using Type = uint4;
|
||||
static constexpr int size = 8;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Qk_vec
|
||||
|
|
@ -163,6 +274,16 @@ struct V_vec_m_<uint16_t, 8> {
|
|||
using Type = uint4;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct V_vec_m_<half, 2> {
|
||||
using Type = half2;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct V_vec_m_<half, 4> {
|
||||
using Type = Half4;
|
||||
};
|
||||
|
||||
template <typename T, int V_VEC_SIZE>
|
||||
struct V_vec_k_ {
|
||||
using Type = typename V_vec_m_<T, V_VEC_SIZE>::Type;
|
||||
|
|
@ -190,18 +311,6 @@ struct V_vec_acum_fp32_<float4> {
|
|||
using Type = float4;
|
||||
};
|
||||
|
||||
struct Float8_ {
|
||||
float2 x;
|
||||
float2 y;
|
||||
float2 z;
|
||||
float2 w;
|
||||
};
|
||||
|
||||
struct Float4_ {
|
||||
float2 x;
|
||||
float2 y;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct V_vec_acum_fp32_<uint32_t> {
|
||||
using Type = float2;
|
||||
|
|
@ -219,7 +328,9 @@ struct V_vec_acum_fp32_<uint4> {
|
|||
// Zero
|
||||
//------------------------------------------------------------
|
||||
|
||||
// TODO: fp16 ?
|
||||
inline __device__ void zero(uint16_t& dst) {
|
||||
dst = uint16_t(0);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline __device__ void zero(T& dst) {
|
||||
|
|
@ -235,6 +346,14 @@ inline __device__ void zero(T& dst) {
|
|||
dst = tmp.raw;
|
||||
}
|
||||
|
||||
inline __device__ uint32_t h0_h0(uint16_t a) {
|
||||
uint32_t b;
|
||||
asm volatile("mov.b32 %0, {%1, %1};"
|
||||
: "=r"(b)
|
||||
: "h"(a));
|
||||
return b;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// vec_conversion
|
||||
//------------------------------------------------------------
|
||||
|
|
@ -268,6 +387,15 @@ inline __device__ float4 add_vec(float4 a, float4 b) {
|
|||
return c;
|
||||
}
|
||||
|
||||
inline __device__ Float8_ add_vec(Float8_ a, Float8_ b) {
|
||||
Float8_ c;
|
||||
c.x = add_vec(a.x, b.x);
|
||||
c.y = add_vec(a.y, b.y);
|
||||
c.z = add_vec(a.z, b.z);
|
||||
c.w = add_vec(a.w, b.w);
|
||||
return c;
|
||||
}
|
||||
|
||||
inline __device__ float HalfToFloat(uint16_t h) {
|
||||
float f;
|
||||
asm volatile("cvt.f32.f16 %0, %1;\n"
|
||||
|
|
@ -378,6 +506,14 @@ inline __device__ float sum(uint4 v) {
|
|||
return sum(c);
|
||||
}
|
||||
|
||||
inline __device__ float sum(Float4_ v) {
|
||||
return v.x.x + v.x.y + v.y.x + v.y.y;
|
||||
}
|
||||
|
||||
inline __device__ float sum(Float8_ v) {
|
||||
return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Mul
|
||||
//------------------------------------------------------------
|
||||
|
|
@ -400,6 +536,14 @@ inline __device__ float2 mul(float2 a, float2 b) {
|
|||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ float2 mul(float a, float2 b) {
|
||||
float2 c;
|
||||
c.x = a * b.x;
|
||||
c.y = a * b.y;
|
||||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ float4 mul(float4 a, float4 b) {
|
||||
float4 c;
|
||||
|
|
@ -410,6 +554,26 @@ inline __device__ float4 mul(float4 a, float4 b) {
|
|||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ float4 mul(float a, float4 b) {
|
||||
float4 c;
|
||||
c.x = a * b.x;
|
||||
c.y = a * b.y;
|
||||
c.z = a * b.z;
|
||||
c.w = a * b.w;
|
||||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ Float8_ mul(float a, Float8_ b) {
|
||||
Float8_ c;
|
||||
c.x = make_float2(a * b.x.x, a * b.x.y);
|
||||
c.y = make_float2(a * b.y.x, a * b.y.y);
|
||||
c.z = make_float2(a * b.z.x, a * b.z.y);
|
||||
c.w = make_float2(a * b.w.x, a * b.w.y);
|
||||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ uint16_t mul(uint16_t a, uint16_t b) {
|
||||
uint16_t c;
|
||||
|
|
@ -428,6 +592,11 @@ inline __device__ uint32_t mul(uint32_t a, uint32_t b) {
|
|||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ uint32_t mul(uint16_t a, uint32_t b) {
|
||||
return mul<uint32_t, uint32_t, uint32_t>(h0_h0(a), b);
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ uint2 mul(uint2 a, uint2 b) {
|
||||
uint2 c;
|
||||
|
|
@ -436,6 +605,15 @@ inline __device__ uint2 mul(uint2 a, uint2 b) {
|
|||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ uint2 mul(uint16_t a, uint2 b) {
|
||||
uint32_t s = h0_h0(a);
|
||||
uint2 c;
|
||||
c.x = mul<uint32_t, uint32_t, uint32_t>(s, b.x);
|
||||
c.y = mul<uint32_t, uint32_t, uint32_t>(s, b.y);
|
||||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ uint4 mul(uint4 a, uint4 b) {
|
||||
uint4 c;
|
||||
|
|
@ -446,6 +624,79 @@ inline __device__ uint4 mul(uint4 a, uint4 b) {
|
|||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ uint4 mul(uint16_t a, uint4 b) {
|
||||
uint32_t s = h0_h0(a);
|
||||
uint4 c;
|
||||
c.x = mul<uint32_t, uint32_t, uint32_t>(s, b.x);
|
||||
c.y = mul<uint32_t, uint32_t, uint32_t>(s, b.y);
|
||||
c.z = mul<uint32_t, uint32_t, uint32_t>(s, b.z);
|
||||
c.w = mul<uint32_t, uint32_t, uint32_t>(s, b.w);
|
||||
return c;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ float mul(uint16_t a, float b) {
|
||||
return HalfToFloat(a) * b;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ float mul(uint16_t a, uint16_t b) {
|
||||
float fa = HalfToFloat(a);
|
||||
float fb = HalfToFloat(b);
|
||||
return fa * fb;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ float2 mul(uint32_t a, uint32_t b) {
|
||||
float2 fa = Half2ToFloat2(a);
|
||||
float2 fb = Half2ToFloat2(b);
|
||||
return mul<float2, float2, float2>(fa, fb);
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ float2 mul(uint16_t a, uint32_t b) {
|
||||
return mul<float2, uint32_t, uint32_t>(h0_h0(a), b);
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ Float4_ mul(uint2 a, uint2 b) {
|
||||
Float4_ fc;
|
||||
fc.x = mul<float2, uint32_t, uint32_t>(a.x, b.x);
|
||||
fc.y = mul<float2, uint32_t, uint32_t>(a.y, b.y);
|
||||
return fc;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ Float4_ mul(uint16_t a, uint2 b) {
|
||||
uint32_t s = h0_h0(a);
|
||||
Float4_ fc;
|
||||
fc.x = mul<float2, uint32_t, uint32_t>(s, b.x);
|
||||
fc.y = mul<float2, uint32_t, uint32_t>(s, b.y);
|
||||
return fc;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ Float8_ mul(uint4 a, uint4 b) {
|
||||
Float8_ fc;
|
||||
fc.x = mul<float2, uint32_t, uint32_t>(a.x, b.x);
|
||||
fc.y = mul<float2, uint32_t, uint32_t>(a.y, b.y);
|
||||
fc.z = mul<float2, uint32_t, uint32_t>(a.z, b.z);
|
||||
fc.w = mul<float2, uint32_t, uint32_t>(a.w, b.w);
|
||||
return fc;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline __device__ Float8_ mul(uint16_t a, uint4 b) {
|
||||
uint32_t s = h0_h0(a);
|
||||
Float8_ fc;
|
||||
fc.x = mul<float2, uint32_t, uint32_t>(s, b.x);
|
||||
fc.y = mul<float2, uint32_t, uint32_t>(s, b.y);
|
||||
fc.z = mul<float2, uint32_t, uint32_t>(s, b.z);
|
||||
fc.w = mul<float2, uint32_t, uint32_t>(s, b.w);
|
||||
return fc;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Fma
|
||||
//------------------------------------------------------------
|
||||
|
|
@ -485,14 +736,6 @@ inline __device__ float4 fma(float a, float4 b, float4 c) {
|
|||
return d;
|
||||
}
|
||||
|
||||
inline __device__ uint32_t h0_h0(uint16_t a) {
|
||||
uint32_t b;
|
||||
asm volatile("mov.b32 %0, {%1, %1};"
|
||||
: "=r"(b)
|
||||
: "h"(a));
|
||||
return b;
|
||||
}
|
||||
|
||||
inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) {
|
||||
uint32_t d;
|
||||
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
|
||||
|
|
@ -555,6 +798,13 @@ inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc) {
|
|||
return fma(h0_h0(a), b, fc);
|
||||
}
|
||||
|
||||
inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc) {
|
||||
Float4_ fd;
|
||||
fd.x = fma(a.x, b.x, fc.x);
|
||||
fd.y = fma(a.y, b.y, fc.y);
|
||||
return fd;
|
||||
}
|
||||
|
||||
inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) {
|
||||
uint32_t s = h0_h0(a);
|
||||
Float4_ fd;
|
||||
|
|
@ -563,6 +813,15 @@ inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) {
|
|||
return fd;
|
||||
}
|
||||
|
||||
inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc) {
|
||||
Float8_ fd;
|
||||
fd.x = fma(a.x, b.x, fc.x);
|
||||
fd.y = fma(a.y, b.y, fc.y);
|
||||
fd.z = fma(a.z, b.z, fc.z);
|
||||
fd.w = fma(a.w, b.w, fc.w);
|
||||
return fd;
|
||||
}
|
||||
|
||||
inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) {
|
||||
uint32_t s = h0_h0(a);
|
||||
Float8_ fd;
|
||||
|
|
@ -573,103 +832,6 @@ inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) {
|
|||
return fd;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Block_sum
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <int WARPS_PER_BLOCK, int WARP_SIZE = 32>
|
||||
inline __device__ float block_sum(float* red_smem, float sum) {
|
||||
// Decompose the thread index into warp / lane.
|
||||
int warp = threadIdx.x / WARP_SIZE;
|
||||
int lane = threadIdx.x % WARP_SIZE;
|
||||
|
||||
// Compute the sum per warp.
|
||||
#pragma unroll
|
||||
for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
|
||||
sum += __shfl_xor_sync(uint32_t(-1), sum, mask);
|
||||
}
|
||||
|
||||
// Warp leaders store the data to shared memory.
|
||||
if (lane == 0) {
|
||||
red_smem[warp] = sum;
|
||||
}
|
||||
|
||||
// Make sure the data is in shared memory.
|
||||
__syncthreads();
|
||||
|
||||
// The warps compute the final sums.
|
||||
if (lane < WARPS_PER_BLOCK) {
|
||||
sum = red_smem[lane];
|
||||
}
|
||||
|
||||
// Parallel reduction inside the warp.
|
||||
#pragma unroll
|
||||
for (int mask = WARPS_PER_BLOCK / 2; mask >= 1; mask /= 2) {
|
||||
sum += __shfl_xor_sync(uint32_t(-1), sum, mask);
|
||||
}
|
||||
|
||||
// Broadcast to other threads.
|
||||
return __shfl_sync(uint32_t(-1), sum, 0);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Shfl_Mask
|
||||
//------------------------------------------------------------
|
||||
|
||||
inline __device__ constexpr uint32_t shfl_mask(int threads) {
|
||||
return threads == 32 ? uint32_t(-1) : (1u << threads) - 1u;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Dot
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <typename A, typename T>
|
||||
inline __device__ float dot(T a, T b) {
|
||||
return sum(mul<A, T, T>(a, b));
|
||||
}
|
||||
|
||||
//------------------------------------------------------------
|
||||
// Qk_Dot
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <int THREADS_PER_KEY, typename K_vec, int N>
|
||||
inline __device__ float qk_dot_(const K_vec (&q)[N], const K_vec (&k)[N]) {
|
||||
using K_vec_acum = K_vec;
|
||||
|
||||
// Compute the parallel products for Q*K^T (treat vector lanes separately).
|
||||
K_vec_acum qk_vec = mul<K_vec_acum, K_vec, K_vec>(q[0], k[0]);
|
||||
#pragma unroll
|
||||
for (int ii = 1; ii < N; ++ii) {
|
||||
qk_vec = fma(q[ii], k[ii], qk_vec);
|
||||
}
|
||||
|
||||
// Finalize the reduction across lanes.
|
||||
float qk = sum(qk_vec);
|
||||
#pragma unroll
|
||||
for (int mask = THREADS_PER_KEY / 2; mask >= 1; mask /= 2) {
|
||||
qk += __shfl_xor_sync(uint32_t(-1), qk, mask);
|
||||
}
|
||||
return qk;
|
||||
}
|
||||
|
||||
template <typename T, int THREADS_PER_KEY>
|
||||
struct Qk_dot {
|
||||
template <typename K_vec, int N>
|
||||
static inline __device__ float dot(const K_vec (&q)[N], const K_vec (&k)[N]) {
|
||||
return qk_dot_<THREADS_PER_KEY>(q, k);
|
||||
}
|
||||
};
|
||||
|
||||
//------------------------------------------------------------
|
||||
// ThreadsPerValue
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <typename T, int head_size>
|
||||
struct ThreadsPerValue {
|
||||
static const int value = head_size * sizeof(T) / 16;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------
|
||||
// ConvertFromFloat
|
||||
//------------------------------------------------------------
|
||||
|
|
@ -691,9 +853,14 @@ inline __device__ uint16_t FloatToHalf(float f) {
|
|||
uint32_t u32;
|
||||
uint16_t u16[2];
|
||||
} tmp;
|
||||
#if 0 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 // Is it better?
|
||||
float zero = 0.f;
|
||||
asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(tmp.u32) : "f"(zero), "f"(f));
|
||||
#else
|
||||
asm volatile("cvt.rn.f16.f32 %0, %1;\n"
|
||||
: "=h"(tmp.u16[0])
|
||||
: "f"(f));
|
||||
#endif
|
||||
return tmp.u16[0];
|
||||
}
|
||||
|
||||
|
|
@ -736,39 +903,8 @@ inline __device__ void ConvertFromFloat(uint4& dst, Float8_ src) {
|
|||
dst.z = Float2ToHalf2(src.z);
|
||||
dst.w = Float2ToHalf2(src.w);
|
||||
}
|
||||
//------------------------------------------------------------
|
||||
// CalcDynamicBlockMemory
|
||||
//------------------------------------------------------------
|
||||
|
||||
template <typename T>
|
||||
inline size_t CalcDynamicBlockMemory(const DecoderMaskedMultiHeadAttentionParams& params,
|
||||
int threads_per_value, int threads_per_block) {
|
||||
// The amount of shared memory needed to store the Q*K^T values in float.
|
||||
#endif
|
||||
|
||||
const int total_sequence_length = params.total_sequence_length;
|
||||
size_t qk_sz = (((total_sequence_length + 3) / 4) * 16);
|
||||
|
||||
// The extra memory needed if we are not using floats for the final logits.
|
||||
size_t logits_sz = 0;
|
||||
|
||||
if (sizeof(T) != 4) {
|
||||
logits_sz = (((total_sequence_length + 3) / 4) * 4 * sizeof(T));
|
||||
}
|
||||
|
||||
// The total size needed during softmax.
|
||||
size_t softmax_sz = qk_sz + logits_sz;
|
||||
|
||||
// The number of partial rows to reduce in the final reduction.
|
||||
int rows_per_red = threads_per_block / threads_per_value;
|
||||
|
||||
// The amount of storage needed to finalize the outputs.
|
||||
size_t red_sz = rows_per_red * params.head_size * sizeof(T) / 2;
|
||||
|
||||
// The max.
|
||||
return std::max(softmax_sz, red_sz);
|
||||
}
|
||||
|
||||
} // namespace decoder_masked_self_attention_details
|
||||
} // namespace cuda
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
Loading…
Reference in a new issue