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

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

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