IQ.Pilot Release Commit @ bec7652
This commit is contained in:
IQ.Lvbs history cleanup
@@ -0,0 +1,400 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Basic operations on generic types.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <limits>
|
||||
#include "base_types.cuh"
|
||||
|
||||
namespace kittens {
|
||||
|
||||
/**
|
||||
* @namespace base_ops
|
||||
*
|
||||
* @brief A namespace for operations on basic data types.
|
||||
*/
|
||||
namespace base_ops {
|
||||
|
||||
/* ---------- CONST OPS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Represents the zero constant operation.
|
||||
*
|
||||
* This operation returns the zero value of the specified type.
|
||||
*
|
||||
* @tparam T The data type for which to return the zero value.
|
||||
* @return The zero value of type T.
|
||||
*/
|
||||
struct zero {
|
||||
template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::zero(); }
|
||||
};
|
||||
/**
|
||||
* @brief Represents the one constant operation.
|
||||
*
|
||||
* This operation returns the one value of the specified type.
|
||||
*
|
||||
* @tparam T The data type for which to return the one value.
|
||||
* @return The one value of type T.
|
||||
*/
|
||||
struct one {
|
||||
template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::one(); }
|
||||
};
|
||||
/**
|
||||
* @brief Represents the positive infinity constant operation.
|
||||
*
|
||||
* This operation returns the positive infinity value of the specified type.
|
||||
*
|
||||
* @tparam T The data type for which to return the positive infinity value.
|
||||
* @return The positive infinity value of type T.
|
||||
*/
|
||||
struct pos_infty {
|
||||
template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::pos_infty(); }
|
||||
};
|
||||
/**
|
||||
* @brief Represents the negative infinity constant operation.
|
||||
*
|
||||
* This operation returns the negative infinity value of the specified type.
|
||||
*
|
||||
* @tparam T The data type for which to return the negative infinity value.
|
||||
* @return The negative infinity value of type T.
|
||||
*/
|
||||
struct neg_infty {
|
||||
template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::neg_infty(); }
|
||||
};
|
||||
|
||||
|
||||
/* ---------- UNARY OPS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Exponential function operation.
|
||||
*
|
||||
* This operation calculates the exponential of the input value.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The exponential of the input value.
|
||||
*/
|
||||
struct exp {
|
||||
template<typename T> static __device__ inline T op(const T &x) { return exp(x); }
|
||||
};
|
||||
template<> __device__ inline float exp::op<float> (const float &x ) { return __expf(x); }
|
||||
template<> __device__ inline float2 exp::op<float2>(const float2 &x) { return float2{__expf(x.x), __expf(x.y)}; }
|
||||
template<> __device__ inline bf16 exp::op<bf16> (const bf16 &x ) { return hexp(x); }
|
||||
template<> __device__ inline bf16_2 exp::op<bf16_2>(const bf16_2 &x) { return h2exp(x); }
|
||||
template<> __device__ inline half exp::op<half> (const half &x ) { return hexp(x); }
|
||||
template<> __device__ inline half_2 exp::op<half_2>(const half_2 &x) { return h2exp(x); }
|
||||
|
||||
/**
|
||||
* @brief Exponential function operation, in base 2
|
||||
*
|
||||
* This operation calculates the exponential of the input value, in base 2.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The exponential of the input value.
|
||||
*/
|
||||
struct exp2 {
|
||||
template<typename T> static __device__ inline T op(const T &x) { return exp2f(x); }
|
||||
};
|
||||
template<> __device__ inline float exp2::op<float> (const float &x ) { return exp2f(x); }
|
||||
template<> __device__ inline float2 exp2::op<float2>(const float2 &x) { return float2{exp2f(x.x), exp2f(x.y)}; }
|
||||
template<> __device__ inline bf16 exp2::op<bf16> (const bf16 &x ) { return hexp2(x); }
|
||||
template<> __device__ inline bf16_2 exp2::op<bf16_2>(const bf16_2 &x) { return h2exp2(x); }
|
||||
template<> __device__ inline half exp2::op<half> (const half &x ) { return hexp2(x); }
|
||||
template<> __device__ inline half_2 exp2::op<half_2>(const half_2 &x) { return h2exp2(x); }
|
||||
/**
|
||||
* @brief Natural log function operation.
|
||||
*
|
||||
* This operation calculates the natural logarithm of the input value.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The natural logarithm of the input value.
|
||||
*/
|
||||
struct log {
|
||||
template<typename T> static __device__ inline T op(const T &x) { return log(x); }
|
||||
};
|
||||
template<> __device__ inline float log::op<float> (const float &x ) { return __logf(x); }
|
||||
template<> __device__ inline float2 log::op<float2>(const float2 &x) { return float2{__logf(x.x), __logf(x.y)}; }
|
||||
template<> __device__ inline bf16 log::op<bf16> (const bf16 &x ) { return hlog(x); }
|
||||
template<> __device__ inline bf16_2 log::op<bf16_2>(const bf16_2 &x) { return h2log(x); }
|
||||
template<> __device__ inline half log::op<half> (const half &x ) { return hlog(x); }
|
||||
template<> __device__ inline half_2 log::op<half_2>(const half_2 &x) { return h2log(x); }
|
||||
/**
|
||||
* @brief Logarithm base 2 operation.
|
||||
*
|
||||
* This operation calculates the logarithm base 2 of the input value.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The logarithm base 2 of the input value.
|
||||
*/
|
||||
struct log2 {
|
||||
template<typename T> static __device__ inline T op(const T &x) { return log2(x); }
|
||||
};
|
||||
template<> __device__ inline float log2::op<float> (const float &x ) { return __log2f(x); }
|
||||
template<> __device__ inline float2 log2::op<float2>(const float2 &x) { return float2{__log2f(x.x), __log2f(x.y)}; }
|
||||
template<> __device__ inline bf16 log2::op<bf16> (const bf16 &x ) { return hlog2(x); }
|
||||
template<> __device__ inline bf16_2 log2::op<bf16_2>(const bf16_2 &x) { return h2log2(x); }
|
||||
template<> __device__ inline half log2::op<half> (const half &x ) { return hlog2(x); }
|
||||
template<> __device__ inline half_2 log2::op<half_2>(const half_2 &x) { return h2log2(x); }
|
||||
/**
|
||||
* @brief Absolute value operation.
|
||||
*
|
||||
* This operation calculates the absolute value of the input.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The absolute value of the input.
|
||||
*/
|
||||
struct abs {
|
||||
template<typename T> static __device__ inline T op(const T &x) { return abs(x); }
|
||||
};
|
||||
template<> __device__ inline float abs::op<float> (const float &x ) { return fabsf(x); }
|
||||
template<> __device__ inline float2 abs::op<float2>(const float2 &x) { return float2{fabsf(x.x), fabsf(x.y)}; }
|
||||
template<> __device__ inline bf16 abs::op<bf16> (const bf16 &x ) { return __habs(x); }
|
||||
template<> __device__ inline bf16_2 abs::op<bf16_2>(const bf16_2 &x) { return __habs2(x); }
|
||||
template<> __device__ inline half abs::op<half> (const half &x ) { return __habs(x); }
|
||||
template<> __device__ inline half_2 abs::op<half_2>(const half_2 &x) { return __habs2(x); }
|
||||
/**
|
||||
* @brief Rectified Linear Unit (ReLU) operation.
|
||||
*
|
||||
* This operation applies the ReLU function to the input, which is the
|
||||
* maximum of zero and the input value.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param x[in] The input value.
|
||||
* @return The result of ReLU function applied to the input.
|
||||
*/
|
||||
struct relu {
|
||||
template<typename T> static __device__ inline T op(const T &x) { return max(x, base_types::constants<T>::zero()); }
|
||||
};
|
||||
template<> __device__ inline float relu::op<float> (const float &x ) { return max(x, 0.f); }
|
||||
template<> __device__ inline float2 relu::op<float2>(const float2 &x) { return float2{max(x.x, 0.f), max(x.y, 0.f)}; }
|
||||
template<> __device__ inline bf16 relu::op<bf16> (const bf16 &x ) { return __hmax(x, base_types::constants<bf16>::zero()); }
|
||||
template<> __device__ inline bf16_2 relu::op<bf16_2>(const bf16_2 &x) { return __hmax2(x, base_types::constants<bf16_2>::zero()); }
|
||||
template<> __device__ inline half relu::op<half> (const half &x ) { return __hmax(x, base_types::constants<half>::zero()); }
|
||||
template<> __device__ inline half_2 relu::op<half_2>(const half_2 &x) { return __hmax2(x, base_types::constants<half_2>::zero()); }
|
||||
/**
|
||||
* @brief Copy operation.
|
||||
*
|
||||
* This operation returns the input value unchanged.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The input value.
|
||||
* @return The same value as the input.
|
||||
*/
|
||||
struct copy { // for non-compile-time setters.
|
||||
template<typename T> static __device__ inline T op(const T &a) { return a; }
|
||||
};
|
||||
|
||||
|
||||
/* ---------- BINARY OPS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Copy2 operation.
|
||||
*
|
||||
* This operation returns the second input value unchanged.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value (ignored).
|
||||
* @param b[in] The second input value.
|
||||
* @return The same value as the second input.
|
||||
*/
|
||||
struct copy2 { // this turns out to be a slightly hacky op that makes some code cleaner :/
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b) { return b; }
|
||||
};
|
||||
/**
|
||||
* @brief Sum operation.
|
||||
*
|
||||
* This operation calculates the sum of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The sum of the input values.
|
||||
*/
|
||||
struct sum {
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b) { return a+b; }
|
||||
};
|
||||
template<> __device__ inline float2 sum::op<float2>(const float2 &a, const float2 &b) {
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
float2 c;
|
||||
asm volatile("add.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
|
||||
return c;
|
||||
#else
|
||||
return float2{a.x+b.x, a.y+b.y};
|
||||
#endif
|
||||
}
|
||||
template<> __device__ inline bf16 sum::op<bf16> (const bf16 &a, const bf16 &b) { return __hadd(a, b); }
|
||||
template<> __device__ inline bf16_2 sum::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hadd2(a, b); }
|
||||
template<> __device__ inline half sum::op<half> (const half &a, const half &b) { return __hadd(a, b); }
|
||||
template<> __device__ inline half_2 sum::op<half_2>(const half_2 &a, const half_2 &b) { return __hadd2(a, b); }
|
||||
/**
|
||||
* @brief Subtraction operation.
|
||||
*
|
||||
* This operation calculates the difference between two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The difference between the input values.
|
||||
*/
|
||||
struct sub {
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b) { return a-b; }
|
||||
};
|
||||
template<> __device__ inline float2 sub::op<float2>(const float2 &a, const float2 &b) {
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
float2 c;
|
||||
asm volatile("sub.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
|
||||
return c;
|
||||
#else
|
||||
return float2{a.x-b.x, a.y-b.y};
|
||||
#endif
|
||||
}
|
||||
template<> __device__ inline bf16 sub::op<bf16> (const bf16 &a, const bf16 &b) { return __hsub(a, b); }
|
||||
template<> __device__ inline bf16_2 sub::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hsub2(a, b); }
|
||||
template<> __device__ inline half sub::op<half> (const half &a, const half &b) { return __hsub(a, b); }
|
||||
template<> __device__ inline half_2 sub::op<half_2>(const half_2 &a, const half_2 &b) { return __hsub2(a, b); }
|
||||
/**
|
||||
* @brief Multiplication operation.
|
||||
*
|
||||
* This operation calculates the product of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The product of the input values.
|
||||
*/
|
||||
struct mul {
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b) { return a*b; }
|
||||
};
|
||||
template<> __device__ inline float2 mul::op<float2>(const float2 &a, const float2 &b) {
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
float2 c;
|
||||
asm volatile("mul.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
|
||||
return c;
|
||||
#else
|
||||
return float2{a.x*b.x, a.y*b.y};
|
||||
#endif
|
||||
}
|
||||
template<> __device__ inline bf16 mul::op<bf16> (const bf16 &a, const bf16 &b) { return __hmul(a, b); }
|
||||
template<> __device__ inline bf16_2 mul::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmul2(a, b); }
|
||||
template<> __device__ inline half mul::op<half> (const half &a, const half &b) { return __hmul(a, b); }
|
||||
template<> __device__ inline half_2 mul::op<half_2>(const half_2 &a, const half_2 &b) { return __hmul2(a, b); }
|
||||
/**
|
||||
* @brief Division operation.
|
||||
*
|
||||
* This operation calculates the quotient of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The quotient of the input values.
|
||||
*/
|
||||
struct div {
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b) { return a/b; }
|
||||
};
|
||||
template<> __device__ inline float2 div::op<float2>(const float2 &a, const float2 &b) { return float2{a.x/b.x, a.y/b.y}; }
|
||||
template<> __device__ inline bf16 div::op<bf16> (const bf16 &a, const bf16 &b) { return __hdiv(a, b); }
|
||||
template<> __device__ inline bf16_2 div::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __h2div(a, b); } // this op is a special snowflake
|
||||
template<> __device__ inline half div::op<half> (const half &a, const half &b) { return __hdiv(a, b); }
|
||||
template<> __device__ inline half_2 div::op<half_2>(const half_2 &a, const half_2 &b) { return __h2div(a, b); }
|
||||
/**
|
||||
* @brief Maximum operation.
|
||||
*
|
||||
* This operation calculates the maximum of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The maximum of the input values.
|
||||
*/
|
||||
struct max {
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::max(a, b); }
|
||||
};
|
||||
template<> __device__ inline float2 max::op<float2>(const float2 &a, const float2 &b) { return float2{::max(a.x, b.x), ::max(a.y, b.y)}; }
|
||||
template<> __device__ inline bf16 max::op<bf16> (const bf16 &a, const bf16 &b) { return __hmax(a, b); }
|
||||
template<> __device__ inline bf16_2 max::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmax2(a, b); }
|
||||
template<> __device__ inline half max::op<half> (const half &a, const half &b) { return __hmax(a, b); }
|
||||
template<> __device__ inline half_2 max::op<half_2>(const half_2 &a, const half_2 &b) { return __hmax2(a, b); }
|
||||
/**
|
||||
* @brief Minimum operation.
|
||||
*
|
||||
* This operation calculates the minimum of two input values.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @return The minimum of the input values.
|
||||
*/
|
||||
struct min {
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::min(a, b); }
|
||||
};
|
||||
template<> __device__ inline float2 min::op<float2>(const float2 &a, const float2 &b) { return float2{::min(a.x, b.x), ::min(a.y, b.y)}; }
|
||||
template<> __device__ inline bf16 min::op<bf16> (const bf16 &a, const bf16 &b) { return __hmin(a, b); }
|
||||
template<> __device__ inline bf16_2 min::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmin2(a, b); }
|
||||
template<> __device__ inline half min::op<half> (const half &a, const half &b) { return __hmin(a, b); }
|
||||
template<> __device__ inline half_2 min::op<half_2>(const half_2 &a, const half_2 &b) { return __hmin2(a, b); }
|
||||
|
||||
|
||||
/* ---------- TERNARY OPS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Fused multiply-add operation A * B + C.
|
||||
*
|
||||
* This operation performs a fused multiply-add, computing (A * B) + C with only one rounding.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The second input value.
|
||||
* @param c[in] The third input value to be added.
|
||||
* @return The result of the fused multiply-add operation.
|
||||
*/
|
||||
struct fma_AxBtC {
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
|
||||
return sum::op<T>(mul::op<T>(a, b), c);
|
||||
}
|
||||
};
|
||||
template<> __device__ inline float2 fma_AxBtC::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
float2 d;
|
||||
asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b), "l"(*(uint64_t*)&c));
|
||||
return d;
|
||||
#else
|
||||
return float2{a.x*b.x+c.x, a.y*b.y+c.y};
|
||||
#endif
|
||||
}
|
||||
/**
|
||||
* @brief Fused multiply-add operation A * C + B.
|
||||
*
|
||||
* This operation performs a fused multiply-add, computing (A * C) + B with only one rounding.
|
||||
* This is particularly useful for attention mechanisms in neural networks.
|
||||
*
|
||||
* @tparam T The data type of the input and output values.
|
||||
* @param a[in] The first input value.
|
||||
* @param b[in] The third input value to be added.
|
||||
* @param c[in] The second input value.
|
||||
* @return The result of the fused multiply-add operation.
|
||||
*/
|
||||
struct fma_AxCtB { // this is the one needed for attention
|
||||
template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
|
||||
return sum::op<T>(mul::op<T>(a, c), b);
|
||||
}
|
||||
};
|
||||
template<> __device__ inline float2 fma_AxCtB::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
float2 d;
|
||||
asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&c), "l"(*(uint64_t*)&b));
|
||||
return d;
|
||||
#else
|
||||
return float2{a.x*c.x+b.x, a.y*c.y+b.y};
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace base_ops
|
||||
|
||||
} // namespace kittens
|
||||
@@ -0,0 +1,519 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief Declarations, manipulations, and wrappers for basic types.
|
||||
*
|
||||
* This file is a bunch of utilities for going back and forth between different types.
|
||||
*
|
||||
* Many of them are for the compiler, so as to clean up the code. It unfortunately
|
||||
* seems necessary when we have types we really care about that are less than word width.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#ifdef KITTENS_HOPPER
|
||||
#include <cuda_fp8.h>
|
||||
#endif
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <string>
|
||||
#include <bit>
|
||||
|
||||
|
||||
namespace kittens {
|
||||
|
||||
/**
|
||||
* @brief Bfloat16 floating-point type.
|
||||
*/
|
||||
using bf16 = __nv_bfloat16;
|
||||
/**
|
||||
* @brief Half-precision floating-point type.
|
||||
*/
|
||||
using half = __half;
|
||||
/**
|
||||
* @brief Packed word of two bfloat16 floating-point values.
|
||||
*/
|
||||
using bf16_2 = __nv_bfloat162;
|
||||
/**
|
||||
* @brief Packed word of two half-precision floating-point values.
|
||||
*/
|
||||
using half_2 = __half2;
|
||||
#ifdef KITTENS_HOPPER
|
||||
/**
|
||||
* @brief float8 floating-point type.
|
||||
*/
|
||||
using fp8e4m3 = __nv_fp8_e4m3;
|
||||
using fp8e5m2 = __nv_fp8_e5m2;
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
using fp8e8m0 = __nv_fp8_e8m0;
|
||||
#endif
|
||||
/**
|
||||
* @brief 2-packed float8 floating-point type.
|
||||
*/
|
||||
using fp8e4m3_2 = __nv_fp8x2_e4m3;
|
||||
using fp8e5m2_2 = __nv_fp8x2_e5m2;
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
using fp8e8m0_2 = __nv_fp8x2_e8m0;
|
||||
#endif
|
||||
/**
|
||||
* @brief 4-packed float8 floating-point type.
|
||||
*/
|
||||
using fp8e4m3_4 = __nv_fp8x4_e4m3;
|
||||
using fp8e5m2_4 = __nv_fp8x4_e5m2;
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
using fp8e8m0_4 = __nv_fp8x4_e8m0;
|
||||
#endif
|
||||
#endif
|
||||
|
||||
namespace ducks {
|
||||
/**
|
||||
* @namespace base_types
|
||||
*
|
||||
* @brief A namespace for concepts for basic data types.
|
||||
*/
|
||||
namespace base_types {
|
||||
|
||||
#ifdef KITTENS_HOPPER
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
template<typename T>
|
||||
concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4> || std::is_same_v<T, fp8e8m0_4>; // could add half_2 later if implemented.
|
||||
template<typename T>
|
||||
concept T1 = std::is_same_v<T, float> || std::is_same_v<T, bf16 > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2> || std::is_same_v<T, fp8e8m0>; // could add half_2 later if implemented.
|
||||
#else
|
||||
template<typename T>
|
||||
concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4>;
|
||||
template<typename T>
|
||||
concept T1 = std::is_same_v<T, float> || std::is_same_v<T, bf16 > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2>;
|
||||
#endif
|
||||
#else
|
||||
template<typename T>
|
||||
concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2>;
|
||||
template<typename T>
|
||||
concept T1 = std::is_same_v<T, float> || std::is_same_v<T, bf16 > || std::is_same_v<T, half>;
|
||||
#endif
|
||||
|
||||
} // namespace base_types
|
||||
} // namespace ducks
|
||||
|
||||
/**
|
||||
* @namespace base_types
|
||||
*
|
||||
* @brief A namespace for ThunderKittens basic data types.
|
||||
*/
|
||||
namespace base_types {
|
||||
|
||||
/**
|
||||
* @brief Provides compile-time constants for different types.
|
||||
*
|
||||
* @tparam T The type for which to provide constants.
|
||||
*/
|
||||
template<typename T> struct constants {
|
||||
/**
|
||||
* @brief Zero
|
||||
* @return Constexpr zero with type T
|
||||
*/
|
||||
static __device__ inline constexpr T zero() { return T{0}; }
|
||||
/**
|
||||
* @brief One
|
||||
* @return Constexpr one with type T
|
||||
*/
|
||||
static __device__ inline constexpr T one() { return T{1}; }
|
||||
/**
|
||||
* @brief Positive infinity. Particularly useful for initializing before a min op.
|
||||
* @return Constexpr positive infinity with type T
|
||||
*/
|
||||
static __device__ inline constexpr T pos_infty() { return T{INFINITY}; } // I'll find a better way at some point but this appears to work.
|
||||
/**
|
||||
* @brief Negative infinity. Particularly useful for initializing before a max op.
|
||||
* @return Constexpr negative infinity with type T
|
||||
*/
|
||||
static __device__ inline constexpr T neg_infty() { return T{-INFINITY}; }
|
||||
};
|
||||
template<> struct constants<float2> {
|
||||
static __device__ inline constexpr float2 zero() { return float2{0.f, 0.f}; }
|
||||
static __device__ inline constexpr float2 one() { return float2{1.f, 1.f}; }
|
||||
static __device__ inline constexpr float2 pos_infty() { return float2{constants<float>::pos_infty(), constants<float>::pos_infty()}; }
|
||||
static __device__ inline constexpr float2 neg_infty() { return float2{constants<float>::neg_infty(), constants<float>::neg_infty()}; }
|
||||
};
|
||||
template<> struct constants<bf16> {
|
||||
static __device__ inline constexpr bf16 zero() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x0000)); } // unfortunately __float2bf16_rn is not constexpr
|
||||
static __device__ inline constexpr bf16 one() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x3F80)); }
|
||||
static __device__ inline constexpr bf16 pos_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x7F80)); }
|
||||
static __device__ inline constexpr bf16 neg_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0xFF80)); }
|
||||
};
|
||||
template<> struct constants<bf16_2> {
|
||||
static __device__ inline constexpr bf16_2 zero() { return bf16_2{constants<bf16>::zero(), constants<bf16>::zero()}; }
|
||||
static __device__ inline constexpr bf16_2 one() { return bf16_2{constants<bf16>::one(), constants<bf16>::one()}; }
|
||||
static __device__ inline constexpr bf16_2 pos_infty() { return bf16_2{constants<bf16>::pos_infty(), constants<bf16>::pos_infty()}; }
|
||||
static __device__ inline constexpr bf16_2 neg_infty() { return bf16_2{constants<bf16>::neg_infty(), constants<bf16>::neg_infty()}; }
|
||||
};
|
||||
template<> struct constants<half> {
|
||||
static __device__ inline constexpr half zero() { return std::bit_cast<__half>(uint16_t(0x0000)); }
|
||||
static __device__ inline constexpr half one() { return std::bit_cast<__half>(uint16_t(0x3C00)); }
|
||||
static __device__ inline constexpr half pos_infty() { return std::bit_cast<__half>(uint16_t(0x7C00)); }
|
||||
static __device__ inline constexpr half neg_infty() { return std::bit_cast<__half>(uint16_t(0xFC00)); }
|
||||
};
|
||||
template<> struct constants<half_2> {
|
||||
static __device__ inline constexpr half_2 zero() { return half_2{constants<half>::zero(), constants<half>::zero()}; }
|
||||
static __device__ inline constexpr half_2 one() { return half_2{constants<half>::one(), constants<half>::one()}; }
|
||||
static __device__ inline constexpr half_2 pos_infty() { return half_2{constants<half>::pos_infty(), constants<half>::pos_infty()}; }
|
||||
static __device__ inline constexpr half_2 neg_infty() { return half_2{constants<half>::neg_infty(), constants<half>::neg_infty()}; }
|
||||
};
|
||||
#ifdef KITTENS_HOPPER
|
||||
template<> struct constants<fp8e4m3> {
|
||||
static __device__ inline constexpr fp8e4m3 zero() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x00)); }
|
||||
static __device__ inline constexpr fp8e4m3 one() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x38)); }
|
||||
};
|
||||
template<> struct constants<fp8e4m3_2> {
|
||||
static __device__ inline constexpr fp8e4m3_2 zero() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x0000)); }
|
||||
static __device__ inline constexpr fp8e4m3_2 one() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x3838)); }
|
||||
};
|
||||
template<> struct constants<fp8e4m3_4> {
|
||||
static __device__ inline constexpr fp8e4m3_4 zero() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x00000000)); }
|
||||
static __device__ inline constexpr fp8e4m3_4 one() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x38383838)); }
|
||||
};
|
||||
template<> struct constants<fp8e5m2> {
|
||||
static __device__ inline constexpr fp8e5m2 zero() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x00)); }
|
||||
static __device__ inline constexpr fp8e5m2 one() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x3C)); }
|
||||
};
|
||||
template<> struct constants<fp8e5m2_2> {
|
||||
static __device__ inline constexpr fp8e5m2_2 zero() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x0000)); }
|
||||
static __device__ inline constexpr fp8e5m2_2 one() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x3C3C)); }
|
||||
};
|
||||
template<> struct constants<fp8e5m2_4> {
|
||||
static __device__ inline constexpr fp8e5m2_4 zero() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x00000000)); }
|
||||
static __device__ inline constexpr fp8e5m2_4 one() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x3C3C3C3C)); }
|
||||
};
|
||||
#endif
|
||||
|
||||
template<> struct constants<int> {
|
||||
static __device__ inline constexpr int zero() { return 0; }
|
||||
static __device__ inline constexpr int one() { return 1; }
|
||||
};
|
||||
template<> struct constants<int2> {
|
||||
static __device__ inline constexpr int2 zero() { return int2{0, 0}; }
|
||||
static __device__ inline constexpr int2 one() { return int2{1, 1}; }
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Provides information about packing of elements for a given type.
|
||||
*
|
||||
* @tparam T The type for which to provide packing information.
|
||||
*/
|
||||
template<typename T> struct packing {
|
||||
/**
|
||||
* @brief The number of elements packed together.
|
||||
*
|
||||
* @return constexpr int representing number of elements within the type.
|
||||
*/
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
/**
|
||||
* @brief Packs a single T element twice (replicated) into its packed type.
|
||||
*
|
||||
* @param i[in] The element to pack.
|
||||
* @return The packed type.
|
||||
*/
|
||||
static __device__ inline constexpr T pack(const bf16 &i);
|
||||
};
|
||||
template<> struct packing<bf16> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = bf16;
|
||||
using packed_type = bf16_2;
|
||||
static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; }
|
||||
};
|
||||
template<> struct packing<bf16_2> {
|
||||
static __device__ inline constexpr int num() { return 2; }
|
||||
using unpacked_type = bf16;
|
||||
using packed_type = bf16_2;
|
||||
static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<half> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = half;
|
||||
using packed_type = half_2;
|
||||
static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; }
|
||||
};
|
||||
template<> struct packing<half_2> {
|
||||
static __device__ inline constexpr int num() { return 2; }
|
||||
using unpacked_type = half;
|
||||
using packed_type = half_2;
|
||||
static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<float> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = float;
|
||||
using packed_type = float2;
|
||||
static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; }
|
||||
};
|
||||
template<> struct packing<float2> {
|
||||
static __device__ inline constexpr int num() { return 2; }
|
||||
using unpacked_type = float;
|
||||
using packed_type = float2;
|
||||
static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<char> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = char;
|
||||
using packed_type = char2;
|
||||
static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<char2> {
|
||||
static __device__ inline constexpr int num() { return 2; }
|
||||
using unpacked_type = char;
|
||||
using packed_type = char2;
|
||||
static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<int> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = int;
|
||||
using packed_type = int2;
|
||||
static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<int2> {
|
||||
static __device__ inline constexpr int num() { return 2; }
|
||||
using unpacked_type = int;
|
||||
using packed_type = int2;
|
||||
static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<uint> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = uint;
|
||||
using packed_type = uint2;
|
||||
static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<uint2> {
|
||||
static __device__ inline constexpr int num() { return 2; }
|
||||
using unpacked_type = uint;
|
||||
using packed_type = uint2;
|
||||
static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
struct uint64_2 { uint64_t x, y; };
|
||||
template<> struct packing<uint64_t> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = uint64_t;
|
||||
using packed_type = uint64_2;
|
||||
static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<uint64_2> {
|
||||
static __device__ inline constexpr int num() { return 2; }
|
||||
using unpacked_type = uint64_t;
|
||||
using packed_type = uint64_2;
|
||||
static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
|
||||
};
|
||||
template<> struct packing<float4> {
|
||||
static __device__ inline constexpr int num() { return 4; }
|
||||
};
|
||||
template<> struct packing<int4> {
|
||||
static __device__ inline constexpr int num() { return 4; }
|
||||
};
|
||||
#ifdef KITTENS_HOPPER
|
||||
template<> struct packing<fp8e4m3> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = fp8e4m3;
|
||||
using packed_type = fp8e4m3_4;
|
||||
};
|
||||
template<> struct packing<fp8e4m3_4> {
|
||||
static __device__ inline constexpr int num() { return 4; }
|
||||
using unpacked_type = fp8e4m3;
|
||||
using packed_type = fp8e4m3_4;
|
||||
};
|
||||
template<> struct packing<fp8e5m2> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = fp8e5m2;
|
||||
using packed_type = fp8e5m2_4;
|
||||
};
|
||||
template<> struct packing<fp8e5m2_4> {
|
||||
static __device__ inline constexpr int num() { return 4; }
|
||||
using unpacked_type = fp8e5m2;
|
||||
using packed_type = fp8e5m2_4;
|
||||
};
|
||||
#ifdef KITTENS_BLACKWELL
|
||||
template<> struct packing<fp8e8m0> {
|
||||
static __device__ inline constexpr int num() { return 1; }
|
||||
using unpacked_type = fp8e8m0;
|
||||
using packed_type = fp8e8m0_4;
|
||||
};
|
||||
template<> struct packing<fp8e8m0_4> {
|
||||
static __device__ inline constexpr int num() { return 4; }
|
||||
using unpacked_type = fp8e8m0;
|
||||
using packed_type = fp8e8m0_4;
|
||||
};
|
||||
#endif
|
||||
#endif
|
||||
|
||||
|
||||
/**
|
||||
* @brief Provides templated functionality to convert between different types.
|
||||
*
|
||||
* @tparam T The target type for conversion.
|
||||
* @tparam U The source type for conversion.
|
||||
*/
|
||||
template<typename T, typename U> struct convertor {
|
||||
/**
|
||||
* @brief Converts a value of type U to type T.
|
||||
*
|
||||
* @param u[in] The value of type U to convert.
|
||||
* @return T The converted value of type T.
|
||||
*/
|
||||
static __host__ __device__ inline T convert(const U & u) {
|
||||
return (T)u;
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float, bf16> {
|
||||
static __host__ __device__ inline float convert(const bf16 & u) {
|
||||
return __bfloat162float(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<bf16, float> {
|
||||
static __host__ __device__ inline bf16 convert(const float & u) {
|
||||
return __float2bfloat16_rn(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float2, bf16_2> {
|
||||
static __host__ __device__ inline float2 convert(const bf16_2 & u) {
|
||||
return __bfloat1622float2(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<bf16_2, float2> {
|
||||
static __host__ __device__ inline bf16_2 convert(const float2 & u) {
|
||||
return __float22bfloat162_rn(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float, half> {
|
||||
static __host__ __device__ inline float convert(const half & u) {
|
||||
return __half2float(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<half, float> {
|
||||
static __host__ __device__ inline half convert(const float & u) {
|
||||
return __float2half(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float2, half_2> {
|
||||
static __host__ __device__ inline float2 convert(const half_2 & u) {
|
||||
return __half22float2(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<half_2, float2> {
|
||||
static __host__ __device__ inline half_2 convert(const float2 & u) {
|
||||
return __float22half2_rn(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<bf16, half> {
|
||||
static __host__ __device__ inline bf16 convert(const half & u) {
|
||||
return __float2bfloat16_rn(__half2float(u));
|
||||
}
|
||||
};
|
||||
template<> struct convertor<half, bf16> {
|
||||
static __host__ __device__ inline half convert(const bf16 & u) {
|
||||
return __float2half(__bfloat162float(u));
|
||||
}
|
||||
};
|
||||
template<> struct convertor<bf16_2, half_2> {
|
||||
static __host__ __device__ inline bf16_2 convert(const half_2 & u) {
|
||||
return __float22bfloat162_rn(__half22float2(u));
|
||||
}
|
||||
};
|
||||
template<> struct convertor<half_2, bf16_2> {
|
||||
static __host__ __device__ inline half_2 convert(const bf16_2 & u) {
|
||||
return __float22half2_rn(__bfloat1622float2(u));
|
||||
}
|
||||
};
|
||||
#ifdef KITTENS_HOPPER
|
||||
// fp8e4m3
|
||||
template<> struct convertor<fp8e4m3_4, float4> {
|
||||
static __host__ __device__ inline fp8e4m3_4 convert(const float4& u) {
|
||||
return __nv_fp8x4_e4m3(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float4, fp8e4m3_4> {
|
||||
static __host__ __device__ inline float4 convert(const fp8e4m3_4& u) {
|
||||
__nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x4_e4m3*>(&u));
|
||||
return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
|
||||
}
|
||||
};
|
||||
template<> struct convertor<fp8e4m3_2, float2> {
|
||||
static __host__ __device__ inline fp8e4m3_2 convert(const float2& u) {
|
||||
return __nv_fp8x2_e4m3(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float2, fp8e4m3_2> {
|
||||
static __host__ __device__ inline float2 convert(const fp8e4m3_2& u) {
|
||||
__nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x2_e4m3*>(&u));
|
||||
return make_float2(float(vals[0]), float(vals[1]));
|
||||
}
|
||||
};
|
||||
template<> struct convertor<fp8e4m3, float> {
|
||||
static __host__ __device__ inline fp8e4m3 convert(const float & u) {
|
||||
return __nv_fp8_e4m3(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float, fp8e4m3> {
|
||||
static __host__ __device__ inline float convert(const fp8e4m3 & u) {
|
||||
return float(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<bf16_2, fp8e4m3_4> {
|
||||
static __host__ __device__ inline bf16_2 convert(const fp8e4m3_4 & u) {
|
||||
float4 f4 = convertor<float4, fp8e4m3_4>::convert(u);
|
||||
float2 f2 = make_float2(f4.x, f4.y);
|
||||
return __float22bfloat162_rn(f2);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<fp8e4m3_4, bf16_2> {
|
||||
static __host__ __device__ inline fp8e4m3_4 convert(const bf16_2 & u) {
|
||||
float2 f2 = __bfloat1622float2(u);
|
||||
float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
|
||||
return __nv_fp8x4_e4m3(f4);
|
||||
}
|
||||
};
|
||||
// fp8e5m2
|
||||
template<> struct convertor<fp8e5m2_4, float4> {
|
||||
static __host__ __device__ inline fp8e5m2_4 convert(const float4& u) {
|
||||
return __nv_fp8x4_e5m2(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float4, fp8e5m2_4> {
|
||||
static __host__ __device__ inline float4 convert(const fp8e5m2_4& u) {
|
||||
__nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x4_e5m2*>(&u));
|
||||
return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
|
||||
}
|
||||
};
|
||||
template<> struct convertor<fp8e5m2_2, float2> {
|
||||
static __host__ __device__ inline fp8e5m2_2 convert(const float2& u) {
|
||||
return __nv_fp8x2_e5m2(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float2, fp8e5m2_2> {
|
||||
static __host__ __device__ inline float2 convert(const fp8e5m2_2& u) {
|
||||
__nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x2_e5m2*>(&u));
|
||||
return make_float2(float(vals[0]), float(vals[1]));
|
||||
}
|
||||
};
|
||||
template<> struct convertor<fp8e5m2, float> {
|
||||
static __host__ __device__ inline fp8e5m2 convert(const float & u) {
|
||||
return __nv_fp8_e5m2(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<float, fp8e5m2> {
|
||||
static __host__ __device__ inline float convert(const fp8e5m2 & u) {
|
||||
return float(u);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<bf16_2, fp8e5m2_4> {
|
||||
static __host__ __device__ inline bf16_2 convert(const fp8e5m2_4 & u) {
|
||||
float4 f4 = convertor<float4, fp8e5m2_4>::convert(u);
|
||||
float2 f2 = make_float2(f4.x, f4.y);
|
||||
return __float22bfloat162_rn(f2);
|
||||
}
|
||||
};
|
||||
template<> struct convertor<fp8e5m2_4, bf16_2> {
|
||||
static __host__ __device__ inline fp8e5m2_4 convert(const bf16_2 & u) {
|
||||
float2 f2 = __bfloat1622float2(u);
|
||||
float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
|
||||
return __nv_fp8x4_e5m2(f4);
|
||||
}
|
||||
};
|
||||
#endif
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief A collection of common resources on which ThunderKittens depends.
|
||||
*/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "util.cuh"
|
||||
#include "base_types.cuh"
|
||||
#include "base_ops.cuh"
|
||||
@@ -0,0 +1,56 @@
|
||||
#pragma once
|
||||
|
||||
// Reset
|
||||
#define TK_RESET "\033[0m"
|
||||
|
||||
// Foreground colors
|
||||
#define TK_FG_BLACK "\033[30m"
|
||||
#define TK_FG_RED "\033[31m"
|
||||
#define TK_FG_GREEN "\033[32m"
|
||||
#define TK_FG_YELLOW "\033[33m"
|
||||
#define TK_FG_BLUE "\033[34m"
|
||||
#define TK_FG_MAGENTA "\033[35m"
|
||||
#define TK_FG_CYAN "\033[36m"
|
||||
#define TK_FG_WHITE "\033[37m"
|
||||
|
||||
// Background colors
|
||||
#define TK_BG_BLACK "\033[40m"
|
||||
#define TK_BG_RED "\033[41m"
|
||||
#define TK_BG_GREEN "\033[42m"
|
||||
#define TK_BG_YELLOW "\033[43m"
|
||||
#define TK_BG_BLUE "\033[44m"
|
||||
#define TK_BG_MAGENTA "\033[45m"
|
||||
#define TK_BG_CYAN "\033[46m"
|
||||
#define TK_BG_WHITE "\033[47m"
|
||||
|
||||
// Bright foreground colors
|
||||
#define TK_FG_BRIGHT_BLACK "\033[90m"
|
||||
#define TK_FG_BRIGHT_RED "\033[91m"
|
||||
#define TK_FG_BRIGHT_GREEN "\033[92m"
|
||||
#define TK_FG_BRIGHT_YELLOW "\033[93m"
|
||||
#define TK_FG_BRIGHT_BLUE "\033[94m"
|
||||
#define TK_FG_BRIGHT_MAGENTA "\033[95m"
|
||||
#define TK_FG_BRIGHT_CYAN "\033[96m"
|
||||
#define TK_FG_BRIGHT_WHITE "\033[97m"
|
||||
|
||||
// Bright background colors
|
||||
#define TK_BG_BRIGHT_BLACK "\033[100m"
|
||||
#define TK_BG_BRIGHT_RED "\033[101m"
|
||||
#define TK_BG_BRIGHT_GREEN "\033[102m"
|
||||
#define TK_BG_BRIGHT_YELLOW "\033[103m"
|
||||
#define TK_BG_BRIGHT_BLUE "\033[104m"
|
||||
#define TK_BG_BRIGHT_MAGENTA "\033[105m"
|
||||
#define TK_BG_BRIGHT_CYAN "\033[106m"
|
||||
#define TK_BG_BRIGHT_WHITE "\033[107m"
|
||||
|
||||
// Text styles
|
||||
#define TK_BOLD "\033[1m"
|
||||
#define TK_DIM "\033[2m"
|
||||
#define TK_ITALIC "\033[3m"
|
||||
#define TK_UNDERLINE "\033[4m"
|
||||
#define TK_BLINK "\033[5m"
|
||||
#define TK_REVERSE "\033[7m"
|
||||
#define TK_HIDDEN "\033[8m"
|
||||
|
||||
// Macro to combine styles
|
||||
#define TK_STYLE(...) "\033[" #__VA_ARGS__ "m"
|
||||
@@ -0,0 +1,314 @@
|
||||
/**
|
||||
* @file
|
||||
* @brief General utilities for ThunderKittens.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <stdint.h>
|
||||
#include <type_traits>
|
||||
#include <concepts>
|
||||
#include <memory>
|
||||
|
||||
// CUDA driver API
|
||||
#define CUCHECK(cmd) do { \
|
||||
CUresult err = cmd; \
|
||||
if (err != CUDA_SUCCESS) { \
|
||||
const char *errStr; \
|
||||
cuGetErrorString(err, &errStr); \
|
||||
fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \
|
||||
__FILE__, __LINE__, errStr); \
|
||||
exit(EXIT_FAILURE); \
|
||||
} \
|
||||
} while(0)
|
||||
|
||||
// CUDA runtime API
|
||||
#define CUDACHECK(cmd) do { \
|
||||
cudaError_t err = cmd; \
|
||||
if (err != cudaSuccess) { \
|
||||
fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \
|
||||
__FILE__, __LINE__, cudaGetErrorString(err)); \
|
||||
exit(EXIT_FAILURE); \
|
||||
} \
|
||||
} while(0)
|
||||
|
||||
/**
|
||||
* @namespace kittens
|
||||
*
|
||||
* @brief The main namespace of ThunderKittens.
|
||||
*/
|
||||
namespace kittens {
|
||||
|
||||
/* ---------- GENERAL CONSTANTS FOR KITTENS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Tile dimension constant.
|
||||
*/
|
||||
template<typename T> constexpr int TILE_COL_DIM = sizeof(T) == 1 ? 32 : 16;
|
||||
template<typename T> constexpr int TILE_ROW_DIM = 16;
|
||||
/**
|
||||
* @brief Tile num elements constant calculated as TILE_DIM squared.
|
||||
*/
|
||||
template<typename T> constexpr int TILE_ELEMENTS{TILE_COL_DIM<T>*TILE_ROW_DIM<T>};
|
||||
/**
|
||||
* @brief Constant representing number of threads in a warp.
|
||||
*/
|
||||
constexpr int WARP_THREADS{32};
|
||||
/**
|
||||
* @brief Constant representing number of threads in a warpgroup of four warps.
|
||||
*/
|
||||
constexpr int WARPGROUP_THREADS{128};
|
||||
/**
|
||||
|
||||
* @brief Constant representing number of warps in a warpgroup of four warps.
|
||||
*/
|
||||
constexpr int WARPGROUP_WARPS{4};
|
||||
/**
|
||||
|
||||
* @brief Get the warp ID of the current thread.
|
||||
* @return The warp ID.
|
||||
*/
|
||||
__device__ static __forceinline__ int warpid() {
|
||||
// uint32_t wid;
|
||||
// asm volatile("mov.u32 %0, %warpid;" : "=r"(wid));
|
||||
// return wid;
|
||||
return threadIdx.x >> 5;
|
||||
}
|
||||
/**
|
||||
* @brief Get the warpgroup ID of the current thread.
|
||||
* @return The warpgroup ID.
|
||||
*/
|
||||
__device__ static __forceinline__ int warpgroupid() { return warpid() >> 2; }
|
||||
/**
|
||||
* @brief Get the lane ID of the current thread within its warp.
|
||||
* @return The lane ID.
|
||||
*/
|
||||
__device__ static __forceinline__ int laneid() {
|
||||
// uint32_t lid;
|
||||
// asm volatile("mov.u32 %0, %laneid;" : "=r"(lid));
|
||||
// return lid;
|
||||
return threadIdx.x & 31;
|
||||
}
|
||||
|
||||
#if defined(KITTENS_HOPPER)
|
||||
constexpr int MAX_SHARED_MEMORY = 227000;
|
||||
#elif defined(KITTENS_A100)
|
||||
constexpr int MAX_SHARED_MEMORY = 164000;
|
||||
#elif defined(KITTENS_4090)
|
||||
constexpr int MAX_SHARED_MEMORY = 100000;
|
||||
#endif
|
||||
|
||||
struct transpose {
|
||||
static constexpr int N = 0; // not transposed
|
||||
static constexpr int T = 1; // transposed
|
||||
};
|
||||
struct axis {
|
||||
static constexpr int ROW = 0; // row axis of a tile
|
||||
static constexpr int COL = 1; // column axis of a tile
|
||||
};
|
||||
|
||||
/* ---------- TYPE HELPERS ---------- */
|
||||
|
||||
/**
|
||||
* @namespace ducks
|
||||
*
|
||||
* @brief ThunderKittens' namespace for template metaprogramming..
|
||||
*
|
||||
* This includes primarily dummy types and concept wrappers, along
|
||||
* with a few additional utilities.
|
||||
*/
|
||||
namespace ducks {
|
||||
|
||||
/**
|
||||
* @brief A type representing an empty default for a template.
|
||||
*/
|
||||
struct default_type {};
|
||||
|
||||
// This macro can't be done as a template, so it doesn't really have a location in kittens.
|
||||
#define typeof(A) typename std::remove_const<typename std::remove_reference<decltype(A)>::type>::type
|
||||
|
||||
}
|
||||
|
||||
/* ---------- SHUFFLE UTILS ---------- */
|
||||
|
||||
/**
|
||||
* @brief Mask constant for all active threads in a warp.
|
||||
*/
|
||||
static constexpr uint32_t MASK_ALL = 0xFFFFFFFF;
|
||||
|
||||
/**
|
||||
* @brief Perform a shuffle down operation on a packed type synchronously across a warp.
|
||||
* @tparam T The type of the value to be shuffled.
|
||||
* @param mask[in] The mask of active threads.
|
||||
* @param f[in] The value to be shuffled.
|
||||
* @param delta[in] The number of positions to shuffle down.
|
||||
* @return The result of the shuffle operation.
|
||||
*/
|
||||
template<typename T>
|
||||
__device__ static inline T packed_shfl_down_sync(uint32_t mask, const T &f, int delta) {
|
||||
return __shfl_down_sync(mask, f, delta);
|
||||
}
|
||||
template<>
|
||||
__device__ inline float2 packed_shfl_down_sync<float2>(uint32_t mask, const float2 &f, int delta) {
|
||||
float2 r;
|
||||
r.x = __shfl_down_sync(mask, f.x, delta);
|
||||
r.y = __shfl_down_sync(mask, f.y, delta);
|
||||
return r;
|
||||
}
|
||||
/**
|
||||
* @brief Perform a packed shuffle operation synchronously across a warp.
|
||||
* @tparam T The type of the value to be shuffled.
|
||||
* @param mask[in] The mask of active threads.
|
||||
* @param f[in] The value to be shuffled.
|
||||
* @param src[in] The source lane from which to shuffle.
|
||||
* @return The result of the shuffle operation.
|
||||
*/
|
||||
template<typename T>
|
||||
__device__ static inline T packed_shfl_sync(uint32_t mask, const T &f, int src) {
|
||||
return __shfl_sync(mask, f, src);
|
||||
}
|
||||
template<>
|
||||
__device__ inline float2 packed_shfl_sync<float2>(uint32_t mask, const float2 &f, int src) {
|
||||
float2 r;
|
||||
r.x = __shfl_sync(mask, f.x, src);
|
||||
r.y = __shfl_sync(mask, f.y, src);
|
||||
return r;
|
||||
}
|
||||
|
||||
/* ---------- SHARED MEMORY UTILS ---------- */
|
||||
|
||||
// namespace ducks {
|
||||
// namespace sb {
|
||||
// struct identifier {};
|
||||
// }
|
||||
// }
|
||||
|
||||
// template<typename Args...>
|
||||
// struct sb {
|
||||
// using identifier = ducks::sb::identifier;
|
||||
// Args... args;
|
||||
// };
|
||||
|
||||
// namespace ducks {
|
||||
// namespace sb {
|
||||
// template<typename T> concept all = requires {
|
||||
// typename T::identifier;
|
||||
// } && std::is_same_v<T::identifier, identifier>;
|
||||
// }
|
||||
// }
|
||||
|
||||
// Joyously stolen from https://github.com/NVIDIA/cutlass/blob/5c447dd84f8ae0e1d48ff9a2eae26ce8c4958101/include/cute/container/alignment.hpp#L51
|
||||
#if defined(__CUDACC__)
|
||||
#define KITTENS_ALIGN_AS(n) __align__(n)
|
||||
#else
|
||||
#define KITTENS_ALIGN_AS(n) alignas(n)
|
||||
#endif
|
||||
|
||||
#ifdef KITTENS_HOPPER
|
||||
#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(128)
|
||||
#else
|
||||
#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(16)
|
||||
#endif
|
||||
|
||||
/**
|
||||
* @brief Dummy structure for alignment purposes. Needed for WGMMA and TMA calls.
|
||||
*/
|
||||
struct KITTENS_DEFAULT_ALIGN alignment_dummy { int dummy; };
|
||||
/**
|
||||
* @brief Very simple allocator for dynamic shared memory. Advances pointer and tracks alignments.
|
||||
* @tparam default_alignment The default alignment this allocator will enforce. If <=0 (default -1) it will not align.
|
||||
*/
|
||||
#ifdef KITTENS_HOPPER
|
||||
template<int default_alignment=1024>
|
||||
#else
|
||||
template<int default_alignment=16>
|
||||
#endif
|
||||
struct shared_allocator {
|
||||
int *ptr;
|
||||
|
||||
private:
|
||||
// Recursive template to generate N-dimensional array type
|
||||
template<typename A, size_t... dims>
|
||||
struct variadic_array;
|
||||
template<typename A, size_t first_dim, size_t... rest_dims>
|
||||
struct variadic_array<A, first_dim, rest_dims...> {
|
||||
using type = typename variadic_array<A, rest_dims...>::type[first_dim];
|
||||
};
|
||||
template<typename A>
|
||||
struct variadic_array<A> {
|
||||
using type = A;
|
||||
};
|
||||
template<typename A, size_t... dims>
|
||||
using variadic_array_t = typename variadic_array<A, dims...>::type;
|
||||
|
||||
template<int alignment>
|
||||
__device__ inline void align_ptr() {
|
||||
if constexpr (alignment > 0) {
|
||||
uint64_t p = reinterpret_cast<uint64_t>(ptr);
|
||||
if(p % alignment != 0) {
|
||||
ptr = (int*)(p + (alignment-(p%alignment)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
/**
|
||||
* @brief Construct a new shared allocator using a pointer to extern shared memory.
|
||||
* @param[in] _ptr Pointer to the start of the extern shared memory.
|
||||
*/
|
||||
__device__ shared_allocator(int *_ptr): ptr(_ptr) {}
|
||||
/**
|
||||
* @brief Allocate shared memory for a single instance or N-dimensional array of type A.
|
||||
* @tparam A The type of the object to allocate.
|
||||
* @tparam dims... A list of dimensions for the N-dimensional array.
|
||||
* @return Reference to the allocated object.
|
||||
*/
|
||||
template<typename A, size_t... dims>
|
||||
__device__ inline variadic_array_t<A, dims...>& allocate() {
|
||||
// static_assert(sizeof(A) % default_alignment == 0, "Type is not aligned properly for array allocation");
|
||||
align_ptr<default_alignment>();
|
||||
using at = variadic_array_t<A, dims...>;
|
||||
at*p = reinterpret_cast<at*>(ptr);
|
||||
ptr += sizeof(at)/sizeof(int);
|
||||
return *p;
|
||||
}
|
||||
/**
|
||||
* @brief Allocate shared memory for a single instance or N-dimensional array of type A.
|
||||
* @tparam alignment An alignment to enforce for this particular object.
|
||||
* @tparam A The type of the object to allocate.
|
||||
* @tparam dims... A list of dimensions for the N-dimensional array.
|
||||
* @return Reference to the allocated object.
|
||||
*/
|
||||
template<int alignment, typename A, size_t... dims>
|
||||
__device__ inline variadic_array_t<A, dims...>& allocate() {
|
||||
// static_assert(sizeof(A) % alignment == 0, "Type is not aligned properly for array allocation");
|
||||
align_ptr<alignment>();
|
||||
using at = variadic_array_t<A, dims...>;
|
||||
at*p = reinterpret_cast<at*>(ptr);
|
||||
ptr += sizeof(at)/sizeof(int);
|
||||
return *p;
|
||||
}
|
||||
};
|
||||
#if (defined(KITTENS_HOPPER) || defined(KITTENS_BLACKWELL))
|
||||
/**
|
||||
* @brief A wrapper for an allocator that enforces sufficient alignment to be used for TMA loads and stores.
|
||||
*/
|
||||
using tma_allocator = shared_allocator<1024>;
|
||||
using tma_swizzle_allocator = tma_allocator; // swizzled TMA modes require up to 1024 byte alignments :/
|
||||
|
||||
/* Get CTA ID within a cluster */
|
||||
__device__ static inline int3 clusterIdx() {
|
||||
int3 cluster_idx;
|
||||
asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(cluster_idx.x));
|
||||
asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(cluster_idx.y));
|
||||
asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(cluster_idx.z));
|
||||
return cluster_idx;
|
||||
}
|
||||
__device__ static inline int cluster_ctarank() {
|
||||
uint32_t ctarank;
|
||||
asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(ctarank));
|
||||
return ctarank;
|
||||
}
|
||||
#endif
|
||||
|
||||
} // namespace kittens
|
||||
Reference in New Issue
Block a user