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

.github/workflows/pr-test-amd.yml

@@ -146,6 +146,7 @@ jobs:
           docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_activation.py
           docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_topk.py
           docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_kvcacheio.py
+          docker exec -w /sglang-checkout/sgl-kernel/tests/sgl_diffusion ci_sglang python3 -m pytest test_timestep_embedding.py
 
   # =============================================== primary ====================================================
 

docker/rocm.Dockerfile

@@ -21,7 +21,6 @@ ENV BUILD_LLVM="0"
 ENV BUILD_AITER_ALL="1"
 ENV BUILD_MOONCAKE="1"
 ENV AITER_COMMIT="v0.1.4"
-ENV NO_DEPS_FLAG=""
 
 # ===============================
 # Base image 942 and args
@@ -32,7 +31,6 @@ ENV BUILD_LLVM="0"
 ENV BUILD_AITER_ALL="1"
 ENV BUILD_MOONCAKE="1"
 ENV AITER_COMMIT="v0.1.9.post1"
-ENV NO_DEPS_FLAG=""
 
 # ===============================
 # Base image 950 and args
@@ -43,7 +41,6 @@ ENV BUILD_LLVM="0"
 ENV BUILD_AITER_ALL="0"
 ENV BUILD_MOONCAKE="1"
 ENV AITER_COMMIT="v0.1.9.post1"
-ENV NO_DEPS_FLAG=""
 # ===============================
 # Chosen arch and args
 FROM ${GPU_ARCH}
@@ -187,9 +184,9 @@ RUN git clone ${SGL_REPO} \
     && cd .. \
     && rm -rf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml \
     && if [ "$BUILD_TYPE" = "srt" ]; then \
-         python -m pip --no-cache-dir install -e "python[srt_hip,diffusion]" ${NO_DEPS_FLAG}; \
+         python -m pip --no-cache-dir install -e "python[srt_hip,diffusion]"; \
        else \
-         python -m pip --no-cache-dir install -e "python[all_hip,diffusion]" ${NO_DEPS_FLAG}; \
+         python -m pip --no-cache-dir install -e "python[all_hip,diffusion]"; \
        fi
 
 RUN python -m pip cache purge

docs/platforms/amd_gpu.md

@@ -52,10 +52,10 @@ pip install --upgrade pip
 cd sgl-kernel
 python setup_rocm.py install
 
-# Install sglang python package
+# Install sglang python package along with diffusion support
 cd ..
 rm -rf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml
-pip install -e "python[all_hip]"
+pip install -e "python[all_hip,diffusion]"
 ```
 
 ### Install Using Docker (Recommended)

python/sglang/multimodal_gen/runtime/layers/visual_embedding.py

@@ -14,11 +14,16 @@ from diffusers.models.embeddings import (
 )
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding
 from diffusers.models.embeddings import Timesteps as _Timesteps
+from diffusers.models.embeddings import (
+    get_timestep_embedding as _get_timestep_embedding,
+)
 
 try:
     from sgl_kernel.elementwise import timestep_embedding as timestep_embedding_cuda
 except Exception as _e:
-    pass
+    # Fallback to diffusers implementation so downstream code can still run
+    # even if `sgl_kernel` is not installed/available.
+    timestep_embedding_cuda = _get_timestep_embedding
 
 from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
 from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear

python/sglang/multimodal_gen/runtime/models/dits/flux.py

@@ -19,10 +19,6 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 from diffusers.models.attention import AttentionModuleMixin, FeedForward
-from diffusers.models.embeddings import (
-    CombinedTimestepGuidanceTextProjEmbeddings,
-    CombinedTimestepTextProjEmbeddings,
-)
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import (
     AdaLayerNormContinuous,
@@ -42,6 +38,10 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     NDRotaryEmbedding,
     apply_flashinfer_rope_qk_inplace,
 )
+from sglang.multimodal_gen.runtime.layers.visual_embedding import (
+    CombinedTimestepGuidanceTextProjEmbeddings,
+    CombinedTimestepTextProjEmbeddings,
+)
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin

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

sgl-kernel/include/hip/hip_math_def.h

@@ -29,7 +29,11 @@ template <typename T>
 __forceinline__ __device__ T shfl_xor_sync(unsigned mask, T var, int laneMask, int width = warpSize);
 
 template <typename srcDtype, typename destDtype>
-__forceinline__ __device__ destDtype cast(srcDtype val);
+__forceinline__ __device__ destDtype cast(srcDtype val) {
+  // Generic fallback used by most scalar types (int/float/double/etc).
+  // Specific types like fp16/bf16 have explicit specializations below.
+  return static_cast<destDtype>(val);
+}
 
 // specialization
 template <>
@@ -43,27 +47,27 @@ __forceinline__ __device__ int shfl_xor_sync(unsigned mask, int var, int laneMas
 }
 
 template <>
-__forceinline__ __device__ float cast(float val) {
+__forceinline__ __device__ float cast<float, float>(float val) {
   return val;
 }
 
 template <>
-__forceinline__ __device__ float cast(__half val) {
+__forceinline__ __device__ float cast<__half, float>(__half val) {
   return __half2float(val);
 }
 
 template <>
-__forceinline__ __device__ float cast(__hip_bfloat16 val) {
+__forceinline__ __device__ float cast<__hip_bfloat16, float>(__hip_bfloat16 val) {
   return __bfloat162float(val);
 }
 
 template <>
-__forceinline__ __device__ __half cast(float fval) {
+__forceinline__ __device__ __half cast<float, __half>(float fval) {
   return __float2half(fval);
 }
 
 template <>
-__forceinline__ __device__ __hip_bfloat16 cast(float fval) {
+__forceinline__ __device__ __hip_bfloat16 cast<float, __hip_bfloat16>(float fval) {
   return __float2bfloat16(fval);
 }
 

sgl-kernel/setup_rocm.py

@@ -54,6 +54,7 @@ sources = [
     "csrc/speculative/eagle_utils.cu",
     "csrc/kvcacheio/transfer.cu",
     "csrc/elementwise/pos_enc.cu",
+    "csrc/sgl_diffusion/elementwise/timestep_embedding.cu",
 ]
 
 cxx_flags = ["-O3"]