sglang/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh

/// \file utils.cuh
/// \brief Core CUDA/device utilities: type aliases, PDL helpers,
///        typed pointer access, kernel launch wrapper, and error checking.
///
/// This header is included (directly or transitively) by nearly every
/// JIT kernel. It provides:
/// - Scalar/packed type aliases (`fp16_t`, `bf16_t`, `fp8_e4m3_t`, ...).
/// - `SGL_DEVICE` macro (forced-inline device function qualifier).
/// - `kWarpThreads` constant (32).
/// - PDL (Programmatic Dependent Launch) helpers for Hopper (sm_90+).
/// - Typed `load_as` / `store_as` for void-pointer access.
/// - `pointer::offset` for safe void-pointer arithmetic.
/// - `host::LaunchKernel` - kernel launcher with optional PDL.
/// - `host::RuntimeDeviceCheck` - CUDA error checking.

#pragma once

#include <sgl_kernel/bits.h>
#include <sgl_kernel/utils.h>

#include <dlpack/dlpack.h>
#include <tvm/ffi/extra/c_env_api.h>

#include <concepts>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <type_traits>
#ifndef USE_ROCM
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#else
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
#ifndef __grid_constant__
#define __grid_constant__
#endif
using cudaError_t = hipError_t;
using cudaStream_t = hipStream_t;
using cudaLaunchConfig_t = hipLaunchConfig_t;
using cudaLaunchAttribute = hipLaunchAttribute;
inline constexpr auto cudaSuccess = hipSuccess;
#define cudaStreamPerThread hipStreamPerThread
#define cudaGetErrorString hipGetErrorString
#define cudaGetLastError hipGetLastError
#define cudaLaunchKernel hipLaunchKernel
#define cudaMemcpyAsync hipMemcpyAsync
#define cudaMemcpyHostToDevice hipMemcpyHostToDevice
#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
#define cudaDeviceGetAttribute hipDeviceGetAttribute
#define cudaDevAttrComputeCapabilityMajor hipDeviceAttributeComputeCapabilityMajor
#define cudaDevAttrComputeCapabilityMinor hipDeviceAttributeComputeCapabilityMinor
#define cudaFuncSetAttribute hipFuncSetAttribute
#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
#endif

