[Feature] Introduce JIT Kernel in sglang (with hicache JIT kernel) (#13453)

python/sglang/jit_kernel/include/sgl_kernel/tensor.h

@@ -0,0 +1,488 @@
+#pragma once
+#include <sgl_kernel/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/dtype.h>
+
+#include <algorithm>
+#include <array>
+#include <concepts>
+#include <cstddef>
+#include <cstdint>
+#include <initializer_list>
+#include <optional>
+#include <ranges>
+#include <source_location>
+#include <span>
+#include <sstream>
+#include <string>
+#include <string_view>
+#include <type_traits>
+#include <utility>
+
+namespace host {
+
+namespace stdr = std::ranges;
+namespace stdv = std::views;
+
+namespace details {
+
+struct SizeRef;
+struct DTypeRef;
+struct DeviceRef;
+
+template <typename T>
+struct dtype_trait {};
+
+template <std::integral T>
+struct dtype_trait<T> {
+  inline static constexpr auto value = DLDataType{
+      .code = std::is_signed_v<T> ? DLDataTypeCode::kDLInt : DLDataTypeCode::kDLUInt,
+      .bits = static_cast<std::uint8_t>(sizeof(T) * 8),
+      .lanes = 1};
+};
+
+template <std::floating_point T>
+struct dtype_trait<T> {
+  inline static constexpr auto value =
+      DLDataType{.code = DLDataTypeCode::kDLFloat, .bits = static_cast<std::uint8_t>(sizeof(T) * 8), .lanes = 1};
+};
+
+inline constexpr auto kAnyDeviceID = -1;
+inline constexpr auto kAnySize = static_cast<int64_t>(-1);
+inline constexpr auto kNullSize = static_cast<int64_t>(0);
+inline constexpr auto kNullDType = static_cast<DLDataTypeCode>(18u);
+inline constexpr auto kNullDevice = static_cast<DLDeviceType>(-1);
+
+template <typename... Ts>
+inline constexpr auto kDTypeList = std::array{dtype_trait<Ts>::value...};
+
+template <auto... Codes>
+inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{
+    DLDevice{.device_type = static_cast<DLDeviceType>(Codes), .device_id = kAnyDeviceID}...};
+
+template <typename T>
+struct PrintAbleSpan {
+  explicit PrintAbleSpan(std::span<const T> data) : data(data) {}
+  std::span<const T> data;
+};
+
+// define DLDataType comparison and printing in root namespace
+template <void* = nullptr>
+inline constexpr auto kDeviceStringMap = [] {
+  constexpr auto map = std::array{
+      std::pair{DLDeviceType::kDLCPU, "cpu"},
+      std::pair{DLDeviceType::kDLCUDA, "cuda"},
+      std::pair{DLDeviceType::kDLCUDAHost, "cuda_host"},
+      std::pair{DLDeviceType::kDLOpenCL, "opencl"},
+      std::pair{DLDeviceType::kDLVulkan, "vulkan"},
+      std::pair{DLDeviceType::kDLMetal, "metal"},
+      std::pair{DLDeviceType::kDLVPI, "vpi"},
+      std::pair{DLDeviceType::kDLROCM, "rocm"},
+      std::pair{DLDeviceType::kDLROCMHost, "rocm_host"},
+      std::pair{DLDeviceType::kDLExtDev, "ext_dev"},
+      std::pair{DLDeviceType::kDLCUDAManaged, "cuda_managed"},
+      std::pair{DLDeviceType::kDLOneAPI, "oneapi"},
+      std::pair{DLDeviceType::kDLWebGPU, "webgpu"},
+      std::pair{DLDeviceType::kDLHexagon, "hexagon"},
+      std::pair{DLDeviceType::kDLMAIA, "maia"},
+      std::pair{DLDeviceType::kDLTrn, "trn"},
+  };
+  constexpr auto max_type = stdr::max(map | stdv::keys);
+  auto result = std::array<std::string_view, max_type + 1>{};
+  for (const auto& [code, name] : map) {
+    result[static_cast<std::size_t>(code)] = name;
+  }
+  return result;
+}();
+
+struct PrintableDevice {
+  DLDevice device;
+};
+
+inline auto& operator<<(std::ostream& os, DLDevice device) {
+  const auto& mapping = kDeviceStringMap<>;
+  const auto entry = static_cast<std::size_t>(device.device_type);
+  host::RuntimeCheck(entry < mapping.size());
+  const auto name = mapping[entry];
+  host::RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type));
+  os << name;
+  if (device.device_id != kAnyDeviceID) os << "[" << device.device_id << "]";
+  return os;
+}
+
+inline auto& operator<<(std::ostream& os, PrintableDevice pd) {
+  return os << pd.device;
+}
+
+template <typename T>
+inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
+  os << "[";
+  for (const auto i : stdv::iota(std::size_t{0}, span.data.size())) {
+    if (i > 0) {
+      os << ", ";
+    }
+    os << span.data[i];
+  }
+  os << "]";
+  return os;
+}
+
+}  // namespace details
+
+struct SymbolicSize {
+ public:
+  SymbolicSize(std::string_view annotation = {}) : m_value(details::kNullSize), m_annotation(annotation) {}
+
+  auto get_name() const -> std::string_view {
+    return m_annotation;
+  }
+  auto set_value(int64_t value) -> void {
+    host::RuntimeCheck(!this->has_value(), "Size value already set");
+    m_value = value;
+  }
+  auto has_value() const -> bool {
+    return m_value != 0;
+  }
+  auto get_value() const -> std::optional<int64_t> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+  auto unwrap() const -> int64_t {
+    host::RuntimeCheck(this->has_value(), "Size value is not set");
+    return m_value;
+  }
+
+  SymbolicSize(const SymbolicSize&) = delete;
+  SymbolicSize& operator=(const SymbolicSize&) = delete;
+
+  auto verify(int64_t dim) -> void {
+    if (this->has_value()) {
+      host::RuntimeCheck(m_value == dim, "Size mismatch: expected ", m_value, " but got ", dim);
+    } else {
+      this->set_value(dim);
+    }
+  }
+
+ private:
+  std::int64_t m_value;
+  std::string_view m_annotation;
+};
+
+inline auto operator==(DLDevice lhs, DLDevice rhs) -> bool {
+  return lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id;
+}
+
+struct SymbolicDType {
+ public:
+  SymbolicDType() : m_value({details::kNullDType, 0, 0}) {}
+
+  auto set_value(DLDataType value) -> void {
+    host::RuntimeCheck(!this->has_value(), "Dtype value already set");
+    host::RuntimeCheck(
+        m_check(value), "Dtype value [", value, "] not in the allowed options: ", details::PrintAbleSpan{m_options});
+    m_value = value;
+  }
+  auto has_value() const -> bool {
+    return m_value.code != details::kNullDType;
+  }
+  auto get_value() const -> std::optional<DLDataType> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+  auto unwrap() const -> DLDataType {
+    host::RuntimeCheck(this->has_value(), "Dtype value is not set");
+    return m_value;
+  }
+
+  auto set_options(std::span<const DLDataType> options) -> void {
+    m_options = options;
+  }
+  template <typename... Ts>
+  auto set_options() -> void {
+    m_options = details::kDTypeList<Ts...>;
+  }
+
+  auto verify(DLDataType dtype) -> void {
+    if (this->has_value()) {
+      host::RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype);
+    } else {
+      this->set_value(dtype);
+    }
+  }
+
+ private:
+  auto m_check(DLDataType value) const -> bool {
+    return stdr::empty(m_options) || (stdr::find(m_options, value) != stdr::end(m_options));
+  }
+
+  std::span<const DLDataType> m_options;
+  DLDataType m_value;
+};
+
+struct SymbolicDevice {
+ public:
+  SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {}
+
+  auto set_value(DLDevice value) -> void {
+    host::RuntimeCheck(!this->has_value(), "Device value already set");
+    host::RuntimeCheck(
+        m_check(value),
+        "Device value [",
+        details::PrintableDevice{value},
+        "] not in the allowed options: ",
+        details::PrintAbleSpan{m_options});
+    m_value = value;
+  }
+  auto has_value() const -> bool {
+    return m_value.device_type != details::kNullDevice;
+  }
+  auto get_value() const -> std::optional<DLDevice> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+  auto unwrap() const -> DLDevice {
+    host::RuntimeCheck(this->has_value(), "Device value is not set");
+    return m_value;
+  }
+
+  auto set_options(std::span<const DLDevice> options) -> void {
+    m_options = options;
+  }
+  template <DLDeviceType... Codes>
+  auto set_options() -> void {
+    m_options = details::kDeviceList<Codes...>;
+  }
+
+  auto verify(DLDevice device) -> void {
+    if (this->has_value()) {
+      host::RuntimeCheck(
+          m_value == device,
+          "Device mismatch: expected ",
+          details::PrintableDevice{m_value},
+          " but got ",
+          details::PrintableDevice{device});
+    } else {
+      this->set_value(device);
+    }
+  }
+
+ private:
+  auto m_check(DLDevice value) const -> bool {
+    return stdr::empty(m_options) || (stdr::any_of(m_options, [value](const DLDevice& opt) {
+             // device type must exactly match
+             if (opt.device_type != value.device_type) return false;
+             // device id can be wildcarded
+             return opt.device_id == details::kAnyDeviceID || opt.device_id == value.device_id;
+           }));
+  }
+
+  std::span<const DLDevice> m_options;
+  DLDevice m_value;
+};
+
+namespace details {
+
+template <typename T>
+struct BaseRef {
+ public:
+  BaseRef(const BaseRef&) = delete;
+  BaseRef& operator=(const BaseRef&) = delete;
+
+  auto operator->() const -> T* {
+    return m_ref;
+  }
+  auto operator*() const -> T& {
+    return *m_ref;
+  }
+  auto rebind(T& other) -> void {
+    m_ref = &other;
+  }
+
+  explicit BaseRef() : m_ref(&m_cache), m_cache() {}
+  BaseRef(T& size) : m_ref(&size), m_cache() {}
+
+ private:
+  T* m_ref;
+  T m_cache;
+};
+
+struct SizeRef : BaseRef<SymbolicSize> {
+  using BaseRef::BaseRef;
+  SizeRef(int64_t value) {
+    if (value != kAnySize) {
+      (**this).set_value(value);
+    } else {
+      // otherwise, we can match any size
+    }
+  }
+
+  auto value_or_name(std::size_t dim) const -> std::string {
+    if (const auto value = (**this).get_value()) {
+      return std::to_string(*value);
+    } else {
+      const auto annotation = (**this).get_name();
+      if (annotation.empty()) {
+        return "dim#" + std::to_string(dim);
+      } else {
+        return static_cast<std::string>(annotation);
+      }
+    }
+  }
+};
+
+struct DTypeRef : BaseRef<SymbolicDType> {
+  using BaseRef::BaseRef;
+  DTypeRef(DLDataType options) {
+    (**this).set_value(options);
+  }
+  DTypeRef(std::initializer_list<DLDataType> options) {
+    (**this).set_options(options);
+  }
+  DTypeRef(std::span<const DLDataType> options) {
+    (**this).set_options(options);
+  }
+};
+
+struct DeviceRef : BaseRef<SymbolicDevice> {
+  using BaseRef::BaseRef;
+  DeviceRef(DLDevice options) {
+    (**this).set_value(options);
+  }
+  DeviceRef(std::initializer_list<DLDevice> options) {
+    (**this).set_options(options);
+  }
+  DeviceRef(std::span<const DLDevice> options) {
+    (**this).set_options(options);
+  }
+};
+
+}  // namespace details
+
+struct TensorMatcher {
+ private:
+  using SizeRef = details::SizeRef;
+  using DTypeRef = details::DTypeRef;
+  using DeviceRef = details::DeviceRef;
+  using Loc_t = std::source_location;
+
+ public:
+  TensorMatcher(const TensorMatcher&) = delete;
+  TensorMatcher& operator=(const TensorMatcher&) = delete;
+
+  explicit TensorMatcher(std::initializer_list<SizeRef> shape) : m_shape(shape), m_strides(), m_dtype() {}
+
+  auto with_strides(std::initializer_list<SizeRef> strides) && -> TensorMatcher&& {
+    // no partial update allowed
+    host::RuntimeCheck(m_strides.size() == 0, "Strides already specified");
+    host::RuntimeCheck(m_shape.size() == strides.size(), "Strides size must match shape size");
+    m_strides = strides;
+    return std::move(*this);
+  }
+
+  template <typename... Ts>
+  auto with_dtype(DTypeRef&& dtype) && -> TensorMatcher&& {
+    m_init_dtype();
+    m_dtype.rebind(*dtype);
+    return std::move(*this);
+  }
+
+  template <typename... Ts>
+  auto with_dtype() && -> TensorMatcher&& {
+    static_assert(sizeof...(Ts) > 0, "At least one dtype option must be specified");
+    m_init_dtype();
+    m_dtype->set_options<Ts...>();
+    return std::move(*this);
+  }
+
+  template <DLDeviceType... Codes>
+  auto with_device(DeviceRef&& device) && -> TensorMatcher&& {
+    m_init_device();
+    m_device.rebind(*device);
+    return std::move(*this);
+  }
+
+  template <DLDeviceType... Codes>
+  auto with_device() && -> TensorMatcher&& {
+    static_assert(sizeof...(Codes) > 0, "At least one device option must be specified");
+    m_init_device();
+    m_device->set_options<Codes...>();
+    return std::move(*this);
+  }
+
+  // once we start verification, we cannot modify anymore
+  auto verify(tvm::ffi::TensorView view, Loc_t loc = Loc_t::current()) const&& -> const TensorMatcher&& {
+    try {
+      this->m_verify_impl(view);
+    } catch (PanicError& e) {
+      auto oss = std::ostringstream{};
+      oss << "Tensor match failed for " << this->debug_str() << " at " << loc.file_name() << ":" << loc.line()
+          << "\n- Root cause:  " << e.detail();
+      throw PanicError(std::move(oss).str());
+    }
+    return std::move(*this);
+  }
+
+  auto debug_str() const -> std::string {
+    auto oss = std::ostringstream{};
+    oss << "Tensor<";
+    std::size_t dim = 0;
+    for (const auto& size_ref : m_shape) {
+      if (dim > 0) {
+        oss << ", ";
+      }
+      oss << size_ref.value_or_name(dim++);
+    }
+    oss << ">";
+    if (m_strides.size() > 0) {
+      oss << " [strides=<";
+      dim = 0;
+      for (const auto& stride_ref : m_strides) {
+        if (dim > 0) {
+          oss << ", ";
+        }
+        oss << stride_ref.value_or_name(dim++);
+      }
+      oss << ">]";
+    }
+    return std::move(oss).str();
+  }
+
+ private:
+  auto m_verify_impl(tvm::ffi::TensorView view) const -> void {
+    const auto dim = static_cast<std::size_t>(view.dim());
+    host::RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ", m_shape.size(), " but got ", dim);
+    for (const auto i : stdv::iota(std::size_t{0}, dim)) {
+      m_shape[i]->verify(view.size(i));
+    }
+    if (this->m_has_strides()) {
+      for (const auto i : stdv::iota(std::size_t{0}, dim)) {
+        m_strides[i]->verify(view.stride(i));
+      }
+    } else {
+      host::RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected");
+    }
+    // since we may double verify, we will force to check
+    m_dtype->verify(view.dtype());
+    m_device->verify(view.device());
+  }
+
+  auto m_init_dtype() -> void {
+    host::RuntimeCheck(!m_has_dtype, "DType already specified");
+    m_has_dtype = true;
+  }
+  auto m_init_device() -> void {
+    host::RuntimeCheck(!m_has_device, "Device already specified");
+    m_has_device = true;
+  }
+  auto m_has_strides() const -> bool {
+    return !m_strides.empty();
+  }
+
+  std::span<const SizeRef> m_shape;
+  std::span<const SizeRef> m_strides;
+  DTypeRef m_dtype;
+  DeviceRef m_device;
+  bool m_has_dtype = false;
+  bool m_has_device = false;
+};
+
+}  // namespace host

