[Feature][JIT Kernel] Fused TP QK norm For Minimax (#20673)

Co-authored-by: Mingyang Jiang <13463932+jmydurant@users.noreply.github.com>

python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh

@@ -48,6 +48,8 @@ struct alignas(128) Semaphore {
 
 struct PullController {
  public:
+  using SignalType = Semaphore;
+
   PullController(void** signals, uint32_t num_gpu) {
     for (uint32_t i = 0; i < num_gpu; ++i) {
       m_signals[i] = static_cast<Semaphore*>(signals[i]);
@@ -90,25 +92,29 @@ struct PullController {
 
 struct PushController {
  public:
+  using SignalType = uint32_t;
   static constexpr int64_t kNumStages = 2;
 
-  PushController(void* ptr) : m_local_signal(static_cast<Semaphore*>(ptr)) {}
+  PushController(void* ptr) : m_local_signal(static_cast<SignalType*>(ptr)) {}
 
-  SGL_DEVICE uint32_t epoch() const {
-    return m_local_signal[blockIdx.x].get_counter();
+  SGL_DEVICE SignalType epoch() const {
+    return m_local_signal[blockIdx.x];
   }
 
   SGL_DEVICE void exit() const {
     __syncthreads();
     if (threadIdx.x == 0) {
-      auto& signal = m_local_signal[blockIdx.x];
-      const auto epoch = signal.get_counter();
-      signal.set_counter((epoch + 1) % kNumStages);
+      this->exit_unsafe(blockIdx.x);
     }
   }
 
+  SGL_DEVICE void exit_unsafe(uint32_t which) const {
+    auto& signal = m_local_signal[which];
+    signal = (signal + 1) % kNumStages;
+  }
+
  private:
-  Semaphore* m_local_signal;
+  SignalType* m_local_signal;
 };
 
 }  // namespace device::distributed

python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh

@@ -93,12 +93,14 @@ struct CustomAllReduceBase : public tvm::ffi::Object {
         // default config for pull kernel, can be updated by `configure()`
         m_num_cta(max_num_cta_pull),
         m_cta_size(256) {
-    RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes()));
+    RuntimeCheck(pull_buffer_size % 128 == 0, "Pull buffer size should be aligned to 128 bytes");
+    RuntimeCheck(push_buffer_size % 128 == 0, "Push buffer size should be aligned to 128 bytes");
     RuntimeCheck(rank < num_gpu, "Invalid rank: ", rank);
     const int64_t kU32Max = static_cast<int64_t>(std::numeric_limits<uint32_t>::max());
     const int64_t push_buffer_size_all = push_all_ranks_bytes();
-    RuntimeCheck(pull_buffer_size <= kU32Max, "Buffer size is too large: ", pull_buffer_size);
+    RuntimeCheck(pull_buffer_size <= kU32Max, "Pull buffer size is too large: ", pull_buffer_size);
     RuntimeCheck(push_buffer_size_all <= kU32Max, "Push buffer size is too large: ", push_buffer_size_all);
+    RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes()));
   }
 
   ExternHandle share_storage() {
@@ -252,19 +254,18 @@ struct CustomAllReduceBase : public tvm::ffi::Object {
     return static_cast<int64_t>(m_graph_capture_inputs.size());
   }
   int64_t pull_signal_bytes() const {
-    return sizeof(device::distributed::Semaphore) * m_max_num_cta_pull;
+    return _align_bytes(sizeof(PullController::SignalType) * m_max_num_cta_pull);
   }
   int64_t push_signal_bytes() const {
-    return sizeof(device::distributed::Semaphore) * m_max_num_cta_push;
+    return _align_bytes(sizeof(PushController::SignalType) * m_max_num_cta_push);
   }
-  int64_t params_bytes() const {
-    return sizeof(AllReduceData) * (1 + m_graph_buffer_count);  // 1 for default
+  int64_t graph_param_bytes() const {
+    return _align_bytes(sizeof(AllReduceData) * (1 + m_graph_buffer_count));  // 1 for default
   }
   int64_t push_all_ranks_bytes() const {
-    return PushController::kNumStages * m_num_gpu * m_push_buffer_bytes;
+    return _align_bytes(PushController::kNumStages * m_num_gpu * m_push_buffer_bytes);
   }
   int64_t storage_bytes() const {
-    // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) |
     return _get_offset_impl(5);
   }
   void* get_pull_signal(void* ptr) const {
@@ -283,16 +284,20 @@ struct CustomAllReduceBase : public tvm::ffi::Object {
     return pointer::offset(ptr, _get_offset_impl(4));
   }
   int64_t _get_offset_impl(int64_t which) const {
+    // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) |
     const int64_t offset_map[5] = {
         /*[0]=*/pull_signal_bytes(),
         /*[1]=*/push_signal_bytes(),
-        /*[2]=*/params_bytes(),
+        /*[2]=*/graph_param_bytes(),
         /*[3]=*/m_pull_buffer_bytes,
         /*[4]=*/push_all_ranks_bytes(),
     };
     RuntimeCheck(which >= 0 && which <= 5, "Invalid offset index: ", which);
     return std::accumulate(offset_map, offset_map + which, int64_t(0));
   }
+  static int64_t _align_bytes(int64_t size) {
+    return div_ceil(size, 128) * 128;
+  }
 
   const int64_t m_pull_buffer_bytes;
   const int64_t m_push_buffer_bytes;

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

@@ -21,10 +21,12 @@ static constexpr uint32_t kFullMask = 0xffffffffu;
  * \param active_mask Bitmask of participating lanes (default: all 32).
  * \return The sum across all active lanes.
  */
-template <typename T>
+template <uint32_t kNumThreads = kWarpThreads, typename T>
 SGL_DEVICE T reduce_sum(T value, uint32_t active_mask = kFullMask) {
+  static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads);
+  static_assert(std::has_single_bit(kNumThreads), "must be pow of 2");
 #pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1)
+  for (int mask = kNumThreads / 2; mask > 0; mask >>= 1)
     value = value + __shfl_xor_sync(active_mask, value, mask, 32);
   return value;
 }