namespace sglang {

#ifndef USE_ROCM
using fp32_t = float;
using fp16_t = __half;
using bf16_t = __nv_bfloat16;
using fp8_e4m3_t = __nv_fp8_e4m3;
using fp8_e5m2_t = __nv_fp8_e5m2;

using fp32x2_t = float2;
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
using fp32_t = float;
using fp16_t = __half;
using bf16_t = __hip_bfloat16;
using fp8_e4m3_t = uint8_t;
using fp8_e5m2_t = uint8_t;
using fp32x2_t = float2;
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

/*
 * LDG Support
 */
#ifndef USE_ROCM
#define SGLANG_LDG(arg) __ldg(arg)
#else
#define SGLANG_LDG(arg) *(arg)
#endif

// DLPack device type for the current platform
#ifndef USE_ROCM
inline constexpr auto kDLGPU = kDLCUDA;
inline constexpr auto kDLGPUHost = kDLCUDAHost;
#else
inline constexpr auto kDLGPU = kDLROCM;
inline constexpr auto kDLGPUHost = kDLROCMHost;
#endif

namespace device {

/// \brief Macro: forced-inline device function qualifier.
#define SGL_DEVICE __forceinline__ __device__
#define SGL_DEVICE_HOST __forceinline__ __device__ __host__

// 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

// Maximum vector size in bytes supported by current architecture.
// Pre-Blackwell / AMD: 128-bit (16 bytes)
// Blackwell or greater: 256-bit (32 bytes)
inline constexpr std::size_t kMaxVecBytes = SGL_ARCH_BLACKWELL_OR_GREATER ? 32 : 16;

/// \brief Number of threads per warp (always 32 on NVIDIA/AMD GPUs).
inline constexpr uint32_t kWarpThreads = 32u;
/// \brief Most implementations prefer this name; keep the alias for them.
inline constexpr uint32_t kWarpSize = kWarpThreads;

/**
 * \brief This thread's index within its logical `kNumThreads` group.
 *
 * \tparam kNumThreads Group width; a power of two, at most 32 on CUDA and at
 * most 64 (the wave) on HIP -- so `64` is a HIP-only instantiation.
 *
 * \note Equals the true in-warp lane only when `blockDim.x` is a multiple of
 * `kNumThreads`; every caller in this tree satisfies that.
 * \note On CUDA prefer this over `threadIdx.x % kNumThreads` when the value
 * feeds an address: `%laneid` is one register read that folds straight into
 * `IMAD.WIDE`, while the modulo makes ptxas re-derive the mask at every address
 * scale. Worth 8 instructions in a two-tile warp copy, measured on sm_100a.
 * That only holds at full width -- a narrower group needs the mask anyway and
 * ties with the modulo.
 */
template <uint32_t kNumThreads = kWarpThreads>
SGL_DEVICE uint32_t get_lane_id() {
#ifndef USE_ROCM
  static_assert(kNumThreads <= 32 && host::is_pow2(kNumThreads));
  uint32_t lane_id;
  asm volatile("mov.u32 %0, %%laneid;" : "=r"(lane_id));
  if constexpr (kNumThreads != 32) lane_id %= kNumThreads;
  return lane_id;
#else
  static_assert(kNumThreads <= 64 && host::is_pow2(kNumThreads));
  // AMD has no lane-id register: `__lane_id()` is computed from the exec mask as
  // a `v_mbcnt_lo`/`v_mbcnt_hi` pair, and the group mask is still needed on top.
  // Masking `threadIdx.x` -- already live in v0 -- is 2 instructions cheaper and
  // yields the same value (measured on gfx950, hipcc 7.0).
  return threadIdx.x % kNumThreads;
#endif
}

/**
 * \brief PDL (Programmatic Dependent Launch): wait for the primary kernel.
 *
 * On Hopper (sm_90+), inserts a `griddepcontrol.wait` instruction to
 * synchronize with a preceding kernel in the same stream. On older
 * architectures or ROCm this is a no-op.
 *
 *\note This is the only thing that orders us against the producer. Per the PTX
 * ISA, `.wait` makes the executing thread wait until every prerequisite grid in
 * flight has COMPLETED and all of its memory operations are performed and made
 * visible to this grid -- so it is what a `PDLTriggerSecondary` upstream does
 * NOT give us. It acts per thread, so every thread that reads producer data has
 * to execute it; put it ahead of the first such load. Stores into our own output
 * buffers depend on nothing upstream and may be issued before it.
 */
template <bool kUsePDL>
SGL_DEVICE void PDLWaitPrimary() {
#if SGL_ARCH_HOPPER_OR_GREATER
  if constexpr (kUsePDL) {
    asm volatile("griddepcontrol.wait;" ::: "memory");
  }
#endif
}

/**
 * \brief PDL: trigger dependent (secondary) kernel launch.
 *
 * On Hopper (sm_90+), inserts a `griddepcontrol.launch_dependents`
 * instruction. On older architectures or ROCm this is a no-op.
 *
 * \note Scheduling only: this carries no memory ordering of its own. The
 * dependent becomes eligible to launch once every CTA in this grid has issued
 * the instruction or has exited, and it may then start before our writes are
 * visible -- making them visible is the job of `PDLWaitPrimary` on the dependent
 * side, which is why the programming guide requires the dependent to call it.
 *
 * Granularity is the CTA: the PTX ISA states that repeated invocations by
 * threads of the same CTA have no side effect past the first, so one thread
 * would do; we call it from all of them because it is free and needs no
 * predication. Leaving it out altogether is safe and merely late, since the
 * trigger is implied once every CTA exits (SASS code `PREEXIT`)
 *
 * Placing it early therefore costs nothing and only buys the dependent a head
 * start on the work that does not depend on us. Even that is opportunistic:
 * concurrent execution is never guaranteed, so nothing may rely on it.
 */
template <bool kUsePDL>
SGL_DEVICE void PDLTriggerSecondary() {
#if SGL_ARCH_HOPPER_OR_GREATER
  if constexpr (kUsePDL) {
    // The "memory" clobber is load-bearing: without it the compiler may sink
    // this kernel's stores past the trigger, and the dependent grid's
    // griddepcontrol.wait only covers writes issued BEFORE launch_dependents.
    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
  }
#endif
}

template <std::integral T, std::integral U>
SGL_DEVICE constexpr auto div_ceil(T a, U b) {
  return (a + b - 1) / b;
}

/**
 * \brief Load data with the specified type and offset from a void pointer.
 * \tparam T The type to load.
 * \param ptr The base pointer.
 * \param offset The offset in number of elements of type T.
 */
template <typename T>
SGL_DEVICE T load_as(const void* ptr, int64_t offset = 0) {
  return static_cast<const T*>(ptr)[offset];
}

/**
 * \brief Store data with the specified type and offset to a void pointer.
 * \tparam T The type to store.
 * \param ptr The base pointer.
 * \param val The value to store.
 * \param offset The offset in number of elements of type T.
 * \note we use type_identity_t to force the caller to explicitly specify
 * the template parameter `T`, which can avoid accidentally using the wrong type.
 */
template <typename T>
SGL_DEVICE void store_as(void* ptr, std::type_identity_t<T> val, int64_t offset = 0) {
  static_cast<T*>(ptr)[offset] = val;
}

/// \brief Safe void-pointer arithmetic (byte-level by default).
namespace pointer {

// we only allow void * pointer arithmetic for safety

template <typename T = char, std::integral... U>
SGL_DEVICE auto offset(void* ptr, U... offset) -> void* {
  return static_cast<T*>(ptr) + (... + offset);
}

template <typename T = char, std::integral... U>
SGL_DEVICE auto offset(const void* ptr, U... offset) -> const void* {
  return static_cast<const T*>(ptr) + (... + offset);
}

}  // namespace pointer

/// 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\";");
#endif
}

template <typename T, std::size_t N>
struct DeviceArray {
 public:
  SGL_DEVICE constexpr static std::size_t size() {
    return N;
  }
  SGL_DEVICE constexpr auto operator[](std::size_t idx) -> T& {
    return m_data[idx];
  }
  SGL_DEVICE constexpr auto operator[](std::size_t idx) const -> const T& {
    return m_data[idx];
  }
  SGL_DEVICE constexpr auto data() const -> const T* {
    return m_data;
  }
  SGL_DEVICE constexpr auto data() -> T* {
    return m_data;
  }