python/sglang/jit_kernel/include/sgl_kernel/utils.cuh

@@ -0,0 +1,101 @@
+#pragma once
+
+#include <sgl_kernel/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/extra/c_env_api.h>
+
+#include <concepts>
+#include <cstddef>
+#include <source_location>
+#include <type_traits>
+
+namespace device {
+
+inline constexpr auto kWarpThreads = 32u;
+
+namespace pointer {
+
+// we only allow void * pointer arithmetic for safety
+
+template <typename T, std::integral... U>
+__always_inline __device__ auto offset(T* ptr, U... offset) -> void* {
+  static_assert(std::is_same_v<T, void>, "Pointer arithmetic is only allowed for void* pointers");
+  return static_cast<char*>(ptr) + (... + offset);
+}
+
+template <typename T, std::integral... U>
+__always_inline __device__ auto offset(const T* ptr, U... offset) -> const void* {
+  static_assert(std::is_same_v<T, void>, "Pointer arithmetic is only allowed for void* pointers");
+  return static_cast<const char*>(ptr) + (... + offset);
+}
+
+}  // namespace pointer
+
+}  // namespace device
+
+namespace host {
+
+inline auto
+RuntimeDeviceCheck(::cudaError_t error, std::source_location location = std::source_location::current()) -> void {
+  if (error != ::cudaSuccess) {
+    [[unlikely]];
+    ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error));
+  }
+}
+
+inline auto RuntimeCudaCheck(std::source_location location = std::source_location::current()) -> void {
+  return RuntimeDeviceCheck(::cudaGetLastError(), location);
+}
+
+template <auto F>
+inline void set_smem_once(std::size_t smem_size) {
+  static const auto last_smem_size = [&] {
+    RuntimeDeviceCheck(::cudaFuncSetAttribute(F, ::cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    return smem_size;
+  }();
+  RuntimeCheck(
+      smem_size <= last_smem_size,
+      "Dynamic shared memory size exceeds the previously set maximum size: ",
+      last_smem_size,
+      " bytes");
+}
+
+struct LaunchKernel {
+ public:
+  explicit LaunchKernel(
+      dim3 grid_dim, dim3 block_dim, DLDevice device, std::size_t dynamic_shared_mem_bytes = 0) noexcept
+      : m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)) {}
+
+  explicit LaunchKernel(
+      dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t dynamic_shared_mem_bytes = 0) noexcept
+      : m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)) {}
+
+  static auto resolve_device(DLDevice device) -> cudaStream_t {
+    return static_cast<cudaStream_t>(::TVMFFIEnvGetStream(device.device_type, device.device_id));
+  }
+
+  LaunchKernel(const LaunchKernel&) = delete;
+  LaunchKernel& operator=(const LaunchKernel&) = delete;
+
+  template <typename T, typename... Args>
+  auto operator()(T&& kernel, Args&&... args) const -> void {
+    host::RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...));
+  }
+
+ private:
+  static auto
+  s_make_config(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;
+  /// TODO: We can add a queue to store the attributes if needed in the future.
+};
+
+}  // namespace host

