Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Kangyan-Zhou merged 4 commits into from
Mar 18, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2a0c584
feat(jit): inject target arch flag into JIT compilation
xingsy97 Mar 7, 2026
61b94bd
add static assert
xingsy97 Mar 12, 2026
bdcd40e
simplify
xingsy97 Mar 12, 2026
455c9fe
fix lint error
xingsy97 Mar 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand All @@ -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");
}
Expand All @@ -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;" :::);
}
Expand Down
12 changes: 12 additions & 0 deletions python/sglang/jit_kernel/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,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:
Expand Down Expand Up @@ -198,6 +202,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."""
Expand Down
Loading