diff --git a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh b/python/sglang/jit_kernel/include/sgl_kernel/math.cuh
index 97e6cf637..2287b31e2 100644
--- a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh
+++ b/python/sglang/jit_kernel/include/sgl_kernel/math.cuh
@@ -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
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh b/python/sglang/jit_kernel/include/sgl_kernel/type.cuh
index 4b6723855..f06bc1407 100644
--- a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh
+++ b/python/sglang/jit_kernel/include/sgl_kernel/type.cuh
@@ -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);