chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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;" :::);
}

12
python/sglang/jit_kernel/utils.py
View File

@@ -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."""