diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 4e939cf4c..dfeebb381 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.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 ==================================================== diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 5b7ba8dbd..6dfd5ec07 100644 --- a/docker/rocm.Dockerfile +++ b/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 diff --git a/docs/platforms/amd_gpu.md b/docs/platforms/amd_gpu.md index f0f869636..068afa997 100644 --- a/docs/platforms/amd_gpu.md +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py index 7f3516a4d..71a88e788 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index b678a633a..7c98211c6 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/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 diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index 94c13fdad..d0147e5f3 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/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, ×tep_embedding); } REGISTER_EXTENSION(common_ops) diff --git a/sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu b/sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu index e239569f5..d241619e1 100644 --- a/sgl-kernel/csrc/sgl_diffusion/elementwise/timestep_embedding.cu +++ b/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 diff --git a/sgl-kernel/include/hip/hip_math_def.h b/sgl-kernel/include/hip/hip_math_def.h index 356ed953f..ac229a4be 100644 --- a/sgl-kernel/include/hip/hip_math_def.h +++ b/sgl-kernel/include/hip/hip_math_def.h @@ -29,7 +29,11 @@ template __forceinline__ __device__ T shfl_xor_sync(unsigned mask, T var, int laneMask, int width = warpSize); template -__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(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 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 fval) { return __float2half(fval); } template <> -__forceinline__ __device__ __hip_bfloat16 cast(float fval) { +__forceinline__ __device__ __hip_bfloat16 cast(float fval) { return __float2bfloat16(fval); } diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index 16a8d596d..a7336a103 100644 --- a/sgl-kernel/setup_rocm.py +++ b/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"]