From 1307968605fc87fae3df6461e29611fbc73830ae Mon Sep 17 00:00:00 2001 From: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Date: Tue, 4 Aug 2026 08:33:54 +0800 Subject: [PATCH] [JIT] Drop redundant per-kernel arch overrides (#32952) Co-authored-by: Claude Fable 5 Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- python/sglang/kernels/jit/__main__.py | 17 ++-- python/sglang/kernels/jit/utils/arch.py | 84 ++++++++++++++++--- .../ops/attention/qprep_bf16_fp8_sm90.py | 33 ++++---- .../sparse_mla_q8kv8_prefill_sm90.py | 29 ++++--- .../kernels/ops/gemm/fp8_blockwise_gemm.py | 36 +++----- 5 files changed, 123 insertions(+), 76 deletions(-) diff --git a/python/sglang/kernels/jit/__main__.py b/python/sglang/kernels/jit/__main__.py index a7c52cf0d..b69b5fef7 100644 --- a/python/sglang/kernels/jit/__main__.py +++ b/python/sglang/kernels/jit/__main__.py @@ -7,8 +7,8 @@ import subprocess from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path -from sglang.kernels.jit.utils import get_jit_cuda_arch, override_jit_cuda_arch -from sglang.kernels.jit.utils.arch import get_default_target_flags +from sglang.kernels.jit.utils import get_jit_cuda_arch +from sglang.kernels.jit.utils.arch import get_default_target_flags, make_jit_cuda_arch from sglang.kernels.jit.utils.compile import DEFAULT_INCLUDE from sglang.kernels.jit.utils.deps import REGISTERED_DEPENDENCIES @@ -69,27 +69,24 @@ def generate_clangd(): if args.cuda_target: assert args.cuda_target.count(".") == 1 major, minor = args.cuda_target.split(".") - major, minor = int(major), int(minor) - context = override_jit_cuda_arch(major, minor) - context.__enter__() + arch = make_jit_cuda_arch(int(major), int(minor)) else: arch = get_jit_cuda_arch() - major, minor = arch.major, f"{arch.minor}{arch.suffix}" assert ( - major > 0 + arch.major > 0 ), "Cannot detect CUDA architecture, please specify --cuda-target explicitly." compile_flags = [ "-xcuda", - f"--cuda-gpu-arch=sm_{major}{minor}", + f"--cuda-gpu-arch=sm_{arch.major}{arch.minor}{arch.suffix}", "-Wall", "-Wextra", - *get_default_target_flags(), + *get_default_target_flags(arch=arch), *[f"-isystem{path}" for path in include_paths], ] # NOTE: for local clangd (fix the missing cluster related macros) - if major >= 9: + if arch.major >= 9: compile_flags.append("-D_CG_LIMIT_INCLUDED_DEPENDENCIES=1") compile_flags.append("-D_CG_HAS_CLUSTER_GROUP=1") diff --git a/python/sglang/kernels/jit/utils/arch.py b/python/sglang/kernels/jit/utils/arch.py index 107dc24f1..9c68aacff 100644 --- a/python/sglang/kernels/jit/utils/arch.py +++ b/python/sglang/kernels/jit/utils/arch.py @@ -3,6 +3,10 @@ from __future__ import annotations import logging +import os +import re +import shutil +import subprocess from contextlib import contextmanager from dataclasses import dataclass from typing import List @@ -14,7 +18,6 @@ from sglang.kernels.jit.utils.common import ( is_hip_runtime, is_musa_runtime, ) -from sglang.srt.utils.common import get_cuda_version logger = logging.getLogger(__name__) @@ -34,18 +37,55 @@ class ArchInfo: return f"-DSGL_CUDA_ARCH={self.major * 100 + self.minor * 10}" +@cache_once +def _jit_cuda_version() -> tuple[int, ...]: + """CUDA version of the nvcc that JIT builds actually run. + + The target has to match the compiler, not the toolkit PyTorch was built + against: a cu129 wheel on a CUDA 12.8 toolkit would otherwise select + `sm_120f`, which nvcc 12.8 rejects. Resolve nvcc the way tvm-ffi does + (`CUDA_HOME` / `CUDA_PATH`, then `$PATH`, then `/usr/local/cuda`) and fall + back to `torch.version.cuda` when it cannot be probed. + """ + cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") + if cuda_home is None: + nvcc_path = shutil.which("nvcc") + cuda_home = ( + os.path.dirname(os.path.dirname(nvcc_path)) + if nvcc_path is not None + else "/usr/local/cuda" + ) + nvcc = os.path.join(cuda_home, "bin", "nvcc") + try: + output = subprocess.check_output([nvcc, "--version"], text=True) + match = re.search(r"release (\d+)\.(\d+)", output) + if match is not None: + return int(match.group(1)), int(match.group(2)) + logger.warning("Cannot parse `%s --version` output: %s", nvcc, output) + except (OSError, subprocess.SubprocessError) as error: + logger.warning("Cannot run `%s --version`: %s", nvcc, error) + from sglang.srt.utils.common import get_cuda_version + + return get_cuda_version() + + def _cuda_arch_suffix(major: int, minor: int) -> str: """Mirror FlashInfer's `_normalize_cuda_arch`: 9.x/10.x+ -> "a"; 12.0 -> "f" and 12.x (x>0) -> "a" (SM120/SM121 need separate cubins to avoid - cudaErrorIllegalInstruction, requires CUDA >= 12.9); below 9.0 -> plain. - Unlike FlashInfer, pre-12.9 CUDA falls back to plain instead of raising. + cudaErrorIllegalInstruction); below 9.0 -> plain. + + The family-specific "f" target needs a CUDA >= 12.9 nvcc; older toolkits + fall back to "a", which SM120 has had since 12.8. Plain sm_120 is never a + valid fallback: without an a/f target CUTLASS's SM120 atoms lose LDSM/STSM + and compile down to trap stubs (verified via SASS), which asserts at launch + instead of failing the build. """ if major == 9: return "a" if major == 12: - if get_cuda_version() < (12, 9): - return "" - return "f" if minor == 0 else "a" + if minor == 0 and _jit_cuda_version() >= (12, 9): + return "f" + return "a" if major >= 10: return "a" return "" @@ -71,7 +111,8 @@ def _init_jit_cuda_arch_once(): _CUDA_ARCH = ArchInfo(major, minor, suffix) -def get_default_target_flags() -> List[str]: +def get_default_target_flags(arch: ArchInfo | None = None) -> List[str]: + """Default compile flags for `arch`, defaulting to the detected local GPU.""" if is_hip_runtime(): flags = ["-DUSE_ROCM", "-std=c++20", "-O3"] # Detect FP8 type based on GPU architecture @@ -86,20 +127,41 @@ def get_default_target_flags() -> List[str]: flags.append("-DHIP_FP8_TYPE_E4M3=1") return flags else: + if arch is None: + arch = get_jit_cuda_arch() return [ - get_jit_cuda_arch().jit_flag, + arch.jit_flag, "-std=c++20", "-O3", "--expt-relaxed-constexpr", ] +def make_jit_cuda_arch(major: int, minor: int) -> ArchInfo: + """Build the JIT target for an explicitly requested capability.""" + return ArchInfo(major, minor, _cuda_arch_suffix(major, minor)) + + @contextmanager -def override_jit_cuda_arch(major: int, minor: int, suffix: str = ""): - """A context manager to temporarily override CUDA architecture.""" +def override_jit_cuda_arch(major: int, minor: int, suffix: str | None = None): + """A context manager to temporarily override CUDA architecture. + + `suffix` defaults to the arch-specific target detection would pick for that + capability; pass it explicitly only to force a different one (an unsuffixed + target loses the arch-only instructions CUTLASS needs, see + `_cuda_arch_suffix`). + + Kernels do not need this to reach an arch-specific target: `get_jit_cuda_arch` + already resolves the local GPU to its a/f target. Reach for it only to compile + for an arch the local GPU is not. + """ global _CUDA_ARCH old_value = get_jit_cuda_arch() - _CUDA_ARCH = ArchInfo(major, minor, suffix) + _CUDA_ARCH = ( + make_jit_cuda_arch(major, minor) + if suffix is None + else ArchInfo(major, minor, suffix) + ) try: yield finally: diff --git a/python/sglang/kernels/ops/attention/qprep_bf16_fp8_sm90.py b/python/sglang/kernels/ops/attention/qprep_bf16_fp8_sm90.py index 91faf9ec7..a8557c1af 100644 --- a/python/sglang/kernels/ops/attention/qprep_bf16_fp8_sm90.py +++ b/python/sglang/kernels/ops/attention/qprep_bf16_fp8_sm90.py @@ -15,7 +15,7 @@ from typing import TYPE_CHECKING import torch from sglang.kernel_api_logging import debug_kernel_api -from sglang.kernels.jit.utils import cache_once, load_jit, override_jit_cuda_arch +from sglang.kernels.jit.utils import cache_once, load_jit if TYPE_CHECKING: from tvm_ffi.module import Module @@ -29,22 +29,21 @@ ROPE_DIM = 64 # qk_rope_head_dim def _jit_qprep_bf16_fp8_module() -> Module: if torch.cuda.get_device_capability()[0] != 9: raise RuntimeError("qprep_bf16_fp8_sm90 requires an SM90 (Hopper) GPU") - with override_jit_cuda_arch(9, 0, "a"): - return load_jit( - "qprep_bf16_fp8_sm90", - cuda_files=["qprep_bf16_fp8_sm90/entry.cuh"], - cuda_wrappers=[("dispatch", "qprep_bf16_fp8_dispatch")], - # Same minimal flag set as the sparse_mla_q8kv8_prefill_sm90 JIT - # build (per-flag ablation there showed the rest are no-ops). - extra_cuda_cflags=[ - "-O3", - "-DNDEBUG", - "-DCUTE_USE_PACKED_TUPLE=1", - "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", - "--use_fast_math", - ], - extra_dependencies=["cutlass"], - ) + return load_jit( + "qprep_bf16_fp8_sm90", + cuda_files=["qprep_bf16_fp8_sm90/entry.cuh"], + cuda_wrappers=[("dispatch", "qprep_bf16_fp8_dispatch")], + # Same minimal flag set as the sparse_mla_q8kv8_prefill_sm90 JIT + # build (per-flag ablation there showed the rest are no-ops). + extra_cuda_cflags=[ + "-O3", + "-DNDEBUG", + "-DCUTE_USE_PACKED_TUPLE=1", + "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", + "--use_fast_math", + ], + extra_dependencies=["cutlass"], + ) # torch._C._cuda_getCurrentRawStream returns the cudaStream_t pointer expected diff --git a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py index e8a8fc10a..355e0dde0 100644 --- a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py +++ b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Optional import torch from sglang.kernel_api_logging import debug_kernel_api -from sglang.kernels.jit.utils import cache_once, load_jit, override_jit_cuda_arch +from sglang.kernels.jit.utils import cache_once, load_jit from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: @@ -59,20 +59,19 @@ def _q8kv8_cuda_flags() -> list[str]: @cache_once def _jit_sparse_mla_q8kv8_prefill_module() -> Module: - with override_jit_cuda_arch(9, 0, "a"): - return load_jit( - "sparse_mla_q8kv8_prefill_sm90", - cuda_files=[ - "sparse_mla_q8kv8_prefill_sm90/entry.cuh", - ], - cuda_wrappers=[ - ("dispatch", "sparse_prefill_q8kv8_dispatch"), - ("dispatch_full", "sparse_prefill_q8kv8_dispatch_full"), - ("dispatch_topk_length", "sparse_prefill_q8kv8_dispatch_topk_length"), - ], - extra_cuda_cflags=_q8kv8_cuda_flags(), - extra_dependencies=["cutlass"], - ) + return load_jit( + "sparse_mla_q8kv8_prefill_sm90", + cuda_files=[ + "sparse_mla_q8kv8_prefill_sm90/entry.cuh", + ], + cuda_wrappers=[ + ("dispatch", "sparse_prefill_q8kv8_dispatch"), + ("dispatch_full", "sparse_prefill_q8kv8_dispatch_full"), + ("dispatch_topk_length", "sparse_prefill_q8kv8_dispatch_topk_length"), + ], + extra_cuda_cflags=_q8kv8_cuda_flags(), + extra_dependencies=["cutlass"], + ) # Pre-resolve entry-point callables on first use to avoid per-call module diff --git a/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py b/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py index 2b883ee97..e7176fa4c 100644 --- a/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py +++ b/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py @@ -1,12 +1,11 @@ from __future__ import annotations -from contextlib import contextmanager from typing import TYPE_CHECKING import torch from sglang.kernel_api_logging import debug_kernel_api -from sglang.kernels.jit.utils import cache_once, load_jit, override_jit_cuda_arch +from sglang.kernels.jit.utils import cache_once, load_jit from sglang.srt.utils.common import is_sm120_supported from sglang.srt.utils.custom_op import register_custom_op @@ -28,31 +27,22 @@ def _fp8_blockwise_cuda_flags() -> list[str]: ] -@contextmanager -def _fp8_blockwise_arch_env(): +@cache_once +def _jit_fp8_blockwise_module() -> Module: + """Compile and cache the SM120 fp8 blockwise GEMM module (handles fp16 + bf16).""" if not is_sm120_supported(): raise RuntimeError( "fp8_blockwise_scaled_mm JIT kernel requires SM120 (Blackwell)." ) - major, minor = torch.cuda.get_device_capability() - # sm_*a target (e.g. sm_120a) required, not plain sm_120. - with override_jit_cuda_arch(major, minor, suffix="a"): - yield - - -@cache_once -def _jit_fp8_blockwise_module() -> Module: - """Compile and cache the SM120 fp8 blockwise GEMM module (handles fp16 + bf16).""" - with _fp8_blockwise_arch_env(): - return load_jit( - "fp8_blockwise_scaled_mm", - cuda_files=["gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh"], - cuda_wrappers=[ - ("fp8_blockwise_scaled_mm", "fp8_blockwise_scaled_mm"), - ], - extra_dependencies=["cutlass"], - extra_cuda_cflags=_fp8_blockwise_cuda_flags(), - ) + return load_jit( + "fp8_blockwise_scaled_mm", + cuda_files=["gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh"], + cuda_wrappers=[ + ("fp8_blockwise_scaled_mm", "fp8_blockwise_scaled_mm"), + ], + extra_dependencies=["cutlass"], + extra_cuda_cflags=_fp8_blockwise_cuda_flags(), + ) @register_custom_op(