python/sglang/jit_kernel/include/sgl_kernel/utils.h

@@ -0,0 +1,88 @@
+#pragma once
+
+#include <dlpack/dlpack.h>
+
+#include <concepts>
+#include <ostream>
+#include <source_location>
+#include <sstream>
+#include <utility>
+
+namespace host {
+
+struct PanicError : public std::runtime_error {
+ public:
+  // copy and move constructors
+  explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {}
+  auto detail() const -> std::string_view {
+    const auto sv = std::string_view{m_message};
+    const auto pos = sv.find(": ");
+    return pos == std::string_view::npos ? sv : sv.substr(pos + 2);
+  }
+
+ private:
+  std::string m_message;
+};
+
+template <typename... Args>
+[[noreturn]]
+inline auto panic(std::source_location location, Args&&... args) -> void {
+  std::ostringstream os;
+  os << "Runtime check failed at " << location.file_name() << ":" << location.line();
+  if constexpr (sizeof...(args) > 0) {
+    os << ": ";
+    (os << ... << std::forward<Args>(args));
+  } else {
+    os << " in " << location.function_name();
+  }
+  throw PanicError(std::move(os).str());
+}
+
+template <typename... Args>
+struct RuntimeCheck {
+  using Loc_t = std::source_location;
+  template <typename Cond>
+  explicit RuntimeCheck(Cond&& condition, Args&&... args, Loc_t location = Loc_t::current()) {
+    if (!condition) {
+      [[unlikely]];
+      ::host::panic(location, std::forward<Args>(args)...);
+    }
+  }
+};
+
+template <typename Cond, typename... Args>
+explicit RuntimeCheck(Cond&&, Args&&...) -> RuntimeCheck<Args...>;
+
+template <std::signed_integral T, std::signed_integral U>
+inline constexpr auto div_ceil(T a, U b) {
+  return (a + b - 1) / b;
+}
+
+template <std::unsigned_integral T, std::unsigned_integral U>
+inline constexpr auto div_ceil(T a, U b) {
+  return (a + b - 1) / b;
+}
+
+inline auto dtype_bytes(DLDataType dtype) -> std::size_t {
+  return static_cast<std::size_t>(dtype.bits / 8);
+}
+
+namespace pointer {
+
+// we only allow void * pointer arithmetic for safety
+
+template <typename T, std::integral... U>
+inline auto offset(T* ptr, U... offset) -> void* {
+  static_assert(std::is_same_v<T, void>, "Pointer arithmetic is only allowed for void* pointers");
+  return static_cast<char*>(ptr) + (... + offset);
+}
+
+template <typename T, std::integral... U>
+inline auto offset(const T* ptr, U... offset) -> const void* {
+  static_assert(std::is_same_v<T, void>, "Pointer arithmetic is only allowed for void* pointers");
+  return static_cast<const char*>(ptr) + (... + offset);
+}
+
+}  // namespace pointer
+
+}  // namespace host

