chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/lora/moe_lora_align_kernel.cu

@@ -0,0 +1,620 @@
+// Temporarily adapted from https://github.com/vllm-project/vllm/blob/main/csrc/moe/moe_align_sum_kernels.cu, will
+// optimize in future refactor
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <cub/cub.cuh>
+#include <tvm/ffi/container/tensor.h>
+
+#include <algorithm>
+
+#ifndef WARP_SIZE
+#define WARP_SIZE 32
+#endif
+
+#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
+
+namespace moe {
+
+template <typename scalar_t>
+SGL_DEVICE void _moe_align_block_size(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t* __restrict__ expert_map,
+    int32_t num_experts,
+    int32_t padded_num_experts,
+    int32_t experts_per_warp,
+    int32_t block_size,
+    size_t numel,
+    int32_t* __restrict__ cumsum,
+    int32_t max_num_tokens_padded,
+    int32_t max_num_m_blocks,
+    int32_t model_offset,
+    int32_t inactive_expert_id,
+    int32_t topk_num,
+    int32_t* token_mask,
+    bool has_expert_map) {
+  extern __shared__ int32_t shared_counts[];
+
+  // Compute input buffer offsets. Typically these will all be 0, except when
+  // using Multi LoRA.
+  int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
+  int expert_ids_offset = max_num_m_blocks * model_offset;
+  int cumsum_offset = (num_experts + 1) * model_offset;
+
+  // Use separate threadblocks to fill sorted_token_ids.
+  // This is safe since the current kernel does not use sorted_token_ids.
+  if (blockIdx.x % 2) {
+    // Initialize sorted_token_ids with numel
+    for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += blockDim.x) {
+      sorted_token_ids[sorted_token_ids_offset + it] = static_cast<int32_t>(numel);
+    }
+    return;
+  }
+
+  const int warp_id = threadIdx.x / WARP_SIZE;
+  const int my_expert_start = warp_id * experts_per_warp;
+
+  for (int i = 0; i < experts_per_warp; ++i) {
+    if (my_expert_start + i < padded_num_experts) {
+      shared_counts[warp_id * experts_per_warp + i] = 0;
+    }
+  }
+
+  __syncthreads();
+
+  const size_t tid = threadIdx.x;
+  const size_t stride = blockDim.x;
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int expert_id = topk_ids[i];
+    if (expert_id < 0 || expert_id >= num_experts) {
+      continue;
+    }
+    if (has_expert_map) {
+      expert_id = expert_map[expert_id];
+      if (expert_id < 0 || expert_id >= num_experts) continue;
+    }
+    int warp_idx = expert_id / experts_per_warp;
+    int expert_offset = expert_id % experts_per_warp;
+    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
+    atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], mask);
+  }
+
+  __syncthreads();
+
+  // Compute prefix sum over token counts per expert
+  using BlockScan = cub::BlockScan<int32_t, 1024>;
+  __shared__ typename BlockScan::TempStorage temp_storage;
+
+  int expert_count = 0;
+  int expert_id = threadIdx.x;
+  if (expert_id < num_experts) {
+    int warp_idx = expert_id / experts_per_warp;
+    int expert_offset = expert_id % experts_per_warp;
+    expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset];
+    expert_count = CEILDIV(expert_count, block_size) * block_size;
+  }
+
+  int cumsum_val;
+  BlockScan(temp_storage).ExclusiveSum(expert_count, cumsum_val);
+  if (expert_id <= num_experts) {
+    cumsum[cumsum_offset + expert_id] = cumsum_val;
+  }
+
+  if (expert_id == num_experts) {
+    total_tokens_post_pad[model_offset] = cumsum_val;
+  }
+
+  __syncthreads();
+
+  if (threadIdx.x < num_experts) {
+    for (int i = cumsum[cumsum_offset + threadIdx.x]; i < cumsum[cumsum_offset + threadIdx.x + 1]; i += block_size) {
+      expert_ids[expert_ids_offset + i / block_size] = threadIdx.x;
+    }
+  }
+
+  // Fill remaining expert_ids with 0
+  const size_t fill_start_idx = cumsum[cumsum_offset + num_experts] / block_size + threadIdx.x;
+  for (size_t i = fill_start_idx; i < max_num_m_blocks; i += blockDim.x) {
+    expert_ids[expert_ids_offset + i] = inactive_expert_id;
+  }
+}
+
+template <typename scalar_t, int32_t fill_threads>
+SGL_DEVICE void _moe_align_block_size_small_batch_expert(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t* __restrict__ expert_map,
+    int32_t num_experts,
+    int32_t block_size,
+    size_t numel,
+    int32_t max_num_tokens_padded,
+    int32_t max_num_m_blocks,
+    int32_t inactive_expert_id,
+    int32_t model_offset,
+    int32_t topk_num,
+    int32_t* token_mask,
+    bool has_expert_map) {
+  // Compute input buffer offsets. Typically these will all be 0, except when
+  // using Multi LoRA.
+  int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
+  int expert_ids_offset = max_num_m_blocks * model_offset;
+
+  // Use an additional group of threads to fill sorted_token_ids.
+  // Since the current kernel will use sorted_token_ids afterward,
+  // we fill sorted_token_ids within the same threadblock to make
+  // synchronization easier.
+  if (threadIdx.x < fill_threads) {
+    // Initialize sorted_token_ids with numel
+    for (size_t it = threadIdx.x; it < max_num_tokens_padded; it += fill_threads) {
+      sorted_token_ids[sorted_token_ids_offset + it] = static_cast<int32_t>(numel);
+    }
+    // Three __syncthreads() corresponding to the other threads
+    __syncthreads();
+    __syncthreads();
+    __syncthreads();
+    return;
+  }
+
+  const size_t tid = threadIdx.x - fill_threads;
+  const size_t stride = blockDim.x - fill_threads;
+
+  extern __shared__ int32_t shared_mem[];
+  int32_t* cumsum = shared_mem;
+  int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1);
+
+  for (int i = 0; i < num_experts; ++i) {
+    tokens_cnts[(tid + 1) * num_experts + i] = 0;
+  }
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (expert_id < 0 || expert_id >= num_experts) continue;
+    if (has_expert_map) {
+      expert_id = expert_map[expert_id];
+      if (expert_id < 0 || expert_id >= num_experts) continue;
+    }
+    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
+    tokens_cnts[(tid + 1) * num_experts + expert_id] += mask;
+  }
+
+  __syncthreads();
+
+  if (tid < num_experts) {
+    tokens_cnts[tid] = 0;
+    for (int i = 1; i <= stride; ++i) {
+      tokens_cnts[i * num_experts + tid] += tokens_cnts[(i - 1) * num_experts + tid];
+    }
+  }
+
+  __syncthreads();
+
+  if (tid == 0) {
+    cumsum[0] = 0;
+    for (int i = 1; i <= num_experts; ++i) {
+      cumsum[i] = cumsum[i - 1] + CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) * block_size;
+    }
+    total_tokens_post_pad[model_offset] = static_cast<int32_t>(cumsum[num_experts]);
+  }
+
+  __syncthreads();
+
+  if (tid < num_experts) {
+    for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) {
+      expert_ids[expert_ids_offset + i / block_size] = tid;
+    }
+  }
+
+  // Fill remaining expert_ids with 0
+  const size_t fill_start_idx = cumsum[num_experts] / block_size + tid;
+  for (size_t i = fill_start_idx; i < max_num_m_blocks; i += stride) {
+    expert_ids[expert_ids_offset + i] = inactive_expert_id;
+  }
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    if (expert_id < 0 || expert_id >= num_experts) continue;
+    if (has_expert_map) {
+      expert_id = expert_map[expert_id];
+      if (expert_id < 0 || expert_id >= num_experts) continue;
+    }
+    int32_t rank_post_pad = tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
+
+    if (token_mask == nullptr || token_mask[i / topk_num]) {
+      sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
+      ++tokens_cnts[tid * num_experts + expert_id];
+    }
+  }
+}
+
+template <typename scalar_t>
+SGL_DEVICE void _count_and_sort_expert_tokens(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ cumsum_buffer,
+    int32_t* __restrict__ expert_map,
+    size_t numel,
+    int32_t num_experts,
+    int32_t max_num_tokens_padded,
+    int32_t* __restrict__ token_mask,
+    int32_t model_offset,
+    int32_t topk_num,
+    bool has_expert_map) {
+  const size_t tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.y;
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    // Under EP, StandardDispatcher writes -1 for experts not owned by this
+    // rank; must filter the sentinel before indexing cumsum/sorted buffers.
+    if (expert_id < 0 || expert_id >= num_experts) {
+      continue;
+    }
+
+    if (has_expert_map) {
+      expert_id = expert_map[expert_id];
+      // filter invalid experts
+      if (expert_id == -1) continue;
+    }
+
+    if (token_mask == nullptr || token_mask[i / topk_num]) {
+      int32_t rank_post_pad = atomicAdd(&cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
+      sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] = i;
+    }
+  }
+}
+
+template <typename scalar_t>
+__global__ void moe_lora_align_block_size_kernel(
+    scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ seg_indptr,
+    int32_t* __restrict__ req_to_lora,
+    int num_reqs,
+    int64_t block_size,
+    int32_t* __restrict__ expert_map,
+    int num_experts,
+    int max_loras,
+    size_t numel,
+    int max_num_tokens_padded,
+    int max_num_m_blocks,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t topk_num,
+    int32_t* total_tokens_post_pad,
+    int32_t* adapter_enabled,
+    int32_t* __restrict__ cumsum,
+    int32_t experts_per_warp,
+    int32_t padded_num_experts,
+    int32_t* lora_ids,
+    int32_t* __restrict__ token_mask,
+    bool has_expert_map) {
+  int lora_idx = blockIdx.x / 2;
+  int lora_id = lora_ids[lora_idx];
+  if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
+    return;
+  }
+
+  int num_tokens = numel / topk_num;
+  int lora_offset = lora_id * num_tokens;
+
+  if (blockIdx.x % 2 == 0) {
+    // 1. Parallel Clear (Reset mask to 0)
+    for (int i = threadIdx.x; i < num_tokens; i += blockDim.x) {
+      token_mask[lora_offset + i] = 0;
+    }
+
+    if (threadIdx.x == 0) {
+      total_tokens_post_pad[lora_id] = 0;
+    }
+
+    __syncthreads();
+
+    // 2. Segment-based Fill
+    for (int r = 0; r < num_reqs; ++r) {
+      if (req_to_lora[r] == lora_id) {
+        int start = seg_indptr[r];
+        int end = seg_indptr[r + 1];
+        for (int i = start + threadIdx.x; i < end; i += blockDim.x) {
+          token_mask[lora_offset + i] = 1;
+        }
+      }
+    }
+
+    __syncthreads();
+  }
+
+  _moe_align_block_size(
+      topk_ids,
+      sorted_token_ids,
+      expert_ids,
+      total_tokens_post_pad,
+      expert_map,
+      num_experts,
+      padded_num_experts,
+      experts_per_warp,
+      block_size,
+      numel,
+      cumsum,
+      max_num_tokens_padded,
+      max_num_m_blocks,
+      lora_id,
+      -1,  // inactive_expert_id padding
+      topk_num,
+      &token_mask[(lora_id * num_tokens)],
+      has_expert_map);
+}
+
+template <typename scalar_t>
+__global__ void lora_count_and_sort_expert_tokens_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ cumsum_buffer,
+    int32_t* __restrict__ expert_map,
+    size_t numel,
+    int32_t num_experts,
+    int32_t max_num_tokens_padded,
+    int32_t topk_num,
+    int32_t* token_mask,
+    int32_t* lora_ids,
+    int32_t* adapter_enabled,
+    bool has_expert_map) {
+  int lora_idx = blockIdx.x;
+  int lora_id = lora_ids[lora_idx];
+  if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
+    return;
+  }
+
+  int num_tokens = numel / topk_num;
+
+  _count_and_sort_expert_tokens(
+      topk_ids,
+      sorted_token_ids,
+      cumsum_buffer,
+      expert_map,
+      numel,
+      num_experts,
+      max_num_tokens_padded,
+      &token_mask[(lora_id * num_tokens)],
+      lora_id,
+      topk_num,
+      has_expert_map);
+}
+
+template <typename scalar_t, int32_t fill_threads>
+__global__ void moe_lora_align_block_size_small_batch_expert_kernel(
+    scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ seg_indptr,
+    int32_t* __restrict__ req_to_lora,
+    int num_reqs,
+    int64_t block_size,
+    int32_t* __restrict__ expert_map,
+    int num_experts,
+    int max_loras,
+    size_t numel,
+    int max_num_tokens_padded,
+    int max_num_m_blocks,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int topk_num,
+    int32_t* total_tokens_post_pad,
+    int32_t* adapter_enabled,
+    int32_t* lora_ids,
+    int32_t* token_mask,
+    bool has_expert_map) {
+  int lora_idx = blockIdx.x;
+  int lora_id = lora_ids[lora_idx];
+  if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
+    return;
+  }
+
+  int num_tokens = numel / topk_num;
+  int lora_offset = lora_id * num_tokens;
+
+  // 1. Parallel Clear (Reset mask to 0)
+  // All threads help clear the mask for this adapter
+  for (int i = threadIdx.x; i < num_tokens; i += blockDim.x) {
+    token_mask[lora_offset + i] = 0;
+  }
+
+  // Initialize output counter
+  if (threadIdx.x == 0) {
+    total_tokens_post_pad[lora_id] = 0;
+  }
+
+  __syncthreads();
+
+  // 2. Segment-based Fill
+  // Iterate over requests. If a request matches this LoRA, fill its range.
+  for (int r = 0; r < num_reqs; ++r) {
+    if (req_to_lora[r] == lora_id) {
+      int start = seg_indptr[r];
+      int end = seg_indptr[r + 1];
+
+      // Parallel Fill: All threads help mark this segment as "1"
+      for (int i = start + threadIdx.x; i < end; i += blockDim.x) {
+        token_mask[lora_offset + i] = 1;
+      }
+    }
+  }
+
+  __syncthreads();
+
+  _moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
+      topk_ids,
+      sorted_token_ids,
+      expert_ids,
+      total_tokens_post_pad,
+      expert_map,
+      num_experts,
+      block_size,
+      numel,
+      max_num_tokens_padded,
+      max_num_m_blocks,
+      -1,  // inactive_expert_id padding
+      lora_id,
+      topk_num,
+      &token_mask[(lora_id * num_tokens)],
+      has_expert_map);
+}
+
+}  // namespace moe
+
+namespace {
+
+template <typename scalar_t>
+struct MoeLoraAlignBlockSizeKernel {
+  static void
+  run(tvm::ffi::TensorView topk_ids,
+      tvm::ffi::TensorView seg_indptr,
+      tvm::ffi::TensorView req_to_lora,
+      int64_t num_experts,
+      int64_t block_size,
+      int64_t max_loras,
+      int64_t max_num_tokens_padded,
+      int64_t max_num_m_blocks,
+      tvm::ffi::TensorView sorted_token_ids,
+      tvm::ffi::TensorView expert_ids,
+      tvm::ffi::TensorView num_tokens_post_pad,
+      tvm::ffi::TensorView adapter_enabled,
+      tvm::ffi::TensorView lora_ids,
+      tvm::ffi::Optional<tvm::ffi::TensorView> maybe_expert_map,
+      tvm::ffi::TensorView cumsum_buffer,
+      tvm::ffi::TensorView token_mask) {
+    using namespace host;
+
+    const int topk_num = topk_ids.size(1);
+
+    RuntimeCheck(block_size > 0, "block_size should be greater than 0. ");
+
+    int device_max_shared_mem;
+    auto device = topk_ids.device();
+    int dev_id = device.device_id;
+    RuntimeDeviceCheck(cudaDeviceGetAttribute(&device_max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev_id));
+    const cudaStream_t stream = LaunchKernel::resolve_device(device);
+
+    int64_t padded_num_experts = ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+
+    // BlockScan uses 1024 threads and assigns one thread per expert.
+    RuntimeCheck(padded_num_experts < 1024, "padded_num_experts must be less than 1024");
+
+    int32_t* token_mask_ptr = static_cast<int32_t*>(token_mask.data_ptr());
+
+    bool has_expert_map = maybe_expert_map.has_value();
+    int32_t* expert_map_ptr = nullptr;
+    if (has_expert_map) {
+      expert_map_ptr = static_cast<int32_t*>(maybe_expert_map.value().data_ptr());
+    }
+    int num_reqs = seg_indptr.size(0) - 1;
+
+    bool small_batch_expert_mode = (topk_ids.numel() < 1024) && (num_experts <= 64);
+
+    if (small_batch_expert_mode) {
+      const int32_t num_thread = std::max((int32_t)num_experts, 128);
+      const int32_t shared_mem = (num_thread + 1) * num_experts * sizeof(int32_t) + (num_experts + 1) * sizeof(int32_t);
+      if (shared_mem > device_max_shared_mem) {
+        RuntimeCheck(false, "Shared memory usage exceeds device limit.");
+      }
+
+      // threadIdx.x >= fill_threads: counting experts and aligning
+      // threadIdx.x < fill_threads: filling sorted_token_ids
+      constexpr int32_t fill_threads = 256;
+
+      dim3 blockDim(num_thread + fill_threads);
+      auto kernel = moe::moe_lora_align_block_size_small_batch_expert_kernel<scalar_t, fill_threads>;
+      RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem));
+
+      LaunchKernel(dim3(max_loras), blockDim, stream, shared_mem)(
+          kernel,
+          static_cast<scalar_t*>(topk_ids.data_ptr()),
+          static_cast<int32_t*>(seg_indptr.data_ptr()),
+          static_cast<int32_t*>(req_to_lora.data_ptr()),
+          num_reqs,
+          block_size,
+          expert_map_ptr,
+          num_experts,
+          max_loras,
+          topk_ids.numel(),
+          max_num_tokens_padded,
+          max_num_m_blocks,
+          static_cast<int32_t*>(sorted_token_ids.data_ptr()),
+          static_cast<int32_t*>(expert_ids.data_ptr()),
+          topk_num,
+          static_cast<int32_t*>(num_tokens_post_pad.data_ptr()),
+          static_cast<int32_t*>(adapter_enabled.data_ptr()),
+          static_cast<int32_t*>(lora_ids.data_ptr()),
+          token_mask_ptr,
+          has_expert_map);
+
+    } else {
+      int num_thread = 1024;
+      dim3 blockDim(num_thread);
+      size_t num_warps = CEILDIV(padded_num_experts, WARP_SIZE);
+
+      size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t);
+
+      auto align_kernel = moe::moe_lora_align_block_size_kernel<scalar_t>;
+
+      // launch two threadblocks for each lora
+      // blockIdx.x % 2 == 0: counting experts and aligning
+      // blockIdx.x % 2 == 1: filling sorted_token_ids
+      LaunchKernel(dim3(max_loras * 2), blockDim, stream, shared_mem_size)(
+          align_kernel,
+          static_cast<scalar_t*>(topk_ids.data_ptr()),
+          static_cast<int32_t*>(seg_indptr.data_ptr()),
+          static_cast<int32_t*>(req_to_lora.data_ptr()),
+          num_reqs,
+          block_size,
+          expert_map_ptr,
+          num_experts,
+          max_loras,
+          topk_ids.numel(),
+          max_num_tokens_padded,
+          max_num_m_blocks,
+          static_cast<int32_t*>(sorted_token_ids.data_ptr()),
+          static_cast<int32_t*>(expert_ids.data_ptr()),
+          topk_num,
+          static_cast<int32_t*>(num_tokens_post_pad.data_ptr()),
+          static_cast<int32_t*>(adapter_enabled.data_ptr()),
+          static_cast<int32_t*>(cumsum_buffer.data_ptr()),
+          WARP_SIZE,
+          padded_num_experts,
+          static_cast<int32_t*>(lora_ids.data_ptr()),
+          token_mask_ptr,
+          has_expert_map);
+
+      const int block_threads = std::min(256, (int)num_thread);
+      const int num_blocks = (topk_ids.numel() + block_threads - 1) / block_threads;
+
+      const int max_blocks = 65535;
+      const int actual_blocks = std::min(num_blocks, max_blocks);
+
+      dim3 gridDims(max_loras, actual_blocks);
+      auto sort_kernel = moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;
+
+      LaunchKernel(gridDims, dim3(block_threads), stream)(
+          sort_kernel,
+          static_cast<scalar_t*>(topk_ids.data_ptr()),
+          static_cast<int32_t*>(sorted_token_ids.data_ptr()),
+          static_cast<int32_t*>(cumsum_buffer.data_ptr()),
+          expert_map_ptr,
+          topk_ids.numel(),
+          num_experts,
+          max_num_tokens_padded,
+          topk_num,
+          token_mask_ptr,
+          static_cast<int32_t*>(lora_ids.data_ptr()),
+          static_cast<int32_t*>(adapter_enabled.data_ptr()),
+          has_expert_map);
+    }
+  }
+};
+
+}  // namespace