[AMD][Diffusion] support timestep embedding kernel for AMD GPUs (#16766)

sgl-kernel/csrc/common_extension_rocm.cc

@@ -219,6 +219,18 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
       "                 Tensor!? key, int head_size,"
       "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   m.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
+  /*
+   * From csrc/sgl_diffusion/elementwise
+   */
+  m.def(
+      "timestep_embedding(Tensor input,"
+      "Tensor output,"
+      "int dim,"
+      "bool flip_sin_to_cos,"
+      "float downscale_freq_shift,"
+      "float scale,"
+      "int max_period) -> Tensor");
+  m.impl("timestep_embedding", torch::kCUDA, &timestep_embedding);
 }
 
 REGISTER_EXTENSION(common_ops)

sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu

@@ -33,7 +33,8 @@ __global__ void timestep_embedding_kernel(
   if (row_idx >= batch_size) {
     return;
   }
-  float t_val = castToFloat(__ldg(&t_ptr[row_idx]));
+  // Use the portable LDG helper (maps to __ldg on CUDA, plain load on ROCm/HIP).
+  float t_val = castToFloat(SGLANG_LDG(&t_ptr[row_idx]));
   float* output_batch_base_ptr = output_ptr + row_idx * dim;
 
   // Calculate half dimension