python/sglang/jit_kernel/include/sgl_kernel/warp.cuh

@@ -0,0 +1,145 @@
+#pragma once
+#include <sgl_kernel/utils.cuh>
+
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+namespace device::warp {
+
+namespace details {
+
+template <std::size_t kUnit>
+inline constexpr auto get_mem_package() {
+  if constexpr (kUnit == 16) {
+    return uint4{};
+  } else if constexpr (kUnit == 8) {
+    return uint2{};
+  } else if constexpr (kUnit == 4) {
+    return uint1{};
+  } else {
+    static_assert(kUnit == 16 || kUnit == 8 || kUnit == 4, "Unsupported memory package size");
+  }
+}
+
+inline constexpr auto default_unit_size(std::size_t x) -> std::size_t {
+  if (x % (16 * kWarpThreads) == 0) return 16;
+  if (x % (8 * kWarpThreads) == 0) return 8;
+  if (x % (4 * kWarpThreads) == 0) return 4;
+  return 0;  // trigger static assert in _get_mem_package
+}
+
+template <std::size_t kBytes, std::size_t kUnit>
+using mem_package_t = decltype(get_mem_package<kUnit>());
+
+template <typename T, std::size_t N>
+struct storage_vec {
+  T data[N];
+};
+
+__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
+  uint32_t tmp;
+  asm volatile("ld.global.cs.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
+  return uint1{tmp};
+}
+
+__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 {
+  uint32_t tmp0, tmp1;
+  asm volatile("ld.global.cs.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
+  return uint2{tmp0, tmp1};
+}
+
+__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 {
+  uint32_t tmp0, tmp1, tmp2, tmp3;
+  asm volatile("ld.global.cs.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src));
+  return uint4{tmp0, tmp1, tmp2, tmp3};
+}
+
+__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) {
+  uint32_t tmp = value.x;
+  asm volatile("st.global.cs.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
+}
+
+__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) {
+  uint32_t tmp0 = value.x;
+  uint32_t tmp1 = value.y;
+  asm volatile("st.global.cs.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
+}
+
+__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) {
+  uint32_t tmp0 = value.x;
+  uint32_t tmp1 = value.y;
+  uint32_t tmp2 = value.z;
+  uint32_t tmp3 = value.w;
+  asm volatile("st.global.cs.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
+}
+
+}  // namespace details
+
+template <
+    std::size_t kBytes,
+    std::size_t kUnit = details::default_unit_size(kBytes),
+    std::size_t kThreads = ::device::kWarpThreads>
+__always_inline __device__ void copy(void* __restrict__ dst, const void* __restrict__ src) {
+  using Package = details::mem_package_t<kBytes, kUnit>;
+  constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
+  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
+  static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
+
+  const auto dst_packed = static_cast<Package*>(dst);
+  const auto src_packed = static_cast<const Package*>(src);
+  const auto lane_id = threadIdx.x % kThreads;
+
+#pragma unroll kLoopCount
+  for (std::size_t i = 0; i < kLoopCount; ++i) {
+    const auto j = i * kThreads + lane_id;
+    dst_packed[j] = src_packed[j];
+  }
+}
+
+template <
+    std::size_t kBytes,
+    std::size_t kUnit = details::default_unit_size(kBytes),
+    std::size_t kThreads = ::device::kWarpThreads>
+__always_inline __device__ auto load_vec(const void* __restrict__ src) {
+  using Package = details::mem_package_t<kBytes, kUnit>;
+  constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
+  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
+  static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
+
+  const auto src_packed = static_cast<const Package*>(src);
+  const auto lane_id = threadIdx.x % kThreads;
+  details::storage_vec<Package, kLoopCount> vec;
+
+#pragma unroll kLoopCount
+  for (std::size_t i = 0; i < kLoopCount; ++i) {
+    const auto j = i * kThreads + lane_id;
+    vec.data[i] = details::load_nc(src_packed + j);
+  }
+
+  return vec;
+}
+
+template <
+    std::size_t kBytes,
+    std::size_t kUnit = details::default_unit_size(kBytes),
+    std::size_t kThreads = ::device::kWarpThreads,
+    typename Tp>
+__always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) {
+  using Package = details::mem_package_t<kBytes, kUnit>;
+  constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
+  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
+  static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
+  static_assert(std::is_same_v<Tp, details::storage_vec<Package, kLoopCount>>);
+
+  const auto dst_packed = static_cast<Package*>(dst);
+  const auto lane_id = threadIdx.x % kThreads;
+
+#pragma unroll kLoopCount
+  for (std::size_t i = 0; i < kLoopCount; ++i) {
+    const auto j = i * kThreads + lane_id;
+    details::store_nc(dst_packed + j, vec.data[i]);
+  }
+}
+
+}  // namespace device::warp