chore: import upstream snapshot with attribution

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

@@ -0,0 +1,30 @@
+#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("get_graph_capture_ptrs", &Class::get_graph_capture_ptrs)
+      .def("get_graph_capture_bases", &Class::get_graph_capture_bases)
+      .def("register_peer_mapped_inputs", &Class::register_peer_mapped_inputs)
+      .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 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);
+    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, input_bytes);
+      /// 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
+      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/distributed/tp_qknorm.cuh

@@ -0,0 +1,325 @@
+// Adapted from https://github.com/NVIDIA/TensorRT-LLM/pull/12163
+// We reuse the custom all reduce push buffer in SGLang
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.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 ParallelQKNormParams {
+  void* __restrict__ buffer[device::distributed::kMaxNumGPU];
+  void* q_ptr;
+  void* k_ptr;
+  const void* __restrict__ q_weight;
+  const void* __restrict__ k_weight;
+  int64_t q_stride_bytes;
+  int64_t k_stride_bytes;
+  float eps;
+  uint32_t rank;
+  uint32_t num_tokens;
+  uint32_t epoch_bytes;
+  uint32_t num_clean_up_count = 0;
+};
+
+template <typename T>
+SGL_DEVICE void ld_global_volatile_8B(T& x, const void* addr, int64_t offset) {
+  static_assert(alignof(T) == 8 && sizeof(T) == 8);
+  addr = device::pointer::offset<T>(addr, offset);
+  uint2 val;
+  asm volatile("ld.volatile.global.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr));
+  x = *reinterpret_cast<const T*>(&val);
+}
+
+template <typename T>
+SGL_DEVICE void st_global_volatile_8B(const T& x, void* addr, int64_t offset) {
+  static_assert(alignof(T) == 8 && sizeof(T) == 8);
+  const uint2 val = *reinterpret_cast<const uint2*>(&x);
+  addr = device::pointer::offset<T>(addr, offset);
+  asm volatile("st.volatile.global.v2.b32 [%2], {%0, %1};" ::"r"(val.x), "r"(val.y), "l"(addr));
+}
+
+[[maybe_unused]]
+SGL_DEVICE float sync_float(float x) {
+  return __shfl_sync(0xffffffffu, x, 0);
+}
+
+[[maybe_unused]]
+constexpr auto next_pow_of_2(uint32_t x) {
+  uint32_t y = 1;
+  while (y < x)
+    y *= 2;
+  return y;
+}
+
+template <typename DType_, uint32_t kNumGPU_, int64_t kQDim_, int64_t kKDim_, bool kUsePDL_>
+struct KernelTrait {
+  // rename the arguments to avoid confusion with the template parameters
+  using DType = DType_;
+  static constexpr uint32_t kNumGPU = kNumGPU_;
+  static constexpr int64_t kQDim = kQDim_;
+  static constexpr int64_t kKDim = kKDim_;
+  static constexpr bool kUsePDL = kUsePDL_;
+
+  static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  static constexpr int64_t kLocalQDim = kQDim / kNumGPU;
+  static constexpr int64_t kLocalKDim = kKDim / kNumGPU;
+  static constexpr uint32_t kNumQThreads = kLocalQDim / (kVecSize * 2);
+  static constexpr uint32_t kNumKThreads = kLocalKDim / (kVecSize * 2);
+  static constexpr uint32_t kNumQWarps = kNumQThreads / device::kWarpThreads;
+  static constexpr uint32_t kNumKWarps = host::div_ceil(kNumKThreads, device::kWarpThreads);
+  static constexpr uint32_t kBlockSize = (kNumQWarps + kNumKWarps) * device::kWarpThreads;
+  static constexpr uint32_t kOccupancy = 2048 / kBlockSize;
+
+  using DType2 = packed_t<DType>;
+  using Storage = device::AlignedVector<DType2, kVecSize>;
+
+  static_assert(std::has_single_bit(kNumGPU), "must be pow of 2");
+  static_assert(kQDim % kNumGPU == 0);
+  static_assert(kKDim % kNumGPU == 0);
+  static_assert(kLocalQDim % (kVecSize * 2) == 0);
+  static_assert(kLocalKDim % (kVecSize * 2) == 0);
+  static_assert(kNumQThreads % device::kWarpThreads == 0);
+  static_assert(kBlockSize <= 1024);
+  static_assert(sizeof(Storage) == 16 && alignof(Storage) == 16);
+  static_assert(kOccupancy * kBlockSize <= 2048);
+};
+
+template <typename Trait>
+__global__ __launch_bounds__(Trait::kBlockSize, Trait::kOccupancy) void parallel_qknorm_across_head(
+    const ParallelQKNormParams __grid_constant__ params, const PushController __grid_constant__ ctrl) {
+  using namespace device;
+
+  // each cta will handle exactly 1 token
+  using Storage = typename Trait::Storage;
+  using DType2 = typename Trait::DType2;
+  const auto &[
+      buffer, q_ptr, k_ptr, q_weight, k_weight, q_stride_bytes, k_stride_bytes, //
+      eps, rank, num_tokens, epoch_bytes, num_clean_up_count
+  ] = params;
+
+  using Package = AlignedVector<float, 2>;
+  constexpr uint32_t kNumGPU = Trait::kNumGPU;
+  constexpr uint32_t kNumQReduce = next_pow_of_2(Trait::kNumQWarps);
+  constexpr uint32_t kNumKReduce = next_pow_of_2(Trait::kNumKWarps);
+  __shared__ float smem_qk[Trait::kNumQWarps + Trait::kNumKWarps];
+  __shared__ float scale_q;
+  __shared__ float scale_k;
+  const auto tx = threadIdx.x;
+  const auto bx = blockIdx.x;
+  /// NOTE: this can hint compiler to optimize `is_valid` out when not needed
+  constexpr uint32_t kActiveThreads = Trait::kNumQThreads + Trait::kNumKThreads;
+  const auto is_valid = Trait::kBlockSize == kActiveThreads || tx < kActiveThreads;
+  const auto smem_q = smem_qk + 0;
+  const auto smem_k = smem_qk + Trait::kNumQWarps;
+  const auto load_q = tx < Trait::kNumQThreads;
+  const auto offset = load_q ? tx : tx - Trait::kNumQThreads;
+  const auto input_ptr = load_q ? q_ptr : k_ptr;
+  const auto weight_ptr = load_q ? q_weight : k_weight;
+  const auto input_stride_bytes = load_q ? q_stride_bytes : k_stride_bytes;
+  PDLWaitPrimary<Trait::kUsePDL>();
+  PDLTriggerSecondary<Trait::kUsePDL>();
+  if (bx >= num_tokens) {
+    [[unlikely]];
+    // In this case, we use the last few blocks to clean up other controllers
+    const auto start = (bx - num_tokens) * blockDim.x + threadIdx.x;
+    const auto stride = (gridDim.x - num_tokens) * blockDim.x;
+    for (uint32_t i = start; i < num_clean_up_count; i += stride)
+      ctrl.exit_unsafe(num_tokens + i);
+    return;
+  }
+  const auto epoch_offset = ctrl.epoch() * epoch_bytes;  // only for comm
+
+  __builtin_assume(bx < num_tokens);  // since we have `bx >= num_tokens`
+  Storage next_input;
+  void* input_i_ptr = pointer::offset(input_ptr, bx * input_stride_bytes);
+  if (is_valid) next_input.load(input_i_ptr, offset);
+
+  for (uint32_t i = bx; i < num_tokens; i += gridDim.x) {
+    // Stage 1. local reduce (warp-level)
+    Storage local_input;
+    {
+      float local_sum = 0.0;
+      if (is_valid) {
+        local_input = next_input;
+#pragma unroll
+        for (uint32_t j = 0; j < Trait::kVecSize; ++j) {
+          const auto [x, y] = cast<fp32x2_t>(local_input[j]);
+          local_sum += x * x + y * y;
+        }
+      }
+      smem_qk[threadIdx.x / kWarpThreads] = warp::reduce_sum(local_sum);
+    }
+
+    // Stage 2. block reduce + push to peer ranks + poll from local rank
+    __syncthreads();
+
+    Storage local_weight;
+    const auto input_next_ptr = pointer::offset(input_i_ptr, gridDim.x * input_stride_bytes);
+    /**
+     * NOTE: Prefetch to hide the latency.
+     * This brings around 20% of performance gain in large batches
+     * The P2P communication is mainly latency bound, so during this waiting period,
+     * We can let some data loading transparently in the background.
+     */
+    if (is_valid) {
+      local_weight.load(weight_ptr, offset);
+      if (i + gridDim.x < num_tokens) next_input.load(input_next_ptr, offset);
+    }
+
+    if (tx < kWarpThreads) {
+      const auto local_sum_q = tx < Trait::kNumQWarps ? smem_q[tx] : 0.0f;
+      const auto local_sum_k = tx < Trait::kNumKWarps ? smem_k[tx] : 0.0f;
+      const auto sum_q = sync_float(warp::reduce_sum<kNumQReduce>(local_sum_q));
+      const auto sum_k = sync_float(warp::reduce_sum<kNumKReduce>(local_sum_k));
+      if (tx < kNumGPU) {  // push a float2 pack to the peer
+        Package sum_q_k;
+        /// NOTE: eps should be scaled down by kNumGPU from host side
+        /// we add here to ensure that the sum is never zero
+        sum_q_k[0] = sum_q + eps;
+        sum_q_k[1] = sum_k + eps;
+        const auto push_ptr = pointer::offset(buffer[tx], epoch_offset);
+        st_global_volatile_8B(sum_q_k, push_ptr, i * kNumGPU + rank);
+        const auto poll_ptr = pointer::offset(buffer[rank], epoch_offset);
+        while (true) {
+          ld_global_volatile_8B(sum_q_k, poll_ptr, i * kNumGPU + tx);
+          if (sum_q_k[0] != 0.0f && sum_q_k[1] != 0.0f) break;
+        }
+        constexpr uint32_t kActiveMask = (1 << kNumGPU) - 1;
+        const auto global_sum_q = warp::reduce_sum<kNumGPU>(sum_q_k[0], kActiveMask);
+        const auto global_sum_k = warp::reduce_sum<kNumGPU>(sum_q_k[1], kActiveMask);
+        scale_q = math::rsqrt(global_sum_q / static_cast<float>(Trait::kQDim));
+        scale_k = math::rsqrt(global_sum_k / static_cast<float>(Trait::kKDim));
+        Package zeros;
+        zeros.fill(0.0f);
+        zeros.store(poll_ptr, i * kNumGPU + tx);
+      }
+    }
+
+    __syncthreads();
+    const auto scale = load_q ? scale_q : scale_k;
+    if (is_valid) {
+#pragma unroll
+      for (uint32_t j = 0; j < Trait::kVecSize; ++j) {
+        const auto fp32_input = cast<fp32x2_t>(local_input[j]);
+        const auto fp32_weight = cast<fp32x2_t>(local_weight[j]);
+        const auto scaled_x = fp32_input.x * scale * fp32_weight.x;
+        const auto scaled_y = fp32_input.y * scale * fp32_weight.y;
+        local_input[j] = cast<DType2>(fp32x2_t{scaled_x, scaled_y});
+      }
+      local_input.store(input_i_ptr, offset);
+    }
+    input_i_ptr = input_next_ptr;
+  }
+  ctrl.exit();
+}
+
+template <typename DType, uint32_t kNumGPU, int64_t kQDim, int64_t kKDim, bool kUsePDL>
+struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase {
+  using Trait = KernelTrait<DType, kNumGPU, kQDim, kKDim, kUsePDL>;
+  static constexpr auto kernel = parallel_qknorm_across_head<Trait>;
+  static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
+
+  void _run(
+      const tvm::ffi::Tensor q,
+      const tvm::ffi::Tensor k,
+      const tvm::ffi::Tensor q_weight,
+      const tvm::ffi::Tensor k_weight,
+      const float eps  // passed in unscaled
+  ) {
+    using namespace host;
+    constexpr auto Q = Trait::kLocalQDim;
+    constexpr auto K = Trait::kLocalKDim;
+    auto N = SymbolicSize{"num_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+    TensorMatcher({N, Q})  // q
+        .with_strides({-1, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q);
+    TensorMatcher({N, K})  // k
+        .with_strides({-1, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(k);
+    TensorMatcher({Q})  // q_weight
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q_weight);
+    TensorMatcher({K})  // k_weight
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(k_weight);
+    const auto device = device_.unwrap();
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    // use at most `world_size` blocks to clean up,
+    // this is based on the observation that occupancy is usually linear
+    // with respect to the world size
+    const bool need_clean = num_tokens < m_max_num_cta_push;
+    const auto num_clean = need_clean ? (m_max_num_cta_push - num_tokens) : 0;
+    const auto num_blocks = need_clean ? num_tokens + div_ceil(num_clean, Trait::kBlockSize)  //
+                                       : m_max_num_cta_push;                                  //
+    const auto num_threads = Trait::kBlockSize;
+    RuntimeCheck(num_blocks <= m_max_num_cta_push, "internal error");
+    ParallelQKNormParams params;
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      params.buffer[i] = get_push_buffer(m_peer_storage[i]);
+    }
+    params.q_ptr = q.data_ptr();
+    params.k_ptr = k.data_ptr();
+    params.q_weight = q_weight.data_ptr();
+    params.k_weight = k_weight.data_ptr();
+    params.q_stride_bytes = q.stride(0) * sizeof(DType);
+    params.k_stride_bytes = k.stride(0) * sizeof(DType);
+    params.eps = eps / kNumGPU;  // scale down eps by number of GPUs
+    params.rank = m_rank;
+    params.num_tokens = num_tokens;
+    params.epoch_bytes = m_push_buffer_bytes;
+    params.num_clean_up_count = num_clean;
+
+    const auto needed_buffer_bytes = static_cast<int64_t>(num_tokens) * 2 * sizeof(float);
+    RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch");
+    RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized");
+    RuntimeCheck(std::bit_cast<intptr_t>(params.q_ptr) % 16 == 0, "q pointer is not properly aligned");
+    RuntimeCheck(std::bit_cast<intptr_t>(params.k_ptr) % 16 == 0, "k pointer is not properly aligned");
+    RuntimeCheck(std::bit_cast<intptr_t>(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned");
+    RuntimeCheck(std::bit_cast<intptr_t>(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned");
+    RuntimeCheck(needed_buffer_bytes <= m_push_buffer_bytes, "Push buffer is too small");
+
+    LaunchKernel(num_blocks, num_threads, device)  //
+        .enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl);
+  }
+
+  static uint32_t get_max_occupancy() {
+    return host::runtime::get_blocks_per_sm(kernel, Trait::kBlockSize);
+  }
+
+  static void
+  run(CustomAllReduceRef obj,
+      const tvm::ffi::Tensor q,
+      const tvm::ffi::Tensor k,
+      const tvm::ffi::Tensor q_weight,
+      const tvm::ffi::Tensor k_weight,
+      const float eps) {
+    using Self = FusedParallelQKNormAcrossHead;
+    return static_cast<Self*>(obj.get())->_run(q, k, q_weight, k_weight, eps);
+  }
+};
+
+}  // namespace