github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

[Inference] Remove unnecessary float4_ and rename float8_ to float8 (#5679)

pull/5660/merge
Steve Luo 7 months ago committed by GitHub
parent
537a3cbc4d
commit
725fbd2ed0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 112 additions and 147 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 9
      extensions/csrc/common/data_type.h
  2. 10
      extensions/csrc/common/vec_type_traits.h
  3. 50
      extensions/csrc/funcs/binary_functor.h
  4. 99
      extensions/csrc/funcs/cast_functor.h
  5. 73
      extensions/csrc/funcs/ternary_functor.h
  6. 8
      extensions/csrc/funcs/unary_functor.h
  7. 10
      extensions/csrc/kernel/cuda/rms_layernorm_kernel.cu

9
extensions/csrc/common/data_type.h
Unescape View File

@ -40,14 +40,7 @@ struct half8 {
#endif #endif
}; };
struct float4_ { struct float8 {
#ifdef COLOSSAL_WITH_CUDA
float2 x;
float2 y;
#endif
};
struct float8_ {
#ifdef COLOSSAL_WITH_CUDA #ifdef COLOSSAL_WITH_CUDA
float2 x; float2 x;
float2 y; float2 y;

10
extensions/csrc/common/vec_type_traits.h
Unescape View File

@ -49,7 +49,7 @@ VEC_TYPE_TRAITS_SPECIALIZATION(half, 4, dtype::half4);
VEC_TYPE_TRAITS_SPECIALIZATION(half, 8, dtype::half8); VEC_TYPE_TRAITS_SPECIALIZATION(half, 8, dtype::half8);
VEC_TYPE_TRAITS_SPECIALIZATION(float, 2, float2) VEC_TYPE_TRAITS_SPECIALIZATION(float, 2, float2)
VEC_TYPE_TRAITS_SPECIALIZATION(float, 4, float4) VEC_TYPE_TRAITS_SPECIALIZATION(float, 4, float4)
VEC_TYPE_TRAITS_SPECIALIZATION(float, 8, dtype::float8_) VEC_TYPE_TRAITS_SPECIALIZATION(float, 8, dtype::float8)
#endif /* defined(COLOSSAL_WITH_CUDA) */ #endif /* defined(COLOSSAL_WITH_CUDA) */
#undef VEC_TYPE_TRAITS_SPECIALIZATION #undef VEC_TYPE_TRAITS_SPECIALIZATION
@ -64,11 +64,11 @@ VEC_TYPE_TRAITS_SPECIALIZATION(float, 8, dtype::float8_)
FLOATVEC_TYPE_TRAITS_SPECIALIZATION(float2, float2) FLOATVEC_TYPE_TRAITS_SPECIALIZATION(float2, float2)
FLOATVEC_TYPE_TRAITS_SPECIALIZATION(float4, float4) FLOATVEC_TYPE_TRAITS_SPECIALIZATION(float4, float4)
FLOATVEC_TYPE_TRAITS_SPECIALIZATION(__nv_bfloat162, float2); FLOATVEC_TYPE_TRAITS_SPECIALIZATION(__nv_bfloat162, float2);
FLOATVEC_TYPE_TRAITS_SPECIALIZATION(dtype::bfloat164, dtype::float4_); FLOATVEC_TYPE_TRAITS_SPECIALIZATION(dtype::bfloat164, float4);
FLOATVEC_TYPE_TRAITS_SPECIALIZATION(dtype::bfloat168, dtype::float8_); FLOATVEC_TYPE_TRAITS_SPECIALIZATION(dtype::bfloat168, dtype::float8);
FLOATVEC_TYPE_TRAITS_SPECIALIZATION(half2, float2); FLOATVEC_TYPE_TRAITS_SPECIALIZATION(half2, float2);
FLOATVEC_TYPE_TRAITS_SPECIALIZATION(dtype::half4, dtype::float4_); FLOATVEC_TYPE_TRAITS_SPECIALIZATION(dtype::half4, float4);
FLOATVEC_TYPE_TRAITS_SPECIALIZATION(dtype::half8, dtype::float8_); FLOATVEC_TYPE_TRAITS_SPECIALIZATION(dtype::half8, dtype::float8);
#endif /* COLOSSAL_WITH_CUDA */ #endif /* COLOSSAL_WITH_CUDA */
#undef FLOATVEC_TYPE_TRAITS_SPECIALIZATION #undef FLOATVEC_TYPE_TRAITS_SPECIALIZATION

50
extensions/csrc/funcs/binary_functor.h
Unescape View File

@ -164,22 +164,22 @@ COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION(
return mul(fa, fb); return mul(fa, fb);
})) }))
COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION( COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION(dtype::bfloat164, dtype::bfloat164,
dtype::bfloat164, dtype::bfloat164, dtype::float4_, BinaryOpType::kMul, float4, BinaryOpType::kMul, DEVICE,
DEVICE, STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float4_ fc; float4 fc;
BinaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2, CastFunctor<__nv_bfloat16, float> cast;
BinaryOpType::kMul> fc.x = cast(lhs.x.x) * cast(rhs.x.x);
mul; fc.y = cast(lhs.x.y) * cast(rhs.x.y);
fc.x = mul(lhs.x, rhs.x); fc.z = cast(lhs.y.x) * cast(rhs.y.x);
fc.y = mul(lhs.y, rhs.y); fc.w = cast(lhs.y.y) * cast(rhs.y.y);
return fc; return fc;
})) }))
COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION( COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION(
dtype::bfloat168, dtype::bfloat168, dtype::float8_, BinaryOpType::kMul, dtype::bfloat168, dtype::bfloat168, dtype::float8, BinaryOpType::kMul,
DEVICE, STMTS_WRAPPER({ DEVICE, STMTS_WRAPPER({
dtype::float8_ fc; dtype::float8 fc;
BinaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2, BinaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2,
BinaryOpType::kMul> BinaryOpType::kMul>
mul; mul;
@ -199,20 +199,22 @@ COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION(
return mul(fa, fb); return mul(fa, fb);
})) }))
COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION( COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION(dtype::half4, dtype::half4, float4,
dtype::half4, dtype::half4, dtype::float4_, BinaryOpType::kMul, DEVICE, BinaryOpType::kMul, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float4_ fc; float4 fc;
BinaryOpFunctor<half2, half2, float2, BinaryOpType::kMul> mul; CastFunctor<half, float> cast;
fc.x = mul(lhs.x, rhs.x); fc.x = cast(lhs.x.x) * cast(rhs.x.x);
fc.y = mul(lhs.y, rhs.y); fc.y = cast(lhs.x.y) * cast(rhs.x.y);
return fc; fc.z = cast(lhs.y.x) * cast(rhs.y.x);
})) fc.w = cast(lhs.y.y) * cast(rhs.y.y);
return fc;
}))
COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION( COLOSSAL_BINARY_FUNCTOR_SPECIALIZATION(
dtype::half8, dtype::half8, dtype::float8_, BinaryOpType::kMul, DEVICE, dtype::half8, dtype::half8, dtype::float8, BinaryOpType::kMul, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float8_ fc; dtype::float8 fc;
BinaryOpFunctor<half2, half2, float2, BinaryOpType::kMul> mul; BinaryOpFunctor<half2, half2, float2, BinaryOpType::kMul> mul;
fc.x = mul(lhs.x, rhs.x); fc.x = mul(lhs.x, rhs.x);
fc.y = mul(lhs.y, rhs.y); fc.y = mul(lhs.y, rhs.y);

99
extensions/csrc/funcs/cast_functor.h
Unescape View File

@ -69,14 +69,16 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float4, dtype::half4, DEVICE,
dst.y = __floats2half2_rn(val.z, val.w); dst.y = __floats2half2_rn(val.z, val.w);
return dst; return dst;
})) }))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::float4_, dtype::half4, DEVICE, COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::half4, float4, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::half4 dst; float4 dst;
dst.x = __float22half2_rn(val.x); dst.x = __half2float(val.x.x);
dst.y = __float22half2_rn(val.y); dst.y = __half2float(val.x.y);
dst.z = __half2float(val.y.x);
dst.w = __half2float(val.y.y);
return dst; return dst;
})) }))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::float8_, dtype::half8, DEVICE, COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::float8, dtype::half8, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::half8 dst; dtype::half8 dst;
dst.x = __float22half2_rn(val.x); dst.x = __float22half2_rn(val.x);
@ -107,6 +109,15 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float4, dtype::bfloat164, DEVICE,
__floats2bfloat162_rn(val.z, val.w); __floats2bfloat162_rn(val.z, val.w);
return dst; return dst;
})) }))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::bfloat164, float4, DEVICE,
STMTS_WRAPPER({
float4 dst;
dst.x = __bfloat162float(val.x.x);
dst.y = __bfloat162float(val.x.y);
dst.z = __bfloat162float(val.y.x);
dst.w = __bfloat162float(val.y.y);
return dst;
}))
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(__nv_bfloat16, __nv_bfloat162, DEVICE, COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(__nv_bfloat16, __nv_bfloat162, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
@ -120,14 +131,7 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float2, __nv_bfloat162, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
return __float22bfloat162_rn(val); return __float22bfloat162_rn(val);
})) }))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::float4_, dtype::bfloat164, DEVICE, COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::float8, dtype::bfloat168, DEVICE,
STMTS_WRAPPER({
dtype::bfloat164 dst;
dst.x = __float22bfloat162_rn(val.x);
dst.y = __float22bfloat162_rn(val.y);
return dst;
}))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::float8_, dtype::bfloat168, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::bfloat168 dst; dtype::bfloat168 dst;
dst.x = __float22bfloat162_rn(val.x); dst.x = __float22bfloat162_rn(val.x);
@ -155,14 +159,7 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float2, __nv_bfloat162, DEVICE,
val.y); val.y);
})) }))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
dtype::float4_, dtype::bfloat164, DEVICE, STMTS_WRAPPER({ dtype::float8, dtype::bfloat168, DEVICE, STMTS_WRAPPER({
dtype::bfloat164 dst;
dst.x = __floats2bfloat162_rn(val.x.x, val.x.y);
dst.y = __floats2bfloat162_rn(val.y.x, val.y.y);
return dst;
}))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
dtype::float8_, dtype::bfloat168, DEVICE, STMTS_WRAPPER({
dtype::bfloat168 dst; dtype::bfloat168 dst;
dst.x = __floats2bfloat162_rn(val.x.x, val.x.y); dst.x = __floats2bfloat162_rn(val.x.x, val.x.y);
dst.y = __floats2bfloat162_rn(val.y.x, val.y.y); dst.y = __floats2bfloat162_rn(val.y.x, val.y.y);
@ -405,35 +402,27 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float4, uint32_t, DEVICE, STMTS_WRAPPER({
(b << 8U) | a; (b << 8U) | a;
})) }))
// fp8x4 -> float4_
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint32_t, dtype::float4_, DEVICE, STMTS_WRAPPER({
dtype::float4_ res;
res.x = CastFunctor<uint16_t, float2>()(static_cast<uint16_t>(val));
res.y =
CastFunctor<uint16_t, float2>()(static_cast<uint16_t>(val >> 16U));
return res;
}))
// fp8x4 -> float4 // fp8x4 -> float4
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint32_t, float4, DEVICE, STMTS_WRAPPER({ uint32_t, float4, DEVICE, STMTS_WRAPPER({
dtype::float4_ tmp = CastFunctor<uint32_t, dtype::float4_>()(val); float4 res;
float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y); res.x = CastFunctor<uint8_t, float>()(static_cast<uint8_t>(val));
res.y = CastFunctor<uint8_t, float>()(static_cast<uint8_t>(val >> 8U));
res.z = CastFunctor<uint8_t, float>()(static_cast<uint8_t>(val >> 16U));
res.w = CastFunctor<uint8_t, float>()(static_cast<uint8_t>(val >> 24U));
return res; return res;
})) }))
// fp8x8 -> float8_ // fp8x8 -> float8
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint2, dtype::float8_, DEVICE, STMTS_WRAPPER({ uint2, dtype::float8, DEVICE, STMTS_WRAPPER({
dtype::float4_ tmp1, tmp2; dtype::float8 res;
tmp1 = CastFunctor<uint32_t, dtype::float4_>()(val.x); res.x = CastFunctor<uint16_t, float2>()(static_cast<uint16_t>(val.x));
tmp2 = CastFunctor<uint32_t, dtype::float4_>()(val.y); res.y =
dtype::float8_ res; CastFunctor<uint16_t, float2>()(static_cast<uint16_t>(val.x >> 16U));
res.x = tmp1.x; res.z = CastFunctor<uint16_t, float2>()(static_cast<uint16_t>(val.y));
res.y = tmp1.y; res.w =
res.z = tmp2.x; CastFunctor<uint16_t, float2>()(static_cast<uint16_t>(val.y >> 16U));
res.w = tmp2.y;
return res; return res;
})) }))
@ -482,34 +471,22 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float2, uint32_t, DEVICE, STMTS_WRAPPER({
return uint32; return uint32;
})) }))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::float4_, uint2, DEVICE, COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float4, uint2, DEVICE, STMTS_WRAPPER({
STMTS_WRAPPER({
uint2 b; uint2 b;
float2 c; float2 c;
c.x = val.x.x; c.x = val.x;
c.y = val.x.y; c.y = val.y;
b.x = CastFunctor<float2, uint32_t>()(c); b.x = CastFunctor<float2, uint32_t>()(c);
c.x = val.y.x; c.x = val.z;
c.y = val.y.y; c.y = val.w;
b.y = CastFunctor<float2, uint32_t>()(c); b.y = CastFunctor<float2, uint32_t>()(c);
return b; return b;
})) }))
// float4_ -> float4
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(dtype::float4_, float4, DEVICE,
STMTS_WRAPPER({
float4 b;
b.x = val.x.x;
b.y = val.x.y;
b.z = val.y.x;
b.w = val.y.y;
return b;
}))
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
dtype::float8_, uint4, DEVICE, STMTS_WRAPPER({ dtype::float8, uint4, DEVICE, STMTS_WRAPPER({
uint4 b; uint4 b;
b.x = CastFunctor<float2, uint32_t>()(val.x); b.x = CastFunctor<float2, uint32_t>()(val.x);
b.y = CastFunctor<float2, uint32_t>()(val.y); b.y = CastFunctor<float2, uint32_t>()(val.y);

73
extensions/csrc/funcs/ternary_functor.h
Unescape View File

@ -94,29 +94,27 @@ COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
return fma(cast(a), b, c); return fma(cast(a), b, c);
})) }))
COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION( COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
dtype::half4, dtype::half4, dtype::float4_, TernaryOpType::kFma, DEVICE, dtype::half4, dtype::half4, float4, TernaryOpType::kFma, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float4_ fd; float4 fd;
TernaryOpFunctor<half2, half2, float2, TernaryOpType::kFma> fma; CastFunctor<dtype::half4, float4> cast;
fd.x = fma(a.x, b.x, c.x); TernaryOpFunctor<float4, float4, float4, TernaryOpType::kFma> fma;
fd.y = fma(a.y, b.y, c.y); fd = fma(cast(a), cast(b), c);
return fd; return fd;
})) }))
COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION( COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
half, dtype::half4, dtype::float4_, TernaryOpType::kFma, DEVICE, half, dtype::half4, float4, TernaryOpType::kFma, DEVICE, STMTS_WRAPPER({
STMTS_WRAPPER({ float4 fd;
dtype::float4_ fd; CastFunctor<half, float> cast0;
CastFunctor<half, half2> cast; CastFunctor<dtype::half4, float4> cast1;
TernaryOpFunctor<half2, half2, float2, TernaryOpType::kFma> fma; TernaryOpFunctor<float, float4, float4, TernaryOpType::kFma> fma;
half2 s = cast(a); fd = fma(cast0(a), cast1(b), c);
fd.x = fma(s, b.x, c.x);
fd.y = fma(s, b.y, c.y);
return fd; return fd;
})) }))
COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION( COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
dtype::half8, dtype::half8, dtype::float8_, TernaryOpType::kFma, DEVICE, dtype::half8, dtype::half8, dtype::float8, TernaryOpType::kFma, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float8_ fd; dtype::float8 fd;
TernaryOpFunctor<half2, half2, float2, TernaryOpType::kFma> fma; TernaryOpFunctor<half2, half2, float2, TernaryOpType::kFma> fma;
fd.x = fma(a.x, b.x, c.x); fd.x = fma(a.x, b.x, c.x);
fd.y = fma(a.y, b.y, c.y); fd.y = fma(a.y, b.y, c.y);
@ -125,9 +123,9 @@ COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
return fd; return fd;
})) }))
COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION( COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
half, dtype::half8, dtype::float8_, TernaryOpType::kFma, DEVICE, half, dtype::half8, dtype::float8, TernaryOpType::kFma, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float8_ fd; dtype::float8 fd;
CastFunctor<half, half2> cast; CastFunctor<half, half2> cast;
TernaryOpFunctor<half2, half2, float2, TernaryOpType::kFma> fma; TernaryOpFunctor<half2, half2, float2, TernaryOpType::kFma> fma;
half2 s = cast(a); half2 s = cast(a);
@ -160,33 +158,28 @@ COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
return fma(cast(a), b, c); return fma(cast(a), b, c);
})) }))
COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION( COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
dtype::bfloat164, dtype::bfloat164, dtype::float4_, TernaryOpType::kFma, dtype::bfloat164, dtype::bfloat164, float4, TernaryOpType::kFma, DEVICE,
DEVICE, STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float4_ fd; float4 fd;
TernaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2, CastFunctor<dtype::bfloat164, float4> cast;
TernaryOpType::kFma> TernaryOpFunctor<float4, float4, float4, TernaryOpType::kFma> fma;
fma; fd = fma(cast(a), cast(b), c);
fd.x = fma(a.x, b.x, c.x);
fd.y = fma(a.y, b.y, c.y);
return fd; return fd;
})) }))
COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION( COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
__nv_bfloat16, dtype::bfloat164, dtype::float4_, TernaryOpType::kFma, __nv_bfloat16, dtype::bfloat164, float4, TernaryOpType::kFma, DEVICE,
DEVICE, STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float4_ fd; float4 fd;
CastFunctor<__nv_bfloat16, __nv_bfloat162> cast; CastFunctor<__nv_bfloat16, float> cast0;
TernaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2, CastFunctor<dtype::bfloat164, float4> cast1;
TernaryOpType::kFma> TernaryOpFunctor<float, float4, float4, TernaryOpType::kFma> fma;
fma; fd = fma(cast0(a), cast1(b), c);
__nv_bfloat162 s = cast(a);
fd.x = fma(s, b.x, c.x);
fd.y = fma(s, b.y, c.y);
return fd; return fd;
})) }))
COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION( COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
dtype::bfloat168, dtype::bfloat168, dtype::float8_, TernaryOpType::kFma, dtype::bfloat168, dtype::bfloat168, dtype::float8, TernaryOpType::kFma,
DEVICE, STMTS_WRAPPER({ DEVICE, STMTS_WRAPPER({
dtype::float8_ fd; dtype::float8 fd;
TernaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2, TernaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2,
TernaryOpType::kFma> TernaryOpType::kFma>
fma; fma;
@ -197,9 +190,9 @@ COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
return fd; return fd;
})) }))
COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION( COLOSSAL_TERNARY_FUNCTOR_SPECIALIZATION(
__nv_bfloat16, dtype::bfloat168, dtype::float8_, TernaryOpType::kFma, __nv_bfloat16, dtype::bfloat168, dtype::float8, TernaryOpType::kFma, DEVICE,
DEVICE, STMTS_WRAPPER({ STMTS_WRAPPER({
dtype::float8_ fd; dtype::float8 fd;
CastFunctor<__nv_bfloat16, __nv_bfloat162> cast; CastFunctor<__nv_bfloat16, __nv_bfloat162> cast;
TernaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2, TernaryOpFunctor<__nv_bfloat162, __nv_bfloat162, float2,
TernaryOpType::kFma> TernaryOpType::kFma>

8
extensions/csrc/funcs/unary_functor.h
Unescape View File

@ -52,13 +52,7 @@ COLOSSAL_UNARY_FUNCTOR_SPECIALIZATION(float2, float, UnaryOpType::kSum, DEVICE,
COLOSSAL_UNARY_FUNCTOR_SPECIALIZATION(float4, float, UnaryOpType::kSum, DEVICE, COLOSSAL_UNARY_FUNCTOR_SPECIALIZATION(float4, float, UnaryOpType::kSum, DEVICE,
{ return val.x + val.y + val.z + val.w; }) { return val.x + val.y + val.z + val.w; })
COLOSSAL_UNARY_FUNCTOR_SPECIALIZATION(dtype::float4_, float, UnaryOpType::kSum, COLOSSAL_UNARY_FUNCTOR_SPECIALIZATION(dtype::float8, float, UnaryOpType::kSum,
DEVICE, {
return val.x.x + val.x.y + val.y.x +
val.y.y;
})
COLOSSAL_UNARY_FUNCTOR_SPECIALIZATION(dtype::float8_, float, UnaryOpType::kSum,
DEVICE, { DEVICE, {
return val.x.x + val.x.y + val.y.x + return val.x.x + val.x.y + val.y.x +
val.y.y + val.z.x + val.z.y + val.y.y + val.z.x + val.z.y +

10
extensions/csrc/kernel/cuda/rms_layernorm_kernel.cu
Unescape View File

@ -283,11 +283,14 @@ void rms_layernorm(
case 4: case 4:
RMSNORM_LAUNCHER(4, block); RMSNORM_LAUNCHER(4, block);
break; break;
case 5:
RMSNORM_LAUNCHER(5, block);
break;
case 8: case 8:
RMSNORM_LAUNCHER(8, block); RMSNORM_LAUNCHER(8, block);
break; break;
default: default:
AT_ERROR("unroll_factor must be 1, 2, 3, 4 or 8"); AT_ERROR("unroll_factor must be 1, 2, 3, 4, 5 or 8");
} }
} }
} }
@ -330,11 +333,14 @@ void fused_add_rms_layernorm(
case 4: case 4:
FUSED_ADD_RMSNORM_LAUNCHER(4, block); FUSED_ADD_RMSNORM_LAUNCHER(4, block);
break; break;
case 5:
FUSED_ADD_RMSNORM_LAUNCHER(5, block);
break;
case 8: case 8:
FUSED_ADD_RMSNORM_LAUNCHER(8, block); FUSED_ADD_RMSNORM_LAUNCHER(8, block);
break; break;
default: default:
AT_ERROR("unroll_factor must be 1, 2, 3, 4 or 8"); AT_ERROR("unroll_factor must be 1, 2, 3, 4, 5 or 8");
} }
} }
} }

Write Preview
Loading…
Cancel
Save