[Diffusion] Add diffusion time embedding to jit kernel (#17658)

python/sglang/jit_kernel/include/sgl_kernel/math.cuh

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

python/sglang/jit_kernel/include/sgl_kernel/type.cuh

@@ -37,6 +37,8 @@ struct dtype_trait {};
   static_assert(true)
 
 SGL_REGISTER_DTYPE_TRAIT(fp32_t, fp32x2_t, SGL_REGISTER_TYPE_END;  //
+                         SGL_REGISTER_FROM_FUNCTION(fp16_t, __half2float);
+                         SGL_REGISTER_FROM_FUNCTION(bf16_t, __bfloat162float);
                          SGL_REGISTER_UNARY_FUNCTION(abs, fabsf);
                          SGL_REGISTER_UNARY_FUNCTION(sqrt, sqrtf);
                          SGL_REGISTER_UNARY_FUNCTION(rsqrt, rsqrtf);