chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Kernel] Dispatch exp/sin/cos through dtype_trait (#19798)

Browse Source
This commit is contained in:
xingsy97
2026-03-06 22:57:52 +08:00
committed by GitHub
parent 2d266c73ea
commit 54634b9a40
2 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/jit_kernel/include/sgl_kernel/math.cuh
View File

@@ -1,8 +1,6 @@
#pragma once
#include <sgl_kernel/type.cuh>
#include <cmath>
namespace device::math {
inline constexpr float log2e = 1.44269504088896340736f;
@@ -35,16 +33,19 @@ SGL_DEVICE T rsqrt(T a) {
return dtype_trait<T>::rsqrt(a);
}
SGL_DEVICE float exp(float a) {
return ::expf(a);
template <typename T>
SGL_DEVICE T exp(T a) {
return dtype_trait<T>::exp(a);
}
SGL_DEVICE float sin(float a) {
return ::sinf(a);
template <typename T>
SGL_DEVICE T sin(T a) {
return dtype_trait<T>::sin(a);
}
SGL_DEVICE float cos(float a) {
return ::cosf(a);
template <typename T>
SGL_DEVICE T cos(T a) {
return dtype_trait<T>::cos(a);
}
} // namespace device::math

3
python/sglang/jit_kernel/include/sgl_kernel/type.cuh
View File

@@ -43,6 +43,9 @@ SGL_REGISTER_DTYPE_TRAIT(
SGL_REGISTER_UNARY_FUNCTION(abs, fabsf);
SGL_REGISTER_UNARY_FUNCTION(sqrt, sqrtf);
SGL_REGISTER_UNARY_FUNCTION(rsqrt, rsqrtf);
SGL_REGISTER_UNARY_FUNCTION(exp, expf);
SGL_REGISTER_UNARY_FUNCTION(sin, sinf);
SGL_REGISTER_UNARY_FUNCTION(cos, cosf);
SGL_REGISTER_BINARY_FUNCTION(max, fmaxf);
SGL_REGISTER_BINARY_FUNCTION(min, fminf););
SGL_REGISTER_DTYPE_TRAIT(fp16_t, fp16x2_t);