[LoRA][I] Add MOE LoRA JIT alignment kernel and tests (#19710)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Jonah Bernard <96398205+Jonahcb@users.noreply.github.com>

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

@@ -0,0 +1,618 @@
+// 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];
+    if (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

python/sglang/jit_kernel/moe_lora_align.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_moe_align_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "moe_lora_align_block_size",
+        *args,
+        cuda_files=["lora/moe_lora_align_kernel.cu"],
+        cuda_wrappers=[
+            ("moe_lora_align_block_size", f"MoeLoraAlignBlockSizeKernel<{args}>::run"),
+        ],
+    )
+
+
+def moe_lora_align_block_size(
+    topk_ids: torch.Tensor,
+    seg_indptr: torch.Tensor,
+    req_to_lora: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    max_loras: int,
+    max_num_tokens_padded: int,
+    max_num_m_blocks: int,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+    adapter_enabled: torch.Tensor,
+    lora_ids: torch.Tensor,
+    maybe_expert_map: Optional[torch.Tensor] = None,
+) -> None:
+    module = _jit_moe_align_module(topk_ids.dtype)
+
+    cumsum_buffer = torch.zeros(
+        max_loras * (num_experts + 1), dtype=torch.int32, device=topk_ids.device
+    )
+    token_mask = torch.empty(
+        (max_loras * topk_ids.shape[0],), dtype=torch.int32, device=topk_ids.device
+    )
+
+    module.moe_lora_align_block_size(
+        topk_ids,
+        seg_indptr,
+        req_to_lora,
+        num_experts,
+        block_size,
+        max_loras,
+        max_num_tokens_padded,
+        max_num_m_blocks,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        adapter_enabled,
+        lora_ids,
+        maybe_expert_map,
+        cumsum_buffer,
+        token_mask,
+    )

python/sglang/jit_kernel/tests/test_moe_lora_align_block_size.py

@@ -0,0 +1,166 @@
+# Temporarily adapted from https://github.com/vllm-project/vllm/blob/main/tests/lora/test_moe_lora_align_sum.py, will optimize in future refactor
+import random
+
+import pytest
+import torch
+
+# ---------------------------------------------------------
+# IMPORT PREBUILT KERNEL
+# ---------------------------------------------------------
+from sglang.jit_kernel.moe_lora_align import moe_lora_align_block_size
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=80, suite="stage-b-test-large-1-gpu")
+
+
+def round_up(x, base):
+    return ((x + base - 1) // base) * base
+
+
+def CEILDIV(x, y):
+    return (x + y - 1) // y
+
+
+def sample_data(num_experts, max_loras, num_tokens, topk_num):
+    # 1. Generate TopK IDs (Flattened tokens)
+    topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32)
+    for i in range(num_tokens):
+        pool = list(range(num_experts))
+        random.shuffle(pool)
+        for j in range(topk_num):
+            topk_ids[i, j] = pool[j]
+
+    # 2. Generate Random Requests (Segments)
+    # We split num_tokens into random chunks to simulate a batch of requests
+    remaining_tokens = num_tokens
+    seg_lens = []
+    while remaining_tokens > 0:
+        # Random length between 1 and remaining
+        length = random.randint(1, min(32, remaining_tokens))
+        if remaining_tokens - length < 0:
+            length = remaining_tokens
+        seg_lens.append(length)
+        remaining_tokens -= length
+
+    # Ensure we cover the full range exactly (cleanup last segment)
+    if sum(seg_lens) < num_tokens:
+        seg_lens.append(num_tokens - sum(seg_lens))
+
+    # 3. Build seg_indptr [0, len1, len1+len2, ...]
+    seg_indptr = torch.cumsum(
+        torch.tensor([0] + seg_lens, dtype=torch.int32), dim=0
+    ).to(dtype=torch.int32)
+
+    # 4. Assign a LoRA ID to each Request
+    num_reqs = len(seg_lens)
+    req_to_lora = torch.randint(0, max_loras, (num_reqs,), dtype=torch.int32)
+
+    return (topk_ids.to("cuda"), seg_indptr.to("cuda"), req_to_lora.to("cuda"))
+
+
+@pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096])
+@pytest.mark.parametrize("topk_num", [6])
+@pytest.mark.parametrize("num_experts", [64, 128, 256, 512])
+@pytest.mark.parametrize("max_loras", [2, 32])
+@pytest.mark.parametrize("block_size", [16])
+def test_moe_lora_align_block_size(
+    num_tokens, topk_num, num_experts, max_loras, block_size
+):
+    # sample data
+    random.seed(1)
+    torch.manual_seed(1)
+
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available, skipping moe_lora_align_block_size test.")
+    # UPDATED: Get the new 3-step mapping tensors
+    topk_ids, seg_indptr, req_to_lora = sample_data(
+        num_experts, max_loras, num_tokens, topk_num
+    )
+
+    # compute paddings
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
+    max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
+
+    # init output tensors
+    sorted_token_ids = torch.full(
+        (max_loras * max_num_tokens_padded,),
+        topk_ids.numel(),
+        dtype=torch.int32,
+        device="cuda",
+    )
+    expert_ids = torch.full(
+        (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda"
+    )
+    num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda")
+    adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda")
+    lora_ids = torch.arange(max_loras, dtype=torch.int32, device="cuda")
+
+    # UPDATED: Call kernel with new signature
+    moe_lora_align_block_size(
+        topk_ids,
+        seg_indptr,  # Arg 2: Pointers
+        req_to_lora,  # Arg 3: Request Map
+        num_experts,
+        block_size,
+        max_loras,
+        max_num_tokens_padded,
+        max_num_m_blocks,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        adapter_enabled,
+        lora_ids,
+        None,
+    )
+
+    # verify values
+    expert_ids = expert_ids.view(max_loras, -1)
+    sorted_token_ids = sorted_token_ids.view(max_loras, -1, block_size)
+
+    # Reconstruct token-level ownership for verification logic
+    # We expand req_to_lora back to [num_tokens] on CPU just to check correctness
+    # This proves the kernel (which used the compressed format) produced the right result
+    cpu_seg_indptr = seg_indptr.cpu()
+    cpu_req_to_lora = req_to_lora.cpu()
+    token_ownership = torch.zeros(num_tokens, dtype=torch.int32)
+
+    for r in range(len(cpu_req_to_lora)):
+        start = cpu_seg_indptr[r]
+        end = cpu_seg_indptr[r + 1]
+        token_ownership[start:end] = cpu_req_to_lora[r]
+
+    token_ownership = token_ownership.to("cuda")
+
+    for lora_idx in range(max_loras):
+        # Count how many tokens actually belong to this LoRA
+        expected_count = (token_ownership == lora_idx).sum().item()
+
+        # Verify the kernel processed a reasonable number of tokens (sanity check)
+        # Note: num_tokens_post_pad includes padding, so it might be larger than expected_count
+        assert num_tokens_post_pad[lora_idx].item() >= expected_count * topk_num
+
+        for token_idx in range(sorted_token_ids.size(1)):
+            block = sorted_token_ids[lora_idx][token_idx]
+            # Valid indices are those less than total numel
+            indices = block[block != topk_ids.numel()]
+
+            if indices.numel() > 0:
+                # 1. Verify routing: Does the token actually route to this expert?
+                expert_id = expert_ids[lora_idx][token_idx]
+                assert torch.all(topk_ids.view(-1)[indices] == expert_id)
+
+                # 2. Verify ownership: Did the kernel grab the correct tokens for this LoRA?
+                # The indices in 'sorted_token_ids' point to the flattened [token, topk] array.
+                # We divide by topk_num to get the original token index.
+                original_token_indices = indices // topk_num
+
+                # Check that all tokens in this block truly belong to 'lora_idx'
+                actual_owners = token_ownership[original_token_indices]
+                assert torch.all(
+                    actual_owners == lora_idx
+                ), f"Kernel put tokens from LoRA {actual_owners} into block for LoRA {lora_idx}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

python/sglang/jit_kernel/utils.py

@@ -94,6 +94,7 @@ CPP_DTYPE_MAP = {
     torch.float8_e4m3fn: "fp8_e4m3_t",
     torch.bfloat16: "bf16_t",
     torch.int8: "int8_t",
+    torch.int32: "int32_t",
     torch.int64: "int64_t",
 }