 private:
  T m_data[N];
};

/**
 * Adapted from
 * https://github.com/deepseek-ai/DeepGEMM/blob/559d79fb6994a58b8a15b4b93bf13ccc16edf247/deep_gemm/include/deep_gemm/common/utils.cuh
 */
SGL_DEVICE_HOST constexpr uint32_t get_tmem_cols(uint32_t num_cols) {
  if (num_cols <= 32) return 32;
  if (num_cols <= 64) return 64;
  if (num_cols <= 128) return 128;
  if (num_cols <= 256) return 256;
  return 512;
}

}  // namespace device

namespace host {

/**
 * \brief Check the CUDA error code and panic with location info on failure.
 */
inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) {
  if (error != ::cudaSuccess) {
    [[unlikely]];
    host::panic(location, "CUDA error: ", ::cudaGetErrorString(error));
  }
}

/// \brief Check the last CUDA error (calls `cudaGetLastError`).
inline void RuntimeDeviceCheck(DebugInfo location = {}) {
  return RuntimeDeviceCheck(::cudaGetLastError(), location);
}

/**
 * \brief Kernel launcher with automatic stream resolution and PDL support.
 *
 * Usage:
 * \code
 *   host::LaunchKernel(grid, block, device)
 *       .enable_pdl(true)(my_kernel, arg0, arg1);
 *   host::LaunchKernel(grid, block, stream)
 *       .config({.use_pdl = true, .cluster_dim = cluster_dim})(my_kernel, arg0);
 * \endcode
 *
 * The constructor resolves the CUDA stream from a `DLDevice` (via `TVMFFIEnvGetStream`)
 * or accepts a raw `cudaStream_t`. The call operator launches the kernel and checks for errors.
 */
