chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[JIT] Inject target architecture flag into JIT compilation (#20103)

Browse Source
This commit is contained in:
xingsy97
2026-03-18 14:16:49 +08:00
committed by GitHub
parent f78d5c3b3c
commit d20e9a20fa
2 changed files with 32 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
View File

@@ -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 <bool kUsePDL>
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 <bool kUsePDL>
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;" :::);
}