diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index ae0577618..48cb52cf3 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -88,6 +88,24 @@ namespace device { /// \brief Macro: forced-inline device function qualifier. #define SGL_DEVICE __forceinline__ __device__ +// Architecture detection: SGL_CUDA_ARCH is injected by load_jit() and is +// available in both host and device compilation passes, whereas __CUDA_ARCH__ +// is only defined by nvcc during the device pass. +#if !defined(USE_ROCM) +#if !defined(SGL_CUDA_ARCH) +#error "SGL_CUDA_ARCH is not defined. JIT compilation must inject -DSGL_CUDA_ARCH via load_jit()." +#endif +#if defined(__CUDA_ARCH__) +static_assert( + __CUDA_ARCH__ == SGL_CUDA_ARCH, "SGL_CUDA_ARCH mismatch: injected arch flag does not match device target"); +#endif +#define SGL_ARCH_HOPPER_OR_GREATER (SGL_CUDA_ARCH >= 900) +#define SGL_ARCH_BLACKWELL_OR_GREATER ((SGL_CUDA_ARCH >= 1000) && (CUDA_VERSION >= 12090)) +#else // USE_ROCM +#define SGL_ARCH_HOPPER_OR_GREATER 0 +#define SGL_ARCH_BLACKWELL_OR_GREATER 0 +#endif + /// \brief Number of threads per warp (always 32 on NVIDIA/AMD GPUs). inline constexpr auto kWarpThreads = 32u; /// \brief Full warp active mask (all 32 lanes). @@ -102,7 +120,7 @@ inline constexpr auto kFullMask = 0xffffffffu; */ template SGL_DEVICE void PDLWaitPrimary() { -#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +#if SGL_ARCH_HOPPER_OR_GREATER if constexpr (kUsePDL) { asm volatile("griddepcontrol.wait;" ::: "memory"); } @@ -117,7 +135,7 @@ SGL_DEVICE void PDLWaitPrimary() { */ template SGL_DEVICE void PDLTriggerSecondary() { -#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +#if SGL_ARCH_HOPPER_OR_GREATER if constexpr (kUsePDL) { asm volatile("griddepcontrol.launch_dependents;" :::); } diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index c9519a267..eb4748d71 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -188,6 +188,10 @@ def load_jit( if is_hip_runtime(): selected_cuda_cflags = DEFAULT_HIP_CFLAGS extra_cuda_cflags = ["-DUSE_ROCM"] + extra_cuda_cflags + else: + extra_cuda_cflags = [ + f"-DSGL_CUDA_ARCH={_get_cuda_arch_value()}" + ] + extra_cuda_cflags if not env_existed: os.environ[env_key] = _get_cuda_arch_list() try: @@ -215,6 +219,14 @@ def is_arch_support_pdl() -> bool: return torch.cuda.get_device_capability(device)[0] >= 9 +@cache_once +def _get_cuda_arch_value() -> int: + """Get CUDA arch value for -DSGL_CUDA_ARCH (e.g. 900 for SM 9.0).""" + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + return major * 100 + minor * 10 + + @cache_once def _get_cuda_arch_list() -> str: """Get the correct CUDA architecture string for TVM_FFI_CUDA_ARCH_LIST."""