diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index 9d6ab1e06..c36f9eeff 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -35,7 +35,8 @@ Add a new operation that scales each element of a tensor by a scalar factor: #include ``` -- **`host::RuntimeCheck(cond, args...)`** — Assert a condition at runtime; throws `PanicError` with file/line info on failure. Prefer this over bare `assert`. +- **`CHECK_HOST(cond) << "msg " << value`** — **Preferred** runtime check: stream-style, throws `PanicError` with file/line info on failure. Zero overhead on the true path — the message expressions are only evaluated when the check fails. +- **`host::RuntimeCheck(cond, args...)`** — Function-style alternative to `CHECK_HOST`. Note its message args are always evaluated (even when the check passes), so prefer `CHECK_HOST` — especially on hot paths. - **`host::Panic(args...)`** — Unconditionally throw a `PanicError` with a descriptive message. - **`host::div_ceil(a, b)`** — Integer ceiling division `(a + b - 1) / b`. - **`host::irange(n)`** / **`host::irange(start, end)`** — Range views for cleaner loops. @@ -57,6 +58,7 @@ Add a new operation that scales each element of a tensor by a scalar factor: - Checks the CUDA error with file/line info after launch via `operator()(kernel, args...)`. - Supports `.enable_pdl(bool)` for PDL (Programmatic Dependent Launch, SM90+). - **`host::RuntimeDeviceCheck(cudaError_t)`** — Check a CUDA error; throw on failure. +- **`CHECK_CUDA(expr) << "context"`** — Stream-style CUDA error check; evaluates `expr` once and throws `PanicError` with `cudaGetErrorString` + file/line info if it is not `cudaSuccess`. Extra streamed context is optional. ### `tensor.h` — Tensor validation (`TensorMatcher`, Symbolic types) @@ -90,18 +92,21 @@ const size_t n = N.unwrap(); const DLDevice dev = device.unwrap(); ``` -### `type.cuh` — `dtype_trait` and `packed_t` +### `type.cuh` — `DTypeTrait`, `packed_t`, and reduction traits ```cpp #include ``` -- **`dtype_trait`** — Static trait struct for each scalar type. Provides: - - `dtype_trait::from(value)` — convert from another type (e.g. `fp32_t` → `fp16_t`) - - `dtype_trait::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math (primarily for `fp32_t`) - - `dtype_trait::max/min(x, y)` — type-dispatched binary math (primarily for `fp32_t`) +- **`DTypeTrait`** — Static trait struct, specialized for integral types, `fp32_t`, `fp16_t`, `bf16_t`, `fp8_e4m3_t`, and their packed x2/x4 variants. Provides: + - `DTypeTrait::from(value)` — convert from another type via the right CUDA intrinsic (e.g. `fp32_t` → `fp16_t`) + - `DTypeTrait::abs/max/min` — type-dispatched math (fp32, fp16/bf16 scalar and x2, integrals) + - `DTypeTrait::sqrt/rsqrt/exp/sin/cos(x)` — `fp32_t` only + - Metadata: `packed_t` / `unpacked_t` / `kVecSize` (packed layout), `kFloatMax` (dtype max as float, e.g. 448.0f for fp8-e4m3), `kZeroBits` - **`packed_t`** — Two-element packed alias: `packed_t` = `fp16x2_t`, `packed_t` = `bf16x2_t`, `packed_t` = `fp32x2_t`. Use for vectorized loads/stores. -- **`device::cast(value)`** — Type-safe cast using `dtype_trait`, e.g. `cast(v)`. +- **`device::cast(value)`** — Type-safe cast using `DTypeTrait`, e.g. `cast(v)`. +- **`device::unpack(value)`** — View a packed value as an `unpacked_t[kVecSize]` array reference (e.g. `fp32x2_t` → `fp32_t[2]`); element writes propagate back to the packed value. +- **`device::ReductionOp` (`SUM`/`MAX`/`MIN`) and `device::ReductionTrait::reduce(x, y)`** — One binary reduction step, dispatched through `DTypeTrait` (packed types reduce elementwise). This is the engine behind `warp::reduce`; use it directly when writing custom reductions. ### `vec.cuh` — Vectorized memory access (`AlignedVector`) @@ -135,8 +140,8 @@ For a **2D tile**, either flatten `(row, col)` into a linear tile index first, o #include ``` -- `device::math::max/min(a, b)` — type-dispatched binary math via `dtype_trait` -- `device::math::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math via `dtype_trait` +- `device::math::max/min(a, b)` — type-dispatched binary math via `DTypeTrait` +- `device::math::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math via `DTypeTrait` ### `warp.cuh` — Warp-level primitives @@ -144,8 +149,8 @@ For a **2D tile**, either flatten `(row, col)` into a linear tile index first, o #include ``` -- `device::warp::reduce_sum(value)` — warp-level sum reduction via `__shfl_xor_sync` -- `device::warp::reduce_max(value)` — warp-level max reduction +- `device::warp::reduce(value, active_mask)` — generic warp reduction via `__shfl_xor_sync`. `Op` is a `device::ReductionOp` (`SUM`/`MAX`/`MIN`); `kNumThreads` is a power-of-two group size (default 32 = full warp); `kInner=true` (default) reduces within each `kNumThreads`-sized group, `kInner=false` reduces across groups (lanes at the same offset in different groups). +- `device::warp::reduce_sum/reduce_max/reduce_min(value)` — convenience wrappers over `reduce`. Work for any type with a `ReductionTrait`: floats, integers, and packed x2 types. ### `cta.cuh` — CTA-level primitives @@ -203,8 +208,8 @@ The implementation fully uses the project abstractions described above: // NOTE: Comments for headers are not common in practice. // It is only shown here for tutorial purposes to highlight the key abstractions. #include // For TensorMatcher, SymbolicSize, SymbolicDevice -#include // For dtype_trait, fp16_t, bf16_t, fp32_t -#include // For RuntimeCheck, div_ceil +#include // For DTypeTrait, fp16_t, bf16_t, fp32_t +#include // For CHECK_HOST, div_ceil #include // For LaunchKernel, SGL_DEVICE #include // For AlignedVector @@ -280,7 +285,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { const uint32_t n = static_cast(N.unwrap()); const DLDevice device = device_.unwrap(); - RuntimeCheck(n > 0, "scale: num_elements must be > 0, got ", n); + CHECK_HOST(n > 0) << "scale: num_elements must be > 0, got " << n; // 2. Choose vector width for 128-bit loads (16 bytes) // fp16/bf16: 8 elements x 2 bytes = 16 bytes @@ -316,10 +321,10 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) { - Use `TensorMatcher` for all tensor validation; never manually check shape/dtype/device - Use `AlignedVector` for vectorised 128-bit loads/stores — significant bandwidth win - Use `LaunchKernel` — it resolves the stream and checks errors automatically -- Use `RuntimeCheck` for runtime assertions with useful error messages +- Use `CHECK_HOST(cond) << ...` for runtime assertions with useful error messages (zero overhead when the check passes) - Prefer passing runtime scalars like `factor` directly unless compile-time specialisation is genuinely required - `fp16_t` / `bf16_t` / `fp32_t` are the project's type aliases (from `utils.cuh`) -- `device::cast` or `dtype_trait::from(val)` for cross-type conversions +- `device::cast` or `DTypeTrait::from(val)` for cross-type conversions - `device::math::` functions for device math instead of bare `__` intrinsics if possible. - Try to use `PDL` feature. In some cases, this will benefit the performance. @@ -626,9 +631,9 @@ cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1 - `python/sglang/jit_kernel/include/sgl_kernel/utils.cuh` — type aliases, `LaunchKernel`, `SGL_DEVICE` - `python/sglang/jit_kernel/include/sgl_kernel/vec.cuh` — `AlignedVector` - `python/sglang/jit_kernel/include/sgl_kernel/tile.cuh` — `tile::Memory` -- `python/sglang/jit_kernel/include/sgl_kernel/type.cuh` — `dtype_trait`, `packed_t`, `device::cast` +- `python/sglang/jit_kernel/include/sgl_kernel/type.cuh` — `DTypeTrait`, `packed_t`, `device::cast`, `device::unpack`, `ReductionTrait` - `python/sglang/jit_kernel/include/sgl_kernel/math.cuh` — `device::math::` -- `python/sglang/jit_kernel/include/sgl_kernel/warp.cuh` — `warp::reduce_sum/max` +- `python/sglang/jit_kernel/include/sgl_kernel/warp.cuh` — `warp::reduce` and `reduce_sum/max/min` wrappers - `python/sglang/jit_kernel/include/sgl_kernel/cta.cuh` — `cta::reduce_max` - `python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh` — `atomic::max` - `python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh` — occupancy / SM count helpers diff --git a/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx b/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx index 8c5623783..28511be09 100644 --- a/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx +++ b/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx @@ -61,19 +61,24 @@ void test() { #### Runtime Checking -`RuntimeCheck` validates conditions at runtime. It accepts optional arguments for error reporting. -If the check fails, these arguments are output to aid debugging. -`RuntimeDeviceCheck` verifies the status of the last kernel launch. +`CHECK_HOST` is the preferred runtime check: stream-style, and zero overhead when the +check passes — the message expressions are only evaluated on failure. +`RuntimeCheck` is the function-style alternative; note its message arguments are always +evaluated, even when the check passes. +`RuntimeDeviceCheck` verifies the status of the last kernel launch, and `CHECK_CUDA` +is its stream-style equivalent for checking a `cudaError_t` with extra context. ```C++ Example #include #include void test() { + CHECK_HOST(1 + 1 == 2) << 1 + 1 << " != " << 2; // preferred host::RuntimeCheck(1 + 1 == 2, 1 + 1, " != ", 2); host::RuntimeDeviceCheck(); // check the provided `cudaError_t` host::RuntimeDeviceCheck(cudaGetLastError()); + CHECK_CUDA(cudaGetLastError()) << "after my_kernel launch"; } ``` @@ -161,7 +166,7 @@ Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](https://github.com/ ```cpp Example #include // For TensorMatcher, SymbolicSize, SymbolicDevice #include // For LaunchKernel -#include // For div_ceil, RuntimeCheck +#include // For div_ceil, CHECK_HOST #include #include @@ -199,8 +204,8 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { const size_t num_elements = N.unwrap(); const size_t grid_size = div_ceil(num_elements, kBlockSize); const DLDevice device = device_.unwrap(); - // some extra runtime checks using host::RuntimeCheck - RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements); + // some extra runtime checks using CHECK_HOST + CHECK_HOST(num_elements > 0) << "We only support non-empty tensors, got num_elements = " << num_elements; // 3. Launch the kernel. Error code will be automatically checked. LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)( @@ -289,12 +294,12 @@ and its key APIs. utils.h host - Host-side essentials: RuntimeCheck, Panic, div_ceil, irange + Host-side essentials: RuntimeCheck, CHECK_HOST(cond) << ..., Panic, div_ceil, irange utils.cuh device / host - Type aliases (fp16_t, bf16_t, ...), SGL_DEVICE macro, PDL helpers, LaunchKernel, RuntimeDeviceCheck + Type aliases (fp16_t, bf16_t, ...), SGL_DEVICE macro, PDL helpers, LaunchKernel, RuntimeDeviceCheck, CHECK_CUDA(expr) << ... source_location.h @@ -347,7 +352,7 @@ and its key APIs. type.cuh (global) / device - dtype_trait<T>, packed_t<T>, device::cast<To>(from) + DTypeTrait<T>, packed_t<T>, device::cast<To>(from) @@ -390,7 +395,7 @@ and its key APIs. warp.cuh device::warp - reduce_sum, reduce_max via __shfl_xor_sync + reduce<Op, kNumThreads, kInner> (SUM/MAX/MIN, grouped or inter-group) and reduce_sum / reduce_max / reduce_min wrappers via __shfl_xor_sync cta.cuh diff --git a/python/sglang/jit_kernel/__main__.py b/python/sglang/jit_kernel/__main__.py index b626fde78..b9c0f9681 100644 --- a/python/sglang/jit_kernel/__main__.py +++ b/python/sglang/jit_kernel/__main__.py @@ -1,16 +1,30 @@ import argparse import logging import os +import re +import shutil +import subprocess from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path -from sglang.jit_kernel.utils import ( - _REGISTERED_DEPENDENCIES, - DEFAULT_INCLUDE, - _get_default_target_flags, - get_jit_cuda_arch, - override_jit_cuda_arch, -) +from sglang.jit_kernel.utils import get_jit_cuda_arch, override_jit_cuda_arch +from sglang.jit_kernel.utils.arch import get_default_target_flags +from sglang.jit_kernel.utils.compile import DEFAULT_INCLUDE +from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES + + +def _clangd_major_version() -> int | None: + clangd = shutil.which("clangd") + if clangd is None: + return None + try: + result = subprocess.run( + [clangd, "--version"], capture_output=True, text=True, check=True + ) + except (OSError, subprocess.CalledProcessError): + return None + match = re.search(r"clangd version (\d+)", result.stdout) + return int(match.group(1)) if match else None def generate_clangd(): @@ -28,7 +42,7 @@ def generate_clangd(): "--dep", nargs="*", default=[], - choices=_REGISTERED_DEPENDENCIES.keys(), + choices=REGISTERED_DEPENDENCIES.keys(), help="Extra dependency libraries to include.", ) parser.add_argument( @@ -42,9 +56,9 @@ def generate_clangd(): dep_include_paths = [] for dep in args.dependencies: - if dep not in _REGISTERED_DEPENDENCIES: + if dep not in REGISTERED_DEPENDENCIES: raise ValueError(f"Dependency {dep} is not registered.") - dep_include_paths += _REGISTERED_DEPENDENCIES[dep]() + dep_include_paths += REGISTERED_DEPENDENCIES[dep]() include_paths = [ *DEFAULT_INCLUDE, @@ -70,9 +84,15 @@ def generate_clangd(): f"--cuda-gpu-arch=sm_{major}{minor}", "-Wall", "-Wextra", - *_get_default_target_flags(), + *get_default_target_flags(), *[f"-isystem{path}" for path in include_paths], ] + + # NOTE: for local clangd (fix the missing cluster related macros) + if major >= 9: + compile_flags.append("-D_CG_LIMIT_INCLUDED_DEPENDENCIES=1") + compile_flags.append("-D_CG_HAS_CLUSTER_GROUP=1") + # NOTE: skip these flags because clangd don't recognize them UNSUPPORTED_FLAGS = {"--expt-relaxed-constexpr"} compile_flags = [flag for flag in compile_flags if flag not in UNSUPPORTED_FLAGS] @@ -83,6 +103,10 @@ CompileFlags: {compile_flags_str} ] """ + # Documentation.CommentFormat lands in clangd 21. + clangd_major = _clangd_major_version() + if clangd_major is not None and clangd_major >= 21: + clangd_content += "Documentation:\n CommentFormat: Doxygen\n" if os.path.exists(".clangd") and not args.overwrite: logger.warning(".clangd file already exists, nothing done.") logger.warning("Use --overwrite to force overwrite the existing .clangd file.") diff --git a/python/sglang/jit_kernel/benchmark/marker.py b/python/sglang/jit_kernel/benchmark/marker.py index 96c6deb34..60e31783b 100644 --- a/python/sglang/jit_kernel/benchmark/marker.py +++ b/python/sglang/jit_kernel/benchmark/marker.py @@ -279,6 +279,9 @@ class Benchmark(Generic[F]): if not DISABLE_LOG_BANDWIDTH: bandwidths.append(float("nan")) continue + except BaseException: + print(f"Benchmark failed at {system}, kwargs =", kwargs) + raise latencies.append(result.times[0] / self._unit_scale) if not DISABLE_LOG_BANDWIDTH and result.memory_footprint is not None: should_log_bandwidth = True diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh index b8ddb3787..c86b92135 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh @@ -114,32 +114,6 @@ SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) { return val; } -/// Warp-wide max/min for integer types. `device::warp::reduce_max` routes through -/// `dtype_trait::max` which is only specialized for FP types. -SGL_DEVICE uint32_t warp_reduce_max_u32(uint32_t val) { -#pragma unroll - for (uint32_t mask = 16; mask > 0; mask >>= 1) { -#ifndef USE_ROCM - val = max(val, __shfl_xor_sync(device::kFullMask, val, mask, 32)); -#else - val = max(val, __shfl_xor(val, mask, 32)); -#endif - } - return val; -} - -SGL_DEVICE uint32_t warp_reduce_min_u32(uint32_t val) { -#pragma unroll - for (uint32_t mask = 16; mask > 0; mask >>= 1) { -#ifndef USE_ROCM - val = min(val, __shfl_xor_sync(device::kFullMask, val, mask, 32)); -#else - val = min(val, __shfl_xor(val, mask, 32)); -#endif - } - return val; -} - __global__ __launch_bounds__(1024, 1) // void plan_compress_prefill_kernel0(const Prefill0Params params) { using namespace device; @@ -185,12 +159,12 @@ __global__ __launch_bounds__(1024, 1) // // For min, treat threads outside `batch_size` as +inf so they don't pull the min down. const uint32_t e_for_max = static_cast(extend_len); const uint32_t e_for_min = (tx < params.batch_size) ? e_for_max : 0xFFFFFFFFu; - warp_max[warp_id] = warp_reduce_max_u32(e_for_max); - warp_min[warp_id] = warp_reduce_min_u32(e_for_min); + warp_max[warp_id] = warp::reduce_max(e_for_max); + warp_min[warp_id] = warp::reduce_min(e_for_min); __syncthreads(); if (warp_id == 0) { - s_max_extend = warp_reduce_max_u32(warp_max[lane_id]); - s_min_extend = warp_reduce_min_u32(warp_min[lane_id]); + s_max_extend = warp::reduce_max(warp_max[lane_id]); + s_min_extend = warp::reduce_min(warp_min[lane_id]); } __syncthreads(); diff --git a/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh b/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh index 25a68e99e..b7db4b4a9 100644 --- a/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh +++ b/python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh @@ -15,7 +15,7 @@ #include // For host dtype helpers and TensorView metadata #include // For RuntimeCheck and div_ceil -#include // For dtype_trait conversions +#include // For DTypeTrait conversions #include // For LaunchKernel and CUDA dtype aliases #include // For device::AlignedVector @@ -99,8 +99,8 @@ __device__ __forceinline__ float to_float(bf16_t v) { template __device__ __forceinline__ T residual_gate_value(T residual, T update, T gate) { - const T product = dtype_trait::from(to_float(update) * to_float(gate)); - return dtype_trait::from(to_float(residual) + to_float(product)); + const T product = DTypeTrait::from(to_float(update) * to_float(gate)); + return DTypeTrait::from(to_float(residual) + to_float(product)); } template diff --git a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh b/python/sglang/jit_kernel/include/sgl_kernel/math.cuh index 2b203e4b1..92c4dad6c 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/math.cuh @@ -2,14 +2,12 @@ /// \brief Device-side math helper functions and constants. /// /// Provides type-generic wrappers around CUDA math intrinsics by -/// dispatching through `dtype_trait`. All functions are forced-inline +/// dispatching through `DTypeTrait`. All functions are forced-inline /// device functions. #pragma once #include -#include - namespace device::math { /// \brief Constant: log2(e) @@ -27,49 +25,49 @@ static_assert(log2e * loge2 == 1.0f, "log2e * loge2 must be 1"); /// \brief Returns the larger of `a` and `b`. template SGL_DEVICE T max(T a, T b) { - return dtype_trait::max(a, b); + return DTypeTrait::max(a, b); } /// \brief Returns the smaller of `a` and `b`. template SGL_DEVICE T min(T a, T b) { - return dtype_trait::min(a, b); + return DTypeTrait::min(a, b); } /// \brief Returns the absolute value of `a`. template SGL_DEVICE T abs(T a) { - return dtype_trait::abs(a); + return DTypeTrait::abs(a); } /// \brief Returns the square root of `a`. template SGL_DEVICE T sqrt(T a) { - return dtype_trait::sqrt(a); + return DTypeTrait::sqrt(a); } /// \brief Returns the reciprocal square root of `a` (i.e. 1 / sqrt(a)). template SGL_DEVICE T rsqrt(T a) { - return dtype_trait::rsqrt(a); + return DTypeTrait::rsqrt(a); } /// \brief Returns e^a. template SGL_DEVICE T exp(T a) { - return dtype_trait::exp(a); + return DTypeTrait::exp(a); } /// \brief Returns sin(a). template SGL_DEVICE T sin(T a) { - return dtype_trait::sin(a); + return DTypeTrait::sin(a); } /// \brief Returns cos(a). template SGL_DEVICE T cos(T a) { - return dtype_trait::cos(a); + return DTypeTrait::cos(a); } } // namespace device::math diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h index 1ae9233a6..915345721 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h +++ b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h @@ -52,10 +52,10 @@ struct DTypeRef; struct DeviceRef; template -struct _dtype_trait {}; +struct DLDataTypeTrait {}; template -struct _dtype_trait { +struct DLDataTypeTrait { inline static constexpr DLDataType value = { .code = std::is_signed_v ? DLDataTypeCode::kDLInt : DLDataTypeCode::kDLUInt, .bits = static_cast(sizeof(T) * 8), @@ -63,45 +63,45 @@ struct _dtype_trait { }; template -struct _dtype_trait { +struct DLDataTypeTrait { inline static constexpr DLDataType value = { .code = DLDataTypeCode::kDLFloat, .bits = static_cast(sizeof(T) * 8), .lanes = 1}; }; #ifdef __CUDACC__ template <> -struct _dtype_trait { +struct DLDataTypeTrait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1}; }; template <> -struct _dtype_trait { +struct DLDataTypeTrait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1}; }; template <> -struct _dtype_trait { +struct DLDataTypeTrait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat8_e4m3fn, .bits = 8, .lanes = 1}; }; #elif defined(__HIPCC__) template <> -struct _dtype_trait { +struct DLDataTypeTrait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1}; }; template <> -struct _dtype_trait { +struct DLDataTypeTrait { inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1}; }; #endif template -struct _device_trait { +struct DLDeviceTrait { inline static constexpr DLDevice value = {.device_type = Code, .device_id = kAnyDeviceID}; }; template -inline constexpr auto kDTypeList = std::array{_dtype_trait::value...}; +inline constexpr auto kDTypeList = std::array{DLDataTypeTrait::value...}; template -inline constexpr auto kDeviceList = std::array{_device_trait::value...}; +inline constexpr auto kDeviceList = std::array{DLDeviceTrait::value...}; template struct PrintAbleSpan { @@ -176,7 +176,7 @@ inline auto& operator<<(std::ostream& os, PrintAbleSpan span) { /// \brief Check whether `dtype` matches the DLDataType for C++ type `T`. template inline bool is_type(DLDataType dtype) { - return dtype == details::_dtype_trait::value; + return dtype == details::DLDataTypeTrait::value; } /** diff --git a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh b/python/sglang/jit_kernel/include/sgl_kernel/type.cuh index a7a534619..3cbd96564 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/type.cuh @@ -1,37 +1,35 @@ /// \file type.cuh /// \brief Dtype trait system for CUDA scalar/packed types. /// -/// `dtype_trait` provides per-type metadata: packed type alias, +/// `DTypeTrait` provides per-type metadata: packed type alias, /// conversion functions (`from`), and unary/binary math operations. /// Use `device::cast(from_value)` for type conversion on device. -/// -/// Registered types: -/// | Scalar | Packed (x2) | Notes | -/// |-----------|-------------|-------------------------------| -/// | `fp32_t` | `fp32x2_t` | Full math ops (abs,sqrt,...) | -/// | `fp16_t` | `fp16x2_t` | Conversion only | -/// | `bf16_t` | `bf16x2_t` | Conversion only | -/// | `fp32x2_t`| `fp32x4_t` | Packed float2 <-> half2/bf162 | #pragma once #include +#include +#include +#include +#include + template -struct dtype_trait {}; +struct DTypeTrait {}; -#define SGL_REGISTER_DTYPE_TRAIT(TYPE, PACK2, ...) \ - template <> \ - struct dtype_trait { \ - using self_t = TYPE; \ - using packed_t = PACK2; \ - template \ - SGL_DEVICE static self_t from(const S& value) { \ - return static_cast(value); \ - } \ - __VA_ARGS__ \ - } +#define SGL_REGISTER_PACKED(SELF, PACKED) \ + using self_t = SELF; \ + using packed_t = PACKED -#define SGL_REGISTER_TYPE_END static_assert(true) +#define SGL_REGISTER_UNPACK(UNPACK, N) \ + using unpacked_t = UNPACK; \ + static constexpr size_t kVecSize = N + +#define SGL_REGISTER_FROM_DEFAULT() \ + template \ + SGL_DEVICE static self_t from(const S& value) { \ + return static_cast(value); \ + } \ + static_assert(true) #define SGL_REGISTER_FROM_FUNCTION(FROM, FN) \ SGL_DEVICE static self_t from(const FROM& x) { \ @@ -45,76 +43,312 @@ struct dtype_trait {}; } \ static_assert(true) +// Also emits a `kHas_` flag so reduction dispatch can detect the op via +// plain member SFINAE (see details::HasMax below) - hipcc mis-evaluates +// requires-expressions that probe device functions, so detection must only +// ever look at data members. #define SGL_REGISTER_BINARY_FUNCTION(NAME, FN) \ + static constexpr bool kHas_##NAME = true; \ SGL_DEVICE static self_t NAME(const self_t& x, const self_t& y) { \ return FN(x, y); \ } \ static_assert(true) -SGL_REGISTER_DTYPE_TRAIT( - fp32_t, fp32x2_t, SGL_REGISTER_TYPE_END; // - SGL_REGISTER_FROM_FUNCTION(fp16_t, __half2float); - SGL_REGISTER_FROM_FUNCTION(bf16_t, __bfloat162float); - SGL_REGISTER_UNARY_FUNCTION(abs, fabsf); - SGL_REGISTER_UNARY_FUNCTION(sqrt, sqrtf); - SGL_REGISTER_UNARY_FUNCTION(rsqrt, rsqrtf); - SGL_REGISTER_UNARY_FUNCTION(exp, expf); - SGL_REGISTER_UNARY_FUNCTION(sin, sinf); - SGL_REGISTER_UNARY_FUNCTION(cos, cosf); - SGL_REGISTER_BINARY_FUNCTION(max, fmaxf); - SGL_REGISTER_BINARY_FUNCTION(min, fminf);); -SGL_REGISTER_DTYPE_TRAIT(fp16_t, fp16x2_t); -SGL_REGISTER_DTYPE_TRAIT(bf16_t, bf16x2_t); +template +struct DTypeTrait { + SGL_REGISTER_PACKED(T, void); + SGL_REGISTER_UNPACK(T, 1); + SGL_REGISTER_FROM_DEFAULT(); + SGL_REGISTER_UNARY_FUNCTION(abs, ::abs); + SGL_REGISTER_BINARY_FUNCTION(max, ::max); + SGL_REGISTER_BINARY_FUNCTION(min, ::min); + static constexpr T kZeroBits = 0; +}; -/// TODO: Add ROCM implementation -SGL_REGISTER_DTYPE_TRAIT( - fp32x2_t, fp32x4_t, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp16x2_t, __half22float2); - SGL_REGISTER_FROM_FUNCTION(bf16x2_t, __bfloat1622float2);); +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(fp32_t, fp32x2_t); + SGL_REGISTER_UNPACK(fp32_t, 1); + SGL_REGISTER_FROM_DEFAULT(); + SGL_REGISTER_FROM_FUNCTION(fp16_t, __half2float); + SGL_REGISTER_FROM_FUNCTION(bf16_t, __bfloat162float); + SGL_REGISTER_UNARY_FUNCTION(abs, fabsf); + SGL_REGISTER_UNARY_FUNCTION(sqrt, sqrtf); + SGL_REGISTER_UNARY_FUNCTION(rsqrt, rsqrtf); + SGL_REGISTER_UNARY_FUNCTION(exp, expf); + SGL_REGISTER_UNARY_FUNCTION(sin, sinf); + SGL_REGISTER_UNARY_FUNCTION(cos, cosf); + SGL_REGISTER_BINARY_FUNCTION(max, fmaxf); + SGL_REGISTER_BINARY_FUNCTION(min, fminf); + static constexpr float kFloatMax = std::numeric_limits::max(); + static constexpr uint32_t kZeroBits = 0x00000000; +}; -SGL_REGISTER_DTYPE_TRAIT( - fp16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22half2_rn);); +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(fp32x2_t, fp32x4_t); + SGL_REGISTER_UNPACK(fp32_t, 2); + SGL_REGISTER_FROM_DEFAULT(); + SGL_REGISTER_FROM_FUNCTION(fp16x2_t, __half22float2); + SGL_REGISTER_FROM_FUNCTION(bf16x2_t, __bfloat1622float2); +}; -SGL_REGISTER_DTYPE_TRAIT( - bf16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22bfloat162_rn);); +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(fp32x4_t, void); + SGL_REGISTER_UNPACK(fp32_t, 4); + SGL_REGISTER_FROM_DEFAULT(); +}; + +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(fp16_t, fp16x2_t); + SGL_REGISTER_UNPACK(fp16_t, 1); + SGL_REGISTER_FROM_DEFAULT(); + SGL_REGISTER_FROM_FUNCTION(fp32_t, __float2half_rn); + SGL_REGISTER_UNARY_FUNCTION(abs, __habs); + SGL_REGISTER_BINARY_FUNCTION(max, __hmax); + SGL_REGISTER_BINARY_FUNCTION(min, __hmin); + // CUDA fp16 max clamp value + static constexpr float kFloatMax = 65504.0f; + static constexpr uint16_t kZeroBits = 0x0000; +}; + +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(fp16x2_t, void); + SGL_REGISTER_UNPACK(fp16_t, 2); + SGL_REGISTER_FROM_DEFAULT(); + SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22half2_rn); + SGL_REGISTER_UNARY_FUNCTION(abs, __habs2); +#ifndef USE_ROCM + SGL_REGISTER_BINARY_FUNCTION(add, __hadd2); + SGL_REGISTER_BINARY_FUNCTION(max, __hmax2); + SGL_REGISTER_BINARY_FUNCTION(min, __hmin2); +#else + // HIP only provides __hmax2/__hmin2 for __hip_bfloat162, not __half2. + // No `add` registered on HIP (packed SUM falls back to lane-wise scalar). + static constexpr bool kHas_max = true; + static constexpr bool kHas_min = true; + SGL_DEVICE static self_t max(const self_t& x, const self_t& y) { + return self_t{__hmax(x.x, y.x), __hmax(x.y, y.y)}; + } + SGL_DEVICE static self_t min(const self_t& x, const self_t& y) { + return self_t{__hmin(x.x, y.x), __hmin(x.y, y.y)}; + } +#endif +}; + +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(bf16_t, bf16x2_t); + SGL_REGISTER_UNPACK(bf16_t, 1); + SGL_REGISTER_FROM_DEFAULT(); +#ifndef USE_ROCM + SGL_REGISTER_FROM_FUNCTION(fp32_t, __float2bfloat16_rn); +#else + // HIP has no _rn-suffixed variant; __float2bfloat16 rounds to nearest. + SGL_REGISTER_FROM_FUNCTION(fp32_t, __float2bfloat16); +#endif + SGL_REGISTER_UNARY_FUNCTION(abs, __habs); + SGL_REGISTER_BINARY_FUNCTION(max, __hmax); + SGL_REGISTER_BINARY_FUNCTION(min, __hmin); + // CUDA bf16 max clamp value + static constexpr float kFloatMax = 3.38953139e38f; + static constexpr uint16_t kZeroBits = 0x0000; +}; + +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(bf16x2_t, void); + SGL_REGISTER_UNPACK(bf16_t, 2); + SGL_REGISTER_FROM_DEFAULT(); + SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22bfloat162_rn); + SGL_REGISTER_UNARY_FUNCTION(abs, __habs2); +#ifndef USE_ROCM + // No `add` on HIP: bf162 __hadd2 is unverified there (packed SUM falls + // back to lane-wise scalar). + SGL_REGISTER_BINARY_FUNCTION(add, __hadd2); +#endif + SGL_REGISTER_BINARY_FUNCTION(max, __hmax2); + SGL_REGISTER_BINARY_FUNCTION(min, __hmin2); +}; #ifndef USE_ROCM -SGL_REGISTER_DTYPE_TRAIT(fp8_e4m3_t, fp8x2_e4m3_t); +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(fp8_e4m3_t, fp8x2_e4m3_t); + SGL_REGISTER_UNPACK(fp8_e4m3_t, 1); + SGL_REGISTER_FROM_DEFAULT(); + // NOTE: CUDA fp8 support explicit cast (i.e. use default from is ok) + + static constexpr float kFloatMax = 448.0f; // CUDA fp8 max clamp value + static constexpr uint8_t kZeroBits = 0x00; +}; + +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(fp8x2_e4m3_t, fp8x4_e4m3_t); + SGL_REGISTER_UNPACK(fp8_e4m3_t, 2); + SGL_REGISTER_FROM_DEFAULT(); + // NOTE: CUDA fp8 support explicit cast (i.e. use default from is ok) +}; + +template <> +struct DTypeTrait { + SGL_REGISTER_PACKED(fp8x4_e4m3_t, void); + SGL_REGISTER_UNPACK(fp8_e4m3_t, 4); + SGL_REGISTER_FROM_DEFAULT(); + // NOTE: CUDA fp8 support explicit cast (i.e. use default from is ok) +}; #endif -#undef SGL_REGISTER_DTYPE_TRAIT +#undef SGL_REGISTER_PACKED +#undef SGL_REGISTER_UNPACK +#undef SGL_REGISTER_FROM_DEFAULT #undef SGL_REGISTER_FROM_FUNCTION +#undef SGL_REGISTER_UNARY_FUNCTION +#undef SGL_REGISTER_BINARY_FUNCTION /// \brief Alias: the packed (x2) type for `T`. template -using packed_t = typename dtype_trait::packed_t; +using packed_t = typename DTypeTrait::packed_t; namespace device { /** * \brief Cast a value from type `From` to type `To` on device. * - * Dispatches through `dtype_trait::from()`, which uses the appropriate + * Dispatches through `DTypeTrait::from()`, which uses the appropriate * CUDA intrinsic (e.g. `__half2float`, `__float22half2_rn`). */ template SGL_DEVICE To cast(const From& value) { - return dtype_trait::from(value); + return DTypeTrait::from(value); } +/** + * \brief View a packed value as an array of its `unpacked_t` elements. + * + * Returns a reference to `value` reinterpreted as `unpacked_t[kVecSize]`, + * so element writes propagate back to the original packed value. + * Constness of `value` is preserved. + */ +template +SGL_DEVICE auto& unpack(T& value) { + using Trait = DTypeTrait>; + using U = typename Trait::unpacked_t; + constexpr size_t kVecSize = Trait::kVecSize; + static_assert(sizeof(T) == sizeof(U) * kVecSize, "packed type must be layout-compatible"); + using A = std::conditional_t, const U, U>; + return reinterpret_cast(value); +} + +enum class ReductionOp : uint8_t { SUM, MAX, MIN }; + +template +struct ReductionTrait {}; + +namespace details { + +// Op detection via the `kHas_*` data members emitted by +// SGL_REGISTER_BINARY_FUNCTION. Deliberately classic void_t member SFINAE: +// hipcc mis-evaluates requires-expressions in device instantiation contexts +// (observed: even `requires { a + b; }` on float came out false), so detection +// must never probe function-call expressions. +template +struct HasAdd : std::false_type {}; +template +struct HasAdd::kHas_add)>> : std::true_type {}; + +template +struct HasMax : std::false_type {}; +template +struct HasMax::kHas_max)>> : std::true_type {}; + +template +struct HasMin : std::false_type {}; +template +struct HasMin::kHas_min)>> : std::true_type {}; + +template +SGL_DEVICE T reduce_recursive(const T& x, const T& y) { + using U = typename DTypeTrait::unpacked_t; + constexpr size_t kVecSize = DTypeTrait::kVecSize; + static_assert(kVecSize > 1, "unsupported scalar type for reduction"); + using Trait = ReductionTrait; + auto& x_unpacked = ::device::unpack(x); + auto& y_unpacked = ::device::unpack(y); + T result{}; + auto& z_unpacked = ::device::unpack(result); +#pragma unroll + for (size_t i = 0; i < kVecSize; ++i) { + z_unpacked[i] = Trait::reduce(x_unpacked[i], y_unpacked[i]); + } + return result; +} + +} // namespace details + +// Dispatch rules, chosen so correctness never depends on detection: +// scalars (kVecSize == 1) call the trait member / operator directly - a +// missing op is a clear compile error at the call line; packed types use the +// native op when the trait registered one and fall back to lane-wise +// recursion otherwise (worst case for a mis-detecting compiler is a slightly +// slower but still correct lane-wise path). +template +struct ReductionTrait { + SGL_DEVICE static T reduce(const T& x, const T& y) { + if constexpr (details::HasAdd::value) { + return DTypeTrait::add(x, y); + } else if constexpr (DTypeTrait::kVecSize == 1) { + return static_cast(x + y); + } else { + return details::reduce_recursive(x, y); + } + } +}; + +template +struct ReductionTrait { + SGL_DEVICE static T reduce(const T& x, const T& y) { + if constexpr (DTypeTrait::kVecSize == 1) { + return DTypeTrait::max(x, y); + } else if constexpr (details::HasMax::value) { + return DTypeTrait::max(x, y); + } else { + return details::reduce_recursive(x, y); + } + } +}; + +template +struct ReductionTrait { + SGL_DEVICE static T reduce(const T& x, const T& y) { + if constexpr (DTypeTrait::kVecSize == 1) { + return DTypeTrait::min(x, y); + } else if constexpr (details::HasMin::value) { + return DTypeTrait::min(x, y); + } else { + return details::reduce_recursive(x, y); + } + } +}; + } // namespace device // --------------------------------------------------------------------------- -// FP8 max clamp value — platform-dependent +// FP8 max clamp value - platform-dependent // CUDA (e4m3fn): 448.0f // AMD FNUZ (e4m3fnuz): 224.0f // AMD E4M3 (e4m3fn): 448.0f // --------------------------------------------------------------------------- #ifndef USE_ROCM -constexpr float kFP8E4M3Max = 448.0f; +inline constexpr float kFP8E4M3Max = 448.0f; #else // USE_ROCM #if HIP_FP8_TYPE_FNUZ -constexpr float kFP8E4M3Max = 224.0f; +inline constexpr float kFP8E4M3Max = 224.0f; #else // HIP_FP8_TYPE_E4M3 -constexpr float kFP8E4M3Max = 448.0f; +inline constexpr float kFP8E4M3Max = 448.0f; #endif // HIP_FP8_TYPE_FNUZ #endif // USE_ROCM diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index 684475117..f54ca9ce8 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -65,6 +65,8 @@ using fp16x2_t = __half2; using bf16x2_t = __nv_bfloat162; using fp8x2_e4m3_t = __nv_fp8x2_e4m3; using fp8x2_e5m2_t = __nv_fp8x2_e5m2; +using fp8x4_e4m3_t = __nv_fp8x4_e4m3; +using fp8x4_e5m2_t = __nv_fp8x4_e5m2; using fp32x4_t = float4; #else @@ -78,6 +80,8 @@ using fp16x2_t = half2; using bf16x2_t = __hip_bfloat162; using fp8x2_e4m3_t = uint16_t; using fp8x2_e5m2_t = uint16_t; +using fp8x4_e4m3_t = uint32_t; +using fp8x4_e5m2_t = uint32_t; using fp32x4_t = float4; #endif @@ -214,9 +218,7 @@ SGL_DEVICE auto offset(const void* ptr, U... offset) -> const void* { } // namespace pointer -/// PTX pragma that lets the compiler spill registers into otherwise-unused -/// shared memory instead of local memory. The radix kernels run at occupancy 2 -/// (32 regs/thread) and rely on this to avoid local-memory traffic. +/// PTX pragma that lets the compiler spill registers into shared memory SGL_DEVICE void enable_smem_spilling() { #if defined(__CUDA_ARCH__) && CUDART_VERSION >= 13000 asm(".pragma \"enable_smem_spilling\";"); @@ -377,4 +379,11 @@ struct LaunchKernel { cudaLaunchAttribute m_attrs[2]; }; +// The empty-true-branch if/else form keeps a trailing `else` in user code +// bound to the user's `if`, not to the macro's. +#define CHECK_CUDA(COND) \ + if (const auto error = (COND); error == ::cudaSuccess) [[likely]] { \ + } else \ + ::host::Error() << "CUDA error: " << ::cudaGetErrorString(error) << ". " + } // namespace host diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.h b/python/sglang/jit_kernel/include/sgl_kernel/utils.h index 3226f79dd..5c9c00495 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.h +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.h @@ -1,14 +1,5 @@ /// \file utils.h /// \brief Host-side C++ utilities used by JIT kernel wrappers. -/// -/// Provides: -/// - `DebugInfo` - wraps `std::source_location` for error reporting. -/// - `RuntimeCheck` - runtime assertion with formatted error messages. -/// - `Panic` - unconditional abort with formatted error messages. -/// - `pointer::offset` - safe void-pointer arithmetic (host side). -/// - `div_ceil` - integer ceiling division. -/// - `dtype_bytes` - byte width of a `DLDataType`. -/// - `irange` - Python-style integer range for range-for loops. #pragma once @@ -83,7 +74,7 @@ template [[noreturn]] inline auto panic(DebugInfo location, Args&&... args) -> void { std::ostringstream os; - os << "Runtime check failed at " << location.file_name() << ":" << location.line(); + os << "Failed at " << location.file_name() << ":" << location.line(); if constexpr (sizeof...(args) > 0) { os << ": "; (os << ... << std::forward(args)); @@ -183,4 +174,38 @@ inline auto irange(T start, T end) { return stdv::iota(start, end); } +/** \brief Error class for stream-style error logging. */ +struct Error { + Error(DebugInfo location = {}) { + m_oss << "Failed at " << location.file_name() << ":" << location.line() << ": "; + } + + template + Error& operator<<(T&& arg) { + m_oss << std::forward(arg); + return *this; + } + + [[noreturn]] + ~Error() noexcept(false) { + throw PanicError(std::move(m_oss).str()); + } + + private: + std::ostringstream m_oss; +}; + +/** + * \brief 0-overhead CHECK macro for host code. This can avoid unnecessary + * instantiation of error messages when the condition is true. + * + * Usage: CHECK_HOST(ptr != nullptr) << "Pointer must not be null"; + */ +// The empty-true-branch if/else form keeps a trailing `else` in user code +// bound to the user's `if`, not to the macro's. +#define CHECK_HOST(COND) \ + if (COND) [[likely]] { \ + } else \ + ::host::Error() + } // namespace host diff --git a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh index 9d82efae1..9fafaf0d8 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh @@ -5,6 +5,9 @@ #include #include +#include +#include + namespace device::warp { /// \brief Full warp active mask. @@ -17,40 +20,103 @@ using mask_t = uint64_t; #endif /** - * \brief Warp-level sum reduction. + * \brief Warp-level reduction. * - * On CUDA: uses __shfl_xor_sync with width=32. + * On CUDA: uses __shfl_xor_sync with width=32. Full-warp reductions + * use a single `redux.sync` instruction when the target supports it. * On HIP: uses __shfl_xor with explicit width parameter (supports wave64 sub-groups). + * \tparam OP Reduction operation to perform (SUM, MAX, MIN). + * \tparam kNumThreads Number of threads as a group. + * \tparam kInner Whether to perform within a group or not. + * \tparam T Type of the value to reduce. + * + * \param value The value to reduce. + * \param active_mask The active mask of threads participating in the reduction. + * + * \note We will divide into groups of `kNumThreads`. + * e.g. kNumThreads = 8, we have 0..7, 8..15, 16..23, 24..31 as groups. + * By reduction is performed within a group. Inter-group reduction will reduce + * over the same offset in different groups. e.g. {0, 8, 16, 24} in the above example. */ -template -SGL_DEVICE T reduce_sum(T value, mask_t active_mask = kFullMask) { +template +SGL_DEVICE T reduce(T value, mask_t active_mask = kFullMask) { static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads); static_assert(std::has_single_bit(kNumThreads), "must be pow of 2"); -#pragma unroll - for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) -#ifndef USE_ROCM - value = value + __shfl_xor_sync(active_mask, value, mask, 32); -#else - value = value + __shfl_xor(value, mask, kNumThreads); + using Trait = ReductionTrait; + +#ifdef SGL_CUDA_ARCH + // CUDA target only + constexpr bool kFullReduction = (kNumThreads == kWarpThreads && kInner) || (kNumThreads == 1 && !kInner); + if constexpr (kFullReduction) { +#if SGL_CUDA_ARCH >= 800 + // 32 bit integer reduction + if constexpr (std::is_integral_v && sizeof(T) <= 4) { + if constexpr (OP == ReductionOp::SUM) { + return __reduce_add_sync(active_mask, value); + } else if constexpr (OP == ReductionOp::MAX) { + return __reduce_max_sync(active_mask, value); + } else if constexpr (OP == ReductionOp::MIN) { + return __reduce_min_sync(active_mask, value); + } + } #endif +#if SGL_CUDA_ARCH >= 1000 && SGL_CUDA_ARCH < 1100 + // 32-bit float reduction + if constexpr (std::is_same_v) { + if constexpr (OP == ReductionOp::MAX) { + float result; + asm("redux.sync.max.f32 %0, %1, %2;" : "=f"(result) : "f"(value), "r"(active_mask)); + return result; + } else if constexpr (OP == ReductionOp::MIN) { + float result; + asm("redux.sync.min.f32 %0, %1, %2;" : "=f"(result) : "f"(value), "r"(active_mask)); + return result; + } + } +#endif + } +#endif // redux.sync for CUDA only + + if constexpr (kInner) { +#pragma unroll + for (uint32_t mask = kNumThreads / 2; mask >= 1; mask >>= 1) { +#ifndef USE_ROCM + value = Trait::reduce(value, __shfl_xor_sync(active_mask, value, mask, 32)); +#else + value = Trait::reduce(value, __shfl_xor(value, mask, kNumThreads)); +#endif + } + } else { +#pragma unroll + for (uint32_t mask = kNumThreads; mask <= kWarpThreads / 2; mask <<= 1) { +#ifndef USE_ROCM + value = Trait::reduce(value, __shfl_xor_sync(active_mask, value, mask, 32)); +#else + // Inter-group shuffle crosses kNumThreads-sized sub-groups, so the + // shuffle width must span the whole warp. + value = Trait::reduce(value, __shfl_xor(value, mask, kWarpThreads)); +#endif + } + } return value; } -/** - * \brief Warp-level max reduction. - */ -template +/** \brief Warp-level sum reduction. */ +template +SGL_DEVICE T reduce_sum(T value, mask_t active_mask = kFullMask) { + return reduce(value, active_mask); +} + +/** \brief Warp-level max reduction. */ +template SGL_DEVICE T reduce_max(T value, mask_t active_mask = kFullMask) { - static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads); - static_assert(std::has_single_bit(kNumThreads), "must be pow of 2"); -#pragma unroll - for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) -#ifndef USE_ROCM - value = math::max(value, __shfl_xor_sync(active_mask, value, mask, 32)); -#else - value = math::max(value, __shfl_xor(value, mask, kNumThreads)); -#endif - return value; + return reduce(value, active_mask); +} + +/** \brief Warp-level min reduction. */ +template +SGL_DEVICE T reduce_min(T value, mask_t active_mask = kFullMask) { + return reduce(value, active_mask); } } // namespace device::warp diff --git a/python/sglang/jit_kernel/per_token_group_quant_8bit.py b/python/sglang/jit_kernel/per_token_group_quant_8bit.py index f2077c802..cf4b753a8 100644 --- a/python/sglang/jit_kernel/per_token_group_quant_8bit.py +++ b/python/sglang/jit_kernel/per_token_group_quant_8bit.py @@ -16,27 +16,26 @@ from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from tvm_ffi.module import Module -from sglang.jit_kernel.utils import CPP_DTYPE_MAP as OUTPUT_DTYPE_MAP - @cache_once def _jit_per_token_group_quant_8bit_module( dtype: torch.dtype, output_type: torch.dtype, group_size: int ) -> Module: dtype_arg = make_cpp_args(dtype) + out_arg = make_cpp_args(output_type) gs_arg = make_cpp_args(group_size) pdl_arg = make_cpp_args(is_arch_support_pdl()) - out_cpp = OUTPUT_DTYPE_MAP[output_type] return load_jit( "per_token_group_quant_8bit", *dtype_arg, + *out_arg, *gs_arg, *pdl_arg, cuda_files=["gemm/per_token_group_quant_8bit.cuh"], cuda_wrappers=[ ( "per_token_group_quant_8bit", - f"per_token_group_quant_8bit<{dtype_arg}, {out_cpp}, {gs_arg}, {pdl_arg}>", + f"per_token_group_quant_8bit<{dtype_arg}, {out_arg}, {gs_arg}, {pdl_arg}>", ) ], ) diff --git a/python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py b/python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py index c18669b35..372f30c82 100644 --- a/python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py +++ b/python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py @@ -40,7 +40,7 @@ def _q8kv8_cuda_flags() -> list[str]: # torch.utils.cpp_extension's AOT path does (COMMON_NVCC_FLAGS). The JIT # toolchain never defines them, so undefining is a no-op. # * --expt-relaxed-constexpr and -O3: already supplied by the JIT default - # target flags (see utils._get_default_target_flags). + # target flags (see utils.arch.get_default_target_flags). # * --expt-extended-lambda, -lineinfo, -D_USE_MATH_DEFINES: not required # by this single-translation-unit kernel. return [ diff --git a/python/sglang/jit_kernel/utils/__init__.py b/python/sglang/jit_kernel/utils/__init__.py new file mode 100644 index 000000000..c321dcdf6 --- /dev/null +++ b/python/sglang/jit_kernel/utils/__init__.py @@ -0,0 +1,31 @@ +"""Public interface of sglang.jit_kernel.utils.""" + +from sglang.jit_kernel.utils.arch import ( + get_jit_cuda_arch, + is_arch_support_pdl, + override_jit_cuda_arch, +) +from sglang.jit_kernel.utils.common import ( + cache_once, + get_ci_test_range, + is_hip_runtime, + is_musa_runtime, + lazy_register_class, + should_run_full_tests, +) +from sglang.jit_kernel.utils.compile import KERNEL_PATH, load_jit, make_cpp_args + +__all__ = [ + "should_run_full_tests", + "get_ci_test_range", + "cache_once", + "lazy_register_class", + "is_hip_runtime", + "is_musa_runtime", + "make_cpp_args", + "load_jit", + "override_jit_cuda_arch", + "get_jit_cuda_arch", + "is_arch_support_pdl", + "KERNEL_PATH", +] diff --git a/python/sglang/jit_kernel/utils/arch.py b/python/sglang/jit_kernel/utils/arch.py new file mode 100644 index 000000000..24c6492ee --- /dev/null +++ b/python/sglang/jit_kernel/utils/arch.py @@ -0,0 +1,119 @@ +"""CUDA/ROCm architecture detection and default compile target flags.""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from dataclasses import dataclass +from typing import List + +import torch + +from sglang.jit_kernel.utils.common import ( + cache_once, + is_hip_runtime, + is_musa_runtime, +) +from sglang.srt.utils.common import get_cuda_version + +logger = logging.getLogger(__name__) + + +@dataclass +class ArchInfo: + major: int + minor: int + suffix: str + + @property + def target_name(self) -> str: + return f"{self.major}.{self.minor}{self.suffix}" + + @property + def jit_flag(self) -> str: + return f"-DSGL_CUDA_ARCH={self.major * 100 + self.minor * 10}" + + +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. + """ + if major == 9: + return "a" + if major == 12: + if get_cuda_version() < (12, 9): + return "" + return "f" if minor == 0 else "a" + if major >= 10: + return "a" + return "" + + +@cache_once +def _init_jit_cuda_arch_once(): + global _CUDA_ARCH + try: + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + except Exception: + logger.warning("Cannot detect CUDA architecture.") + major, minor = 0, 0 # invalid value to trigger compile error if used + # JIT builds target the exact local GPU, so the arch-specific target is + # always correct on Hopper+ and unlocks arch-only instructions (redux.f32). + # HIP/MUSA capability numbers aren't CUDA SM versions and stay unsuffixed. + suffix = ( + "" + if (is_hip_runtime() or is_musa_runtime()) + else _cuda_arch_suffix(major, minor) + ) + _CUDA_ARCH = ArchInfo(major, minor, suffix) + + +def get_default_target_flags() -> List[str]: + if is_hip_runtime(): + flags = ["-DUSE_ROCM", "-std=c++20", "-O3"] + # Detect FP8 type based on GPU architecture + try: + device = torch.cuda.current_device() + gcn_arch = torch.cuda.get_device_properties(device).gcnArchName + if "gfx942" in gcn_arch: + flags.append("-DHIP_FP8_TYPE_FNUZ=1") + else: + flags.append("-DHIP_FP8_TYPE_E4M3=1") + except Exception: + flags.append("-DHIP_FP8_TYPE_E4M3=1") + return flags + else: + return [ + get_jit_cuda_arch().jit_flag, + "-std=c++20", + "-O3", + "--expt-relaxed-constexpr", + ] + + +@contextmanager +def override_jit_cuda_arch(major: int, minor: int, suffix: str = ""): + """A context manager to temporarily override CUDA architecture.""" + global _CUDA_ARCH + old_value = get_jit_cuda_arch() + _CUDA_ARCH = ArchInfo(major, minor, suffix) + try: + yield + finally: + _CUDA_ARCH = old_value + + +def get_jit_cuda_arch() -> ArchInfo: + """Get the current CUDA architecture info.""" + _init_jit_cuda_arch_once() + return _CUDA_ARCH + + +@cache_once +def is_arch_support_pdl() -> bool: + if is_hip_runtime() or is_musa_runtime(): + return False + return get_jit_cuda_arch().major >= 9 diff --git a/python/sglang/jit_kernel/utils/common.py b/python/sglang/jit_kernel/utils/common.py new file mode 100644 index 000000000..63437b532 --- /dev/null +++ b/python/sglang/jit_kernel/utils/common.py @@ -0,0 +1,79 @@ +"""Shared helpers: caching decorator, CI test gating, and runtime detection.""" + +from __future__ import annotations + +import functools +from typing import Any, Callable, Dict, List, TypeVar + +import torch + +from sglang.srt.environ import envs +from sglang.utils import is_in_ci + +F = TypeVar("F", bound=Callable[..., Any]) +T = TypeVar("T") + + +def should_run_full_tests() -> bool: + return envs.SGLANG_JIT_KERNEL_RUN_FULL_TESTS.get() + + +def get_ci_test_range(full_range: List[Any], ci_range: List[Any]) -> List[Any]: + if should_run_full_tests(): + return full_range + return ci_range if is_in_ci() else full_range + + +def cache_once(fn: F) -> F: + """ + NOTE: `functools.lru_cache` is not compatible with `torch.compile` + So we manually implement a simple cache_once decorator to replace it. + """ + result_map = {} + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + key = (args, tuple(sorted(kwargs.items()))) + if key not in result_map: + result_map[key] = fn(*args, **kwargs) + return result_map[key] + + return wrapper # type: ignore + + +@cache_once +def is_hip_runtime() -> bool: + return bool(torch.version.hip) + + +@cache_once +def is_musa_runtime() -> bool: + return hasattr(torch.version, "musa") and torch.version.musa is not None + + +_REGISTERED_CLASSES: Dict[type, type] = {} + + +def lazy_register_class(name: str, init_fn: Callable[[], None]) -> Callable[[T], T]: + """A decorator to lazily register a tvm-ffi object class on first use. + + `init_fn` runs once (typically JIT-compiling and registering the C++ + reflection) right before the class is registered under the FFI type key + `name`; afterwards instantiation proceeds normally. + """ + + def decorator(cls: T) -> T: + def __new__(cls, *args, **kwargs): + import tvm_ffi + + if cls not in _REGISTERED_CLASSES: + init_fn() # lazy initialization before registration once + _REGISTERED_CLASSES[cls] = tvm_ffi.register_object(name)(cls) + cls = _REGISTERED_CLASSES[cls] + return original_new(cls, *args, **kwargs) + + original_new = cls.__new__ + cls.__new__ = __new__ + return cls + + return decorator diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils/compile.py similarity index 50% rename from python/sglang/jit_kernel/utils.py rename to python/sglang/jit_kernel/utils/compile.py index 66ab63a58..1196d8559 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils/compile.py @@ -1,6 +1,7 @@ +"""JIT compilation: load_jit, the build cache, and C++ template arguments.""" + from __future__ import annotations -import functools import hashlib import importlib.util import logging @@ -8,89 +9,20 @@ import os import pathlib import re from contextlib import contextmanager -from dataclasses import dataclass -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - List, - Optional, - Tuple, - TypeAlias, - TypeVar, - Union, -) +from typing import TYPE_CHECKING, List, Tuple, TypeAlias, Union import torch -from sglang.srt.environ import envs -from sglang.utils import is_in_ci +from sglang.jit_kernel.utils.arch import get_default_target_flags, get_jit_cuda_arch +from sglang.jit_kernel.utils.common import cache_once, is_hip_runtime +from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES if TYPE_CHECKING: from tvm_ffi import Module -F = TypeVar("F", bound=Callable[..., Any]) - logger = logging.getLogger(__name__) -def should_run_full_tests() -> bool: - return envs.SGLANG_JIT_KERNEL_RUN_FULL_TESTS.get() - - -def get_ci_test_range(full_range: List[Any], ci_range: List[Any]) -> List[Any]: - if should_run_full_tests(): - return full_range - return ci_range if is_in_ci() else full_range - - -def cache_once(fn: F) -> F: - """ - NOTE: `functools.lru_cache` is not compatible with `torch.compile` - So we manually implement a simple cache_once decorator to replace it. - """ - result_map = {} - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - key = (args, tuple(sorted(kwargs.items()))) - if key not in result_map: - result_map[key] = fn(*args, **kwargs) - return result_map[key] - - return wrapper # type: ignore - - -_REGISTERED_CLASSES: Dict[type, type] = {} -T = TypeVar("T") - - -def lazy_register_class(name: str, init_fn: Callable[[], None]) -> Callable[[T], T]: - """A decorator to lazily register a tvm-ffi object class on first use. - - `init_fn` runs once (typically JIT-compiling and registering the C++ - reflection) right before the class is registered under the FFI type key - `name`; afterwards instantiation proceeds normally. - """ - - def decorator(cls: T) -> T: - def __new__(cls, *args, **kwargs): - import tvm_ffi - - if cls not in _REGISTERED_CLASSES: - init_fn() # lazy initialization before registration once - _REGISTERED_CLASSES[cls] = tvm_ffi.register_object(name)(cls) - cls = _REGISTERED_CLASSES[cls] - return original_new(cls, *args, **kwargs) - - original_new = cls.__new__ - cls.__new__ = __new__ - return cls - - return decorator - - def _make_wrapper(tup: Tuple[str, str]) -> str: export_name, kernel_name = tup return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));" @@ -140,7 +72,10 @@ def _local_jit_source_hash(source_files: List[str]) -> str: @cache_once def _resolve_kernel_path() -> pathlib.Path: - cur_dir = pathlib.Path(__file__).parent.resolve() + # Resolve via the package spec so the lookup is location-independent. + spec = importlib.util.find_spec("sglang.jit_kernel") + assert spec is not None and spec.origin is not None + cur_dir = pathlib.Path(spec.origin).parent.resolve() # first, try this directory structure def _environment_install(): @@ -172,28 +107,28 @@ class CPPArgList(list[str]): CPP_DTYPE_MAP = { - torch.float: "fp32_t", + torch.float64: "double", + torch.float32: "fp32_t", torch.float16: "fp16_t", - torch.float8_e4m3fn: "fp8_e4m3_t", torch.bfloat16: "bf16_t", + # The fnuz variants are the ROCm-side torch dtypes; fp8_*_t resolves to + # the matching HIP type there (see HIP_FP8_TYPE_* in utils.cuh). + torch.float8_e4m3fn: "fp8_e4m3_t", + torch.float8_e4m3fnuz: "fp8_e4m3_t", + torch.float8_e5m2: "fp8_e5m2_t", + torch.float8_e5m2fnuz: "fp8_e5m2_t", torch.int8: "int8_t", + torch.int16: "int16_t", torch.int32: "int32_t", torch.int64: "int64_t", + torch.uint8: "uint8_t", + torch.uint16: "uint16_t", + torch.uint32: "uint32_t", + torch.uint64: "uint64_t", + torch.bool: "bool", } -# AMD/ROCm note: -@cache_once -def is_hip_runtime() -> bool: - return bool(torch.version.hip) - - -# MThreads/MUSA note: -@cache_once -def is_musa_runtime() -> bool: - return hasattr(torch.version, "musa") and torch.version.musa is not None - - def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList: def _convert(arg: CPP_TEMPLATE_TYPE) -> str: if isinstance(arg, bool): @@ -294,9 +229,9 @@ def load_jit( cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files] for dep in set(extra_dependencies or []): - if dep not in _REGISTERED_DEPENDENCIES: + if dep not in REGISTERED_DEPENDENCIES: raise ValueError(f"Dependency {dep} is not registered.") - extra_include_paths += _REGISTERED_DEPENDENCIES[dep]() + extra_include_paths += REGISTERED_DEPENDENCIES[dep]() module_name = "sgl_kernel_jit_" + "_".join(str(arg) for arg in args) if cpp_files or cuda_files: @@ -338,7 +273,7 @@ def load_jit( cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_get_default_target_flags() + extra_cuda_cflags, + extra_cuda_cflags=get_default_target_flags() + extra_cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, @@ -351,40 +286,13 @@ def load_jit( cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_get_default_target_flags() + extra_cuda_cflags, + extra_cuda_cflags=get_default_target_flags() + extra_cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) -@dataclass -class ArchInfo: - major: int - minor: int - suffix: str - - @property - def target_name(self) -> str: - return f"{self.major}.{self.minor}{self.suffix}" - - @property - def jit_flag(self) -> str: - return f"-DSGL_CUDA_ARCH={self.major * 100 + self.minor * 10}" - - -@cache_once -def _init_jit_cuda_arch_once(): - global _CUDA_ARCH - try: - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) - except Exception: - logger.warning("Cannot detect CUDA architecture.") - major, minor = 0, 0 # invalid value to trigger compile error if used - _CUDA_ARCH = ArchInfo(major, minor, "") - - @contextmanager def _jit_compile_context(): if is_hip_runtime(): @@ -400,200 +308,3 @@ def _jit_compile_context(): os.environ.pop(env_key, None) else: os.environ[env_key] = old_value - - -# NOTE: this might also be used in __main__.py for compile flags export -def _get_default_target_flags() -> List[str]: - if is_hip_runtime(): - flags = ["-DUSE_ROCM", "-std=c++20", "-O3"] - # Detect FP8 type based on GPU architecture - try: - device = torch.cuda.current_device() - gcn_arch = torch.cuda.get_device_properties(device).gcnArchName - if "gfx942" in gcn_arch: - flags.append("-DHIP_FP8_TYPE_FNUZ=1") - else: - flags.append("-DHIP_FP8_TYPE_E4M3=1") - except Exception: - flags.append("-DHIP_FP8_TYPE_E4M3=1") - return flags - else: - return [ - get_jit_cuda_arch().jit_flag, - "-std=c++20", - "-O3", - "--expt-relaxed-constexpr", - ] - - -@contextmanager -def override_jit_cuda_arch(major: int, minor: int, suffix: str = ""): - """A context manager to temporarily override CUDA architecture.""" - global _CUDA_ARCH - old_value = get_jit_cuda_arch() - _CUDA_ARCH = ArchInfo(major, minor, suffix) - try: - yield - finally: - _CUDA_ARCH = old_value - - -def get_jit_cuda_arch() -> ArchInfo: - """Get the current CUDA architecture info.""" - _init_jit_cuda_arch_once() - return _CUDA_ARCH - - -@cache_once -def is_arch_support_pdl() -> bool: - if is_hip_runtime() or is_musa_runtime(): - return False - return get_jit_cuda_arch().major >= 9 - - -def _find_package_root(package: str) -> Optional[pathlib.Path]: - spec = importlib.util.find_spec(package) - if spec is None or spec.origin is None: - return None - return pathlib.Path(spec.origin).resolve().parent - - -# NOTE: this might also be used in __main__.py for compile flags export -_REGISTERED_DEPENDENCIES: Dict[str, Callable[[], List[str]]] = {} - - -def register_dependency(name: str): - def decorator(f: Callable[[], List[str]]) -> Callable[[], List[str]]: - if name in _REGISTERED_DEPENDENCIES: - raise ValueError(f"Dependency {name} already registered") - _REGISTERED_DEPENDENCIES[name] = f - return f - - return decorator - - -@register_dependency("flashinfer") -def get_flashinfer_include_paths() -> List[str]: - include_paths: List[str] = [] - flashinfer_root = _find_package_root("flashinfer") - if flashinfer_root is None: - raise RuntimeError( - "Cannot find flashinfer package. Please install flashinfer to get" - "the required headers for JIT compilation." - ) - - flashinfer_data = flashinfer_root / "data" - candidates = [ - flashinfer_data / "include", - flashinfer_data / "csrc", - flashinfer_data / "cutlass" / "include", - flashinfer_data / "cutlass" / "tools" / "util" / "include", - flashinfer_data / "spdlog" / "include", - ] - - for path in candidates: - if not path.exists(): - raise RuntimeError( - f"Required header path {path} for flashinfer dependency not found." - " Please check your flashinfer installation." - ) - include_paths.append(str(path)) - return include_paths - - -def get_mathdx_root() -> Optional[pathlib.Path]: - """Locate the NVIDIA Math-DX install (cuBLASDx headers). - - Searches in order: - 1. ``$MATHDX_HOME`` env var (extracted Math-DX archive root). - 2. The ``nvidia-mathdx`` PyPI package, if installed. - """ - env_home = os.environ.get("MATHDX_HOME") - if env_home: - candidate = pathlib.Path(env_home).expanduser().resolve() - if (candidate / "include").exists(): - return candidate - - # The ``nvidia-mathdx`` wheel installs as the namespace package - # ``nvidia.mathdx`` (no __init__, so spec.origin is None); resolve it via - # submodule_search_locations rather than _find_package_root, which only - # handles regular packages. - spec = importlib.util.find_spec("nvidia.mathdx") - if spec is not None: - roots = list(spec.submodule_search_locations or []) - if spec.origin is not None: - roots.append(str(pathlib.Path(spec.origin).parent)) - for root in roots: - candidate = pathlib.Path(root).resolve() - if (candidate / "include").exists(): - return candidate - - return None - - -@register_dependency("mathdx") -def get_mathdx_include_paths() -> List[str]: - root = get_mathdx_root() - if root is None: - raise RuntimeError( - "Cannot find NVIDIA Math-DX (cuBLASDx) headers. " - "Install the `nvidia-mathdx` package " - "(`pip install nvidia-mathdx`) or set MATHDX_HOME to an " - "extracted Math-DX archive root." - ) - candidates = [root / "include"] - cutlass = root / "external" / "cutlass" / "include" - if cutlass.exists(): - candidates.append(cutlass) - return [str(p) for p in candidates] - - -@register_dependency("cutlass") -def get_cutlass_include_paths() -> List[str]: - include_paths: List[str] = [] - - flashinfer_root = _find_package_root("flashinfer") - if flashinfer_root is not None: - candidates = [ - flashinfer_root / "data" / "cutlass" / "include", - flashinfer_root / "data" / "cutlass" / "tools" / "util" / "include", - ] - for path in candidates: - if path.exists(): - include_paths.append(str(path)) - - deep_gemm_root = _find_package_root("deep_gemm") - if deep_gemm_root is not None: - candidate = deep_gemm_root / "include" - if candidate.exists(): - include_paths.append(str(candidate)) - - # De-duplicate while preserving order. - unique_paths = [] - seen = set() - for path in include_paths: - if path in seen: - continue - seen.add(path) - unique_paths.append(path) - - if not unique_paths: - raise RuntimeError( - "Cannot find CUTLASS headers required for JIT compilation. " - "Please install flashinfer or deep_gemm with CUTLASS headers." - ) - return unique_paths - - -__all__ = [ - "should_run_full_tests", - "get_ci_test_range", - "cache_once", - "is_hip_runtime", - "make_cpp_args", - "load_jit", - "override_jit_cuda_arch", - "get_jit_cuda_arch", - "is_arch_support_pdl", - "register_dependency", -] diff --git a/python/sglang/jit_kernel/utils/deps.py b/python/sglang/jit_kernel/utils/deps.py new file mode 100644 index 000000000..16ce9b22f --- /dev/null +++ b/python/sglang/jit_kernel/utils/deps.py @@ -0,0 +1,141 @@ +"""Header-only dependency registration (flashinfer, cutlass, mathdx, ...).""" + +from __future__ import annotations + +import importlib.util +import os +import pathlib +from typing import Callable, Dict, List, Optional + + +def _find_package_root(package: str) -> Optional[pathlib.Path]: + spec = importlib.util.find_spec(package) + if spec is None or spec.origin is None: + return None + return pathlib.Path(spec.origin).resolve().parent + + +REGISTERED_DEPENDENCIES: Dict[str, Callable[[], List[str]]] = {} + + +def register_dependency(name: str): + def decorator(f: Callable[[], List[str]]) -> Callable[[], List[str]]: + if name in REGISTERED_DEPENDENCIES: + raise ValueError(f"Dependency {name} already registered") + REGISTERED_DEPENDENCIES[name] = f + return f + + return decorator + + +@register_dependency("flashinfer") +def get_flashinfer_include_paths() -> List[str]: + include_paths: List[str] = [] + flashinfer_root = _find_package_root("flashinfer") + if flashinfer_root is None: + raise RuntimeError( + "Cannot find flashinfer package. Please install flashinfer to get" + "the required headers for JIT compilation." + ) + + flashinfer_data = flashinfer_root / "data" + candidates = [ + flashinfer_data / "include", + flashinfer_data / "csrc", + flashinfer_data / "cutlass" / "include", + flashinfer_data / "cutlass" / "tools" / "util" / "include", + flashinfer_data / "spdlog" / "include", + ] + + for path in candidates: + if not path.exists(): + raise RuntimeError( + f"Required header path {path} for flashinfer dependency not found." + " Please check your flashinfer installation." + ) + include_paths.append(str(path)) + return include_paths + + +def get_mathdx_root() -> Optional[pathlib.Path]: + """Locate the NVIDIA Math-DX install (cuBLASDx headers). + + Searches in order: + 1. ``$MATHDX_HOME`` env var (extracted Math-DX archive root). + 2. The ``nvidia-mathdx`` PyPI package, if installed. + """ + env_home = os.environ.get("MATHDX_HOME") + if env_home: + candidate = pathlib.Path(env_home).expanduser().resolve() + if (candidate / "include").exists(): + return candidate + + # The ``nvidia-mathdx`` wheel installs as the namespace package + # ``nvidia.mathdx`` (no __init__, so spec.origin is None); resolve it via + # submodule_search_locations rather than _find_package_root, which only + # handles regular packages. + spec = importlib.util.find_spec("nvidia.mathdx") + if spec is not None: + roots = list(spec.submodule_search_locations or []) + if spec.origin is not None: + roots.append(str(pathlib.Path(spec.origin).parent)) + for root in roots: + candidate = pathlib.Path(root).resolve() + if (candidate / "include").exists(): + return candidate + + return None + + +@register_dependency("mathdx") +def get_mathdx_include_paths() -> List[str]: + root = get_mathdx_root() + if root is None: + raise RuntimeError( + "Cannot find NVIDIA Math-DX (cuBLASDx) headers. " + "Install the `nvidia-mathdx` package " + "(`pip install nvidia-mathdx`) or set MATHDX_HOME to an " + "extracted Math-DX archive root." + ) + candidates = [root / "include"] + cutlass = root / "external" / "cutlass" / "include" + if cutlass.exists(): + candidates.append(cutlass) + return [str(p) for p in candidates] + + +@register_dependency("cutlass") +def get_cutlass_include_paths() -> List[str]: + include_paths: List[str] = [] + + flashinfer_root = _find_package_root("flashinfer") + if flashinfer_root is not None: + candidates = [ + flashinfer_root / "data" / "cutlass" / "include", + flashinfer_root / "data" / "cutlass" / "tools" / "util" / "include", + ] + for path in candidates: + if path.exists(): + include_paths.append(str(path)) + + deep_gemm_root = _find_package_root("deep_gemm") + if deep_gemm_root is not None: + candidate = deep_gemm_root / "include" + if candidate.exists(): + include_paths.append(str(candidate)) + + # De-duplicate while preserving order. + unique_paths = [] + seen = set() + for path in include_paths: + if path in seen: + continue + seen.add(path) + unique_paths.append(path) + + if not unique_paths: + raise RuntimeError( + "Cannot find CUTLASS headers required for JIT compilation. " + "Please install flashinfer or deep_gemm with CUTLASS headers." + ) + return unique_paths