struct LaunchKernel {
 private:
  struct KernelConfig {
    bool use_pdl = false;
    std::optional<dim3> cluster_dim = std::nullopt;
  };

 public:
  explicit LaunchKernel(
      dim3 grid_dim,
      dim3 block_dim,
      DLDevice device,
      std::size_t dynamic_shared_mem_bytes = 0,
      DebugInfo location = {}) noexcept
      : m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)),
        m_location(location) {}

  explicit LaunchKernel(
      dim3 grid_dim,
      dim3 block_dim,
      cudaStream_t stream,
      std::size_t dynamic_shared_mem_bytes = 0,
      DebugInfo location = {}) noexcept
      : m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)), m_location(location) {}

  LaunchKernel(const LaunchKernel&) = delete;
  LaunchKernel& operator=(const LaunchKernel&) = delete;

  static auto resolve_device(DLDevice device) -> cudaStream_t {
    return static_cast<cudaStream_t>(::TVMFFIEnvGetStream(device.device_type, device.device_id));
  }

  auto enable_pdl(bool enabled = true) -> LaunchKernel& {
#ifdef USE_ROCM
    (void)enabled;
    m_config.numAttrs = 0;
#else
    if (enabled) {
      auto& attr = m_attrs[m_config.numAttrs++];
      attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
      attr.val.programmaticStreamSerializationAllowed = true;
      m_config.attrs = m_attrs;
    }
#endif
    return *this;
  }

  auto enable_cluster(dim3 cluster_dim) -> LaunchKernel& {
#ifdef USE_ROCM
    (void)cluster_dim;
#else
    auto& attr = m_attrs[m_config.numAttrs++];
    attr.id = cudaLaunchAttributeClusterDimension;
    attr.val.clusterDim = {cluster_dim.x, cluster_dim.y, cluster_dim.z};
    m_config.attrs = m_attrs;
#endif
    return *this;
  }

  /**
   * \brief Configure the kernel launch with the given options.
   * \param config The kernel configuration options.
   * \return A reference to this `LaunchKernel` for chaining.
   * \note This is a convenience method that applies multiple configurations at once.
   * We are in favor of this instead of `enable_pdl` and `enable_cluster`.
   * We enforce use of designated initializers for better readability.
   */
  auto config(const KernelConfig& config) -> LaunchKernel& {
    if (config.use_pdl) this->enable_pdl(true);
    if (config.cluster_dim) this->enable_cluster(*config.cluster_dim);
    return *this;
  }

  template <typename T, typename... Args>
  auto operator()(T&& kernel, Args&&... args) const -> void {
#ifdef USE_ROCM
    hipLaunchKernelGGL(
        std::forward<T>(kernel),
        m_config.gridDim,
        m_config.blockDim,
        m_config.dynamicSmemBytes,
        m_config.stream,
        std::forward<Args>(args)...);
    RuntimeDeviceCheck(m_location);
#else
    RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...), m_location);
#endif
  }

  template <typename T, typename... Args>
  auto launch(T&& kernel, Args&&... args) const -> void {
    return (*this)(std::forward<T>(kernel), std::forward<Args>(args)...);
  }

 private:
  static auto s_make_config(  // Make a config for kernel launch
      dim3 grid_dim,
      dim3 block_dim,
      cudaStream_t stream,
      std::size_t smem) -> cudaLaunchConfig_t {
    auto config = ::cudaLaunchConfig_t{};
    config.gridDim = grid_dim;
    config.blockDim = block_dim;
    config.dynamicSmemBytes = smem;
    config.stream = stream;
    config.numAttrs = 0;
    return config;
  }

  cudaLaunchConfig_t m_config;
  const DebugInfo m_location;
  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

}  // namespace sglang