From d3b972cbf0e349c1934afb4c64c2e058a9a2dcbd Mon Sep 17 00:00:00 2001 From: Yikai Zhang Date: Fri, 28 Aug 2026 01:18:42 -0700 Subject: [PATCH] fix(lora): build the MoE LoRA align JIT kernel on ROCm (#36379) Co-authored-by: amdpilot-upstream-sync --- .../jit/csrc/lora/moe_lora_align_kernel.cu | 19 ++++++++++++++++++- .../ops/moe/test_moe_lora_align_block_size.py | 3 ++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/lora/moe_lora_align_kernel.cu b/python/sglang/kernels/jit/csrc/lora/moe_lora_align_kernel.cu index e6a3e6365..0fbdb3256 100644 --- a/python/sglang/kernels/jit/csrc/lora/moe_lora_align_kernel.cu +++ b/python/sglang/kernels/jit/csrc/lora/moe_lora_align_kernel.cu @@ -6,10 +6,26 @@ #include +#ifndef USE_ROCM #include +#else +#include +#endif #include #include +#include + +#ifdef USE_ROCM +// The JIT build does not run hipify, so this file maps the CUDA spellings it +// uses onto HIP itself. Deliberately local rather than in utils.cuh: that +// header sits in nearly every kernel's dependency closure, which the build +// cache is keyed on, so editing it rebuilds every JIT module. +namespace cub = hipcub; +#define cudaDevAttrMaxSharedMemoryPerBlockOptin hipDeviceAttributeSharedMemPerBlockOptin +#define cudaFuncSetAttribute hipFuncSetAttribute +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize +#endif #ifndef WARP_SIZE #define WARP_SIZE 32 @@ -529,7 +545,8 @@ struct MoeLoraAlignBlockSizeKernel { dim3 blockDim(num_thread + fill_threads); auto kernel = moe::moe_lora_align_block_size_small_batch_expert_kernel; - RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem)); + const auto fptr = std::bit_cast(kernel); + RuntimeDeviceCheck(cudaFuncSetAttribute(fptr, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem)); LaunchKernel(dim3(max_loras), blockDim, stream, shared_mem)( kernel, diff --git a/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py b/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py index 66ef8b0fd..b0962d36d 100644 --- a/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py +++ b/test/registered/kernels/ops/moe/test_moe_lora_align_block_size.py @@ -9,9 +9,10 @@ import torch # IMPORT PREBUILT KERNEL # --------------------------------------------------------- from sglang.kernels.ops.moe.moe_lora_align import moe_lora_align_block_size -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=28, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd") def round_up(x, base):