[JIT Kernel][Feature] Support JIT custom all reduce (rewrite as v2) (#19880)

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh

@@ -0,0 +1,27 @@
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/custom_all_reduce.cuh>
+
+#include <cstdint>
+#include <cstring>
+
+inline void register_custom_all_reduce() {
+  namespace refl = tvm::ffi::reflection;
+  using Class = host::distributed::CustomAllReduceBase;
+  refl::ObjectDef<Class>()
+      .def(refl::init<uint32_t, uint32_t, uint32_t, uint32_t, int64_t, int64_t, int64_t>(), "__init__")
+      .def("share_storage", &Class::share_storage)
+      .def("share_graph_inputs", &Class::share_graph_inputs)
+      .def("post_init", &Class::post_init)
+      .def("register_inputs", &Class::register_inputs)
+      .def("set_cuda_graph_capture", &Class::set_cuda_graph_capture)
+      .def("free_ipc_handles", &Class::free_ipc_handles)
+      .def("free_storage", &Class::free_storage)
+      .def("configure_pull", &Class::configure_pull);
+}

python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh

@@ -0,0 +1,205 @@
+// Partially migrated from AOT kernel:
+// https://github.com/sgl-project/sglang/blob/v0.5.9/sgl-kernel/csrc/allreduce/custom_all_reduce.cu
+// Which was originally adapted from:
+// https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cu
+// We redesign the controller interface to minimize control plane traffic,
+// and fuse the reduce-scatter and broadcast in the 2-shot all reduce
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/common.cuh>
+#include <sgl_kernel/distributed/custom_all_reduce.cuh>
+
+#include <bit>
+#include <cstdint>
+#include <cstring>
+
+namespace {
+
+using device::distributed::PullController;
+using host::distributed::AllReduceData;
+using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
+
+struct AllReduceParams {
+  void* __restrict__ output;
+  uint32_t rank;
+  uint32_t num_items;  // NOTE: support at most 4G, but that's too much
+};
+
+[[maybe_unused]]
+SGL_DEVICE void prefetch_uniform_ptr(const void* ptr) {
+  asm volatile("prefetchu.L1 [%0];" ::"l"(ptr) : "memory");
+}
+
+#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1)
+
+template <bool kBroadcast, typename DType, uint32_t kNumGPU>
+SGL_DEVICE void all_reduce_impl(const AllReduceParams& params, DType* (&input)[kNumGPU]) {
+  using namespace device;
+
+  constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  using DType2 = packed_t<DType>;
+  using Storage = AlignedVector<DType2, kVecSize>;
+  const auto& [output, rank, num_items] = params;
+
+  for (auto i = blockIdx.x;; i += gridDim.x) {
+    const auto offset = i * blockDim.x + threadIdx.x;
+    if (offset * kVecSize * 2 >= num_items) break;
+    Storage storage[kNumGPU];
+
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      storage[i].load(input[i], offset);
+    }
+    const Storage result = distributed::reduce_impl(storage);
+    if constexpr (kBroadcast) {
+#pragma unroll
+      for (uint32_t i = 0; i < kNumGPU; ++i) {
+        result.store(input[i], offset);
+      }
+    } else {
+      result.store(output, offset);
+    }
+  }
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+CUSTOM_AR_KERNEL void all_reduce_one_shot_kernel(
+    const AllReduceData* __restrict__ data,
+    const AllReduceParams __grid_constant__ params,
+    const PullController __grid_constant__ ctrl) {
+  /// NOTE: we assume the data array is ready before the previous kernel
+  DType* input[kNumGPU];
+  prefetch_uniform_ptr(data);
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i)
+    input[i] = static_cast<DType*>(data->input[i]);
+  device::PDLWaitPrimary<kUsePDL>();
+
+  ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
+  all_reduce_impl</*kBroadcast=*/false>(params, input);
+
+  device::PDLTriggerSecondary<kUsePDL>();
+  ctrl.sync</*kFence=*/0, /*kStart=*/0>(params.rank, kNumGPU);
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+CUSTOM_AR_KERNEL void all_reduce_two_shot_kernel(
+    const AllReduceData* __restrict__ data,
+    const AllReduceParams __grid_constant__ params,
+    const PullController __grid_constant__ ctrl) {
+  // get the range of this rank
+  using device::kWarpThreads, device::div_ceil;
+
+  prefetch_uniform_ptr(data);
+  DType* input[kNumGPU];
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i)
+    input[i] = static_cast<DType*>(data->input[i]);
+
+  constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  const uint32_t num_items = params.num_items;
+  const uint32_t total_vec = num_items / (kVecSize * 2);  // must be divisible here
+  const uint32_t vec_per_rank = div_ceil(div_ceil(total_vec, kNumGPU), kWarpThreads) * kWarpThreads;
+  const uint32_t local_vec_start = min(params.rank * vec_per_rank, total_vec);
+  const uint32_t local_vec_finish = min(local_vec_start + vec_per_rank, total_vec);
+  const uint32_t local_start = local_vec_start * kVecSize * 2;
+  const uint32_t local_length = (local_vec_finish - local_vec_start) * kVecSize * 2;
+  const auto local_params = AllReduceParams{
+      .output = nullptr,  // this is not used for 2-shot all reduce
+      .rank = params.rank,
+      .num_items = local_length,
+  };
+
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i)
+    input[i] += local_start;
+
+  device::PDLWaitPrimary<kUsePDL>();
+
+  ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
+  all_reduce_impl</*kBroadcast=*/true>(local_params, input);
+
+  device::PDLTriggerSecondary<kUsePDL>();
+  ctrl.sync</*kFence=*/1, /*kStart=*/0>(params.rank, kNumGPU);
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+struct CustomAllReducePull : public CustomAllReduceBase {
+  static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  static constexpr auto one_shot_kernel = all_reduce_one_shot_kernel<DType, kNumGPU, kUsePDL>;
+  static constexpr auto two_shot_kernel = all_reduce_two_shot_kernel<DType, kNumGPU, kUsePDL>;
+  static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
+
+  tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) {
+    using namespace host;
+    const bool use_2shot = (shot == 2);
+    const auto device = input.device();
+    const auto input_ptr = input.data_ptr();
+    const auto buffer_ptr = get_pull_buffer(m_storage);
+    const auto num_items_int64 = input.numel();
+    const auto num_items = static_cast<uint32_t>(num_items_int64);
+    const auto items_per_block = m_cta_size * kVecSize * 2;
+    const auto needed_blocks = div_ceil(num_items, items_per_block);
+    const auto num_blocks = std::min(needed_blocks, m_num_cta);
+    const auto kernel = use_2shot ? two_shot_kernel : one_shot_kernel;
+    // only 1-shot + graph capture need extra output buffer
+    const auto output = (m_is_graph_capturing && !use_2shot) ? ffi::empty_like(input) : input;
+    const auto params = AllReduceParams{
+        .output = use_2shot ? nullptr : output.data_ptr(),
+        .rank = m_rank,
+        .num_items = num_items,
+    };
+
+    RuntimeCheck(input.IsContiguous(), "Input tensor must be contiguous");
+    RuntimeCheck(m_num_gpu == kNumGPU, "Mismatch GPU count");
+    RuntimeCheck(shot == 1 || shot == 2, "Invalid shot count: ", shot);
+    RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported");
+    RuntimeCheck(is_type<DType>(input.dtype()), "Input dtype mismatch");
+    RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
+    RuntimeCheck(m_pull_ctrl.has_value(), "Controller is not initialized");
+    RuntimeCheck(static_cast<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
+
+    const auto& ctrl = *m_pull_ctrl;
+    const auto stream = LaunchKernel::resolve_device(device);
+    auto launch = LaunchKernel{num_blocks, m_cta_size, stream};
+    launch.enable_pdl(kUsePDL);
+    const auto check_capturing = [&] {
+      if (!m_is_graph_capturing) return false;  // override to avoid cudaRT call overhead
+      cudaStreamCaptureStatus status;
+      RuntimeDeviceCheck(cudaStreamIsCapturing(stream, &status));
+      return status == cudaStreamCaptureStatusActive;
+    };
+    if (check_capturing()) {
+      // no-op if not really capturing, we're in a dummy run
+      const auto data_ptr = allocate_graph_capture_input(input_ptr);
+      /// NOTE: we assume when the graph is replayed, the data_ptr should be ready
+      launch(kernel, data_ptr, params, ctrl);
+    } else {
+      // 1.copy the input to the buffer
+      const auto input_bytes = static_cast<int64_t>(sizeof(DType) * num_items);
+      RuntimeCheck(input_bytes <= m_pull_buffer_bytes, "Input is too large, num items: ", num_items);
+      RuntimeDeviceCheck(cudaMemcpyAsync(buffer_ptr, input_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream));
+      // 2. launch the all reduce kernel
+      const auto data_ptr = get_data_ptr();  // use default buffer
+      launch(kernel, data_ptr, params, ctrl);
+      if (use_2shot) {  // 3. copy the reduced result back to the output, because 2-shot doesn't write to output
+        RuntimeDeviceCheck(cudaMemcpyAsync(input_ptr, buffer_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream));
+      }
+    }
+    return output;
+  }
+};
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
+  using Impl = CustomAllReducePull<DType, kNumGPU, kUsePDL>;
+  return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh

@@ -0,0 +1,253 @@
+// Partially adapted from:
+// https://github.com/flashinfer-ai/flashinfer/blob/v0.6.4/include/flashinfer/comm/trtllm_allreduce_fusion.cuh
+// We simplify the lamport design and minimize the ring buffer count (from 3 -> 2)
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/common.cuh>
+#include <sgl_kernel/distributed/custom_all_reduce.cuh>
+
+#include <cstdint>
+#include <cstring>
+
+namespace {
+
+using device::distributed::PushController;
+using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
+
+struct AllReducePushData {
+  void* __restrict__ buffer[device::distributed::kMaxNumGPU];
+  const void* input;
+  void* output;
+  uint32_t rank;
+  uint32_t num_items;
+  uint32_t buffer_bytes;
+  uint32_t epoch_bytes;
+};
+
+#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1)
+
+template <typename T>
+struct fp_trait {};
+
+// TODO: support more dtypes
+template <>
+struct fp_trait<bf16_t> {
+  using type = uint16_t;
+  [[maybe_unused]]
+  static constexpr uint16_t pos_zero = 0x0000u;
+  [[maybe_unused]]
+  static constexpr uint16_t neg_zero = 0x8000u;
+};
+
+template <>
+struct fp_trait<fp16_t> {
+  using type = uint16_t;
+  [[maybe_unused]]
+  static constexpr uint16_t pos_zero = 0x0000u;
+  [[maybe_unused]]
+  static constexpr uint16_t neg_zero = 0x8000u;
+};
+
+template <>
+struct fp_trait<float> {
+  using type = uint32_t;
+  [[maybe_unused]]
+  static constexpr uint32_t pos_zero = 0x00000000u;
+  [[maybe_unused]]
+  static constexpr uint32_t neg_zero = 0x80000000u;
+};
+
+template <typename DType>
+SGL_DEVICE void clear_pos_zero(DType& val) {
+  using Trait = fp_trait<DType>;
+  const auto ptr = reinterpret_cast<typename Trait::type*>(&val);
+  if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero;
+}
+
+template <typename DType>
+SGL_DEVICE bool is_pos_zero(const DType& val) {
+  using Trait = fp_trait<DType>;
+  const auto ptr = reinterpret_cast<const typename Trait::type*>(&val);
+  return *ptr == Trait::pos_zero;
+}
+
+template <typename DType>
+SGL_DEVICE DType get_pos_zero() {
+  using Trait = fp_trait<DType>;
+  const auto value = Trait::pos_zero;
+  return *reinterpret_cast<const DType*>(&value);
+}
+
+template <typename T>
+SGL_DEVICE void ld_global_volatile_16B(T& x, const void* addr, int64_t offset) {
+  static_assert(alignof(T) == 16 && sizeof(T) == 16);
+  addr = device::pointer::offset<T>(addr, offset);
+  uint4 val;
+  asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];"
+               : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
+               : "l"(addr));
+  x = *reinterpret_cast<const T*>(&val);
+}
+
+template <typename T>
+SGL_DEVICE void st_global_volatile_16B(const T& x, void* addr, int64_t offset) {
+  static_assert(alignof(T) == 16 && sizeof(T) == 16);
+  const uint4 val = *reinterpret_cast<const uint4*>(&x);
+  addr = device::pointer::offset<T>(addr, offset);
+  asm volatile(
+      "st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr));
+}
+
+template <typename DType, uint32_t kNumGPU>
+SGL_DEVICE void push_impl(DType* (&push_buf)[kNumGPU], const void* data, uint32_t num_items) {
+  using namespace device;
+  constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  using Storage = AlignedVector<packed_t<DType>, kVecSize>;
+
+  for (auto i = blockIdx.x;; i += gridDim.x) {
+    const auto offset = i * blockDim.x + threadIdx.x;
+    if (offset * kVecSize * 2 >= num_items) break;
+    Storage vec;
+    vec.load(data, offset);
+#pragma unroll
+    for (uint32_t j = 0; j < kVecSize; ++j) {
+      clear_pos_zero(vec[j].x);
+      clear_pos_zero(vec[j].y);
+    }
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      st_global_volatile_16B(vec, push_buf[i], offset);
+    }
+  }
+}
+
+template <typename DType, uint32_t kNumGPU>
+SGL_DEVICE void poll_impl(DType* (&poll_buf)[kNumGPU], void* data, uint32_t num_items) {
+  using namespace device;
+  constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  using Storage = AlignedVector<packed_t<DType>, kVecSize>;
+
+  for (auto i = blockIdx.x;; i += gridDim.x) {
+    const auto offset = i * blockDim.x + threadIdx.x;
+    if (offset * kVecSize * 2 >= num_items) break;
+    Storage storage[kNumGPU];
+
+    while (true) {
+      bool has_pos_zero = false;
+#pragma unroll
+      for (uint32_t i = 0; i < kNumGPU; ++i) {
+        ld_global_volatile_16B(storage[i], poll_buf[i], offset);
+#pragma unroll
+        for (auto j = 0; j < kVecSize; ++j) {
+          has_pos_zero |= is_pos_zero(storage[i][j].x);
+          has_pos_zero |= is_pos_zero(storage[i][j].y);
+        }
+      }
+      if (!has_pos_zero) break;
+    }
+
+    const Storage result = distributed::reduce_impl(storage);
+    result.store(data, offset);
+
+    Storage pos_zeros;
+    pos_zeros.fill({get_pos_zero<DType>(), get_pos_zero<DType>()});
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      pos_zeros.store(poll_buf[i], offset);
+    }
+  }
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+CUSTOM_AR_KERNEL void all_reduce_one_shot_push_kernel(
+    const AllReducePushData __grid_constant__ params,  //
+    const PushController __grid_constant__ ctrl) {
+  using namespace device;
+
+  const auto [buffer, input, output, rank, num_items, buffer_bytes, epoch_bytes] = params;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // Phase 1: Push data from input to all ranks' buffers
+  const auto epoch_offset = ctrl.epoch() * epoch_bytes;
+  DType* push_buf[kNumGPU];
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i) {
+    push_buf[i] = static_cast<DType*>(pointer::offset(buffer[i], rank * buffer_bytes, epoch_offset));
+  }
+  push_impl(push_buf, input, num_items);
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // Phase 2: Poll local data
+  DType* poll_buf[kNumGPU];
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i) {
+    poll_buf[i] = static_cast<DType*>(pointer::offset(buffer[rank], i * buffer_bytes, epoch_offset));
+  }
+  poll_impl(poll_buf, output, num_items);
+  ctrl.exit();
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+struct CustomAllReducePush : public CustomAllReduceBase {
+  static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
+
+  tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) {
+    using namespace host;
+    const auto device = input.device();
+    const auto input_ptr = input.data_ptr();
+    const auto num_items_int64 = input.numel();
+    const auto num_items = static_cast<uint32_t>(num_items_int64);
+    const auto num_blocks = m_max_num_cta_push;  // must be constant to ensure correctness
+    const auto num_threads = [&] {
+      for (const auto t : {128u, 256u, 512u}) {
+        if (t * num_blocks * 2 * kVecSize >= num_items) return t;
+      }
+      return 1024u;
+    }();
+    const auto output = input;
+    AllReducePushData params;
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      params.buffer[i] = get_push_buffer(m_peer_storage[i]);
+    }
+    params.input = input_ptr;
+    params.output = input_ptr;
+    params.rank = m_rank;
+    params.num_items = num_items;
+    params.buffer_bytes = m_push_buffer_bytes;
+    params.epoch_bytes = kNumGPU * params.buffer_bytes;
+
+    RuntimeCheck(input.IsContiguous(), "Input must be contiguous");
+    RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch");
+    RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported");
+    RuntimeCheck(is_type<DType>(input.dtype()), "Input dtype mismatch");
+    RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
+    RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized");
+    RuntimeCheck(shot == 1, "Push all-reduce only supports 1-shot, got: ", shot);
+    RuntimeCheck(static_cast<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
+
+    const auto input_bytes = static_cast<int64_t>(sizeof(DType) * num_items_int64);
+    RuntimeCheck(input_bytes <= m_push_buffer_bytes, "Input is too large, num items: ", num_items);
+
+    const auto kernel = all_reduce_one_shot_push_kernel<DType, kNumGPU, kUsePDL>;
+    LaunchKernel(num_blocks, num_threads, device)  //
+        .enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl);
+    return output;
+  }
+};
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
+  using Impl = CustomAllReducePush<DType, kNumGPU, kUsePDL>;
+  return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh

@@ -40,8 +40,6 @@ struct Vec {
 
 using I4 = Vec<int, 4>;
 
-using host::div_ceil;
-
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
 // No support for async
 #else